diff --git a/comfy/hooks.py b/comfy/hooks.py index 3fc55a8d6..80af32ffd 100644 --- a/comfy/hooks.py +++ b/comfy/hooks.py @@ -1,8 +1,14 @@ from typing import TYPE_CHECKING, List, Dict, Tuple import enum +import torch if TYPE_CHECKING: + from comfy.model_patcher import ModelPatcher from comfy.model_base import BaseModel + from comfy.sd import CLIP +import comfy.lora +import comfy.model_management +from node_helpers import conditioning_set_values class EnumHookMode(enum.Enum): MinVram = "minvram" @@ -11,7 +17,7 @@ class EnumHookMode(enum.Enum): class HookRef: pass -class HookWeight: +class Hook: def __init__(self): self.hook_ref = HookRef() self.hook_keyframe = HookWeightKeyframeGroup() @@ -28,36 +34,36 @@ class HookWeight: self.hook_keyframe.reset() def clone(self): - c = HookWeight() + c = Hook() c.hook_ref = self.hook_ref c.hook_keyframe = self.hook_keyframe return c - def __eq__(self, other: 'HookWeight'): + def __eq__(self, other: 'Hook'): return self.__class__ == other.__class__ and self.hook_ref == other.hook_ref def __hash__(self): return hash(self.hook_ref) -class HookWeightGroup: +class HookGroup: def __init__(self): - self.hooks: List[HookWeight] = [] + self.hooks: List[Hook] = [] - def add(self, hook: HookWeight): + def add(self, hook: Hook): if hook not in self.hooks: self.hooks.append(hook) - def contains(self, hook: HookWeight): + def contains(self, hook: Hook): return hook in self.hooks def clone(self): - c = HookWeightGroup() + c = HookGroup() # TODO: review if clone is necessary for hook in self.hooks: c.add(hook.clone()) return c - def clone_and_combine(self, other: 'HookWeightGroup'): + def clone_and_combine(self, other: 'HookGroup'): c = self.clone() for hook in other.hooks: c.add(hook.clone()) @@ -69,8 +75,8 @@ class HookWeightGroup: hook.hook_keyframe = hook_kf @staticmethod - def combine_all_hooks(hooks_list: List['HookWeightGroup'], require_count=1) -> 'HookWeightGroup': - actual: List[HookWeightGroup] = [] + def combine_all_hooks(hooks_list: List['HookGroup'], require_count=1) -> 'HookGroup': + actual: List[HookGroup] = [] for group in hooks_list: if group is not None: actual.append(group) @@ -79,7 +85,7 @@ class HookWeightGroup: # if only 1 hook, just reutnr itself without cloning if len(actual) == 1: return actual[0] - final_hook: HookWeightGroup = None + final_hook: HookGroup = None for hook in actual: if final_hook is None: final_hook = hook.clone() @@ -203,3 +209,149 @@ def get_sorted_list_via_attr(objects: List, attr: str) -> List: for object_list in sorted_attrs.values(): sorted_list.extend(object_list) return sorted_list + + +def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: Dict[str, torch.Tensor], + hook: Hook, strength_model: float, strength_clip: float): + key_map = {} + if model is not None: + key_map = comfy.lora.model_lora_keys_unet(model.model, key_map) + if clip is not None: + key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map) + + loaded: Dict[str] = comfy.lora.load_lora(lora, key_map) + if model is not None: + new_modelpatcher = model.clone() + k = new_modelpatcher.add_hook_patches(hook=hook, patches=loaded, strength_patch=strength_model) + else: + 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) + else: + k1 = () + new_clip = None + k = set(k) + k1 = set(k1) + for x in loaded: + if (x not in k) and (x not in k1): + print(f"NOT LOADED {x}") + return (new_modelpatcher, new_clip) + +def load_hook_model_as_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', + model_loaded: 'ModelPatcher', clip_loaded: '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 + for key in list(patches_model.keys()): + if key.startswith("model_sampling"): + expected_model_keys.discard(key) + patches_model.pop(key, None) + k = new_modelpatcher.add_hook_patches(hook=hook, patches=patches_model, strength_patch=strength_model, is_diff=True) + else: + 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) + else: + k1 = () + new_clip = None + + k = set(k) + k1 = set(k1) + if model is not None and model_loaded is not None: + for key in expected_model_keys: + if key not in k: + print(f"MODEL-AS-LORA NOT LOADED {key}") + if clip is not None and clip_loaded is not None: + for key in expected_clip_keys: + if key not in k1: + print(f"CLIP-AS-LORA NOT LOADED {key}") + + return (new_modelpatcher, new_clip) + +def set_hooks_for_conditioning(cond, hooks: HookGroup): + if hooks is None: + return cond + return conditioning_set_values(cond, {'hooks': hooks}) + +def set_timesteps_for_conditioning(cond, timestep_range: Tuple[float,float]): + if timestep_range is None: + return cond + return conditioning_set_values(cond, {"start_percent": timestep_range[0], + "end_percent": timestep_range[1]}) + +def set_mask_for_conditioning(cond, mask: torch.Tensor, set_cond_area: str, strength: float): + if mask is None: + return cond + set_area_to_bounds = False + if set_cond_area != 'default': + set_area_to_bounds = True + if len(mask.shape) < 3: + mask = mask.unsqueeze(0) + return conditioning_set_values(cond, {'mask': mask, + 'set_area_to_bounds': set_area_to_bounds, + 'mask_strength': strength}) + +def combine_conditioning(conds: List): + combined_conds = [] + for cond in conds: + combined_conds.extend(cond) + return combined_conds + +def set_mask_conds(conds: List, strength: float, set_cond_area: str, + opt_mask: torch.Tensor=None, opt_hooks: HookGroup=None, opt_timestep_range: Tuple[float,float]=None): + masked_conds = [] + for c in conds: + # first, apply lora_hook to conditioning, if provided + c = set_hooks_for_conditioning(c, opt_hooks) + # next, apply mask to conditioning + c = set_mask_for_conditioning(cond=c, mask=opt_mask, strength=strength, set_cond_area=set_cond_area) + # apply timesteps, if present + c = set_timesteps_for_conditioning(cond=c, timestep_range=opt_timestep_range) + # finally, apply mask to conditioning and store + masked_conds.append(c) + return masked_conds + +def set_mask_and_combine_conds(conds: List, new_conds: List, strength: float=1.0, set_cond_area: str="default", + opt_mask: torch.Tensor=None, opt_hooks: HookGroup=None, opt_timestep_range: Tuple[float,float]=None): + combined_conds = [] + for c, masked_c in zip(conds, new_conds): + # first, apply lora_hook to new conditioning, if provided + masked_c = set_hooks_for_conditioning(masked_c, opt_hooks) + # next, apply mask to new conditioning, if provided + masked_c = set_mask_for_conditioning(cond=masked_c, mask=opt_mask, set_cond_area=set_cond_area, strength=strength) + # apply timesteps, if present + masked_c = set_timesteps_for_conditioning(cond=masked_c, timestep_range=opt_timestep_range) + # finally, combine with existing conditioning and store + combined_conds.append(combine_conditioning([c, masked_c])) + return combined_conds + +def set_default_and_combine_conds(conds: list, new_conds: list, + opt_hooks: HookGroup=None, opt_timestep_range: Tuple[float,float]=None): + combined_conds = [] + for c, new_c in zip(conds, new_conds): + # first, apply lora_hook to new conditioning, if provided + new_c = set_hooks_for_conditioning(new_c, opt_hooks) + # next, add default_cond key to cond so that during sampling, it can be identified + new_c = conditioning_set_values(new_c, {'default': True}) + # apply timesteps, if present + new_c = set_timesteps_for_conditioning(cond=new_c, timestep_range=opt_timestep_range) + # finally, combine with existing conditioning and store + combined_conds.append(combine_conditioning([c, new_c])) + return combined_conds diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 14cc7f46a..cc71c7dca 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -115,8 +115,8 @@ class ModelPatcher: self.hook_patches: Dict[comfy.hooks.HookRef] = {} self.hook_backup: Dict[str, Tuple[torch.Tensor, torch.device]] = {} - self.cached_hook_patches: Dict[comfy.hooks.HookWeightGroup, Dict[str, torch.Tensor]] = {} - self.current_hooks: Optional[comfy.hooks.HookWeightGroup] = None + self.cached_hook_patches: Dict[comfy.hooks.HookGroup, Dict[str, torch.Tensor]] = {} + self.current_hooks: Optional[comfy.hooks.HookGroup] = None # TODO: hook_mode should be entirely removed; behavior should be determined by remaining VRAM/memory self.hook_mode = comfy.hooks.EnumHookMode.MaxSpeed @@ -567,7 +567,7 @@ class ModelPatcher: def set_hook_mode(self, hook_mode: comfy.hooks.EnumHookMode): self.hook_mode = hook_mode - def prepare_hook_patches_current_keyframe(self, t: torch.Tensor, hook_group: comfy.hooks.HookWeightGroup): + def prepare_hook_patches_current_keyframe(self, t: torch.Tensor, hook_group: comfy.hooks.HookGroup): curr_t = t[0] for hook in hook_group.hooks: changed = hook.hook_keyframe.prepare_current_keyframe(curr_t=curr_t) @@ -584,7 +584,7 @@ class ModelPatcher: if cached_group.contains(hook): self.cached_hook_patches.pop(cached_group) - def add_hook_patches(self, hook: comfy.hooks.HookWeight, patches, strength_patch=1.0, strength_model=1.0, is_diff=False): + def add_hook_patches(self, hook: comfy.hooks.Hook, patches, strength_patch=1.0, strength_model=1.0, is_diff=False): # NOTE: this mirrors behavior of add_patches func current_hook_patches: Dict[str,List] = self.hook_patches.get(hook.hook_ref, {}) p = set() @@ -620,7 +620,7 @@ class ModelPatcher: self.patches_uuid = uuid.uuid4() return list(p) - def get_combined_hook_patches(self, hooks: comfy.hooks.HookWeightGroup): + def get_combined_hook_patches(self, hooks: comfy.hooks.HookGroup): # combined_patches will contain weights of all relevant hooks, per key combined_patches = {} if hooks is not None: @@ -633,18 +633,18 @@ class ModelPatcher: else: # patches are stored as tuples: (strength_patch, (tuple_with_weights,), strength_model) for patch in hook_patches[key]: - new_patch = List(patch) + new_patch = list(patch) new_patch[0] *= hook.strength - current_patches.append(Tuple(new_patch)) + current_patches.append(tuple(new_patch)) combined_patches[key] = current_patches return combined_patches - def apply_hooks(self, hooks: comfy.hooks.HookWeightGroup): + def apply_hooks(self, hooks: comfy.hooks.HookGroup): if self.current_hooks == hooks: return self.patch_hooks(hooks=hooks) - def patch_hooks(self, hooks: comfy.hooks.HookWeightGroup): + def patch_hooks(self, hooks: comfy.hooks.HookGroup): self.unpatch_hooks() model_sd = self.model_state_dict() # if have cached weights for hooks, use it @@ -680,7 +680,7 @@ class ModelPatcher: self.cached_hook_patches.clear() self.current_hooks = None - def patch_hook_weight_to_device(self, hooks: comfy.hooks.HookWeightGroup, combined_patches: dict, key: str): + def patch_hook_weight_to_device(self, hooks: comfy.hooks.HookGroup, combined_patches: dict, key: str): if key not in combined_patches: return weight: torch.Tensor = comfy.utils.get_attr(self.model, key) @@ -719,7 +719,7 @@ class ModelPatcher: comfy.utils.copy_to_param(self.model, k, self.hook_backup[k][0].to(device=self.hook_backup[k][1])) self.hook_backup.clear() - self.current_hooks = None # TODO: should this be clear_cached_hook_weights instead? + self.current_hooks = None def clean_hooks(self): self.unpatch_hooks() diff --git a/comfy/samplers.py b/comfy/samplers.py index 4d1ecef16..24ff8d7eb 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -140,7 +140,7 @@ def cond_cat(c_list): return out -def finalize_default_conds(hooked_to_run: Dict[comfy.hooks.HookWeightGroup,List[Tuple[Tuple,int]]], default_conds: List[List[Dict]], x_in, timestep): +def finalize_default_conds(hooked_to_run: Dict[comfy.hooks.HookGroup,List[Tuple[Tuple,int]]], default_conds: List[List[Dict]], x_in, timestep): # need to figure out remaining unmasked area for conds default_mults = [] for _ in default_conds: @@ -171,7 +171,7 @@ def finalize_default_conds(hooked_to_run: Dict[comfy.hooks.HookWeightGroup,List[ continue # replace p's mult with calculated mult p = p._replace(mult=mult) - hook: comfy.hooks.HookWeightGroup = x.get('hook', None) + hook: comfy.hooks.HookGroup = x.get('hook', None) hooked_to_run.setdefault(hook, list()) hooked_to_run[hook] += [(p, i)] @@ -180,7 +180,7 @@ def calc_cond_batch(model, conds: List[List[Dict]], x_in, timestep, model_option out_counts = [] # separate conds by matching hooks # TODO: implement default_conds support - hooked_to_run: Dict[comfy.hooks.HookWeightGroup,List[Tuple[Tuple,int]]] = {} + hooked_to_run: Dict[comfy.hooks.HookGroup,List[Tuple[Tuple,int]]] = {} default_conds = [] has_default_conds = False @@ -199,7 +199,7 @@ def calc_cond_batch(model, conds: List[List[Dict]], x_in, timestep, model_option p = comfy.samplers.get_area_and_mult(x, x_in, timestep) if p is None: continue - hooks: comfy.hooks.HookWeightGroup = x.get('hooks', None) + hooks: comfy.hooks.HookGroup = x.get('hooks', None) if hooks is not None: model.current_patcher.prepare_hook_patches_current_keyframe(timestep, hooks) hooked_to_run.setdefault(hooks, list()) diff --git a/comfy_extras/nodes_hooks.py b/comfy_extras/nodes_hooks.py new file mode 100644 index 000000000..febb47486 --- /dev/null +++ b/comfy_extras/nodes_hooks.py @@ -0,0 +1,432 @@ +from typing import TYPE_CHECKING, Dict, List, Tuple +import torch + +if TYPE_CHECKING: + from comfy.model_patcher import ModelPatcher + from comfy.sd import CLIP + +import comfy.hooks +import comfy.sd +import folder_paths + +########################################### +# Mask, Combine, and Hook Conditioning +#------------------------------------------ +class PairConditioningSetProperties: + NodeId = 'PairConditioningSetProperties' + NodeName = 'Pair Cond Set Props' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "positive_NEW": ("CONDITIONING", ), + "negative_NEW": ("CONDITIONING", ), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "set_cond_area": (["default", "mask bounds"],), + }, + "optional": { + "opt_mask": ("MASK", ), + "opt_hooks": ("HOOKS",), + "opt_timesteps": ("TIMESTEPS_RANGE",), + } + } + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING") + RETURN_NAMES = ("positive", "negative") + CATEGORY = "advanced/hooks/cond pair" + FUNCTION = "set_properties" + + def set_properties(self, positive_NEW, negative_NEW, + strength: float, set_cond_area: str, + opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.Hook=None, opt_timesteps: Tuple=None): + final_positive, final_negative = comfy.hooks.set_mask_conds(conds=[positive_NEW, negative_NEW], + strength=strength, set_cond_area=set_cond_area, + opt_mask=opt_mask, opt_hooks=opt_hooks, opt_timestep_range=opt_timesteps) + return (final_positive, final_negative) + +class ConditioningSetProperties: + NodeId = 'ConditioningSetProperties' + NodeName = 'Cond Set Props' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "cond_NEW": ("CONDITIONING", ), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "set_cond_area": (["default", "mask bounds"],), + }, + "optional": { + "opt_mask": ("MASK", ), + "opt_hooks": ("HOOKS",), + "opt_timesteps": ("TIMESTEPS_RANGE",), + } + } + + RETURN_TYPES = ("CONDITIONING",) + RETURN_NAMES = ("positive", "negative") + CATEGORY = "advanced/hooks/cond single" + FUNCTION = "set_properties" + + def set_properties(self, cond_NEW, + strength: float, set_cond_area: str, + opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.Hook=None, opt_timesteps: Tuple=None): + (final_cond,) = comfy.hooks.set_mask_conds(conds=[cond_NEW], + strength=strength, set_cond_area=set_cond_area, + opt_mask=opt_mask, opt_hooks=opt_hooks, opt_timestep_range=opt_timesteps) + return (final_cond,) + +class PairConditioningCombine: + NodeId = 'PairConditioningCombine' + NodeName = 'Pair Cond Combine' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "positive_A": ("CONDITIONING",), + "negative_A": ("CONDITIONING",), + "positive_B": ("CONDITIONING",), + "negative_B": ("CONDITIONING",), + }, + } + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING") + RETURN_NAMES = ("positive", "negative") + CATEGORY = "advanced/hooks/cond pair" + FUNCTION = "combine" + + def combine(self, positive_A, negative_A, positive_B, negative_B): + final_positive, final_negative = comfy.hooks.set_mask_and_combine_conds(conds=[positive_A, negative_A], new_conds=[positive_B, negative_B],) + return (final_positive, final_negative,) + +class PairConditioningSetDefaultAndCombine: + NodeId = 'PairConditioningSetDefaultCombine' + NodeName = 'Pair Cond Set Default Combine' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "positive_DEFAULT": ("CONDITIONING",), + "negative_DEFAULT": ("CONDITIONING",), + }, + "optional": { + "opt_hooks": ("HOOKS",), + } + } + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING") + RETURN_NAMES = ("positive", "negative") + CATEGORY = "advanced/hooks/cond pair" + FUNCTION = "set_default_and_combine" + + def set_default_and_combine(self, positive, negative, positive_DEFAULT, negative_DEFAULT, + opt_hooks: comfy.hooks.HookGroup=None): + final_positive, final_negative = comfy.hooks.set_default_and_combine_conds(conds=[positive, negative], new_conds=[positive_DEFAULT, negative_DEFAULT], + opt_hooks=opt_hooks) + return (final_positive, final_negative) + +class ConditioningSetDefaultAndCombine: + NodeId = 'ConditioningSetDefaultCombine' + NodeName = 'Cond Set Default Combine' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "cond": ("CONDITIONING",), + "cond_DEFAULT": ("CONDITIONING",), + }, + "optional": { + "opt_hooks": ("HOOKS",), + } + } + + RETURN_TYPES = ("CONDITIONING",) + CATEGORY = "advanced/hooks/cond single" + FUNCTION = "set_default_and_combine" + + def append_and_combine(self, cond, cond_DEFAULT, + opt_hooks: comfy.hooks.HookGroup=None): + (final_conditioning,) = comfy.hooks.set_default_and_combine_conds(conds=[cond], new_conds=[cond_DEFAULT], + opt_hooks=opt_hooks) + return (final_conditioning,) +#------------------------------------------ +########################################### + + +########################################### +# Register Hooks +#------------------------------------------ +class RegisterHookLora: + NodeId = 'RegisterHookLora' + NodeName = 'Register Hook LoRA' + def __init__(self): + self.loaded_lora = None + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "clip": ("CLIP",), + "lora_name": (folder_paths.get_filename_list("loras"), ), + "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + "strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("MODEL", "CLIP", "HOOKS") + CATEGORY = "advanced/hooks/register" + FUNCTION = "register_lora" + + def register_lora(self, model: 'ModelPatcher', clip: 'CLIP', lora_name: str, + strength_model: float, strength_clip: float): + if strength_model == 0 and strength_clip == 0: + return (model, clip, None) + + lora_path = folder_paths.get_full_path("loras", lora_name) + lora = None + if self.loaded_lora is not None: + if self.loaded_lora[0] == lora_path: + lora = self.loaded_lora[1] + else: + temp = self.loaded_lora + self.loaded_lora = None + del temp + + if lora is None: + lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + self.loaded_lora = (lora_path, lora) + + hook = comfy.hooks.Hook() + hook_group = comfy.hooks.HookGroup() + hook_group.add(hook) + model_lora, clip_lora = comfy.hooks.load_hook_lora_for_models(model=model, clip=clip, lora=lora, hook=hook, + strength_model=strength_model, strength_clip=strength_clip) + return (model_lora, clip_lora, hook_group) + +class RegisterHookLoraModelOnly(RegisterHookLora): + NodeId = 'RegisterHookLoraModelOnly' + NodeName = 'Register Hook LoRA (MO)' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "clip": ("CLIP",), + "lora_name": (folder_paths.get_filename_list("loras"), ), + "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + "strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("MODEL", "HOOKS") + CATEGORY = "advanced/hooks/register" + FUNCTION = "register_lora_model_only" + + def register_lora_model_only(self, model: 'ModelPatcher', lora_name: str, strength_model: float): + model_lora, _, hooks = self.register_lora(model=model, clip=None, lora_name=lora_name, + strength_model=strength_model, strength_clip=0) + return (model_lora, hooks) + +class RegisterHookModelAsLora: + NodeId = 'RegisterHookModelAsLora' + NodeName = 'Register Hook Model as LoRA' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "clip": ("CLIP",), + "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), + "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + "strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("MODEL", "CLIP", "HOOKS") + CATEGORY = "advanced/hooks/register" + FUNCTION = "register_model_as_lora" + + def register_model_as_lora(self, model: 'ModelPatcher', clip: 'CLIP', ckpt_name: str, + strength_model: float, strength_clip: float): + ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) + model_loaded = out[0] + clip_loaded = out[1] + + hook = comfy.hooks.Hook() + hook_group = comfy.hooks.HookGroup() + hook_group.add(hook) + model_lora, clip_lora = comfy.hooks.load_hook_model_as_lora_for_models(model=model, clip=clip, + model_loaded=model_loaded, clip_loaded=clip_loaded, + hook=hook, + strength_model=strength_model, strength_clip=strength_clip) + return (model_lora, clip_lora, hook_group) + +class RegisterHookModelAsLoraModelOnly: + NodeId = 'RegisterHookModelAsLoraModelOnly' + NodeName = 'Register Hook Model as LoRA (MO)' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), + "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("MODEL", "HOOKS") + CATEGORY = "advanced/hooks/register" + FUNCTION = "register_model_as_lora_model_only" + + def register_model_as_lora_model_only(self, model: 'ModelPatcher', ckpt_name: str, strength_model: float): + model_lora, _, hooks = RegisterHookModelAsLora.register_model_as_lora(self, model=model, clip=None, ckpt_name=ckpt_name, + strength_model=strength_model, strength_clip=0) + return (model_lora, hooks) +#------------------------------------------ +########################################### + + +########################################### +# Schedule Hooks +#------------------------------------------ +#------------------------------------------ +########################################### + + +class SetModelHooksOnCond: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "conditioning": ("CONDITIONING",), + "hooks": ("HOOKS",), + }, + } + + RETURN_TYPES = ("CONDITIONING",) + CATEGORY = "advanced/hooks/manual" + FUNCTION = "attach_hook" + + def attach_hook(self, conditioning, hooks: comfy.hooks.HookGroup): + return (comfy.hooks.set_hooks_for_conditioning(conditioning, hooks),) + + +########################################### +# Combine Hooks +#------------------------------------------ +class CombineHooks: + NodeId = 'CombineHooks2' + NodeName = 'Combine Hooks [2]' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + }, + "optional": { + "hooks_A": ("HOOKS",), + "hooks_B": ("HOOKS",), + } + } + + RETURN_TYPES = ("HOOKS",) + CATEGORY = "advanced/hooks/combine" + FUNCTION = "combine_hooks" + + def combine_hooks(self, + hooks_A: comfy.hooks.HookGroup=None, + hooks_B: comfy.hooks.HookGroup=None): + candidates = [hooks_A, hooks_B] + return (comfy.hooks.HookGroup.combine_all_hooks(candidates),) + +class CombineHooksFour: + NodeId = 'CombineHooks4' + NodeName = 'Combine Hooks [4]' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + }, + "optional": { + "hooks_A": ("HOOKS",), + "hooks_B": ("HOOKS",), + "hooks_C": ("HOOKS",), + "hooks_D": ("HOOKS",), + } + } + + RETURN_TYPES = ("HOOKS",) + CATEGORY = "advanced/hooks/combine" + FUNCTION = "combine_hooks" + + def combine_hooks(self, + hooks_A: comfy.hooks.HookGroup=None, + hooks_B: comfy.hooks.HookGroup=None, + hooks_C: comfy.hooks.HookGroup=None, + hooks_D: comfy.hooks.HookGroup=None): + candidates = [hooks_A, hooks_B, hooks_C, hooks_D] + return (comfy.hooks.HookGroup.combine_all_hooks(candidates),) + +class CombineHooksEight: + NodeId = 'CombineHooks8' + NodeName = 'Combine Hooks [8]' + @classmethod + def INPUT_TYPES(s): + return { + "required": { + }, + "optional": { + "hooks_A": ("HOOKS",), + "hooks_B": ("HOOKS",), + "hooks_C": ("HOOKS",), + "hooks_D": ("HOOKS",), + "hooks_E": ("HOOKS",), + "hooks_F": ("HOOKS",), + "hooks_G": ("HOOKS",), + "hooks_H": ("HOOKS",), + } + } + + RETURN_TYPES = ("HOOKS",) + CATEGORY = "advanced/hooks/combine" + FUNCTION = "combine_hooks" + + def combine_hooks(self, + hooks_A: comfy.hooks.HookGroup=None, + hooks_B: comfy.hooks.HookGroup=None, + hooks_C: comfy.hooks.HookGroup=None, + hooks_D: comfy.hooks.HookGroup=None, + hooks_E: comfy.hooks.HookGroup=None, + hooks_F: comfy.hooks.HookGroup=None, + hooks_G: comfy.hooks.HookGroup=None, + hooks_H: comfy.hooks.HookGroup=None): + candidates = [hooks_A, hooks_B, hooks_C, hooks_D, hooks_E, hooks_F, hooks_G, hooks_H] + return (comfy.hooks.HookGroup.combine_all_hooks(candidates),) +#------------------------------------------ +########################################### + +node_list = [ + # Register + RegisterHookLora, + RegisterHookLoraModelOnly, + RegisterHookModelAsLora, + RegisterHookModelAsLoraModelOnly, + # Combine + CombineHooks, + CombineHooksFour, + CombineHooksEight, + # Attach + ConditioningSetProperties, + PairConditioningSetProperties, + ConditioningSetDefaultAndCombine, + PairConditioningSetDefaultAndCombine, + PairConditioningCombine +] +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +for node in node_list: + NODE_CLASS_MAPPINGS[node.NodeId] = node + NODE_DISPLAY_NAME_MAPPINGS[node.NodeId] = node.NodeName diff --git a/nodes.py b/nodes.py index 0d3749d42..36022f731 100644 --- a/nodes.py +++ b/nodes.py @@ -2104,6 +2104,7 @@ def init_builtin_extra_nodes(): "nodes_flux.py", "nodes_lora_extract.py", "nodes_torch_compile.py", + "nodes_hooks.py", ] import_failed = []