openvinotoolkit/stable-diffusion-webui
1
1
Fork 0
mirror of https://github.com/openvinotoolkit/stable-diffusion-webui.git synced 2024-12-15 07:03:06 +03:00
Code Issues Packages Projects Releases Wiki Activity

Add polling callback

Browse Source
This commit is contained in:
rucadi 2022-12-16 17:10:13 +01:00
parent 226d840e84
commit 3662a274e2
2 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
modules/script_callbacks.py
View File

@ -74,6 +74,7 @@ callback_map = dict(
callbacks_infotext_pasted=[], callbacks_infotext_pasted=[],
callbacks_script_unloaded=[], callbacks_script_unloaded=[],
callbacks_before_ui=[], callbacks_before_ui=[],
callbacks_on_polling=[],
) )
@ -89,6 +90,12 @@ def app_started_callback(demo: Optional[Blocks], app: FastAPI):
except Exception: except Exception:
report_exception(c, 'app_started_callback') 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): def model_loaded_callback(sd_model):
for c in callback_map['callbacks_model_loaded']: 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) 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): def on_model_loaded(callback):
"""register a function to be called when the stable diffusion model is created; the model is """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. """ passed as an argument; this function is also called when the script is reloaded. """

1
webui.py
View File

@ -171,6 +171,7 @@ def create_api(app):
def wait_on_server(demo=None): def wait_on_server(demo=None):
while 1: while 1:
time.sleep(0.5) time.sleep(0.5)
modules.script_callbacks.app_polling_callback(None, demo)
if shared.state.need_restart: if shared.state.need_restart:
shared.state.need_restart = False shared.state.need_restart = False
time.sleep(0.5) time.sleep(0.5)