From a20be20ac72ce89e616e66d200df017bdd853ad7 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 16 Nov 2024 17:37:57 -0600 Subject: [PATCH] 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 | 11 +++++++++++ comfy_extras/nodes_clip_sdxl.py | 4 ++-- comfy_extras/nodes_hooks.py | 9 ++++++--- comfy_extras/nodes_sd3.py | 2 +- 4 files changed, 20 insertions(+), 6 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index 1c86aea8d..6c2139a1e 100644 --- a/comfy/sd.py +++ b/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: diff --git a/comfy_extras/nodes_clip_sdxl.py b/comfy_extras/nodes_clip_sdxl.py index 9dcfcacb8..4ed4352b2 100644 --- a/comfy_extras/nodes_clip_sdxl.py +++ b/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, diff --git a/comfy_extras/nodes_hooks.py b/comfy_extras/nodes_hooks.py index 2cf2a679c..daa62d500 100644 --- a/comfy_extras/nodes_hooks.py +++ b/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) diff --git a/comfy_extras/nodes_sd3.py b/comfy_extras/nodes_sd3.py index 2d8113dac..99279196e 100644 --- a/comfy_extras/nodes_sd3.py +++ b/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):