From 1742c04bab5ff88b53dee60cfb9b90076dd98512 Mon Sep 17 00:00:00 2001 From: rucadi Date: Fri, 16 Dec 2022 17:10:13 +0100 Subject: [PATCH 1/2] 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 8e22f8755..715e1830d 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 c2d0c6be8..5f7a53f9c 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) From 0c8825b2bec3a68836eacf37718306c9c78554a0 Mon Sep 17 00:00:00 2001 From: rucadi Date: Fri, 16 Dec 2022 18:31:20 +0100 Subject: [PATCH 2/2] Add a callback called before reloading the server --- modules/script_callbacks.py | 13 ++++++++++++- webui.py | 1 + 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index 715e1830d..b646b0f97 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -64,6 +64,7 @@ callback_map = dict( callbacks_before_component=[], callbacks_after_component=[], callbacks_on_polling=[], + callbacks_on_reload=[], ) @@ -71,7 +72,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: @@ -86,6 +86,14 @@ def app_polling_callback(demo: Optional[Blocks], app: FastAPI): 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']: try: @@ -195,6 +203,9 @@ 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 diff --git a/webui.py b/webui.py index 5f7a53f9c..dda34249b 100644 --- a/webui.py +++ b/webui.py @@ -108,6 +108,7 @@ def wait_on_server(demo=None): time.sleep(0.5) modules.script_callbacks.app_polling_callback(None, demo) if shared.state.need_restart: + modules.script_callbacks.app_reload_callback(None, demo) shared.state.need_restart = False time.sleep(0.5) demo.close()