From 3662a274e2b6482c4ad831cc2d7976d919b40212 Mon Sep 17 00:00:00 2001 From: rucadi Date: Fri, 16 Dec 2022 17:10:13 +0100 Subject: [PATCH] Add polling callback --- modules/script_callbacks.py | 12 ++++++++++++ webui.py | 1 + 2 files changed, 13 insertions(+) diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index 4bb45ec7..7763936f 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -74,6 +74,7 @@ callback_map = dict( callbacks_infotext_pasted=[], callbacks_script_unloaded=[], callbacks_before_ui=[], + callbacks_on_polling=[], ) @@ -89,6 +90,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']: @@ -227,6 +234,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; this function is also called when the script is reloaded. """ diff --git a/webui.py b/webui.py index 5b5c2139..6c2b511c 100644 --- a/webui.py +++ b/webui.py @@ -171,6 +171,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)