Merge pull request #8789 from Rucadi/master

Add polling and reload callback for extensions.

1 files changed, 24 insertions, 1 deletions

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 07911876..cede1cb0 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -93,6 +93,8 @@ callback_map = dict(
     callbacks_infotext_pasted=[],

     callbacks_script_unloaded=[],

     callbacks_before_ui=[],

+    callbacks_on_reload=[],

+    callbacks_on_polling=[],

 )

 

 

@@ -100,7 +102,6 @@ def clear_callbacks():
     for callback_list in callback_map.values():

         callback_list.clear()

 

-

 def app_started_callback(demo: Optional[Blocks], app: FastAPI):

     for c in callback_map['callbacks_app_started']:

         try:

@@ -108,6 +109,20 @@ def app_started_callback(demo: Optional[Blocks], app: FastAPI):
         except Exception:

             report_exception(c, 'app_started_callback')

 

+def app_polling_callback(demo: Optional[Blocks], app: FastAPI):

+    for c in callback_map['callbacks_on_polling']:

+        try:

+            c.callback()

+        except Exception:

+            report_exception(c, 'callbacks_on_polling')

+

+def app_reload_callback(demo: Optional[Blocks], app: FastAPI):

+    for c in callback_map['callbacks_on_reload']:

+        try:

+            c.callback()

+        except Exception:

+            report_exception(c, 'callbacks_on_reload')

+

 

 def model_loaded_callback(sd_model):

     for c in callback_map['callbacks_model_loaded']:

@@ -254,6 +269,14 @@ def on_app_started(callback):
     add_callback(callback_map['callbacks_app_started'], callback)

 

 

+def on_polling(callback):

+    """register a function to be called on each polling of the server."""

+    add_callback(callback_map['callbacks_on_polling'], callback)

+

+def on_before_reload(callback):

+    """register a function to be called just before the server reloads."""

+    add_callback(callback_map['callbacks_on_reload'], 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; this function is also called when the script is reloaded. """