Made CLIP work with hook patches

This commit is contained in:
kosinkadink1@gmail.com 2024-09-15 15:47:09 +09:00
parent 9ded65a616
commit a5034df6db
3 changed files with 35 additions and 11 deletions

View File

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

View File

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

View File

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