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: [PATCH] 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(-)

diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index 08b9e7408..08866e412 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 c1f534edf..670e2151b 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 fc0de61fa..79cbb92e2 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 da0e48aa5..e36bd3c94 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)),