Add polling callback

2 files changed, 13 insertions, 0 deletions

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 8e22f875..715e1830 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -63,6 +63,7 @@ callback_map = dict(
     callbacks_cfg_denoiser=[],

     callbacks_before_component=[],

     callbacks_after_component=[],

+    callbacks_on_polling=[],

 )

 

 

@@ -78,6 +79,12 @@ 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 model_loaded_callback(sd_model):

     for c in callback_map['callbacks_model_loaded']:

@@ -184,6 +191,11 @@ 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_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 c2d0c6be..5f7a53f9 100644
--- a/webui.py
+++ b/webui.py
@@ -106,6 +106,7 @@ def create_api(app):
 def wait_on_server(demo=None):

     while 1:

         time.sleep(0.5)

+        modules.script_callbacks.app_polling_callback(None, demo)

         if shared.state.need_restart:

             shared.state.need_restart = False

             time.sleep(0.5)