Merge pull request #4021 from AUTOMATIC1111/add-kdiff-cfgdenoiser-callback

Add mid-kdiffusion cfgdenoiser script callback - access latents, conditionings and sigmas mid-sampling

1 files changed, 37 insertions, 1 deletions

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index ce264690..da88635b 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -26,6 +26,24 @@ class ImageSaveParams:
         """dictionary with parameters for image's PNG info data; infotext will have the key 'parameters'"""

 

 

+class CFGDenoiserParams:

+    def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps):

+        self.x = x

+        """Latent image representation in the process of being denoised"""

+        

+        self.image_cond = image_cond

+        """Conditioning image"""

+        

+        self.sigma = sigma

+        """Current sigma noise step value"""

+        

+        self.sampling_step = sampling_step

+        """Current Sampling step number"""

+        

+        self.total_sampling_steps = total_sampling_steps

+        """Total number of sampling steps planned"""

+

+

 ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])

 callbacks_app_started = []

 callbacks_model_loaded = []

@@ -33,6 +51,7 @@ callbacks_ui_tabs = []
 callbacks_ui_settings = []

 callbacks_before_image_saved = []

 callbacks_image_saved = []

+callbacks_cfg_denoiser = []

 

 

 def clear_callbacks():

@@ -41,7 +60,7 @@ def clear_callbacks():
     callbacks_ui_settings.clear()

     callbacks_before_image_saved.clear()

     callbacks_image_saved.clear()

-

+    callbacks_cfg_denoiser.clear()

 

 def app_started_callback(demo: Blocks, app: FastAPI):

     for c in callbacks_app_started:

@@ -95,6 +114,14 @@ def image_saved_callback(params: ImageSaveParams):
             report_exception(c, 'image_saved_callback')

 

 

+def cfg_denoiser_callback(params: CFGDenoiserParams):

+    for c in callbacks_cfg_denoiser:

+        try:

+            c.callback(params)

+        except Exception:

+            report_exception(c, 'cfg_denoiser_callback')

+

+

 def add_callback(callbacks, fun):

     stack = [x for x in inspect.stack() if x.filename != __file__]

     filename = stack[0].filename if len(stack) > 0 else 'unknown file'

@@ -147,3 +174,12 @@ def on_image_saved(callback):
         - params: ImageSaveParams - parameters the image was saved with. Changing fields in this object does nothing.

     """

     add_callback(callbacks_image_saved, 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:

+        - params: CFGDenoiserParams - parameters to be passed to the inner model and sampling state details.

+    """

+    add_callback(callbacks_cfg_denoiser, callback)

+