Made WrapperHook functional, added another wrapper/callback getter, added ON_DETACH callback to ModelPatcher

comfy/hooks.py

@@ -24,8 +24,8 @@ class EnumHookType(enum.Enum):
     Patch = "patch"
     ObjectPatch = "object_patch"
     AddModels = "add_models"
-    Callbacks = "add_callback"
-    Wrappers = "add_wrapper"
+    Callbacks = "callbacks"
+    Wrappers = "wrappers"
     SetInjections = "add_injections"
 
 class EnumWeightTarget(enum.Enum):
@@ -131,6 +131,7 @@ class WeightHook(Hook):
             else:
                 weights = self.weights_clip
         k = model.add_hook_patches(hook=self, patches=weights, strength_patch=strength)
+        registered.append(self)
         return True
         # TODO: add logs about any keys that were not applied
 
@@ -204,7 +205,7 @@ class CallbackHook(Hook):
     # TODO: add functionality
 
 class WrapperHook(Hook):
-    def __init__(self, wrappers_dict: dict[str, dict[str, list[Callable]]]=None):
+    def __init__(self, wrappers_dict: dict[str, dict[str, dict[str, list[Callable]]]]=None):
         super().__init__(hook_type=EnumHookType.Wrappers)
         self.wrappers_dict = wrappers_dict
 
@@ -215,11 +216,13 @@ class WrapperHook(Hook):
         c.wrappers_dict = self.wrappers_dict
         return c
     
-    def add_hook_wrapper(self, model: 'ModelPatcher', model_options: dict, target: EnumWeightTarget, registered: list[Hook]):
+    def add_hook_patches(self, model: 'ModelPatcher', model_options: dict, target: EnumWeightTarget, registered: list[Hook]):
         if not self.should_register(model, model_options, target, registered):
             return False
-        add_model_options = {"transformer_options": {"wrappers": self.wrappers_dict}}
+        add_model_options = {"transformer_options": self.wrappers_dict}
         comfy.patcher_extension.merge_nested_dicts(model_options, add_model_options, copy_dict1=False)
+        registered.append(self)
+        return True
 
 class SetInjectionsHook(Hook):
     def __init__(self, key: str=None, injections: list['PatcherInjection']=None):

comfy/model_patcher.py

@@ -753,6 +753,8 @@ class ModelPatcher:
         self.model_patches_to(self.offload_device)
         if unpatch_all:
             self.unpatch_model(self.offload_device, unpatch_weights=unpatch_all)
+        for callback in self.get_all_callbacks(CallbacksMP.ON_DETACH):
+            callback(self, unpatch_all)
         return self.model
 
     def current_loaded_device(self):

comfy/patcher_extension.py

@@ -4,6 +4,7 @@ from typing import Callable
 class CallbacksMP:
     ON_CLONE = "on_clone"
     ON_LOAD = "on_load_after"
+    ON_DETACH = "on_detach_after"
     ON_CLEANUP = "on_cleanup"
     ON_PRE_RUN = "on_pre_run"
     ON_PREPARE_STATE = "on_prepare_state"
@@ -23,11 +24,19 @@ def add_callback(call_type: str, callback: Callable, transformer_options: dict,
 
 def add_callback_with_key(call_type: str, key: str, callback: Callable, transformer_options: dict, is_model_options=False):
     if is_model_options:
-        transformer_options = transformer_options.get("transformer_options", {})
-    callbacks: dict[str, dict[str, list]] = transformer_options.get("callbacks", {})
+        transformer_options = transformer_options.setdefault("transformer_options", {})
+    callbacks: dict[str, dict[str, list]] = transformer_options.setdefault("callbacks", {})
     c = callbacks.setdefault(call_type, {}).setdefault(key, [])
     c.append(callback)
 
+def get_callbacks_with_key(call_type: str, key: str, transformer_options: dict, is_model_options=False):
+    if is_model_options:
+        transformer_options = transformer_options.get("transformer_options", {})
+    c_list = []
+    callbacks: dict[str, list] = transformer_options.get("callbacks", {})
+    c_list.extend(callbacks.get(call_type, {}).get(key, []))
+    return c_list
+
 def get_all_callbacks(call_type: str, transformer_options: dict, is_model_options=False):
     if is_model_options:
         transformer_options = transformer_options.get("transformer_options", {})
@@ -55,11 +64,19 @@ def add_wrapper(wrapper_type: str, wrapper: Callable, transformer_options: dict,
 
 def add_wrapper_with_key(wrapper_type: str, key: str, wrapper: Callable, transformer_options: dict, is_model_options=False):
     if is_model_options:
-        transformer_options = transformer_options.get("transformer_options", {})
-    wrappers: dict[str, dict[str, list]] = transformer_options.get("wrappers", {})
+        transformer_options = transformer_options.setdefault("transformer_options", {})
+    wrappers: dict[str, dict[str, list]] = transformer_options.setdefault("wrappers", {})
     w = wrappers.setdefault(wrapper_type, {}).setdefault(key, [])
     w.append(wrapper)
 
+def get_wrappers_with_key(wrapper_type: str, key: str, transformer_options: dict, is_model_options=False):
+    if is_model_options:
+        transformer_options = transformer_options.get("transformer_options", {})
+    w_list = []
+    wrappers: dict[str, list] = transformer_options.get("wrappers", {})
+    w_list.extend(wrappers.get(wrapper_type, {}).get(key, []))
+    return w_list
+
 def get_all_wrappers(wrapper_type: str, transformer_options: dict, is_model_options=False):
     if is_model_options:
         transformer_options = transformer_options.get("transformer_options", {})