From 549b475be9b5cf1dca0a7bad2b6a6381e50e2b37 Mon Sep 17 00:00:00 2001
From: catboxanon <122327233+catboxanon@users.noreply.github.com>
Date: Tue, 29 Aug 2023 14:22:04 -0400
Subject: [PATCH] Add noisy latent to ExtraNoiseParams for callback

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

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index fab23551..c99695eb 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -29,12 +29,15 @@ class ImageSaveParams:
 
 
 class ExtraNoiseParams:
-    def __init__(self, noise, x):
+    def __init__(self, noise, x, xi):
         self.noise = noise
         """Random noise generated by the seed"""
 
         self.x = x
-        """Latent image representation of the image"""
+        """Latent representation of the image"""
+
+        self.xi = xi
+        """Noisy latent representation of the image"""
 
 
 class CFGDenoiserParams:
diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index 72c352a6..8a8c87e0 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -148,7 +148,7 @@ 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_params = ExtraNoiseParams(noise, x, xi)
             extra_noise_callback(extra_noise_params)
             noise = extra_noise_params.noise
             xi += noise * opts.img2img_extra_noise
diff --git a/modules/sd_samplers_timesteps.py b/modules/sd_samplers_timesteps.py
index 7a6cbd46..b17a8f93 100644
--- a/modules/sd_samplers_timesteps.py
+++ b/modules/sd_samplers_timesteps.py
@@ -107,7 +107,7 @@ 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_params = ExtraNoiseParams(noise, x, xi)
             extra_noise_callback(extra_noise_params)
             noise = extra_noise_params.noise
             xi += noise * opts.img2img_extra_noise * sqrt_alpha_cumprod