From 8285a149d8c488ae6c7a566eb85fb5e825145464 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Tue, 8 Aug 2023 19:20:11 +0300
Subject: add CFG denoiser implementation for DDIM, PLMS and UniPC (this is the
 commit when you can run both old and new implementations to compare them)

---
 modules/sd_samplers_timesteps.py | 147 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 147 insertions(+)
 create mode 100644 modules/sd_samplers_timesteps.py

(limited to 'modules/sd_samplers_timesteps.py')

diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
new file mode 100644
index 00000000..8560d009
--- /dev/null
+++ b/modules/sd_samplers_timesteps.py
@@ -0,0 +1,147 @@
+import torch
+import inspect
+from modules import devices, sd_samplers_common, sd_samplers_timesteps_impl
+from modules.sd_samplers_cfg_denoiser import CFGDenoiser
+
+from modules.shared import opts
+import modules.shared as shared
+
+samplers_timesteps = [
+    ('k_DDIM', sd_samplers_timesteps_impl.ddim, ['k_ddim'], {}),
+    ('k_PLMS', sd_samplers_timesteps_impl.plms, ['k_plms'], {}),
+    ('k_UniPC', sd_samplers_timesteps_impl.unipc, ['k_unipc'], {}),
+]
+
+
+samplers_data_timesteps = [
+    sd_samplers_common.SamplerData(label, lambda model, funcname=funcname: CompVisSampler(funcname, model), aliases, options)
+    for label, funcname, aliases, options in samplers_timesteps
+]
+
+
+class CompVisTimestepsDenoiser(torch.nn.Module):
+    def __init__(self, model, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.inner_model = model
+
+    def forward(self, input, timesteps, **kwargs):
+        return self.inner_model.apply_model(input, timesteps, **kwargs)
+
+
+class CompVisTimestepsVDenoiser(torch.nn.Module):
+    def __init__(self, model, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.inner_model = model
+
+    def predict_eps_from_z_and_v(self, x_t, t, v):
+        return self.inner_model.sqrt_alphas_cumprod[t.to(torch.int), None, None, None] * v + self.inner_model.sqrt_one_minus_alphas_cumprod[t.to(torch.int), None, None, None] * x_t
+
+    def forward(self, input, timesteps, **kwargs):
+        model_output = self.inner_model.apply_model(input, timesteps, **kwargs)
+        e_t = self.predict_eps_from_z_and_v(input, timesteps, model_output)
+        return e_t
+
+
+class CFGDenoiserTimesteps(CFGDenoiser):
+
+    def __init__(self, model, sampler):
+        super().__init__(model, sampler)
+
+        self.alphas = model.inner_model.alphas_cumprod
+
+    def get_pred_x0(self, x_in, x_out, sigma):
+        ts = int(sigma.item())
+
+        s_in = x_in.new_ones([x_in.shape[0]])
+        a_t = self.alphas[ts].item() * s_in
+        sqrt_one_minus_at = (1 - a_t).sqrt()
+
+        pred_x0 = (x_in - sqrt_one_minus_at * x_out) / a_t.sqrt()
+
+        return pred_x0
+
+
+class CompVisSampler(sd_samplers_common.Sampler):
+    def __init__(self, funcname, sd_model):
+        super().__init__(funcname)
+
+        self.eta_option_field = 'eta_ddim'
+        self.eta_infotext_field = 'Eta DDIM'
+
+        denoiser = CompVisTimestepsVDenoiser if sd_model.parameterization == "v" else CompVisTimestepsDenoiser
+        self.model_wrap = denoiser(sd_model)
+        self.model_wrap_cfg = CFGDenoiserTimesteps(self.model_wrap, self)
+
+    def get_timesteps(self, p, steps):
+        discard_next_to_last_sigma = self.config is not None and self.config.options.get('discard_next_to_last_sigma', False)
+        if opts.always_discard_next_to_last_sigma and not discard_next_to_last_sigma:
+            discard_next_to_last_sigma = True
+            p.extra_generation_params["Discard penultimate sigma"] = True
+
+        steps += 1 if discard_next_to_last_sigma else 0
+
+        timesteps = torch.clip(torch.asarray(list(range(0, 1000, 1000 // steps)), device=devices.device) + 1, 0, 999)
+
+        return timesteps
+
+    def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
+        steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps)
+
+        timesteps = self.get_timesteps(p, steps)
+        timesteps_sched = timesteps[:t_enc]
+
+        alphas_cumprod = shared.sd_model.alphas_cumprod
+        sqrt_alpha_cumprod = torch.sqrt(alphas_cumprod[timesteps[t_enc]])
+        sqrt_one_minus_alpha_cumprod = torch.sqrt(1 - alphas_cumprod[timesteps[t_enc]])
+
+        xi = x * sqrt_alpha_cumprod + noise * sqrt_one_minus_alpha_cumprod
+
+        extra_params_kwargs = self.initialize(p)
+        parameters = inspect.signature(self.func).parameters
+
+        if 'timesteps' in parameters:
+            extra_params_kwargs['timesteps'] = timesteps_sched
+        if 'is_img2img' in parameters:
+            extra_params_kwargs['is_img2img'] = True
+
+        self.model_wrap_cfg.init_latent = x
+        self.last_latent = x
+        extra_args = {
+            'cond': conditioning,
+            'image_cond': image_conditioning,
+            'uncond': unconditional_conditioning,
+            'cond_scale': p.cfg_scale,
+            's_min_uncond': self.s_min_uncond
+        }
+
+        samples = self.launch_sampling(t_enc + 1, lambda: self.func(self.model_wrap_cfg, xi, extra_args=extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs))
+
+        if self.model_wrap_cfg.padded_cond_uncond:
+            p.extra_generation_params["Pad conds"] = True
+
+        return samples
+
+    def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
+        steps = steps or p.steps
+        timesteps = self.get_timesteps(p, steps)
+
+        extra_params_kwargs = self.initialize(p)
+        parameters = inspect.signature(self.func).parameters
+
+        if 'timesteps' in parameters:
+            extra_params_kwargs['timesteps'] = timesteps
+
+        self.last_latent = x
+        samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={
+            'cond': conditioning,
+            'image_cond': image_conditioning,
+            'uncond': unconditional_conditioning,
+            'cond_scale': p.cfg_scale,
+            's_min_uncond': self.s_min_uncond
+        }, disable=False, callback=self.callback_state, **extra_params_kwargs))
+
+        if self.model_wrap_cfg.padded_cond_uncond:
+            p.extra_generation_params["Pad conds"] = True
+
+        return samples
+
-- 
cgit v1.2.3


From a8a256f9b5b445206818bfc8a363ed5a1ba50c86 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Tue, 8 Aug 2023 21:07:18 +0300
Subject: REMOVE

---
 modules/processing.py               |   3 -
 modules/sd_hijack.py                |   4 +-
 modules/sd_hijack_inpainting.py     |  95 ---------------
 modules/sd_samplers.py              |   7 +-
 modules/sd_samplers_cfg_denoiser.py |   1 -
 modules/sd_samplers_compvis.py      | 224 ------------------------------------
 modules/sd_samplers_kdiffusion.py   |   3 +-
 modules/sd_samplers_timesteps.py    |   6 +-
 8 files changed, 7 insertions(+), 336 deletions(-)
 delete mode 100644 modules/sd_hijack_inpainting.py
 delete mode 100644 modules/sd_samplers_compvis.py

(limited to 'modules/sd_samplers_timesteps.py')

diff --git a/modules/processing.py b/modules/processing.py
index 31745006..61ba5f11 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -1112,9 +1112,6 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 
         img2img_sampler_name = self.hr_sampler_name or self.sampler_name
 
-        if self.sampler_name in ['PLMS', 'UniPC']:  # PLMS/UniPC do not support img2img so we just silently switch to DDIM
-            img2img_sampler_name = 'DDIM'
-
         self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model)
 
         if self.latent_scale_mode is not None:
diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 9ad98199..46652fbd 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -5,7 +5,7 @@ from types import MethodType
 from modules import devices, sd_hijack_optimizations, shared, script_callbacks, errors, sd_unet
 from modules.hypernetworks import hypernetwork
 from modules.shared import cmd_opts
-from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr, sd_hijack_inpainting
+from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr
 
 import ldm.modules.attention
 import ldm.modules.diffusionmodules.model
@@ -34,8 +34,6 @@ ldm.modules.diffusionmodules.model.print = shared.ldm_print
 ldm.util.print = shared.ldm_print
 ldm.models.diffusion.ddpm.print = shared.ldm_print
 
-sd_hijack_inpainting.do_inpainting_hijack()
-
 optimizers = []
 current_optimizer: sd_hijack_optimizations.SdOptimization = None
 
diff --git a/modules/sd_hijack_inpainting.py b/modules/sd_hijack_inpainting.py
deleted file mode 100644
index 2d44b856..00000000
--- a/modules/sd_hijack_inpainting.py
+++ /dev/null
@@ -1,95 +0,0 @@
-import torch
-
-import ldm.models.diffusion.ddpm
-import ldm.models.diffusion.ddim
-import ldm.models.diffusion.plms
-
-from ldm.models.diffusion.ddim import noise_like
-from ldm.models.diffusion.sampling_util import norm_thresholding
-
-
-@torch.no_grad()
-def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
-                  temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
-                  unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None, dynamic_threshold=None):
-    b, *_, device = *x.shape, x.device
-
-    def get_model_output(x, t):
-        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
-            e_t = self.model.apply_model(x, t, c)
-        else:
-            x_in = torch.cat([x] * 2)
-            t_in = torch.cat([t] * 2)
-
-            if isinstance(c, dict):
-                assert isinstance(unconditional_conditioning, dict)
-                c_in = {}
-                for k in c:
-                    if isinstance(c[k], list):
-                        c_in[k] = [
-                            torch.cat([unconditional_conditioning[k][i], c[k][i]])
-                            for i in range(len(c[k]))
-                        ]
-                    else:
-                        c_in[k] = torch.cat([unconditional_conditioning[k], c[k]])
-            else:
-                c_in = torch.cat([unconditional_conditioning, c])
-
-            e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
-            e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
-
-        if score_corrector is not None:
-            assert self.model.parameterization == "eps"
-            e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
-
-        return e_t
-
-    alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
-    alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
-    sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
-    sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
-
-    def get_x_prev_and_pred_x0(e_t, index):
-        # select parameters corresponding to the currently considered timestep
-        a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
-        a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
-        sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
-        sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
-
-        # current prediction for x_0
-        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
-        if quantize_denoised:
-            pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
-        if dynamic_threshold is not None:
-            pred_x0 = norm_thresholding(pred_x0, dynamic_threshold)
-        # direction pointing to x_t
-        dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
-        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
-        if noise_dropout > 0.:
-            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
-        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
-        return x_prev, pred_x0
-
-    e_t = get_model_output(x, t)
-    if len(old_eps) == 0:
-        # Pseudo Improved Euler (2nd order)
-        x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
-        e_t_next = get_model_output(x_prev, t_next)
-        e_t_prime = (e_t + e_t_next) / 2
-    elif len(old_eps) == 1:
-        # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
-        e_t_prime = (3 * e_t - old_eps[-1]) / 2
-    elif len(old_eps) == 2:
-        # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
-        e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
-    elif len(old_eps) >= 3:
-        # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
-        e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
-
-    x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
-
-    return x_prev, pred_x0, e_t
-
-
-def do_inpainting_hijack():
-    ldm.models.diffusion.plms.PLMSSampler.p_sample_plms = p_sample_plms
diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index fe206894..05dbe2b5 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -1,11 +1,10 @@
-from modules import sd_samplers_compvis, sd_samplers_kdiffusion, sd_samplers_timesteps, shared
+from modules import sd_samplers_kdiffusion, sd_samplers_timesteps, shared
 
 # imports for functions that previously were here and are used by other modules
 from modules.sd_samplers_common import samples_to_image_grid, sample_to_image  # noqa: F401
 
 all_samplers = [
     *sd_samplers_kdiffusion.samplers_data_k_diffusion,
-    *sd_samplers_compvis.samplers_data_compvis,
     *sd_samplers_timesteps.samplers_data_timesteps,
 ]
 all_samplers_map = {x.name: x for x in all_samplers}
@@ -42,10 +41,8 @@ def set_samplers():
     global samplers, samplers_for_img2img
 
     hidden = set(shared.opts.hide_samplers)
-    hidden_img2img = set(shared.opts.hide_samplers + ['PLMS', 'UniPC'])
-
     samplers = [x for x in all_samplers if x.name not in hidden]
-    samplers_for_img2img = [x for x in all_samplers if x.name not in hidden_img2img]
+    samplers_for_img2img = [x for x in all_samplers if x.name not in hidden]
 
     samplers_map.clear()
     for sampler in all_samplers:
diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py
index 166a00c7..d826222c 100644
--- a/modules/sd_samplers_cfg_denoiser.py
+++ b/modules/sd_samplers_cfg_denoiser.py
@@ -1,4 +1,3 @@
-from collections import deque
 import torch
 from modules import prompt_parser, devices, sd_samplers_common
 
diff --git a/modules/sd_samplers_compvis.py b/modules/sd_samplers_compvis.py
deleted file mode 100644
index 4a8396f9..00000000
--- a/modules/sd_samplers_compvis.py
+++ /dev/null
@@ -1,224 +0,0 @@
-import math
-import ldm.models.diffusion.ddim
-import ldm.models.diffusion.plms
-
-import numpy as np
-import torch
-
-from modules.shared import state
-from modules import sd_samplers_common, prompt_parser, shared
-import modules.models.diffusion.uni_pc
-
-
-samplers_data_compvis = [
-    sd_samplers_common.SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), [], {"default_eta_is_0": True, "uses_ensd": True, "no_sdxl": True}),
-    sd_samplers_common.SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), [], {"no_sdxl": True}),
-    sd_samplers_common.SamplerData('UniPC', lambda model: VanillaStableDiffusionSampler(modules.models.diffusion.uni_pc.UniPCSampler, model), [], {"no_sdxl": True}),
-]
-
-
-class VanillaStableDiffusionSampler:
-    def __init__(self, constructor, sd_model):
-        self.sampler = constructor(sd_model)
-        self.is_ddim = hasattr(self.sampler, 'p_sample_ddim')
-        self.is_plms = hasattr(self.sampler, 'p_sample_plms')
-        self.is_unipc = isinstance(self.sampler, modules.models.diffusion.uni_pc.UniPCSampler)
-        self.orig_p_sample_ddim = None
-        if self.is_plms:
-            self.orig_p_sample_ddim = self.sampler.p_sample_plms
-        elif self.is_ddim:
-            self.orig_p_sample_ddim = self.sampler.p_sample_ddim
-        self.mask = None
-        self.nmask = None
-        self.init_latent = None
-        self.sampler_noises = None
-        self.step = 0
-        self.stop_at = None
-        self.eta = None
-        self.config = None
-        self.last_latent = None
-
-        self.conditioning_key = sd_model.model.conditioning_key
-
-    def number_of_needed_noises(self, p):
-        return 0
-
-    def launch_sampling(self, steps, func):
-        state.sampling_steps = steps
-        state.sampling_step = 0
-
-        try:
-            return func()
-        except sd_samplers_common.InterruptedException:
-            return self.last_latent
-
-    def p_sample_ddim_hook(self, x_dec, cond, ts, unconditional_conditioning, *args, **kwargs):
-        x_dec, ts, cond, unconditional_conditioning = self.before_sample(x_dec, ts, cond, unconditional_conditioning)
-
-        res = self.orig_p_sample_ddim(x_dec, cond, ts, *args, unconditional_conditioning=unconditional_conditioning, **kwargs)
-
-        x_dec, ts, cond, unconditional_conditioning, res = self.after_sample(x_dec, ts, cond, unconditional_conditioning, res)
-
-        return res
-
-    def before_sample(self, x, ts, cond, unconditional_conditioning):
-        if state.interrupted or state.skipped:
-            raise sd_samplers_common.InterruptedException
-
-        if self.stop_at is not None and self.step > self.stop_at:
-            raise sd_samplers_common.InterruptedException
-
-        # Have to unwrap the inpainting conditioning here to perform pre-processing
-        image_conditioning = None
-        uc_image_conditioning = None
-        if isinstance(cond, dict):
-            if self.conditioning_key == "crossattn-adm":
-                image_conditioning = cond["c_adm"]
-                uc_image_conditioning = unconditional_conditioning["c_adm"]
-            else:
-                image_conditioning = cond["c_concat"][0]
-            cond = cond["c_crossattn"][0]
-            unconditional_conditioning = unconditional_conditioning["c_crossattn"][0]
-
-        conds_list, tensor = prompt_parser.reconstruct_multicond_batch(cond, self.step)
-        unconditional_conditioning = prompt_parser.reconstruct_cond_batch(unconditional_conditioning, self.step)
-
-        assert all(len(conds) == 1 for conds in conds_list), 'composition via AND is not supported for DDIM/PLMS samplers'
-        cond = tensor
-
-        # for DDIM, shapes must match, we can't just process cond and uncond independently;
-        # filling unconditional_conditioning with repeats of the last vector to match length is
-        # not 100% correct but should work well enough
-        if unconditional_conditioning.shape[1] < cond.shape[1]:
-            last_vector = unconditional_conditioning[:, -1:]
-            last_vector_repeated = last_vector.repeat([1, cond.shape[1] - unconditional_conditioning.shape[1], 1])
-            unconditional_conditioning = torch.hstack([unconditional_conditioning, last_vector_repeated])
-        elif unconditional_conditioning.shape[1] > cond.shape[1]:
-            unconditional_conditioning = unconditional_conditioning[:, :cond.shape[1]]
-
-        if self.mask is not None:
-            img_orig = self.sampler.model.q_sample(self.init_latent, ts)
-            x = img_orig * self.mask + self.nmask * x
-
-        # Wrap the image conditioning back up since the DDIM code can accept the dict directly.
-        # Note that they need to be lists because it just concatenates them later.
-        if image_conditioning is not None:
-            if self.conditioning_key == "crossattn-adm":
-                cond = {"c_adm": image_conditioning, "c_crossattn": [cond]}
-                unconditional_conditioning = {"c_adm": uc_image_conditioning, "c_crossattn": [unconditional_conditioning]}
-            else:
-                cond = {"c_concat": [image_conditioning], "c_crossattn": [cond]}
-                unconditional_conditioning = {"c_concat": [image_conditioning], "c_crossattn": [unconditional_conditioning]}
-
-        return x, ts, cond, unconditional_conditioning
-
-    def update_step(self, last_latent):
-        if self.mask is not None:
-            self.last_latent = self.init_latent * self.mask + self.nmask * last_latent
-        else:
-            self.last_latent = last_latent
-
-        sd_samplers_common.store_latent(self.last_latent)
-
-        self.step += 1
-        state.sampling_step = self.step
-        shared.total_tqdm.update()
-
-    def after_sample(self, x, ts, cond, uncond, res):
-        if not self.is_unipc:
-            self.update_step(res[1])
-
-        return x, ts, cond, uncond, res
-
-    def unipc_after_update(self, x, model_x):
-        self.update_step(x)
-
-    def initialize(self, p):
-        if self.is_ddim:
-            self.eta = p.eta if p.eta is not None else shared.opts.eta_ddim
-        else:
-            self.eta = 0.0
-
-        if self.eta != 0.0:
-            p.extra_generation_params["Eta DDIM"] = self.eta
-
-        if self.is_unipc:
-            keys = [
-                ('UniPC variant', 'uni_pc_variant'),
-                ('UniPC skip type', 'uni_pc_skip_type'),
-                ('UniPC order', 'uni_pc_order'),
-                ('UniPC lower order final', 'uni_pc_lower_order_final'),
-            ]
-
-            for name, key in keys:
-                v = getattr(shared.opts, key)
-                if v != shared.opts.get_default(key):
-                    p.extra_generation_params[name] = v
-
-        for fieldname in ['p_sample_ddim', 'p_sample_plms']:
-            if hasattr(self.sampler, fieldname):
-                setattr(self.sampler, fieldname, self.p_sample_ddim_hook)
-        if self.is_unipc:
-            self.sampler.set_hooks(lambda x, t, c, u: self.before_sample(x, t, c, u), lambda x, t, c, u, r: self.after_sample(x, t, c, u, r), lambda x, mx: self.unipc_after_update(x, mx))
-
-        self.mask = p.mask if hasattr(p, 'mask') else None
-        self.nmask = p.nmask if hasattr(p, 'nmask') else None
-
-
-    def adjust_steps_if_invalid(self, p, num_steps):
-        if ((self.config.name == 'DDIM') and p.ddim_discretize == 'uniform') or (self.config.name == 'PLMS') or (self.config.name == 'UniPC'):
-            if self.config.name == 'UniPC' and num_steps < shared.opts.uni_pc_order:
-                num_steps = shared.opts.uni_pc_order
-            valid_step = 999 / (1000 // num_steps)
-            if valid_step == math.floor(valid_step):
-                return int(valid_step) + 1
-
-        return num_steps
-
-    def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
-        steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps)
-        steps = self.adjust_steps_if_invalid(p, steps)
-        self.initialize(p)
-
-        self.sampler.make_schedule(ddim_num_steps=steps, ddim_eta=self.eta, ddim_discretize=p.ddim_discretize, verbose=False)
-        x1 = self.sampler.stochastic_encode(x, torch.tensor([t_enc] * int(x.shape[0])).to(shared.device), noise=noise)
-
-        self.init_latent = x
-        self.last_latent = x
-        self.step = 0
-
-        # Wrap the conditioning models with additional image conditioning for inpainting model
-        if image_conditioning is not None:
-            if self.conditioning_key == "crossattn-adm":
-                conditioning = {"c_adm": image_conditioning, "c_crossattn": [conditioning]}
-                unconditional_conditioning = {"c_adm": torch.zeros_like(image_conditioning), "c_crossattn": [unconditional_conditioning]}
-            else:
-                conditioning = {"c_concat": [image_conditioning], "c_crossattn": [conditioning]}
-                unconditional_conditioning = {"c_concat": [image_conditioning], "c_crossattn": [unconditional_conditioning]}
-
-        samples = self.launch_sampling(t_enc + 1, lambda: self.sampler.decode(x1, conditioning, t_enc, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning))
-
-        return samples
-
-    def sample(self, p, x, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
-        self.initialize(p)
-
-        self.init_latent = None
-        self.last_latent = x
-        self.step = 0
-
-        steps = self.adjust_steps_if_invalid(p, steps or p.steps)
-
-        # Wrap the conditioning models with additional image conditioning for inpainting model
-        # dummy_for_plms is needed because PLMS code checks the first item in the dict to have the right shape
-        if image_conditioning is not None:
-            if self.conditioning_key == "crossattn-adm":
-                conditioning = {"dummy_for_plms": np.zeros((conditioning.shape[0],)), "c_crossattn": [conditioning], "c_adm": image_conditioning}
-                unconditional_conditioning = {"c_crossattn": [unconditional_conditioning], "c_adm": torch.zeros_like(image_conditioning)}
-            else:
-                conditioning = {"dummy_for_plms": np.zeros((conditioning.shape[0],)), "c_crossattn": [conditioning], "c_concat": [image_conditioning]}
-                unconditional_conditioning = {"c_crossattn": [unconditional_conditioning], "c_concat": [image_conditioning]}
-
-        samples_ddim = self.launch_sampling(steps, lambda: self.sampler.sample(S=steps, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning, x_T=x, eta=self.eta)[0])
-
-        return samples_ddim
diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index 3a2e01b7..27a73486 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -1,8 +1,7 @@
-from collections import deque
 import torch
 import inspect
 import k_diffusion.sampling
-from modules import devices, sd_samplers_common, sd_samplers_extra, sd_samplers_cfg_denoiser
+from modules import sd_samplers_common, sd_samplers_extra, sd_samplers_cfg_denoiser
 
 from modules.shared import opts
 import modules.shared as shared
diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
index 8560d009..d89d0efb 100644
--- a/modules/sd_samplers_timesteps.py
+++ b/modules/sd_samplers_timesteps.py
@@ -7,9 +7,9 @@ from modules.shared import opts
 import modules.shared as shared
 
 samplers_timesteps = [
-    ('k_DDIM', sd_samplers_timesteps_impl.ddim, ['k_ddim'], {}),
-    ('k_PLMS', sd_samplers_timesteps_impl.plms, ['k_plms'], {}),
-    ('k_UniPC', sd_samplers_timesteps_impl.unipc, ['k_unipc'], {}),
+    ('DDIM', sd_samplers_timesteps_impl.ddim, ['ddim'], {}),
+    ('PLMS', sd_samplers_timesteps_impl.plms, ['plms'], {}),
+    ('UniPC', sd_samplers_timesteps_impl.unipc, ['unipc'], {}),
 ]
 
 
-- 
cgit v1.2.3


From f8ff8c0638997fd0aef217db1505598846f14782 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Tue, 8 Aug 2023 22:09:40 +0300
Subject: merge errors

---
 modules/sd_samplers_cfg_denoiser.py | 23 +++++++++++++++++++++--
 modules/sd_samplers_common.py       |  6 +++++-
 modules/sd_samplers_kdiffusion.py   | 17 ++++++++++++-----
 modules/sd_samplers_timesteps.py    | 27 +++++++++++++++++----------
 4 files changed, 55 insertions(+), 18 deletions(-)

(limited to 'modules/sd_samplers_timesteps.py')

diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py
index d826222c..a532e013 100644
--- a/modules/sd_samplers_cfg_denoiser.py
+++ b/modules/sd_samplers_cfg_denoiser.py
@@ -38,16 +38,24 @@ class CFGDenoiser(torch.nn.Module):
     negative prompt.
     """
 
-    def __init__(self, model, sampler):
+    def __init__(self, sampler):
         super().__init__()
-        self.inner_model = model
+        self.model_wrap = None
         self.mask = None
         self.nmask = None
         self.init_latent = None
+        self.steps = None
         self.step = 0
         self.image_cfg_scale = None
         self.padded_cond_uncond = False
         self.sampler = sampler
+        self.model_wrap = None
+        self.p = None
+
+    @property
+    def inner_model(self):
+        raise NotImplementedError()
+
 
     def combine_denoised(self, x_out, conds_list, uncond, cond_scale):
         denoised_uncond = x_out[-uncond.shape[0]:]
@@ -68,10 +76,21 @@ class CFGDenoiser(torch.nn.Module):
     def get_pred_x0(self, x_in, x_out, sigma):
         return x_out
 
+    def update_inner_model(self):
+        self.model_wrap = None
+
+        c, uc = self.p.get_conds()
+        self.sampler.sampler_extra_args['cond'] = c
+        self.sampler.sampler_extra_args['uncond'] = uc
+
     def forward(self, x, sigma, uncond, cond, cond_scale, s_min_uncond, image_cond):
         if state.interrupted or state.skipped:
             raise sd_samplers_common.InterruptedException
 
+        if sd_samplers_common.apply_refiner(self):
+            cond = self.sampler.sampler_extra_args['cond']
+            uncond = self.sampler.sampler_extra_args['uncond']
+
         # at self.image_cfg_scale == 1.0 produced results for edit model are the same as with normal sampling,
         # so is_edit_model is set to False to support AND composition.
         is_edit_model = shared.sd_model.cond_stage_key == "edit" and self.image_cfg_scale is not None and self.image_cfg_scale != 1.0
diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index 15f27970..fa3614ff 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -202,8 +202,9 @@ class Sampler:
 
         self.conditioning_key = shared.sd_model.model.conditioning_key
 
-        self.model_wrap = None
+        self.p = None
         self.model_wrap_cfg = None
+        self.sampler_extra_args = None
 
     def callback_state(self, d):
         step = d['i']
@@ -215,6 +216,7 @@ class Sampler:
         shared.total_tqdm.update()
 
     def launch_sampling(self, steps, func):
+        self.model_wrap_cfg.steps = steps
         state.sampling_steps = steps
         state.sampling_step = 0
 
@@ -234,6 +236,8 @@ class Sampler:
         return p.steps
 
     def initialize(self, p) -> dict:
+        self.p = p
+        self.model_wrap_cfg.p = p
         self.model_wrap_cfg.mask = p.mask if hasattr(p, 'mask') else None
         self.model_wrap_cfg.nmask = p.nmask if hasattr(p, 'nmask') else None
         self.model_wrap_cfg.step = 0
diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index 3ff4b634..95a43cef 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -52,17 +52,24 @@ k_diffusion_scheduler = {
 }
 
 
+class CFGDenoiserKDiffusion(sd_samplers_cfg_denoiser.CFGDenoiser):
+    @property
+    def inner_model(self):
+        if self.model_wrap is None:
+            denoiser = k_diffusion.external.CompVisVDenoiser if shared.sd_model.parameterization == "v" else k_diffusion.external.CompVisDenoiser
+            self.model_wrap = denoiser(shared.sd_model, quantize=shared.opts.enable_quantization)
+
+        return self.model_wrap
+
+
 class KDiffusionSampler(sd_samplers_common.Sampler):
     def __init__(self, funcname, sd_model):
-
         super().__init__(funcname)
 
-        self.extra_params = sampler_extra_params.get(funcname, [])
         self.func = funcname if callable(funcname) else getattr(k_diffusion.sampling, self.funcname)
 
-        denoiser = k_diffusion.external.CompVisVDenoiser if sd_model.parameterization == "v" else k_diffusion.external.CompVisDenoiser
-        self.model_wrap = denoiser(sd_model, quantize=shared.opts.enable_quantization)
-        self.model_wrap_cfg = sd_samplers_cfg_denoiser.CFGDenoiser(self.model_wrap, self)
+        self.model_wrap_cfg = CFGDenoiserKDiffusion(self)
+        self.model_wrap = self.model_wrap_cfg.inner_model
 
     def get_sigmas(self, p, steps):
         discard_next_to_last_sigma = self.config is not None and self.config.options.get('discard_next_to_last_sigma', False)
diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
index d89d0efb..965e61c6 100644
--- a/modules/sd_samplers_timesteps.py
+++ b/modules/sd_samplers_timesteps.py
@@ -44,10 +44,10 @@ class CompVisTimestepsVDenoiser(torch.nn.Module):
 
 class CFGDenoiserTimesteps(CFGDenoiser):
 
-    def __init__(self, model, sampler):
-        super().__init__(model, sampler)
+    def __init__(self, sampler):
+        super().__init__(sampler)
 
-        self.alphas = model.inner_model.alphas_cumprod
+        self.alphas = shared.sd_model.alphas_cumprod
 
     def get_pred_x0(self, x_in, x_out, sigma):
         ts = int(sigma.item())
@@ -60,6 +60,14 @@ class CFGDenoiserTimesteps(CFGDenoiser):
 
         return pred_x0
 
+    @property
+    def inner_model(self):
+        if self.model_wrap is None:
+            denoiser = CompVisTimestepsVDenoiser if shared.sd_model.parameterization == "v" else CompVisTimestepsDenoiser
+            self.model_wrap = denoiser(shared.sd_model)
+
+        return self.model_wrap
+
 
 class CompVisSampler(sd_samplers_common.Sampler):
     def __init__(self, funcname, sd_model):
@@ -68,9 +76,7 @@ class CompVisSampler(sd_samplers_common.Sampler):
         self.eta_option_field = 'eta_ddim'
         self.eta_infotext_field = 'Eta DDIM'
 
-        denoiser = CompVisTimestepsVDenoiser if sd_model.parameterization == "v" else CompVisTimestepsDenoiser
-        self.model_wrap = denoiser(sd_model)
-        self.model_wrap_cfg = CFGDenoiserTimesteps(self.model_wrap, self)
+        self.model_wrap_cfg = CFGDenoiserTimesteps(self)
 
     def get_timesteps(self, p, steps):
         discard_next_to_last_sigma = self.config is not None and self.config.options.get('discard_next_to_last_sigma', False)
@@ -106,7 +112,7 @@ class CompVisSampler(sd_samplers_common.Sampler):
 
         self.model_wrap_cfg.init_latent = x
         self.last_latent = x
-        extra_args = {
+        self.sampler_extra_args = {
             'cond': conditioning,
             'image_cond': image_conditioning,
             'uncond': unconditional_conditioning,
@@ -114,7 +120,7 @@ class CompVisSampler(sd_samplers_common.Sampler):
             's_min_uncond': self.s_min_uncond
         }
 
-        samples = self.launch_sampling(t_enc + 1, lambda: self.func(self.model_wrap_cfg, xi, extra_args=extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs))
+        samples = self.launch_sampling(t_enc + 1, lambda: self.func(self.model_wrap_cfg, xi, extra_args=self.sampler_extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs))
 
         if self.model_wrap_cfg.padded_cond_uncond:
             p.extra_generation_params["Pad conds"] = True
@@ -132,13 +138,14 @@ class CompVisSampler(sd_samplers_common.Sampler):
             extra_params_kwargs['timesteps'] = timesteps
 
         self.last_latent = x
-        samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args={
+        self.sampler_extra_args = {
             'cond': conditioning,
             'image_cond': image_conditioning,
             'uncond': unconditional_conditioning,
             'cond_scale': p.cfg_scale,
             's_min_uncond': self.s_min_uncond
-        }, disable=False, callback=self.callback_state, **extra_params_kwargs))
+        }
+        samples = self.launch_sampling(steps, lambda: self.func(self.model_wrap_cfg, x, extra_args=self.sampler_extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs))
 
         if self.model_wrap_cfg.padded_cond_uncond:
             p.extra_generation_params["Pad conds"] = True
-- 
cgit v1.2.3


From 0e83c675257f473e024511845e7940802333fd5f Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Tue, 8 Aug 2023 22:27:32 +0300
Subject: by request: fix tiled vae extension

---
 modules/sd_samplers_kdiffusion.py | 5 +++--
 modules/sd_samplers_timesteps.py  | 4 ++++
 2 files changed, 7 insertions(+), 2 deletions(-)

(limited to 'modules/sd_samplers_timesteps.py')

diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index f47431af..5613b8c1 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -1,7 +1,8 @@
 import torch
 import inspect
 import k_diffusion.sampling
-from modules import sd_samplers_common, sd_samplers_extra, sd_samplers_cfg_denoiser
+from modules import sd_samplers_common, sd_samplers_extra
+from modules.sd_samplers_cfg_denoiser import CFGDenoiser
 
 from modules.shared import opts
 import modules.shared as shared
@@ -62,7 +63,7 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
 
         denoiser = k_diffusion.external.CompVisVDenoiser if sd_model.parameterization == "v" else k_diffusion.external.CompVisDenoiser
         self.model_wrap = denoiser(sd_model, quantize=shared.opts.enable_quantization)
-        self.model_wrap_cfg = sd_samplers_cfg_denoiser.CFGDenoiser(self.model_wrap, self)
+        self.model_wrap_cfg = CFGDenoiser(self.model_wrap, self)
 
     def get_sigmas(self, p, steps):
         discard_next_to_last_sigma = self.config is not None and self.config.options.get('discard_next_to_last_sigma', False)
diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
index d89d0efb..f61799a8 100644
--- a/modules/sd_samplers_timesteps.py
+++ b/modules/sd_samplers_timesteps.py
@@ -1,5 +1,6 @@
 import torch
 import inspect
+import sys
 from modules import devices, sd_samplers_common, sd_samplers_timesteps_impl
 from modules.sd_samplers_cfg_denoiser import CFGDenoiser
 
@@ -145,3 +146,6 @@ class CompVisSampler(sd_samplers_common.Sampler):
 
         return samples
 
+
+sys.modules['modules.sd_samplers_compvis'] = sys.modules[__name__]
+VanillaStableDiffusionSampler = CompVisSampler  # temp. compatibility with older extensions
-- 
cgit v1.2.3


From aeb76ef174bc8a1904b25ca0b0b5009395f07d96 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Mon, 14 Aug 2023 08:49:02 +0300
Subject: repair DDIM/PLMS/UniPC batches

---
 modules/sd_samplers_timesteps.py      |  5 ++---
 modules/sd_samplers_timesteps_impl.py | 18 ++++++++++--------
 2 files changed, 12 insertions(+), 11 deletions(-)

(limited to 'modules/sd_samplers_timesteps.py')

diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
index 16572c7e..6aed2974 100644
--- a/modules/sd_samplers_timesteps.py
+++ b/modules/sd_samplers_timesteps.py
@@ -51,10 +51,9 @@ class CFGDenoiserTimesteps(CFGDenoiser):
         self.alphas = shared.sd_model.alphas_cumprod
 
     def get_pred_x0(self, x_in, x_out, sigma):
-        ts = int(sigma.item())
+        ts = sigma.to(dtype=int)
 
-        s_in = x_in.new_ones([x_in.shape[0]])
-        a_t = self.alphas[ts].item() * s_in
+        a_t = self.alphas[ts][:, None, None, None]
         sqrt_one_minus_at = (1 - a_t).sqrt()
 
         pred_x0 = (x_in - sqrt_one_minus_at * x_out) / a_t.sqrt()
diff --git a/modules/sd_samplers_timesteps_impl.py b/modules/sd_samplers_timesteps_impl.py
index d32e3521..a72daafd 100644
--- a/modules/sd_samplers_timesteps_impl.py
+++ b/modules/sd_samplers_timesteps_impl.py
@@ -16,16 +16,17 @@ def ddim(model, x, timesteps, extra_args=None, callback=None, disable=None, eta=
     sigmas = eta * np.sqrt((1 - alphas_prev.cpu().numpy()) / (1 - alphas.cpu()) * (1 - alphas.cpu() / alphas_prev.cpu().numpy()))
 
     extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
+    s_in = x.new_ones((x.shape[0]))
+    s_x = x.new_ones((x.shape[0], 1, 1, 1))
     for i in tqdm.trange(len(timesteps) - 1, disable=disable):
         index = len(timesteps) - 1 - i
 
         e_t = model(x, timesteps[index].item() * s_in, **extra_args)
 
-        a_t = alphas[index].item() * s_in
-        a_prev = alphas_prev[index].item() * s_in
-        sigma_t = sigmas[index].item() * s_in
-        sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_in
+        a_t = alphas[index].item() * s_x
+        a_prev = alphas_prev[index].item() * s_x
+        sigma_t = sigmas[index].item() * s_x
+        sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_x
 
         pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
         dir_xt = (1. - a_prev - sigma_t ** 2).sqrt() * e_t
@@ -47,13 +48,14 @@ def plms(model, x, timesteps, extra_args=None, callback=None, disable=None):
 
     extra_args = {} if extra_args is None else extra_args
     s_in = x.new_ones([x.shape[0]])
+    s_x = x.new_ones((x.shape[0], 1, 1, 1))
     old_eps = []
 
     def get_x_prev_and_pred_x0(e_t, index):
         # select parameters corresponding to the currently considered timestep
-        a_t = alphas[index].item() * s_in
-        a_prev = alphas_prev[index].item() * s_in
-        sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_in
+        a_t = alphas[index].item() * s_x
+        a_prev = alphas_prev[index].item() * s_x
+        sqrt_one_minus_at = sqrt_one_minus_alphas[index].item() * s_x
 
         # current prediction for x_0
         pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
-- 
cgit v1.2.3


From c1a31ec9f75c8dfe4ddcb0061f06e2704db98359 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Mon, 14 Aug 2023 08:59:15 +0300
Subject: revert to applying mask before denoising for k-diffusion, like it was
 before

---
 modules/sd_samplers_cfg_denoiser.py | 6 +++++-
 modules/sd_samplers_timesteps.py    | 1 +
 2 files changed, 6 insertions(+), 1 deletion(-)

(limited to 'modules/sd_samplers_timesteps.py')

diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py
index 113425b2..bc9b97e4 100644
--- a/modules/sd_samplers_cfg_denoiser.py
+++ b/modules/sd_samplers_cfg_denoiser.py
@@ -56,6 +56,7 @@ class CFGDenoiser(torch.nn.Module):
         self.sampler = sampler
         self.model_wrap = None
         self.p = None
+        self.mask_before_denoising = False
 
     @property
     def inner_model(self):
@@ -104,7 +105,7 @@ class CFGDenoiser(torch.nn.Module):
 
         assert not is_edit_model or all(len(conds) == 1 for conds in conds_list), "AND is not supported for InstructPix2Pix checkpoint (unless using Image CFG scale = 1.0)"
 
-        if self.mask is not None:
+        if self.mask_before_denoising and self.mask is not None:
             x = self.init_latent * self.mask + self.nmask * x
 
         batch_size = len(conds_list)
@@ -206,6 +207,9 @@ class CFGDenoiser(torch.nn.Module):
         else:
             denoised = self.combine_denoised(x_out, conds_list, uncond, cond_scale)
 
+        if not self.mask_before_denoising and self.mask is not None:
+            denoised = self.init_latent * self.mask + self.nmask * denoised
+
         self.sampler.last_latent = self.get_pred_x0(torch.cat([x_in[i:i + 1] for i in denoised_image_indexes]), torch.cat([x_out[i:i + 1] for i in denoised_image_indexes]), sigma)
 
         if opts.live_preview_content == "Prompt":
diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
index 6aed2974..c1f534ed 100644
--- a/modules/sd_samplers_timesteps.py
+++ b/modules/sd_samplers_timesteps.py
@@ -49,6 +49,7 @@ class CFGDenoiserTimesteps(CFGDenoiser):
         super().__init__(sampler)
 
         self.alphas = shared.sd_model.alphas_cumprod
+        self.mask_before_denoising = True
 
     def get_pred_x0(self, x_in, x_out, sigma):
         ts = sigma.to(dtype=int)
-- 
cgit v1.2.3


From 45be87afc6b173ebda94f427b8d396a2842ddf3c Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Mon, 14 Aug 2023 21:48:05 +0300
Subject: correctly add Eta DDIM to infotext when it's 1.0 and do not add it
 when it's 0.0.

---
 modules/sd_samplers_common.py    | 3 ++-
 modules/sd_samplers_timesteps.py | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

(limited to 'modules/sd_samplers_timesteps.py')

diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index 07fc4434..8886de7e 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -217,6 +217,7 @@ class Sampler:
 
         self.eta_option_field = 'eta_ancestral'
         self.eta_infotext_field = 'Eta'
+        self.eta_default = 1.0
 
         self.conditioning_key = shared.sd_model.model.conditioning_key
 
@@ -273,7 +274,7 @@ class Sampler:
                 extra_params_kwargs[param_name] = getattr(p, param_name)
 
         if 'eta' in inspect.signature(self.func).parameters:
-            if self.eta != 1.0:
+            if self.eta != self.eta_default:
                 p.extra_generation_params[self.eta_infotext_field] = self.eta
 
             extra_params_kwargs['eta'] = self.eta
diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
index c1f534ed..66e83ff7 100644
--- a/modules/sd_samplers_timesteps.py
+++ b/modules/sd_samplers_timesteps.py
@@ -76,6 +76,7 @@ class CompVisSampler(sd_samplers_common.Sampler):
 
         self.eta_option_field = 'eta_ddim'
         self.eta_infotext_field = 'Eta DDIM'
+        self.eta_default = 0.0
 
         self.model_wrap_cfg = CFGDenoiserTimesteps(self)
 
-- 
cgit v1.2.3


From 371b24b17c1cf98c9068a4b585b93cc1610702dc Mon Sep 17 00:00:00 2001
From: catboxanon <122327233+catboxanon@users.noreply.github.com>
Date: Tue, 15 Aug 2023 02:19:19 -0400
Subject: Add extra img2img noise

---
 modules/sd_samplers_kdiffusion.py | 4 ++++
 modules/sd_samplers_timesteps.py  | 4 ++++
 modules/shared_options.py         | 3 ++-
 scripts/xyz_grid.py               | 2 ++
 4 files changed, 12 insertions(+), 1 deletion(-)

(limited to 'modules/sd_samplers_timesteps.py')

diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index 08b9e740..08866e41 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -145,6 +145,10 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
 
         xi = x + noise * sigma_sched[0]
 
+        if opts.img2img_extra_noise > 0:
+            p.extra_generation_params["Extra noise"] = opts.img2img_extra_noise
+            xi += noise * opts.img2img_extra_noise
+
         extra_params_kwargs = self.initialize(p)
         parameters = inspect.signature(self.func).parameters
 
diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
index c1f534ed..670e2151 100644
--- a/modules/sd_samplers_timesteps.py
+++ b/modules/sd_samplers_timesteps.py
@@ -103,6 +103,10 @@ class CompVisSampler(sd_samplers_common.Sampler):
 
         xi = x * sqrt_alpha_cumprod + noise * sqrt_one_minus_alpha_cumprod
 
+        if opts.img2img_extra_noise > 0:
+            p.extra_generation_params["Extra noise"] = opts.img2img_extra_noise
+            xi += noise * opts.img2img_extra_noise * sqrt_alpha_cumprod
+
         extra_params_kwargs = self.initialize(p)
         parameters = inspect.signature(self.func).parameters
 
diff --git a/modules/shared_options.py b/modules/shared_options.py
index fc0de61f..79cbb92e 100644
--- a/modules/shared_options.py
+++ b/modules/shared_options.py
@@ -166,7 +166,8 @@ For img2img, VAE is used to process user's input image before the sampling, and
 
 options_templates.update(options_section(('img2img', "img2img"), {
     "inpainting_mask_weight": OptionInfo(1.0, "Inpainting conditioning mask strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext='Conditional mask weight'),
-    "initial_noise_multiplier": OptionInfo(1.0, "Noise multiplier for img2img", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01}, infotext='Noise multiplier'),
+    "initial_noise_multiplier": OptionInfo(1.0, "Noise multiplier for img2img", gr.Slider, {"minimum": 0.0, "maximum": 1.5, "step": 0.001}, infotext='Noise multiplier'),
+    "img2img_extra_noise": OptionInfo(0.0, "Extra noise multiplier for img2img and hires fix", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext='Extra noise').info("0 = disabled (default); should be lower than denoising strength"),
     "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."),
     "img2img_fix_steps": OptionInfo(False, "With img2img, do exactly the amount of steps the slider specifies.").info("normally you'd do less with less denoising"),
     "img2img_background_color": OptionInfo("#ffffff", "With img2img, fill transparent parts of the input image with this color.", ui_components.FormColorPicker, {}),
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index da0e48aa..e36bd3c9 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -241,6 +241,8 @@ axis_options = [
     AxisOption("Eta", float, apply_field("eta")),
     AxisOption("Clip skip", int, apply_clip_skip),
     AxisOption("Denoising", float, apply_field("denoising_strength")),
+    AxisOption("Initial noise multiplier", float, apply_field("initial_noise_multiplier")),
+    AxisOption("Extra noise", float, apply_override("img2img_extra_noise")),
     AxisOptionTxt2Img("Hires upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]),
     AxisOptionImg2Img("Cond. Image Mask Weight", float, apply_field("inpainting_mask_weight")),
     AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)),
-- 
cgit v1.2.3


From 254be4eeb259fd3bd2452250fca1bd278fa248ff Mon Sep 17 00:00:00 2001
From: catboxanon <122327233+catboxanon@users.noreply.github.com>
Date: Wed, 16 Aug 2023 21:45:19 -0400
Subject: Add extra noise callback

---
 modules/script_callbacks.py       | 26 ++++++++++++++++++++++++++
 modules/sd_samplers_kdiffusion.py |  4 ++++
 modules/sd_samplers_timesteps.py  |  4 ++++
 3 files changed, 34 insertions(+)

(limited to 'modules/sd_samplers_timesteps.py')

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 77ee55ee..fab23551 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -28,6 +28,15 @@ class ImageSaveParams:
         """dictionary with parameters for image's PNG info data; infotext will have the key 'parameters'"""
 
 
+class ExtraNoiseParams:
+    def __init__(self, noise, x):
+        self.noise = noise
+        """Random noise generated by the seed"""
+
+        self.x = x
+        """Latent image representation of the image"""
+
+
 class CFGDenoiserParams:
     def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps, text_cond, text_uncond):
         self.x = x
