Use Torch for pyramid RNG, ensure seed is set in samplers even when add_noise is disabled

extra_samplers.py

@@ -8,7 +8,6 @@ from tqdm.auto import trange, tqdm
 import comfy.sample
 
 from comfy.k_diffusion.sampling import BrownianTreeNoiseSampler, PIDStepSizeController, get_ancestral_step, to_d, default_noise_sampler
-import random
 
 # The following function adds the samplers during initialization, in __init__.py
 def add_samplers():
@@ -264,7 +263,7 @@ def highres_pyramid_noise_like(x, discount=0.7):
     u = torch.nn.Upsample(size=(orig_h, orig_w), mode='bilinear')
     noise = (torch.rand_like(x) - 0.5) * 2 * 1.73 # Start with scaled uniform noise
     for i in range(4):
-        r = random.random()*2+2 # Rather than always going 2x,
+        r = torch.rand(1).item() * 2 + 2 # Rather than always going 2x,
         h, w = min(orig_h*15, int(h*(r**i))), min(orig_w*15, int(w*(r**i)))
         noise += u(torch.randn(b, c, h, w).to(x)) * discount**i
         if h>=orig_h*15 or w>=orig_w*15: break # Lowest resolution is 1x1
@@ -794,4 +793,4 @@ discard_penultimate_sigma_samplers = set((
     "clyb_4m_sde_momentumized"
 ))
 
-extra_schedulers = {}
+extra_schedulers = {}

nodes.py

@@ -7,7 +7,6 @@ import torch
 import numpy as np
 from tqdm.auto import trange
 
-import random
 def pyramid_noise_like(size, dtype, layout, generator, device="cpu", discount=0.8):
     b, c, h, w = size
     orig_h = h
@@ -15,7 +14,7 @@ def pyramid_noise_like(size, dtype, layout, generator, device="cpu", discount=0.
     noise = torch.zeros(size=size, dtype=dtype, layout=layout, device=device)
     r = 1
     for i in range(5):
-        r *= 2 # Rather than always going 2x, 
+        r *= 2 # Rather than always going 2x,
         #w, h = max(1, int(w/(r**i))), max(1, int(h/(r**i)))
         noise += torch.nn.functional.interpolate((torch.normal(mean=0, std=0.5 ** i, size=(b, c, h * r, w * r), dtype=dtype, layout=layout, generator=generator, device=device)), size=(orig_h, orig_w), mode='nearest-exact') * discount**i
         #if w>=orig_w*16 or h>=orig_h*16: break
@@ -63,7 +62,7 @@ def prepare_noise(latent_image, seed, noise_type, noise_inds=None): # From `samp
             noise_func = torch.randn
     if noise_inds is None:
         return noise_func(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)
     noises = []
     for i in range(unique_inds[-1]+1):
@@ -320,11 +319,12 @@ class SamplerCustomNoise:
         latent = latent_image
         latent_image = latent["samples"]
         if not add_noise:
+            torch.manual_seed(noise_seed)
             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 = prepare_noise(latent_image, noise_seed, noise_type, batch_inds)
-        
+
         if noise_is_latent:
             noise += latent_image.cpu()# * noise.std()
             noise.sub_(noise.mean()).div_(noise.std())
@@ -383,6 +383,7 @@ class SamplerCustomNoiseDuo:
         latent = latent_image
         latent_image = latent["samples"]
         if not add_noise:
+            torch.manual_seed(noise_seed)
             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
@@ -455,6 +456,7 @@ class SamplerCustomModelMixtureDuo:
         latent = latent_image
         latent_image = latent["samples"]
         if not add_noise:
+            torch.manual_seed(noise_seed)
             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
@@ -491,4 +493,4 @@ class SamplerCustomModelMixtureDuo:
             out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
         else:
             out_denoised = out
-        return (out, out_denoised)
+        return (out, out_denoised)