From 9d40212485febe05a662dd0346e6def83e456288 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Tue, 13 Sep 2022 21:49:58 +0300 Subject: [PATCH 1/2] 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) From 87e8b9a2ab3f033e7fdadbb2fe258857915980ac Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Fri, 16 Sep 2022 09:47:03 +0300 Subject: [PATCH 2/2] prevent replacing torch_randn globally (instead replacing k_diffusion.sampling.torch) and add a setting to disable this all --- modules/processing.py | 2 +- modules/sd_samplers.py | 25 ++++++++++++++++++++----- modules/shared.py | 3 ++- 3 files changed, 23 insertions(+), 7 deletions(-) diff --git a/modules/processing.py b/modules/processing.py index aab72903..5abdfd7c 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -122,7 +122,7 @@ def slerp(val, low, high): 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: + if p is not None and p.sampler is not None and len(seeds) > 1 and opts.enable_batch_seeds: sampler_noises = [[] for _ in range(p.sampler.number_of_needed_noises(p))] else: sampler_noises = None diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index f77fe43f..d478c5bc 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -175,7 +175,19 @@ def extended_trange(count, *args, **kwargs): shared.total_tqdm.update() -original_randn_like = torch.randn_like +class TorchHijack: + def __init__(self, kdiff_sampler): + self.kdiff_sampler = kdiff_sampler + + def __getattr__(self, item): + if item == 'randn_like': + return self.kdiff_sampler.randn_like + + if hasattr(torch, item): + return getattr(torch, item) + + raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, item)) + class KDiffusionSampler: def __init__(self, funcname, sd_model): @@ -186,8 +198,6 @@ class KDiffusionSampler: 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"]) @@ -200,8 +210,7 @@ class KDiffusionSampler: if noise is not None and x.shape == noise.shape: res = noise else: - print('generating') - res = original_randn_like(x) + res = torch.randn_like(x) self.sampler_noise_index += 1 return res @@ -223,6 +232,9 @@ class KDiffusionSampler: if hasattr(k_diffusion.sampling, 'trange'): k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs) + if self.sampler_noises is not None: + k_diffusion.sampling.torch = TorchHijack(self) + return self.func(self.model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state) def sample(self, p, x, conditioning, unconditional_conditioning): @@ -232,6 +244,9 @@ class KDiffusionSampler: if hasattr(k_diffusion.sampling, 'trange'): k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs) + if self.sampler_noises is not None: + k_diffusion.sampling.torch = TorchHijack(self) + samples_ddim = self.func(self.model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state) return samples_ddim diff --git a/modules/shared.py b/modules/shared.py index bc39ad1c..ac870ec4 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -124,7 +124,8 @@ class Options: "add_model_hash_to_info": OptionInfo(False, "Add model hash to generation information"), "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."), "font": OptionInfo("", "Font for image grids that have text"), - "enable_emphasis": OptionInfo(True, "Use (text) to make model pay more attention to text text and [text] to make it pay less attention"), + "enable_emphasis": OptionInfo(True, "Use (text) to make model pay more attention to text and [text] to make it pay less attention"), + "enable_batch_seeds": OptionInfo(True, "Make K-diffusion samplers produce same images in a batch as when making a single image"), "save_txt": OptionInfo(False, "Create a text file next to every image with generation parameters."), "ESRGAN_tile": OptionInfo(192, "Tile size for upscaling. 0 = no tiling.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}), "ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for upscaling. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),