make it possible for scripts to add cross attention optimizations

add UI selection for cross attention optimization

1 files changed, 21 insertions, 0 deletions

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 3c21a362..40f388a5 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -110,6 +110,7 @@ callback_map = dict(
     callbacks_script_unloaded=[],

     callbacks_before_ui=[],

     callbacks_on_reload=[],

+    callbacks_list_optimizers=[],

 )

 

 

@@ -258,6 +259,18 @@ 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'

@@ -409,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)