From bcc6a22178a17d5b77e0cbd16dbcf9c3e69ca89b Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Thu, 14 Nov 2024 08:06:02 -0600 Subject: [PATCH] Made WrapperHook functional, added another wrapper/callback getter, added ON_DETACH callback to ModelPatcher --- comfy/hooks.py | 13 ++++++++----- comfy/model_patcher.py | 2 ++ comfy/patcher_extension.py | 25 +++++++++++++++++++++---- 3 files changed, 31 insertions(+), 9 deletions(-) diff --git a/comfy/hooks.py b/comfy/hooks.py index 7530523ba..f0f112c63 100644 --- a/comfy/hooks.py +++ b/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): diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 50fa977f7..f482dcfd7 100644 --- a/comfy/model_patcher.py +++ b/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): diff --git a/comfy/patcher_extension.py b/comfy/patcher_extension.py index a95cd8746..514469185 100644 --- a/comfy/patcher_extension.py +++ b/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", {})