siliconflow/onediff

[Bug] KeyError: '_attn2' with IPAdapter

kvochko opened this issue · 2 comments

Your current environment information

PyTorch version: 2.1.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OneFlow version: path: ['.venv/lib/python3.10/site-packages/oneflow'], version: 0.9.1.dev20240803+cu122, git_commit: d23c061, cmake_build_type: Release, rdma: True, mlir: True, enterprise: False
Nexfort version: none
OneDiff version: 1.2.1.dev24+g9231f556
OneDiffX version: none

OS: Debian GNU/Linux 11 (bullseye) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: Could not collect
CMake version: version 3.18.4
Libc version: glibc-2.31

Python version: 3.10.12 (main, Jul  5 2023, 18:54:27) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.10.0-32-cloud-amd64-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 550.54.15
cuDNN version: Could not collect
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
Byte Order:                           Little Endian
Address sizes:                        46 bits physical, 48 bits virtual
CPU(s):                               12
On-line CPU(s) list:                  0-11
Thread(s) per core:                   2
Core(s) per socket:                   6
Socket(s):                            1
NUMA node(s):                         1
Vendor ID:                            GenuineIntel
CPU family:                           6
Model:                                85
Model name:                           Intel(R) Xeon(R) CPU @ 2.20GHz
Stepping:                             7
CPU MHz:                              2200.192
BogoMIPS:                             4400.38
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            192 KiB
L1i cache:                            192 KiB
L2 cache:                             6 MiB
L3 cache:                             38.5 MiB
NUMA node0 CPU(s):                    0-11
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:        Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Enhanced / Automatic IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
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 tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat avx512_vnni md_clear arch_capabilities

Versions of relevant libraries:
[pip3] diffusers==0.28.2
[pip3] ltx-pytorchnvcodec==0.1.2
[pip3] mypy==1.9.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] onnx==1.16.0
[pip3] onnx2torch==1.5.15
[pip3] onnxruntime==1.17.1
[pip3] onnxruntime-gpu==1.19.2
[pip3] open-clip-torch==2.24.0
[pip3] pytorch-lightning==2.2.1
[pip3] taming-transformers==0.0.1
[pip3] torch==2.1.2
[pip3] torchaudio==2.1.2
[pip3] torchmetrics==1.3.2
[pip3] torchsde==0.2.6
[pip3] torchvision==0.16.2
[pip3] transformers==4.39.3
[pip3] triton==2.1.0
[conda] No relevant packages

🐛 Describe the bug

When I use onediff with IPAdapter from ComfyUI, I get the following error. Minimal workflow to reproduce:

_attn2_key_error.json

[ERROR](GRAPH:OneflowGraph_1:OneflowGraph) building graph got error.
  0%|                                                                                                                                                         | 0/5 [00:13<?, ?it/s]
