[Bug] SDXL with Self-Attention Guidance Compiled Error

[blacklong28]

Your current environment information

Collecting environment information...
PyTorch version: 2.4.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OneFlow version: path: ['/usr/local/lib/python3.10/dist-packages/oneflow'], version: 0.9.1.dev20241202+cu122, git_commit: 4710a9b, cmake_build_type: Release, rdma: True, mlir: True, enterprise: False
Nexfort version: 0.1.dev299
OneDiff version: 1.2.1.dev7
OneDiffX version: 1.2.1.dev7

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.30.0
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 6 2024, 20:22:13) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-101-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA L20
Nvidia driver version: 535.183.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 52 bits physical, 57 bits virtual
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) Gold 6462C
CPU family: 6
Model: 143
Thread(s) per core: 2
Core(s) per socket: 8
Socket(s): 1
Stepping: 8
CPU max MHz: 3900.0000
CPU min MHz: 800.0000
BogoMIPS: 6600.00
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ibrs_enhanced fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 wbnoinvd ida arat hwp hwp_notify hwp_act_window hwp_epp hwp_pkg_req avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid bus_lock_detect cldemote movdiri movdir64b enqcmd fsrm md_clear serialize tsxldtrk amx_bf16 avx512_fp16 amx_tile amx_int8 arch_capabilities
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 384 KiB (8 instances)
L1i cache: 256 KiB (8 instances)
L2 cache: 16 MiB (8 instances)
L3 cache: 60 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Unknown: No mitigations
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced IBRS, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected

Versions of relevant libraries:
[pip3] clip-anytorch==2.6.0
[pip3] dctorch==0.1.2
[pip3] diffusers==0.28.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.0
[pip3] onnx==1.15.0
[pip3] onnx-graphsurgeon==0.5.2
[pip3] onnx2torch==1.5.15
[pip3] onnxconverter-common==1.14.0
[pip3] onnxmltools==1.12.0
[pip3] onnxruntime==1.17.3
[pip3] onnxruntime_extensions==0.12.0
[pip3] onnxruntime-gpu==1.18.0
[pip3] open-clip-torch==2.24.0
[pip3] optree==0.13.0
[pip3] pytorch-lightning==1.2.9
[pip3] pytorch-triton==3.1.0+cf34004b8a
[pip3] sentence-transformers==3.1.0
[pip3] stable-fast==1.0.5+torch230cu121
[pip3] taming-transformers==0.0.1
[pip3] torch==2.4.0
[pip3] torch_tensorrt==2.3.0
[pip3] torchao==0.1
[pip3] torchaudio==2.4.0
[pip3] torchdiffeq==0.2.4
[pip3] torchmetrics==0.2.0
[pip3] torchscale==0.3.0
[pip3] torchsde==0.2.6
[pip3] torchvision==0.19.0
[pip3] transformers==4.38.2
[pip3] triton==3.0.0
[conda] Could not collect

🐛 Describe the bug

When I use onediff to accelerate this comfyui workflow(
sdxl_sampler_debug.json
), I get the following error:

Cache lookup: Key='5d34d8fb-52f4-4203-84aa-46e62fe4433b', Cached Model Type='<class 'NoneType'>'
#12 [ModelSpeedup]: 0.00s
#11 [SelfAttentionGuidance]: 0.00s
Requested to load SDXL
Loading 1 new model
  0%|                                                                                                                                                                                               | 0/10 [00:00<?, ?it/s][ERROR](GRAPH:OneflowGraph_0:OneflowGraph) building graph got error.
ERROR [2024-12-04 10:32:08] /usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py:44 - Exception in forward: e=TypeError('expected Tensor as element 0 in argument 1, but got Tensor')
WARNING [2024-12-04 10:32:08] /usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py:45 - Recompile oneflow module ...
[ERROR](GRAPH:OneflowGraph_1:OneflowGraph) building graph got error.
  0%|                                                                                                                                                                                               | 0/10 [00:07<?, ?it/s]
