Add polling callback

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

View File

@@ -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"""