!!! Exception during processing !!! '_attn2'
Traceback (most recent call last):
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py", line 42, in wrapper
    return func(self, *args, **kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/online_quantization_utils.py", line 65, in wrapper
    output = func(self, *args, **kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/graph_management_utils.py", line 120, in wrapper
    ret = func(self, *args, **kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/args_tree_util.py", line 74, in wrapper
    else func(self, *mapped_args, **mapped_kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py", line 154, in forward
    output = dpl_graph(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 295, in __call__
    self._compile(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 861, in _compile
    return self._dynamic_input_graph_cache._compile(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/cache.py", line 121, in _compile
    return graph._compile(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 865, in _compile
    return self._compile_new(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 884, in _compile_new
    _, eager_outputs = self.build_graph(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 1429, in build_graph
    outputs = self.__build_graph(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 1577, in __build_graph
    outputs = self.build(*lazy_args, **lazy_kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/graph.py", line 19, in build
    return self.model(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "src/comfyui_custom_nodes/onediff_comfy_nodes/modules/oneflow/infer_compiler_registry/register_comfy/openaimodel.py", line 96, in forward
    h = forward_timestep_embed(
  File "src/comfyui/comfy/ldm/modules/diffusionmodules/openaimodel.py", line 44, in forward_timestep_embed
    x = layer(x, context, transformer_options)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "src/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 ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "src/comfyui/comfy/ldm/modules/attention.py", line 618, in forward
    n = attn2_replace_patch[block_attn2](n, context_attn2, value_attn2, extra_options)
  File "src/comfyui_custom_nodes/onediff_comfy_nodes/modules/oneflow/infer_compiler_registry/register_comfy/CrossAttentionPatch.py", line 59, in __call__
    patch_kwargs = extra_options["_attn2"].get(self.forward_patch_key)
KeyError: '_attn2'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "src/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 "src/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 "src/comfyui/execution.py", line 169, in _map_node_over_list
    process_inputs(input_dict, i)
  File "src/comfyui/execution.py", line 158, in process_inputs
    results.append(getattr(obj, func)(**inputs))
  File "src/comfyui/nodes.py", line 1430, in sample
    return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
  File "src/comfyui/nodes.py", line 1397, in common_ksampler
    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
  File "src/comfyui/custom_nodes/ComfyUI-Advanced-ControlNet/adv_control/control_reference.py", line 47, in refcn_sample
    return orig_comfy_sample(model, *args, **kwargs)
  File "src/comfyui/custom_nodes/ComfyUI-Advanced-ControlNet/adv_control/utils.py", line 112, in uncond_multiplier_check_cn_sample
    return orig_comfy_sample(model, *args, **kwargs)
  File "src/comfyui/custom_nodes/ComfyUI-AnimateDiff-Evolved/animatediff/sampling.py", line 420, in motion_sample
    return orig_comfy_sample(model, noise, *args, **kwargs)
  File "src/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 "src/comfyui/comfy/sample.py", line 56, 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 "src/comfyui/custom_nodes/ComfyUI_smZNodes/smZNodes.py", line 117, in KSampler_sample
    return orig_fn(*args, **kwargs)
  File "src/comfyui/comfy/samplers.py", line 829, 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 "src/comfyui/custom_nodes/ComfyUI_smZNodes/smZNodes.py", line 134, in sample
    return orig_fn(*args, **kwargs)
  File "src/comfyui/comfy/samplers.py", line 729, in sample
    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
  File "src/comfyui/comfy/samplers.py", line 716, in sample
    output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
  File "src/comfyui/comfy/samplers.py", line 695, in inner_sample
    samples = sampler.sample(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
  File "src/comfyui/comfy/samplers.py", line 600, in sample
    samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
  File ".venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "src/comfyui/comfy/k_diffusion/sampling.py", line 619, in sample_dpmpp_sde
    denoised = model(x, sigmas[i] * s_in, **extra_args)
  File "src/comfyui/custom_nodes/ComfyUI_smZNodes/smZNodes.py", line 105, in KSamplerX0Inpaint___call__
    return store.KSamplerX0Inpaint___call__(*args, **kwargs)
  File "src/comfyui/comfy/samplers.py", line 299, in __call__
    out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
  File "src/comfyui/comfy/samplers.py", line 682, in __call__
    return self.predict_noise(*args, **kwargs)
  File "src/comfyui/comfy/samplers.py", line 685, 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 "src/comfyui/custom_nodes/ComfyUI_smZNodes/smZNodes.py", line 175, in sampling_function
    out = orig_fn(*args, **kwargs)
  File "src/comfyui/comfy/samplers.py", line 279, in sampling_function
    out = calc_cond_batch(model, conds, x, timestep, model_options)
  File "src/comfyui/custom_nodes/ComfyUI_smZNodes/smZNodes.py", line 85, in calc_cond_batch
    return store.calc_cond_batch(*args, **kwargs)
  File "src/comfyui/comfy/samplers.py", line 228, in calc_cond_batch
    output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
  File "src/comfyui/custom_nodes/ComfyUI-Advanced-ControlNet/adv_control/utils.py", line 64, in apply_model_uncond_cleanup_wrapper
    return orig_apply_model(self, *args, **kwargs)
  File "src/comfyui/comfy/model_base.py", line 142, in apply_model
    model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds).float()
  File ".venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py", line 48, in wrapper
    return func(self, *args, **kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/online_quantization_utils.py", line 65, in wrapper
    output = func(self, *args, **kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/graph_management_utils.py", line 123, in wrapper
    ret = func(self, *args, **kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/args_tree_util.py", line 74, in wrapper
    else func(self, *mapped_args, **mapped_kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/deployable_module.py", line 154, in forward
    output = dpl_graph(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 295, in __call__
    self._compile(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 861, in _compile
    return self._dynamic_input_graph_cache._compile(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/cache.py", line 121, in _compile
    return graph._compile(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 865, in _compile
    return self._compile_new(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 884, in _compile_new
    _, eager_outputs = self.build_graph(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 1429, in build_graph
    outputs = self.__build_graph(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/graph.py", line 1577, in __build_graph
    outputs = self.build(*lazy_args, **lazy_kwargs)
  File ".venv/lib/python3.10/site-packages/onediff/infer_compiler/backends/oneflow/graph.py", line 19, in build
    return self.model(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "src/comfyui_custom_nodes/onediff_comfy_nodes/modules/oneflow/infer_compiler_registry/register_comfy/openaimodel.py", line 96, in forward
    h = forward_timestep_embed(
  File "src/comfyui/comfy/ldm/modules/diffusionmodules/openaimodel.py", line 44, in forward_timestep_embed
    x = layer(x, context, transformer_options)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "src/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 ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 188, in __call__
    result = self.__block_forward(*args, **kwargs)
  File ".venv/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 239, in __block_forward
    result = unbound_forward_of_module_instance(self, *args, **kwargs)
  File "src/comfyui/comfy/ldm/modules/attention.py", line 618, in forward
    n = attn2_replace_patch[block_attn2](n, context_attn2, value_attn2, extra_options)
  File "src/comfyui_custom_nodes/onediff_comfy_nodes/modules/oneflow/infer_compiler_registry/register_comfy/CrossAttentionPatch.py", line 59, in __call__
    patch_kwargs = extra_options["_attn2"].get(self.forward_patch_key)
KeyError: '_attn2'

ComfyUI commit: ca08597670c180554ab494e9452e12132d9b346a
ComfyUI_IPAdapter_plus commit: 88a71407c545e4eb0f223294f5b56302ef8696f3

The culprit is ComfyUI-smZNodes, which monkey-patches some functions in the sampler: https://github.com/shiimizu/ComfyUI_smZNodes/blob/5e584280a59e4fce69718db90e9b004664bfd48f/smZNodes.py#L29-L44