xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-13 10:57:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added initial set of hook-related nodes, added code to register hooks for loras/model-as-loras, small renaming/refactoring

Browse Source
This commit is contained in:
kosinkadink1@gmail.com 2024-09-15 08:33:17 +09:00
parent f5abdc6f86
commit 9ded65a616
5 changed files with 612 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

176
comfy/hooks.py
View File

@ -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

22
comfy/model_patcher.py
View File

@ -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()

8
comfy/samplers.py
View File

@ -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())

432
comfy_extras/nodes_hooks.py Normal file
View File

@ -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

1
nodes.py
View File

@ -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 = []