From f04fa75eaee262011817bab4a6e4ffb052c7e22c Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Mon, 29 Jan 2024 10:10:10 -0500
Subject: [PATCH] add on_after_ui callback

---
 modules/script_callbacks.py | 16 ++++++++++++++++
 modules/ui.py               |  1 +
 webui.py                    |  1 +
 3 files changed, 18 insertions(+)

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 0d0809573..d96baf733 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -123,6 +123,7 @@ callback_map = dict(
     callbacks_infotext_pasted=[],
     callbacks_script_unloaded=[],
     callbacks_before_ui=[],
+    callbacks_after_ui=[],
     callbacks_on_reload=[],
 )
 
@@ -359,6 +360,16 @@ def before_ui_callback():
             report_exception(e, c, 'before_ui')
 
 
+def after_ui_callback():
+    for c in reversed(callback_map['callbacks_after_ui']):
+        try:
+            t0 = time.time()
+            c.callback()
+            timer(t0, c.script, 'after_ui')
+        except Exception as e:
+            report_exception(e, c, 'after_ui')
+
+
 def add_callback(callbacks, fun):
     # stack = [x for x in inspect.stack(0) if x.filename != __file__]
     # filename = stack[0].filename if len(stack) > 0 else 'unknown file'
@@ -536,3 +547,8 @@ def on_script_unloaded(callback):
 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_after_ui(callback):
+    """register a function to be called before the UI is created."""
+    add_callback(callback_map['callbacks_after_ui'], callback)
diff --git a/modules/ui.py b/modules/ui.py
index 22ee5a296..b6d99ed78 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -405,6 +405,7 @@ def create_ui(startup_timer = None):
                 loadsave.add_block(interface, ifid)
             loadsave.add_component(f"webui/Tabs@{tabs.elem_id}", tabs)
             loadsave.setup_ui()
+
         if opts.notification_audio_enable and os.path.exists(os.path.join(script_path, opts.notification_audio_path)):
             gr.Audio(interactive=False, value=os.path.join(script_path, opts.notification_audio_path), elem_id="audio_notification", visible=False)
 
diff --git a/webui.py b/webui.py
index f89e20560..53ca3730c 100644
--- a/webui.py
+++ b/webui.py
@@ -304,6 +304,7 @@ def webui(restart=False):
 
     start_common()
     start_ui()
+    modules.script_callbacks.after_ui_callback()
     modules.sd_models.write_metadata()
     load_model()
     shared.opts.save(shared.config_filename)