Added apply_to_conds to Set CLIP Hooks, modified relevant code to allow text encoding to automatically apply hooks to output conds when apply_to_conds is set to True

comfy/sd.py

@@ -97,6 +97,7 @@ class CLIP:
         self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
         self.patcher.hook_mode = comfy.hooks.EnumHookMode.MinVram
         self.patcher.is_clip = True
+        self.apply_hooks_to_conds = None
         if params['device'] == load_device:
             model_management.load_models_gpu([self.patcher], force_full_load=True)
         self.layer_idx = None
@@ -110,6 +111,7 @@ class CLIP:
         n.tokenizer = self.tokenizer
         n.layer_idx = self.layer_idx
         n.use_clip_schedule = self.use_clip_schedule
+        n.apply_hooks_to_conds = self.apply_hooks_to_conds
         return n
 
     def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
@@ -121,6 +123,11 @@ class CLIP:
     def tokenize(self, text, return_word_ids=False):
         return self.tokenizer.tokenize_with_weights(text, return_word_ids)
 
+    def add_hooks_to_dict(self, pooled_dict: dict[str]):
+        if self.apply_hooks_to_conds:
+            pooled_dict["hooks"] = self.apply_hooks_to_conds
+        return pooled_dict
+
     def encode_from_tokens_scheduled(self, tokens, unprojected=False, add_dict: dict[str]=None, show_pbar=True):
         all_cond_pooled: list[tuple[torch.Tensor, dict[str]]] = []
         all_hooks = self.patcher.forced_hooks
@@ -155,6 +162,8 @@ class CLIP:
             # add/update any keys with the provided add_dict
             if add_dict is not None:
                 pooled_dict.update(add_dict)
+            # add hooks stored on clip
+            self.add_hooks_to_dict(pooled_dict)
             all_cond_pooled.append([cond, pooled_dict])
             if show_pbar:
                 pbar.update(1)
@@ -179,6 +188,8 @@ class CLIP:
             if len(o) > 2:
                 for k in o[2]:
                     out[k] = o[2][k]
+            if self.apply_hooks_to_conds:
+                out["hooks"] = self.apply_hooks_to_conds
             return out
 
         if return_pooled:

comfy_extras/nodes_clip_sdxl.py

@@ -20,7 +20,7 @@ class CLIPTextEncodeSDXLRefiner:
         if clip.use_clip_schedule:
             return (clip.encode_from_tokens_scheduled(tokens, add_dict={"aesthetic_score": ascore, "width": width, "height": height}), )
         cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
-        return ([[cond, {"pooled_output": pooled, "aesthetic_score": ascore, "width": width,"height": height}]], )
+        return ([[cond, clip.add_hooks_to_dict({"pooled_output": pooled, "aesthetic_score": ascore, "width": width,"height": height})]], )
 
 class CLIPTextEncodeSDXL:
     @classmethod
@@ -52,7 +52,7 @@ class CLIPTextEncodeSDXL:
         if clip.use_clip_schedule:
             return (clip.encode_from_tokens_scheduled(tokens, add_dict={"width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}), )
         cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
-        return ([[cond, {"pooled_output": pooled, "width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}]], )
+        return ([[cond, clip.add_hooks_to_dict({"pooled_output": pooled, "width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height})]], )
 
 NODE_CLASS_MAPPINGS = {
     "CLIPTextEncodeSDXLRefiner": CLIPTextEncodeSDXLRefiner,

comfy_extras/nodes_hooks.py

@@ -226,7 +226,8 @@ class SetClipHooks:
         return {
             "required": {
                 "clip": ("CLIP",),
-                "schedule_clip": ("BOOLEAN", {"default": True})
+                "apply_to_conds": ("BOOLEAN", {"default": True}),
+                "schedule_clip": ("BOOLEAN", {"default": False})
             },
             "optional": {
                 "hooks": ("HOOKS",)
@@ -237,11 +238,13 @@ class SetClipHooks:
     CATEGORY = "advanced/hooks/clip"
     FUNCTION = "apply_hooks"
 
-    def apply_hooks(self, clip: 'CLIP', schedule_clip: bool, hooks: comfy.hooks.HookGroup=None):
+    def apply_hooks(self, clip: 'CLIP', schedule_clip: bool, apply_to_conds: bool, hooks: comfy.hooks.HookGroup=None):
         if hooks is not None:
             clip = clip.clone()
-            clip.use_clip_schedule = schedule_clip
+            if apply_to_conds:
+                clip.apply_hooks_to_conds = hooks
             clip.patcher.forced_hooks = hooks.clone()
+            clip.use_clip_schedule = schedule_clip
             if not clip.use_clip_schedule:
                 clip.patcher.forced_hooks.set_keyframes_on_hooks(None)
             clip.patcher.register_all_hook_patches(hooks.get_dict_repr(), comfy.hooks.EnumWeightTarget.Clip)

comfy_extras/nodes_sd3.py

@@ -81,7 +81,7 @@ class CLIPTextEncodeSD3:
         if clip.use_clip_schedule:
             return (clip.encode_from_tokens_scheduled(tokens), )
         cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
-        return ([[cond, {"pooled_output": pooled}]], )
+        return ([[cond, clip.add_hooks_to_dict({"pooled_output": pooled})]], )
 
 
 class ControlNetApplySD3(nodes.ControlNetApplyAdvanced):