Made CLIP work with hook patches

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.patcher.add_hook_patches(hook=hook, patches=loaded, strength_patch=strength_clip)
-        #k1 = new_clip.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.patcher.add_hook_patches(hook=hook, patches=patches_clip, strength_patch=strength_clip, is_diff=True)
-        #k1 = new_clip.add_hook_patches(hook=hook, patches=patches_clip, strength_patch=strength_clip, is_diff=True)
     else:
         k1 = ()
         new_clip = None

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)

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 = {}