mirror/stable-diffusion-webui
2
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2024-12-21 07:30:02 +08: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 685f9631b5
commit 1742c04bab
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

@ -63,6 +63,7 @@ callback_map = dict(
callbacks_cfg_denoiser=[], callbacks_cfg_denoiser=[],
callbacks_before_component=[], callbacks_before_component=[],
callbacks_after_component=[], callbacks_after_component=[],
callbacks_on_polling=[],
) )
@ -78,6 +79,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']:
@ -184,6 +191,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""" passed as an argument"""

1
webui.py
View File

@ -106,6 +106,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)