From 423f22228306ae72d0480e25add9777c3c5d8fdf Mon Sep 17 00:00:00 2001 From: Maiko Tan Date: Sun, 30 Oct 2022 22:46:43 +0800 Subject: [PATCH 1/2] feat: add app started callback --- modules/script_callbacks.py | 15 +++++++++++++++ webui.py | 3 +++ 2 files changed, 18 insertions(+) diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index 6ea58d616..f5509629d 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,10 @@ def add_callback(callbacks, fun): callbacks.append(ScriptCallback(filename, fun)) +def on_app_started(callback): + 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""" diff --git a/webui.py b/webui.py index 29530872c..13407e429 100644 --- a/webui.py +++ b/webui.py @@ -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 @@ -135,6 +136,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() From 081df45da47feadfb055552c0ad5c6e6ecdd9f28 Mon Sep 17 00:00:00 2001 From: Maiko Sinkyaet Tan Date: Mon, 31 Oct 2022 08:47:43 +0800 Subject: [PATCH 2/2] docs: add python doc (?) not sure if this available... --- modules/script_callbacks.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index f5509629d..d8aa6f008 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -103,6 +103,8 @@ def add_callback(callbacks, 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)