Merge branch 'master' into hn-activation

1 files changed, 94 insertions, 8 deletions

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index dc520abc..c28e220e 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -2,23 +2,73 @@ import sys
 import traceback

 from collections import namedtuple

 import inspect

+from typing import Optional

 

+from fastapi import FastAPI

+from gradio import Blocks

 

 def report_exception(c, job):

     print(f"Error executing callback {job} for {c.script}", file=sys.stderr)

     print(traceback.format_exc(), file=sys.stderr)

 

 

+class ImageSaveParams:

+    def __init__(self, image, p, filename, pnginfo):

+        self.image = image

+        """the PIL image itself"""

+

+        self.p = p

+        """p object with processing parameters; either StableDiffusionProcessing or an object with same fields"""

+

+        self.filename = filename

+        """name of file that the image would be saved to"""

+

+        self.pnginfo = pnginfo

+        """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 = []

 callbacks_ui_tabs = []

 callbacks_ui_settings = []

+callbacks_before_image_saved = []

 callbacks_image_saved = []

+callbacks_cfg_denoiser = []

+

 

 def clear_callbacks():

     callbacks_model_loaded.clear()

     callbacks_ui_tabs.clear()

+    callbacks_ui_settings.clear()

+    callbacks_before_image_saved.clear()

     callbacks_image_saved.clear()

+    callbacks_cfg_denoiser.clear()

+

+def app_started_callback(demo: Optional[Blocks], app: FastAPI):

+    for c in callbacks_app_started:

+        try:

+            c.callback(demo, app)

+        except Exception:

+            report_exception(c, 'app_started_callback')

 

 

 def model_loaded_callback(sd_model):

@@ -49,14 +99,30 @@ def ui_settings_callback():
             report_exception(c, 'ui_settings_callback')

 

 

-def image_saved_callback(image, p, fullfn, txt_fullfn):

+def before_image_saved_callback(params: ImageSaveParams):

+    for c in callbacks_before_image_saved:

+        try:

+            c.callback(params)

+        except Exception:

+            report_exception(c, 'before_image_saved_callback')

+

+

+def image_saved_callback(params: ImageSaveParams):

     for c in callbacks_image_saved:

         try:

-            c.callback(image, p, fullfn, txt_fullfn)

+            c.callback(params)

         except Exception:

             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'

@@ -64,6 +130,11 @@ def add_callback(callbacks, fun):
     callbacks.append(ScriptCallback(filename, fun))

 

 

+def on_app_started(callback):

+    """register a function to be called when the webui started, the gradio `Block` component and

+    fastapi `FastAPI` object are passed as the arguments"""

+    add_callback(callbacks_app_started, callback)

+

 

 def on_model_loaded(callback):

     """register a function to be called when the stable diffusion model is created; the model is

@@ -90,11 +161,26 @@ def on_ui_settings(callback):
     add_callback(callbacks_ui_settings, callback)

 

 

-def on_save_imaged(callback):

-    """register a function to be called after modules.images.save_image is called.

-    The callback is called with three arguments:

-        - p - procesing object (or a dummy object with same fields if the image is saved using save button)

-        - fullfn - image filename

-        - txt_fullfn - text file with parameters; may be None

+def on_before_image_saved(callback):

+    """register a function to be called before an image is saved to a file.

+    The callback is called with one argument:

+        - params: ImageSaveParams - parameters the image is to be saved with. You can change fields in this object.

+    """

+    add_callback(callbacks_before_image_saved, callback)

+

+

+def on_image_saved(callback):

+    """register a function to be called after an image is saved to a file.

+    The callback is called with one argument:

+        - 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)

+