Uniform, Scheduled and randomized seeds depending on list or INT input

comfy/sample.py

@@ -3,25 +3,42 @@
 import comfy.samplers
 import comfy.conds
 import comfy.utils
+import random
 import math
 import numpy as np
 
-def prepare_noise(latent_image, seed, noise_inds=None):
+def prepare_noise(latent_image, seeds, batch_behavior = "randomize"):
     """
-    creates random noise given a latent image and a seed.
-    optional arg skip can be used to skip and discard x number of noise generations for a given seed
+    Creates noise based on the batch behavior, a latent image and a seed or a list of seeds.
     """
-    generator = torch.manual_seed(seed)
-    if noise_inds is None:
-        return torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
-
-    unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
+    num_latents = latent_image.size(0)
+
+    if not isinstance(seeds, list):
+        seeds = [seeds]
+
+    generator = torch.Generator()
+    generator.manual_seed(seeds[0])  # Use the first seed as the default generator seed
+
     noises = []
-    for i in range(unique_inds[-1]+1):
-        noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
-        if i in unique_inds:
-            noises.append(noise)
-    noises = [noises[i] for i in inverse]
+
+    for i in range(num_latents):
+        if i < len(seeds):  # Use the provided seeds if available then follow behavior
+            seed = seeds[i]
+        elif batch_behavior == "randomize" :
+            seed = int(random.random() * 99999999999999)
+        elif batch_behavior == "fixed":
+            seed = seeds[-1]
+        elif batch_behavior == "increment":
+            seed = seeds[-1] + i
+        else :
+            seed = seeds[-1] - i
+
+        generator.manual_seed(seed)
+        print("seed:", seed)
+        noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, device="cpu", generator=generator)
+        noises.append(noise)
+
+    noises = [noises[i] for i in range(num_latents)]
     noises = torch.cat(noises, axis=0)
     return noises
 

nodes.py

@@ -1308,13 +1308,12 @@ def set_mask(self, samples, mask):
         s["noise_mask"] = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
         return (s,)
 
-def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
+def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, batch_behavior = "randomize", disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
     latent_image = latent["samples"]
     if disable_noise:
         noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
     else:
-        batch_inds = latent["batch_index"] if "batch_index" in latent else None
-        noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
+        noise = comfy.sample.prepare_noise(latent_image, seed, batch_behavior)
 
     noise_mask = None
     if "noise_mask" in latent:
@@ -1343,6 +1342,7 @@ def INPUT_TYPES(s):
                     "negative": ("CONDITIONING", ),
                     "latent_image": ("LATENT", ),
                     "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                    "batch_behavior": (["randomize", "fixed", "increment", "decrement"],),
                      }
                 }
 
@@ -1351,8 +1351,8 @@ def INPUT_TYPES(s):
 
     CATEGORY = "sampling"
 
-    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
-        return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
+    def sample(self, model, seed, batch_behavior, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
+        return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, batch_behavior=batch_behavior)
 
 class KSamplerAdvanced:
     @classmethod
@@ -1371,6 +1371,7 @@ def INPUT_TYPES(s):
                     "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
                     "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
                     "return_with_leftover_noise": (["disable", "enable"], ),
+                    "batch_behavior": (["randomize", "fixed", "increment", "decrement"],),
                      }
                 }
 
@@ -1379,14 +1380,14 @@ def INPUT_TYPES(s):
 
     CATEGORY = "sampling"
 
-    def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0):
+    def sample(self, model, add_noise, noise_seed, batch_behavior, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0):
         force_full_denoise = True
         if return_with_leftover_noise == "enable":
             force_full_denoise = False
         disable_noise = False
         if add_noise == "disable":
             disable_noise = True
-        return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
+        return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, batch_behavior=batch_behavior, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
 
 class SaveImage:
     def __init__(self):