Add extra noise callback

1 files changed, 4 insertions, 0 deletions

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)