AddNoise node does not reshape latent dimensions before applying noise
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)