From a5034df6db06b646fb11c5b4c8fb947ba4fadff7 Mon Sep 17 00:00:00 2001 From: "kosinkadink1@gmail.com" Date: Sun, 15 Sep 2024 15:47:09 +0900 Subject: [PATCH] Made CLIP work with hook patches --- comfy/hooks.py | 12 +++--------- comfy/model_patcher.py | 10 +++++++++- comfy_extras/nodes_hooks.py | 24 +++++++++++++++++++++++- 3 files changed, 35 insertions(+), 11 deletions(-) diff --git a/comfy/hooks.py b/comfy/hooks.py index 80af32ffd..8eb9e1e59 100644 --- a/comfy/hooks.py +++ b/comfy/hooks.py @@ -58,7 +58,6 @@ class HookGroup: def clone(self): c = HookGroup() - # TODO: review if clone is necessary for hook in self.hooks: c.add(hook.clone()) return c @@ -90,7 +89,7 @@ class HookGroup: if final_hook is None: final_hook = hook.clone() else: - final_hook - final_hook.clone_and_combine() + final_hook = final_hook.clone_and_combine(hook) return final_hook class HookWeightKeyframe: @@ -227,11 +226,9 @@ def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: Dict[st k = () new_modelpatcher = None - # TODO: make hooks work with clip if clip is not None: new_clip = clip.clone() - k1 = [] - #k1 = new_clip.add_hook_patches(hook=hook, patches=loaded, strength_patch=strength_clip) + k1 = new_clip.patcher.add_hook_patches(hook=hook, patches=loaded, strength_patch=strength_clip) else: k1 = () new_clip = None @@ -247,7 +244,6 @@ def load_hook_model_as_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', hook: Hook, strength_model: float, strength_clip: float): if model is not None and model_loaded is not None: new_modelpatcher = model.clone() - comfy.model_management.unload_model_clones(new_modelpatcher) expected_model_keys = set(model_loaded.model.state_dict().keys()) patches_model: Dict[str, torch.Tensor] = model_loaded.model.state_dict() # do not include ANY model_sampling components of the model that should act as a patch @@ -260,14 +256,12 @@ def load_hook_model_as_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', k = () new_modelpatcher = None - # TODO: make hooks work with clip if clip is not None and clip_loaded is not None: new_clip = clip.clone() comfy.model_management.unload_model_clones(new_clip.patcher) expected_clip_keys = clip_loaded.patcher.model.state_dict().copy() patches_clip: Dict[str, torch.Tensor] = clip_loaded.cond_stage_model.state_dict() - k1 = [] - #k1 = new_clip.add_hook_patches(hook=hook, patches=patches_clip, strength_patch=strength_clip, is_diff=True) + k1 = new_clip.patcher.add_hook_patches(hook=hook, patches=patches_clip, strength_patch=strength_clip, is_diff=True) else: k1 = () new_clip = None diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index cc71c7dca..587080f44 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -117,6 +117,7 @@ class ModelPatcher: self.hook_backup: Dict[str, Tuple[torch.Tensor, torch.device]] = {} self.cached_hook_patches: Dict[comfy.hooks.HookGroup, Dict[str, torch.Tensor]] = {} self.current_hooks: Optional[comfy.hooks.HookGroup] = None + self.forced_hooks: Optional[comfy.hooks.HookGroup] = None # NOTE: only used for CLIP # TODO: hook_mode should be entirely removed; behavior should be determined by remaining VRAM/memory self.hook_mode = comfy.hooks.EnumHookMode.MaxSpeed @@ -164,7 +165,8 @@ class ModelPatcher: for k in self.cached_hook_patches[group]: n.cached_hook_patches[group][k] = self.cached_hook_patches[group][k] n.hook_backup = self.hook_backup - n.current_hooks = self.current_hooks + n.current_hooks = self.current_hooks.clone() if self.current_hooks else self.current_hooks + n.forced_hooks = self.forced_hooks.clone() if self.forced_hooks else self.forced_hooks n.hook_mode = self.hook_mode return n @@ -181,6 +183,8 @@ class ModelPatcher: return False if self.current_hooks != clone.current_hooks: return False + if self.forced_hooks != clone.forced_hooks: + return False if self.hook_patches.keys() != clone.hook_patches.keys(): return False @@ -358,6 +362,7 @@ class ModelPatcher: comfy.utils.set_attr_param(self.model, key, out_weight) def load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False): + self.unpatch_hooks() mem_counter = 0 patch_counter = 0 lowvram_counter = 0 @@ -441,6 +446,7 @@ class ModelPatcher: self.model.lowvram_patch_counter += patch_counter self.model.device = device_to self.model.model_loaded_weight_memory = mem_counter + self.apply_hooks(self.forced_hooks) def patch_model(self, device_to=None, lowvram_model_memory=0, load_weights=True, force_patch_weights=False): for k in self.object_patches: @@ -459,6 +465,7 @@ class ModelPatcher: def unpatch_model(self, device_to=None, unpatch_weights=True): if unpatch_weights: + self.unpatch_hooks() if self.model.model_lowvram: for m in self.model.modules(): wipe_lowvram_weight(m) @@ -606,6 +613,7 @@ class ModelPatcher: if is_diff: # take difference between desired weight and existing weight to get diff # TODO: try to implement diff via strength_path/strength_model diff + comfy.model_management.unload_model_clones(self) model_dtype = comfy.utils.get_attr(self.model, key).dtype if model_dtype in [torch.float8_e5m2, torch.float8_e4m3fn]: diff_weight = (patches[k].to(torch.float32)-comfy.utils.get_attr(self.model, key).to(torch.float32)).to(model_dtype) diff --git a/comfy_extras/nodes_hooks.py b/comfy_extras/nodes_hooks.py index febb47486..9e91db5ff 100644 --- a/comfy_extras/nodes_hooks.py +++ b/comfy_extras/nodes_hooks.py @@ -150,6 +150,27 @@ class ConditioningSetDefaultAndCombine: (final_conditioning,) = comfy.hooks.set_default_and_combine_conds(conds=[cond], new_conds=[cond_DEFAULT], opt_hooks=opt_hooks) return (final_conditioning,) + +class SetClipHooks: + NodeId = 'SetClipHooks' + NodeName = 'Set CLIP Hooks' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "clip": ("CLIP",), + "hooks": ("HOOKS",), + } + } + + RETURN_TYPES = ("CLIP",) + CATEGORY = "advanced/hooks/clip" + FUNCTION = "apply_hooks" + + def apply_hooks(self, clip: 'CLIP', hooks: comfy.hooks.HookGroup): + clip = clip.clone() + clip.patcher.forced_hooks = hooks + return (clip,) #------------------------------------------ ########################################### @@ -422,7 +443,8 @@ node_list = [ PairConditioningSetProperties, ConditioningSetDefaultAndCombine, PairConditioningSetDefaultAndCombine, - PairConditioningCombine + PairConditioningCombine, + SetClipHooks, ] NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {}