From 9d40212485febe05a662dd0346e6def83e456288 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Tue, 13 Sep 2022 21:49:58 +0300
Subject: first attempt to produce crrect seeds in batch

---
 modules/devices.py     | 10 ++++++++++
 modules/processing.py  | 18 ++++++++++++++++--
 modules/sd_samplers.py | 25 +++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/modules/devices.py b/modules/devices.py
index e4430e1a..07bb2339 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -48,3 +48,13 @@ def randn(seed, shape):
     torch.manual_seed(seed)
     return torch.randn(shape, device=device)
 
+
+def randn_without_seed(shape):
+    # Pytorch currently doesn't handle setting randomness correctly when the metal backend is used.
+    if device.type == 'mps':
+        generator = torch.Generator(device=cpu)
+        noise = torch.randn(shape, generator=generator, device=cpu).to(device)
+        return noise
+
+    return torch.randn(shape, device=device)
+
diff --git a/modules/processing.py b/modules/processing.py
index f33560ee..aab72903 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -119,8 +119,14 @@ def slerp(val, low, high):
     return res
 
 
-def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, seed_resize_from_h=0, seed_resize_from_w=0):
+def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, seed_resize_from_h=0, seed_resize_from_w=0, p=None):
     xs = []
+
+    if p is not None and p.sampler is not None and len(seeds) > 1:
+        sampler_noises = [[] for _ in range(p.sampler.number_of_needed_noises(p))]
+    else:
+        sampler_noises = None
+
     for i, seed in enumerate(seeds):
         noise_shape = shape if seed_resize_from_h <= 0 or seed_resize_from_w <= 0 else (shape[0], seed_resize_from_h//8, seed_resize_from_w//8)
 
@@ -155,9 +161,17 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
             x[:, ty:ty+h, tx:tx+w] = noise[:, dy:dy+h, dx:dx+w]
             noise = x
 
+        if sampler_noises is not None:
+            cnt = p.sampler.number_of_needed_noises(p)
 
+            for j in range(cnt):
+                sampler_noises[j].append(devices.randn_without_seed(tuple(noise_shape)))
 
         xs.append(noise)
+
+    if sampler_noises is not None:
+        p.sampler.sampler_noises = [torch.stack(n).to(shared.device) for n in sampler_noises]
+
     x = torch.stack(xs).to(shared.device)
     return x
 
@@ -254,7 +268,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                 comments += model_hijack.comments
 
             # we manually generate all input noises because each one should have a specific seed
-            x = create_random_tensors([opt_C, p.height // opt_f, p.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w)
+            x = create_random_tensors([opt_C, p.height // opt_f, p.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w, p=p)
 
             if p.n_iter > 1:
                 shared.state.job = f"Batch {n+1} out of {p.n_iter}"
diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 7ef507f1..f77fe43f 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -93,6 +93,10 @@ class VanillaStableDiffusionSampler:
         self.mask = None
         self.nmask = None
         self.init_latent = None
+        self.sampler_noises = None
+
+    def number_of_needed_noises(self, p):
+        return 0
 
     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning):
         t_enc = int(min(p.denoising_strength, 0.999) * p.steps)
@@ -171,16 +175,37 @@ def extended_trange(count, *args, **kwargs):
         shared.total_tqdm.update()
 
 
+original_randn_like = torch.randn_like
+
 class KDiffusionSampler:
     def __init__(self, funcname, sd_model):
         self.model_wrap = k_diffusion.external.CompVisDenoiser(sd_model)
         self.funcname = funcname
         self.func = getattr(k_diffusion.sampling, self.funcname)
         self.model_wrap_cfg = CFGDenoiser(self.model_wrap)
+        self.sampler_noises = None
+        self.sampler_noise_index = 0
+
+        k_diffusion.sampling.torch.randn_like = self.randn_like
 
     def callback_state(self, d):
         store_latent(d["denoised"])
 
+    def number_of_needed_noises(self, p):
+        return p.steps
+
+    def randn_like(self, x):
+        noise = self.sampler_noises[self.sampler_noise_index] if self.sampler_noises is not None and self.sampler_noise_index < len(self.sampler_noises) else None
+
+        if noise is not None and x.shape == noise.shape:
+            res = noise
+        else:
+            print('generating')
+            res = original_randn_like(x)
+
+        self.sampler_noise_index += 1
+        return res
+
     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning):
         t_enc = int(min(p.denoising_strength, 0.999) * p.steps)
         sigmas = self.model_wrap.get_sigmas(p.steps)
-- 
cgit v1.2.1