Convert callbacks into a private map, add utility functions for removing callbacks

Author: dfaker

modules/script_callbacks.py

@@ -46,33 +46,31 @@ def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps):
 
 
 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:
             report_exception(c, 'app_started_callback')
 
 
 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,31 +90,31 @@ 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:
             report_exception(c, '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:
             report_exception(c, 'before_image_saved_callback')
 
 
 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:
             report_exception(c, 'image_saved_callback')
 
 
 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,35 +166,35 @@ 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):
     """register a function to be called before an image is saved to a file.
     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):
     """register a function to be called after an image is saved to a file.
     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):
     """register a function to be called in the kdiffussion cfg_denoiser method after building the inner model inputs.
     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)