Merge branch 'release_candidate'

1 files changed, 30 insertions, 1 deletions

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 9ed7ad21..08bc5256 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -1,3 +1,4 @@
+import dataclasses

 import inspect

 import os

 from collections import namedtuple

@@ -41,7 +42,7 @@ class ExtraNoiseParams:
 

 

 class CFGDenoiserParams:

-    def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps, text_cond, text_uncond):

+    def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps, text_cond, text_uncond, denoiser=None):

         self.x = x

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

 

@@ -63,6 +64,9 @@ class CFGDenoiserParams:
         self.text_uncond = text_uncond

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

 

+        self.denoiser = denoiser

+        """Current CFGDenoiser object with processing parameters"""

+

 

 class CFGDenoisedParams:

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

@@ -103,6 +107,15 @@ class ImageGridLoopParams:
         self.rows = rows

 

 

+@dataclasses.dataclass

+class BeforeTokenCounterParams:

+    prompt: str

+    steps: int

+    styles: list

+

+    is_positive: bool = True

+

+

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

 callback_map = dict(

     callbacks_app_started=[],

@@ -125,6 +138,7 @@ callback_map = dict(
     callbacks_on_reload=[],

     callbacks_list_optimizers=[],

     callbacks_list_unets=[],

+    callbacks_before_token_counter=[],

 )

 

 

@@ -306,6 +320,14 @@ def list_unets_callback():
     return res

 

 

+def before_token_counter_callback(params: BeforeTokenCounterParams):

+    for c in callback_map['callbacks_before_token_counter']:

+        try:

+            c.callback(params)

+        except Exception:

+            report_exception(c, 'before_token_counter')

+

+

 def add_callback(callbacks, fun):

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

     filename = stack[0].filename if stack else 'unknown file'

@@ -480,3 +502,10 @@ def on_list_unets(callback):
     The function will be called with one argument, a list, and shall add objects of type modules.sd_unet.SdUnetOption to it."""

 

     add_callback(callback_map['callbacks_list_unets'], callback)

+

+

+def on_before_token_counter(callback):

+    """register a function to be called when UI is counting tokens for a prompt.

+    The function will be called with one argument of type BeforeTokenCounterParams, and should modify its fields if necessary."""

+

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