!!! Exception during processing !!! expected Tensor as element 0 in argument 1, but got Tensor
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py", line 42, in wrapper
    return func(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/online_quantization_utils.py", line 65, in wrapper
    output = func(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/graph_management_utils.py", line 120, in wrapper
    ret = func(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/args_tree_util.py", line 70, in wrapper
    output = func(self, *mapped_args, **mapped_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py", line 154, in forward
    output = dpl_graph(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 295, in __call__
    self._compile(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 861, in _compile
    return self._dynamic_input_graph_cache._compile(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/cache.py", line 121, in _compile
    return graph._compile(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 865, in _compile
    return self._compile_new(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 884, in _compile_new
    _, eager_outputs = self.build_graph(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 1429, in build_graph
    outputs = self.__build_graph(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 1577, in __build_graph
    outputs = self.build(*lazy_args, **lazy_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/graph.py", line 19, in build
    return self.model(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "/root/project/ComfyUI/custom_nodes/onediff_comfy_nodes/modules/oneflow/infer_compiler_registry/register_comfy/openaimodel.py", line 119, in forward
    h = forward_timestep_embed(
  File "/root/project/ComfyUI/comfy/ldm/modules/diffusionmodules/openaimodel.py", line 44, in forward_timestep_embed
    x = layer(x, context, transformer_options)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "/root/project/ComfyUI/custom_nodes/onediff_comfy_nodes/modules/oneflow/infer_compiler_registry/register_comfy/attention.py", line 130, in forward
    x = block(x, context=context[i], transformer_options=transformer_options)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "/root/project/ComfyUI/comfy/ldm/modules/attention.py", line 593, in forward
    n = attn1_replace_patch[block_attn1](n, context_attn1, value_attn1, extra_options)
  File "/root/project/ComfyUI/comfy_extras/nodes_sag.py", line 138, in attn_and_record
    (out, sim) = attention_basic_with_sim(q, k, v, heads=heads, attn_precision=extra_options["attn_precision"])
  File "/root/project/ComfyUI/comfy_extras/nodes_sag.py", line 31, in attention_basic_with_sim
    sim = einsum('b i d, b j d -> b i j', q, k) * scale
  File "/usr/local/lib/python3.10/dist-packages/torch/functional.py", line 386, in einsum
    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
TypeError: expected Tensor as element 0 in argument 1, but got Tensor

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/root/project/ComfyUI/execution.py", line 323, in execute
    output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
  File "/root/project/ComfyUI/execution.py", line 198, in get_output_data
    return_values = _map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
  File "/root/project/ComfyUI/execution.py", line 169, in _map_node_over_list
    process_inputs(input_dict, i)
  File "/root/project/ComfyUI/execution.py", line 158, in process_inputs
    results.append(getattr(obj, func)(**inputs))
  File "/root/project/ComfyUI/nodes.py", line 1454, in sample
    return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
  File "/root/project/ComfyUI/nodes.py", line 1421, in common_ksampler
    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
  File "/root/project/ComfyUI/custom_nodes/ComfyUI-Impact-Pack/modules/impact/sample_error_enhancer.py", line 9, in informative_sample
    return original_sample(*args, **kwargs)  # This code helps interpret error messages that occur within exceptions but does not have any impact on other operations.
  File "/root/project/ComfyUI/custom_nodes/ComfyUI-Advanced-ControlNet/adv_control/sampling.py", line 116, in acn_sample
    return orig_comfy_sample(model, *args, **kwargs)
  File "/root/project/ComfyUI/custom_nodes/ComfyUI-Advanced-ControlNet/adv_control/utils.py", line 117, in uncond_multiplier_check_cn_sample
    return orig_comfy_sample(model, *args, **kwargs)
  File "/root/project/ComfyUI/comfy/sample.py", line 43, in sample
    samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
  File "/root/project/ComfyUI/custom_nodes/ComfyUI-TiledDiffusion/utils.py", line 51, in KSampler_sample
    return orig_fn(*args, **kwargs)
  File "/root/project/ComfyUI/comfy/samplers.py", line 877, in sample
    return sample(self.model, noise, positive, negative, cfg, self.device, sampler, sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
  File "/root/project/ComfyUI/comfy/samplers.py", line 775, in sample
    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
  File "/root/project/ComfyUI/comfy/samplers.py", line 762, in sample
    output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
  File "/root/project/ComfyUI/comfy/samplers.py", line 719, in inner_sample
    samples = sampler.sample(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
  File "/root/project/ComfyUI/custom_nodes/ComfyUI-TiledDiffusion/utils.py", line 34, in KSAMPLER_sample
    return orig_fn(*args, **kwargs)
  File "/root/project/ComfyUI/comfy/samplers.py", line 624, in sample
    samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/root/project/ComfyUI/comfy/k_diffusion/sampling.py", line 155, in sample_euler
    denoised = model(x, sigma_hat * s_in, **extra_args)
  File "/root/project/ComfyUI/comfy/samplers.py", line 299, in __call__
    out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
  File "/root/project/ComfyUI/comfy/samplers.py", line 706, in __call__
    return self.predict_noise(*args, **kwargs)
  File "/root/project/ComfyUI/comfy/samplers.py", line 709, in predict_noise
    return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)
  File "/root/project/ComfyUI/comfy/samplers.py", line 279, in sampling_function
    out = calc_cond_batch(model, conds, x, timestep, model_options)
  File "/root/project/ComfyUI/custom_nodes/onediff_comfy_nodes/modules/sd_hijack_utils.py", line 58, in hijacked_method
    return self(*args, **kwargs)
  File "/root/project/ComfyUI/custom_nodes/onediff_comfy_nodes/modules/sd_hijack_utils.py", line 95, in __call__
    return sub_func(self._orig_func, *args, **kwargs)
  File "/root/project/ComfyUI/custom_nodes/onediff_comfy_nodes/modules/oneflow/hijack_samplers.py", line 122, in calc_cond_batch_of
    output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
  File "/root/project/ComfyUI/custom_nodes/ComfyUI-Advanced-ControlNet/adv_control/utils.py", line 69, in apply_model_uncond_cleanup_wrapper
    return orig_apply_model(self, *args, **kwargs)
  File "/root/project/ComfyUI/comfy/model_base.py", line 187, in apply_model
    model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds)#.float()
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1603, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py", line 48, in wrapper
    return func(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/online_quantization_utils.py", line 65, in wrapper
    output = func(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/graph_management_utils.py", line 123, in wrapper
    ret = func(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/args_tree_util.py", line 70, in wrapper
    output = func(self, *mapped_args, **mapped_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py", line 154, in forward
    output = dpl_graph(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 295, in __call__
    self._compile(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 861, in _compile
    return self._dynamic_input_graph_cache._compile(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/cache.py", line 121, in _compile
    return graph._compile(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 865, in _compile
    return self._compile_new(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 884, in _compile_new
    _, eager_outputs = self.build_graph(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 1429, in build_graph
    outputs = self.__build_graph(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/graph.py", line 1577, in __build_graph
    outputs = self.build(*lazy_args, **lazy_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onediff/infer_compiler/backends/oneflow/graph.py", line 19, in build
    return self.model(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "/root/project/ComfyUI/custom_nodes/onediff_comfy_nodes/modules/oneflow/infer_compiler_registry/register_comfy/openaimodel.py", line 119, in forward
    h = forward_timestep_embed(
  File "/root/project/ComfyUI/comfy/ldm/modules/diffusionmodules/openaimodel.py", line 44, in forward_timestep_embed
    x = layer(x, context, transformer_options)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "/root/project/ComfyUI/custom_nodes/onediff_comfy_nodes/modules/oneflow/infer_compiler_registry/register_comfy/attention.py", line 130, in forward
    x = block(x, context=context[i], transformer_options=transformer_options)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "/root/project/ComfyUI/comfy/ldm/modules/attention.py", line 593, in forward
    n = attn1_replace_patch[block_attn1](n, context_attn1, value_attn1, extra_options)
  File "/root/project/ComfyUI/comfy_extras/nodes_sag.py", line 138, in attn_and_record
    (out, sim) = attention_basic_with_sim(q, k, v, heads=heads, attn_precision=extra_options["attn_precision"])
  File "/root/project/ComfyUI/comfy_extras/nodes_sag.py", line 31, in attention_basic_with_sim
    sim = einsum('b i d, b j d -> b i j', q, k) * scale
  File "/usr/local/lib/python3.10/dist-packages/torch/functional.py", line 386, in einsum
    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
TypeError: expected Tensor as element 0 in argument 1, but got Tensor