diff options
| -rw-r--r-- | modules/script_callbacks.py | 17 | ||||
| -rw-r--r-- | webui.py | 3 |
2 files changed, 20 insertions, 0 deletions
diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index 6ea58d61..d8aa6f00 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -3,6 +3,8 @@ import traceback from collections import namedtuple
import inspect
+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)
@@ -25,6 +27,7 @@ class ImageSaveParams: ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])
+callbacks_app_started = []
callbacks_model_loaded = []
callbacks_ui_tabs = []
callbacks_ui_settings = []
@@ -40,6 +43,14 @@ def clear_callbacks(): callbacks_image_saved.clear()
+def app_started_callback(demo: 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):
for c in callbacks_model_loaded:
try:
@@ -91,6 +102,12 @@ 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
passed as an argument"""
@@ -23,6 +23,7 @@ import modules.sd_hijack import modules.sd_models
import modules.shared as shared
import modules.txt2img
+import modules.script_callbacks
import modules.ui
from modules import devices
@@ -140,6 +141,8 @@ def webui(): if launch_api:
create_api(app)
+ modules.script_callbacks.app_started_callback(demo, app)
+
wait_on_server(demo)
sd_samplers.set_samplers()
|