编译diffusion模型失败

[zhangvia]

Describe the bug

A clear and concise description of what the bug is.

Your environment

ubuntu 20.04

OneDiff git commit id

onediff-0.13.0.dev202403200123

OneFlow version info

version: 0.9.1.dev20240319+cu118
git_commit: a130148
cmake_build_type: Release
rdma: True
mlir: True

How To Reproduce

我在这个模型上尝试了编译repo.该模型也是基于diffusion方法 unet迭代生成图像的。转换时成功的，但在编译静态图时出现了无效输入参数的问题：

[ERROR](GRAPH:OneflowGraph_0:OneflowGraph) building graph got error.
2024-03-21 05:53:54,210 [ERROR] with_oneflow_compile.py:187 - Exception in __call__: e=TypeError('\x1b[1m\x1b[38;2;255;000;000mError\x1b[0m: less(): received an invalid combination of arguments. The valid signatures are:\n\t*0: Tensor (Tensor input, Tensor other)\n\t*1: Tensor (Tensor input, Scalar other)\n\t*2: Tensor (Scalar input, Tensor other)\n')
2024-03-21 05:53:54,210 [WARNING] with_oneflow_compile.py:188 - Recompile oneflow module ...
2024-03-21 05:53:54,210 [DEBUG] with_oneflow_compile.py:34 - Convert <class 'models.unet_autoenc.BeatGANsAutoencModel'> ...
2024-03-21 05:53:56,086 [DEBUG] with_oneflow_compile.py:36 - Convert <class 'models.unet_autoenc.BeatGANsAutoencModel'> done!
2024-03-21 05:53:56,122 [INFO] with_oneflow_compile.py:336 - Building a graph for <class 'models.unet_autoenc.BeatGANsAutoencModel'> ...
[ERROR](GRAPH:OneflowGraph_1:OneflowGraph) building graph got error.
2024-03-21 05:53:56,159 [INFO] core.py:570 - virtual tryon:  :-: 6.81
Traceback (most recent call last):
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/onediff/infer_compiler/with_oneflow_compile.py", line 185, in wrapper
return func(self, *args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/onediff/infer_compiler/utils/graph_management_utils.py", line 82, in wrapper
ret = func(self, *args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/onediff/infer_compiler/with_oneflow_compile.py", line 272, in __call__
output = dpl_graph(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 281, in __call__
return self._dynamic_input_graph_cache(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/cache.py", line 115, in __call__
return graph(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 284, in __call__
self._compile(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 852, in _compile
return self._compile_new(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 871, in _compile_new
_, eager_outputs = self.build_graph(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 1374, in build_graph
outputs = self.__build_graph(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 1491, in __build_graph
outputs = self.build(*lazy_args, **lazy_kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/onediff/infer_compiler/with_oneflow_compile.py", line 341, in build
return self.model(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 186, in __call__
result = self.__block_forward(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 236, in __block_forward
result = unbound_forward_of_module_instance(self, *args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/tryon/VirtualTryon/models/unet_autoenc.py", line 158, in forward
cond_mask = prob_mask_like((x.shape[0],), prob = prob, device = x.device)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/tryon/VirtualTryon/models/unet_autoenc.py", line 17, in prob_mask_like
return torch.zeros(shape, device = device).float().uniform_(0, 1) < prob
TypeError: Error: less(): received an invalid combination of arguments. The valid signatures are:
*0: Tensor (Tensor input, Tensor other)
*1: Tensor (Tensor input, Scalar other)
*2: Tensor (Scalar input, Tensor other)
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/cog/server/worker.py", line 221, in _predict
result = predict(**payload)
File "/media/74nvme/research/tryon/predict.py", line 337, in predict
sr_try, mask_nd_by, hair_nd, input_try, output_try, person_ary, portrait_mask = self.process(data, cloth_path, mote_params, gender,
File "/media/74nvme/research/tryon/predict.py", line 190, in process
rst = self.ArcVTO.generate(mode=t_model, random_seed=-1)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/tryon/VirtualTryon.py", line 178, in generate
rst = model.predict(garmentIn, personIn, poseIn, sample_algorithm=sample_algorithm, cond_scale=cond_scale, noise=noise, nsteps=num_steps)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/tryon/VirtualTryon/predict.py", line 56, in predict
xs, x0_preds = ddim_steps(noise, seq, self.model, self.betas.cuda(), [src, tgt_pose], cond_scale = cond_scale)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/tryon/VirtualTryon/diffusion.py", line 36, in ddim_steps
et = model.forward_with_cond_scale(x = torch.cat([xt, target_pose],1), t = t, cond = cond, cond_scale = cond_scale)[0]
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/tryon/VirtualTryon/models/unet_autoenc.py", line 123, in forward_with_cond_scale
logits = self.__call__(x, t, cond=cond, prob = 1)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/onediff/infer_compiler/utils/args_tree_util.py", line 50, in wrapper
output = func(self, *mapped_args, **mapped_kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/onediff/infer_compiler/with_oneflow_compile.py", line 191, in wrapper
return func(self, *args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/onediff/infer_compiler/utils/graph_management_utils.py", line 82, in wrapper
ret = func(self, *args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/onediff/infer_compiler/with_oneflow_compile.py", line 272, in __call__
output = dpl_graph(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 281, in __call__
return self._dynamic_input_graph_cache(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/cache.py", line 115, in __call__
return graph(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 284, in __call__
self._compile(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 852, in _compile
return self._compile_new(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 871, in _compile_new
_, eager_outputs = self.build_graph(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 1374, in build_graph
outputs = self.__build_graph(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 1491, in __build_graph
outputs = self.build(*lazy_args, **lazy_kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/onediff/infer_compiler/with_oneflow_compile.py", line 341, in build
return self.model(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 186, in __call__
result = self.__block_forward(*args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 236, in __block_forward
result = unbound_forward_of_module_instance(self, *args, **kwargs)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/tryon/VirtualTryon/models/unet_autoenc.py", line 158, in forward
cond_mask = prob_mask_like((x.shape[0],), prob = prob, device = x.device)
File "/media/74nvme/software/miniconda3/envs/stable-fast/lib/python3.10/site-packages/tryon/VirtualTryon/models/unet_autoenc.py", line 17, in prob_mask_like
return torch.zeros(shape, device = device).float().uniform_(0, 1) < prob
TypeError: Error: less(): received an invalid combination of arguments. The valid signatures are:
*0: Tensor (Tensor input, Tensor other)
*1: Tensor (Tensor input, Scalar other)
*2: Tensor (Scalar input, Tensor other)```