@@ -100,6 +109,7 @@ callback_map = dict(
     callbacks_ui_settings=[],
     callbacks_before_image_saved=[],
     callbacks_image_saved=[],
+    callbacks_extra_noise=[],
     callbacks_cfg_denoiser=[],
     callbacks_cfg_denoised=[],
     callbacks_cfg_after_cfg=[],
@@ -189,6 +199,14 @@ def image_saved_callback(params: ImageSaveParams):
             report_exception(c, 'image_saved_callback')
 
 
+def extra_noise_callback(params: ExtraNoiseParams):
+    for c in callback_map['callbacks_extra_noise']:
+        try:
+            c.callback(params)
+        except Exception:
+            report_exception(c, 'callbacks_extra_noise')
+
+
 def cfg_denoiser_callback(params: CFGDenoiserParams):
     for c in callback_map['callbacks_cfg_denoiser']:
         try:
@@ -367,6 +385,14 @@ def on_image_saved(callback):
     add_callback(callback_map['callbacks_image_saved'], callback)
 
 
+def on_extra_noise(callback):
+    """register a function to be called before adding extra noise in img2img or hires fix;
+    The callback is called with one argument:
+        - params: ExtraNoiseParams - contains noise determined by seed and latent representation of image
+    """
+    add_callback(callback_map['callbacks_extra_noise'], callback)
+
+
 def on_cfg_denoiser(callback):
     """register a function to be called in the kdiffussion cfg_denoiser method after building the inner model inputs.
     The callback is called with one argument:
diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index 08866e41..b9e0d577 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -3,6 +3,7 @@ import inspect
 import k_diffusion.sampling
 from modules import sd_samplers_common, sd_samplers_extra, sd_samplers_cfg_denoiser
 from modules.sd_samplers_cfg_denoiser import CFGDenoiser  # noqa: F401
+from modules.script_callbacks import ExtraNoiseParams, extra_noise_callback
 
 from modules.shared import opts
 import modules.shared as shared
@@ -147,6 +148,9 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
 
         if opts.img2img_extra_noise > 0:
             p.extra_generation_params["Extra noise"] = opts.img2img_extra_noise
+            extra_noise_params = ExtraNoiseParams(noise, x)
+            extra_noise_callback(extra_noise_params)
+            noise = extra_noise_params.noise
             xi += noise * opts.img2img_extra_noise
 
         extra_params_kwargs = self.initialize(p)
diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
index 68aea454..7a6cbd46 100644
--- a/modules/sd_samplers_timesteps.py
+++ b/modules/sd_samplers_timesteps.py
@@ -3,6 +3,7 @@ import inspect
 import sys
 from modules import devices, sd_samplers_common, sd_samplers_timesteps_impl
 from modules.sd_samplers_cfg_denoiser import CFGDenoiser
+from modules.script_callbacks import ExtraNoiseParams, extra_noise_callback
 
 from modules.shared import opts
 import modules.shared as shared
@@ -106,6 +107,9 @@ class CompVisSampler(sd_samplers_common.Sampler):
 
         if opts.img2img_extra_noise > 0:
             p.extra_generation_params["Extra noise"] = opts.img2img_extra_noise
+            extra_noise_params = ExtraNoiseParams(noise, x)
+            extra_noise_callback(extra_noise_params)
+            noise = extra_noise_params.noise
             xi += noise * opts.img2img_extra_noise * sqrt_alpha_cumprod
 
         extra_params_kwargs = self.initialize(p)
-- 
cgit v1.2.3