Merge pull request #3982 from MaikoTan/on-started-callback

feat: add app started callback

1 files changed, 17 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"""