diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index c28e220e..4a7fb944 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -46,25 +46,23 @@ class CFGDenoiserParams: ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"]) -callbacks_app_started = [] -callbacks_model_loaded = [] -callbacks_ui_tabs = [] -callbacks_ui_settings = [] -callbacks_before_image_saved = [] -callbacks_image_saved = [] -callbacks_cfg_denoiser = [] +__callback_map = dict( + callbacks_app_started=[], + callbacks_model_loaded=[], + callbacks_ui_tabs=[], + callbacks_ui_settings=[], + callbacks_before_image_saved=[], + callbacks_image_saved=[], + callbacks_cfg_denoiser=[] +) def clear_callbacks(): - callbacks_model_loaded.clear() - callbacks_ui_tabs.clear() - callbacks_ui_settings.clear() - callbacks_before_image_saved.clear() - callbacks_image_saved.clear() - callbacks_cfg_denoiser.clear() + for callback_list in __callback_map.values(): + callback_list.clear() def app_started_callback(demo: Optional[Blocks], app: FastAPI): - for c in callbacks_app_started: + for c in __callback_map['callbacks_app_started']: try: c.callback(demo, app) except Exception: @@ -72,7 +70,7 @@ def app_started_callback(demo: Optional[Blocks], app: FastAPI): def model_loaded_callback(sd_model): - for c in callbacks_model_loaded: + for c in __callback_map['callbacks_model_loaded']: try: c.callback(sd_model) except Exception: @@ -82,7 +80,7 @@ def model_loaded_callback(sd_model): def ui_tabs_callback(): res = [] - for c in callbacks_ui_tabs: + for c in __callback_map['callbacks_ui_tabs']: try: res += c.callback() or [] except Exception: @@ -92,7 +90,7 @@ def ui_tabs_callback(): def ui_settings_callback(): - for c in callbacks_ui_settings: + for c in __callback_map['callbacks_ui_settings']: try: c.callback() except Exception: @@ -100,7 +98,7 @@ def ui_settings_callback(): def before_image_saved_callback(params: ImageSaveParams): - for c in callbacks_before_image_saved: + for c in __callback_map['callbacks_before_image_saved']: try: c.callback(params) except Exception: @@ -108,7 +106,7 @@ def before_image_saved_callback(params: ImageSaveParams): def image_saved_callback(params: ImageSaveParams): - for c in callbacks_image_saved: + for c in __callback_map['callbacks_image_saved']: try: c.callback(params) except Exception: @@ -116,7 +114,7 @@ def image_saved_callback(params: ImageSaveParams): def cfg_denoiser_callback(params: CFGDenoiserParams): - for c in callbacks_cfg_denoiser: + for c in __callback_map['callbacks_cfg_denoiser']: try: c.callback(params) except Exception: @@ -129,17 +127,33 @@ def add_callback(callbacks, fun): callbacks.append(ScriptCallback(filename, fun)) + +def remove_current_script_callbacks(): + stack = [x for x in inspect.stack() if x.filename != __file__] + filename = stack[0].filename if len(stack) > 0 else 'unknown file' + if filename == 'unknown file': + return + for callback_list in __callback_map.values(): + for callback_to_remove in [cb for cb in callback_list if cb.script == filename]: + callback_list.remove(callback_to_remove) + + +def remove_callbacks_for_function(callback_func): + for callback_list in __callback_map.values(): + for callback_to_remove in [cb for cb in callback_list if cb.callback == callback_func]: + callback_list.remove(callback_to_remove) + def on_app_started(callback): """register a function to be called when the webui started, the gradio `Block` component and fastapi `FastAPI` object are passed as the arguments""" - add_callback(callbacks_app_started, callback) + add_callback(__callback_map['callbacks_app_started'], 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""" - add_callback(callbacks_model_loaded, callback) + add_callback(__callback_map['callbacks_model_loaded'], callback) def on_ui_tabs(callback): @@ -152,13 +166,13 @@ def on_ui_tabs(callback): title is tab text displayed to user in the UI elem_id is HTML id for the tab """ - add_callback(callbacks_ui_tabs, callback) + add_callback(__callback_map['callbacks_ui_tabs'], callback) def on_ui_settings(callback): """register a function to be called before UI settings are populated; add your settings by using shared.opts.add_option(shared.OptionInfo(...)) """ - add_callback(callbacks_ui_settings, callback) + add_callback(__callback_map['callbacks_ui_settings'], callback) def on_before_image_saved(callback): @@ -166,7 +180,7 @@ def on_before_image_saved(callback): The callback is called with one argument: - params: ImageSaveParams - parameters the image is to be saved with. You can change fields in this object. """ - add_callback(callbacks_before_image_saved, callback) + add_callback(__callback_map['callbacks_before_image_saved'], callback) def on_image_saved(callback): @@ -174,7 +188,7 @@ def on_image_saved(callback): The callback is called with one argument: - params: ImageSaveParams - parameters the image was saved with. Changing fields in this object does nothing. """ - add_callback(callbacks_image_saved, callback) + add_callback(__callback_map['callbacks_image_saved'], callback) def on_cfg_denoiser(callback): @@ -182,5 +196,5 @@ def on_cfg_denoiser(callback): The callback is called with one argument: - params: CFGDenoiserParams - parameters to be passed to the inner model and sampling state details. """ - add_callback(callbacks_cfg_denoiser, callback) + add_callback(__callback_map['callbacks_cfg_denoiser'], callback)