Merge branch 'dev' into reorder-hotkeys

1 files changed, 60 insertions, 7 deletions

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 17109732..40f388a5 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -32,27 +32,42 @@ class CFGDenoiserParams:
     def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps, text_cond, text_uncond):

         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"""

-        

+

         self.text_cond = text_cond

         """ Encoder hidden states of text conditioning from prompt"""

-        

+

         self.text_uncond = text_uncond

         """ Encoder hidden states of text conditioning from negative prompt"""

 

 

 class CFGDenoisedParams:

+    def __init__(self, x, sampling_step, total_sampling_steps, inner_model):

+        self.x = x

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

+

+        self.sampling_step = sampling_step

+        """Current Sampling step number"""

+

+        self.total_sampling_steps = total_sampling_steps

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

+

+        self.inner_model = inner_model

+        """Inner model reference used for denoising"""

+

+

+class AfterCFGCallbackParams:

     def __init__(self, x, sampling_step, total_sampling_steps):

         self.x = x

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

@@ -87,6 +102,7 @@ callback_map = dict(
     callbacks_image_saved=[],

     callbacks_cfg_denoiser=[],

     callbacks_cfg_denoised=[],

+    callbacks_cfg_after_cfg=[],

     callbacks_before_component=[],

     callbacks_after_component=[],

     callbacks_image_grid=[],

@@ -94,6 +110,7 @@ callback_map = dict(
     callbacks_script_unloaded=[],

     callbacks_before_ui=[],

     callbacks_on_reload=[],

+    callbacks_list_optimizers=[],

 )

 

 

@@ -186,6 +203,14 @@ def cfg_denoised_callback(params: CFGDenoisedParams):
             report_exception(c, 'cfg_denoised_callback')

 

 

+def cfg_after_cfg_callback(params: AfterCFGCallbackParams):

+    for c in callback_map['callbacks_cfg_after_cfg']:

+        try:

+            c.callback(params)

+        except Exception:

+            report_exception(c, 'cfg_after_cfg_callback')

+

+

 def before_component_callback(component, **kwargs):

     for c in callback_map['callbacks_before_component']:

         try:

@@ -234,13 +259,25 @@ def before_ui_callback():
             report_exception(c, 'before_ui')

 

 

+def list_optimizers_callback():

+    res = []

+

+    for c in callback_map['callbacks_list_optimizers']:

+        try:

+            c.callback(res)

+        except Exception:

+            report_exception(c, 'list_optimizers')

+

+    return res

+

+

 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'

 

     callbacks.append(ScriptCallback(filename, fun))

 

-    

+

 def remove_current_script_callbacks():

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

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

@@ -332,6 +369,14 @@ def on_cfg_denoised(callback):
     add_callback(callback_map['callbacks_cfg_denoised'], callback)

 

 

+def on_cfg_after_cfg(callback):

+    """register a function to be called in the kdiffussion cfg_denoiser method after cfg calculations are completed.

+    The callback is called with one argument:

+        - params: AfterCFGCallbackParams - parameters to be passed to the script for post-processing after cfg calculation.

+    """

+    add_callback(callback_map['callbacks_cfg_after_cfg'], callback)

+

+

 def on_before_component(callback):

     """register a function to be called before a component is created.

     The callback is called with arguments:

@@ -377,3 +422,11 @@ def on_before_ui(callback):
     """register a function to be called before the UI is created."""

 

     add_callback(callback_map['callbacks_before_ui'], callback)

+

+

+def on_list_optimizers(callback):

+    """register a function to be called when UI is making a list of cross attention optimization options.

+    The function will be called with one argument, a list, and shall add objects of type modules.sd_hijack_optimizations.SdOptimization

+    to it."""

+

+    add_callback(callback_map['callbacks_list_optimizers'], callback)