ComfyUI icon indicating copy to clipboard operation
ComfyUI copied to clipboard

AddNoise node does not reshape latent dimensions before applying noise

Open drozbay opened this issue 1 year ago • 0 comments

Expected Behavior

When using the AddNoise node before a sampling process that requires a latent with more than 4 channels and when inputting a 4 channel latent, I expect the latent to be reshaped to the proper number of channels before the noise is added to the latent. Example input latent shape: (1, 4, 128, 128) output latent shape: (1, 16, 128, 128)

Actual Behavior

When using the AddNoise node with a 4 channel latent input, the noise is added to the 4 channel latent without reshaping it. Even if a follow-up node attempts to run the function: comfy.sample.fix_empty_latent_channels(model, latent_image), it will not be able to reshape the latent because the noise has already been applied. This will result in a tensor shape error error during sampling, because of the mismatched latent and model shapes. With SD3, for example, the error is: RuntimeError: Given groups=1, weight of size [1536, 16, 2, 2], expected input[2, 4, 128, 128] to have 16 channels, but got 4 channels instead

Steps to Reproduce

addnoise_latentshape_test.json

Debug Logs

# ComfyUI Error Report
## Error Details
- **Node Type:** SamplerCustom
- **Exception Type:** RuntimeError
- **Exception Message:** Given groups=1, weight of size [1536, 16, 2, 2], expected input[2, 4, 128, 128] to have 16 channels, but got 4 channels instead
## Stack Trace

  File "H:\AppsDir\git\ComfyBase\ComfyUI\execution.py", line 317, 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 "H:\AppsDir\git\ComfyBase\ComfyUI\execution.py", line 192, 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 "H:\AppsDir\git\ComfyBase\ComfyUI\execution.py", line 169, in _map_node_over_list
    process_inputs(input_dict, i)

  File "H:\AppsDir\git\ComfyBase\ComfyUI\execution.py", line 158, in process_inputs
    results.append(getattr(obj, func)(**inputs))
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy_extras\nodes_custom_sampler.py", line 455, in sample
    samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\custom_nodes\ComfyUI-AnimateDiff-Evolved\animatediff\sampling.py", line 420, in motion_sample
    return orig_comfy_sample(model, noise, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\custom_nodes\ComfyUI-Advanced-ControlNet\adv_control\sampling.py", line 116, in acn_sample
    return orig_comfy_sample(model, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\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 "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\sample.py", line 48, in sample_custom
    samples = comfy.samplers.sample(model, noise, positive, negative, cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\custom_nodes\ComfyUI_smZNodes\smZNodes.py", line 134, in sample
    return orig_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\samplers.py", line 729, in sample
    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\custom_nodes\UltraCascade\modules\stage_up.py", line 32, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\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 "H:\AppsDir\git\ComfyBase\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 "H:\AppsDir\git\ComfyBase\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 "h:\AppsDir\git\ComfyBase\python_embeded\Lib\site-packages\torch\utils\_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^

  File "h:\AppsDir\git\ComfyBase\ComfyUI\comfy\k_diffusion\sampling.py", line 144, in sample_euler
    denoised = model(x, sigma_hat * s_in, **extra_args)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\custom_nodes\ComfyUI_smZNodes\smZNodes.py", line 105, in KSamplerX0Inpaint___call__
    return store.KSamplerX0Inpaint___call__(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\samplers.py", line 299, in __call__
    out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\samplers.py", line 682, in __call__
    return self.predict_noise(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\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 "H:\AppsDir\git\ComfyBase\ComfyUI\custom_nodes\UltraCascade\modules\stage_up.py", line 16, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\custom_nodes\ComfyUI_smZNodes\smZNodes.py", line 175, in sampling_function
    out = orig_fn(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\samplers.py", line 279, in sampling_function
    out = calc_cond_batch(model, conds, x, timestep, model_options)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\custom_nodes\ComfyUI_smZNodes\smZNodes.py", line 85, in calc_cond_batch
    return store.calc_cond_batch(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\samplers.py", line 228, in calc_cond_batch
    output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\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 "H:\AppsDir\git\ComfyBase\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 "h:\AppsDir\git\ComfyBase\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "h:\AppsDir\git\ComfyBase\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\ldm\modules\diffusionmodules\mmdit.py", line 934, in forward
    return super().forward(x, timesteps, context=context, y=y, control=control)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\ldm\modules\diffusionmodules\mmdit.py", line 909, in forward
    x = self.x_embedder(x) + comfy.ops.cast_to_input(self.cropped_pos_embed(hw, device=x.device), x)
        ^^^^^^^^^^^^^^^^^^

  File "h:\AppsDir\git\ComfyBase\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "h:\AppsDir\git\ComfyBase\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\ldm\modules\diffusionmodules\mmdit.py", line 116, in forward
    x = self.proj(x)
        ^^^^^^^^^^^^

  File "h:\AppsDir\git\ComfyBase\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "h:\AppsDir\git\ComfyBase\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "H:\AppsDir\git\ComfyBase\ComfyUI\comfy\ops.py", line 106, in forward
    return super().forward(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "h:\AppsDir\git\ComfyBase\python_embeded\Lib\site-packages\torch\nn\modules\conv.py", line 458, in forward
    return self._conv_forward(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "h:\AppsDir\git\ComfyBase\python_embeded\Lib\site-packages\torch\nn\modules\conv.py", line 454, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

System Information

  • ComfyUI Version: v0.1.3-24-gb643eae
  • Arguments: H:\AppsDir\git\ComfyBase\ComfyUI/main.py --windows-standalone-build --disable-auto-launch --fast
  • OS: nt
  • Python Version: 3.11.8 (tags/v3.11.8:db85d51, Feb 6 2024, 22:03:32) [MSC v.1937 64 bit (AMD64)]
  • Embedded Python: true
  • PyTorch Version: 2.4.0+cu121

Devices

  • Name: cuda:0 NVIDIA GeForce RTX 4080 SUPER : cudaMallocAsync
    • Type: cuda
    • VRAM Total: 17170956288
    • VRAM Free: 5244308026
    • Torch VRAM Total: 10468982784
    • Torch VRAM Free: 29739578


### Other

This is the simple fix that works for me, in .\comfy_extras\nodes_custom_sampler.py:
```diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py
index 219975e..b2c3635 100644
--- a/comfy_extras/nodes_custom_sampler.py
+++ b/comfy_extras/nodes_custom_sampler.py
@@ -644,6 +644,9 @@ class AddNoise:
 
         latent = latent_image
         latent_image = latent["samples"]
+        latent = latent.copy()
+        latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
+        latent["samples"] = latent_image
 
         noisy = noise.generate_noise(latent)
 

drozbay avatar Sep 01 '24 18:09 drozbay