From 2582a0fd3b3e91c5fba9e5e561cbdf5fee835063 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Thu, 18 May 2023 22:48:28 +0300
Subject: make it possible for scripts to add cross attention optimizations add
 UI selection for cross attention optimization

---
 modules/script_callbacks.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

(limited to 'modules/script_callbacks.py')

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)
-- 
cgit v1.2.3