Add extra noise callback

1 files changed, 26 insertions, 0 deletions

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: