diff --git a/comfy/sample.py b/comfy/sample.py index 2be0cae5f..617816882 100644 --- a/comfy/sample.py +++ b/comfy/sample.py @@ -37,6 +37,11 @@ def prepare_noise(latent_image, seed, noise_inds=None): return noises +def prepare_empty_noise(latent_image): + if latent_image.is_nested: + return comfy.nested_tensor.NestedTensor([torch.zeros_like(t, device="cpu") for t in latent_image.unbind()]) + return torch.zeros_like(latent_image, device="cpu") + def fix_empty_latent_channels(model, latent_image, downscale_ratio_spacial=None, downscale_ratio_temporal=None): if latent_image.is_nested: return latent_image diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index d5aa730d2..c73a8f6dc 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -718,15 +718,7 @@ class Noise_EmptyNoise: self.seed = 0 def generate_noise(self, input_latent): - latent_image = input_latent["samples"] - if latent_image.is_nested: - tensors = latent_image.unbind() - zeros = [] - for t in tensors: - zeros.append(torch.zeros(t.shape, dtype=t.dtype, layout=t.layout, device="cpu")) - return comfy.nested_tensor.NestedTensor(zeros) - else: - return torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") + return comfy.sample.prepare_empty_noise(input_latent["samples"]) class Noise_RandomNoise: diff --git a/nodes.py b/nodes.py index a7f91720f..ec298e1de 100644 --- a/nodes.py +++ b/nodes.py @@ -1570,7 +1570,7 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None), latent.get("downscale_ratio_temporal", None)) if disable_noise: - noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") + noise = comfy.sample.prepare_empty_noise(latent_image) else: batch_inds = latent["batch_index"] if "batch_index" in latent else None noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)