mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-10-03 06:37:04 +08:00
Added support for adding weight hooks that aren't registered on the ModelPatcher at sampling time
This commit is contained in:
parent
f5c899f42a
commit
4b472ba44c
209
comfy/hooks.py
209
comfy/hooks.py
@ -1,5 +1,6 @@
|
|||||||
from typing import TYPE_CHECKING, List, Dict, Tuple
|
from typing import TYPE_CHECKING, List, Dict, Tuple, Callable
|
||||||
import enum
|
import enum
|
||||||
|
import math
|
||||||
import torch
|
import torch
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
@ -15,41 +16,23 @@ class EnumHookMode(enum.Enum):
|
|||||||
MinVram = "minvram"
|
MinVram = "minvram"
|
||||||
MaxSpeed = "maxspeed"
|
MaxSpeed = "maxspeed"
|
||||||
|
|
||||||
class InterpolationMethod:
|
class EnumHookType(enum.Enum):
|
||||||
LINEAR = "linear"
|
Weight = "weight"
|
||||||
EASE_IN = "ease_in"
|
Patch = "patch"
|
||||||
EASE_OUT = "ease_out"
|
|
||||||
EASE_IN_OUT = "ease_in_out"
|
|
||||||
|
|
||||||
_LIST = [LINEAR, EASE_IN, EASE_OUT, EASE_IN_OUT]
|
class EnumWeightTarget(enum.Enum):
|
||||||
|
Model = "model"
|
||||||
|
Clip = "clip"
|
||||||
|
|
||||||
@classmethod
|
class _HookRef:
|
||||||
def get_weights(cls, num_from: float, num_to: float, length: int, method: str, reverse=False):
|
|
||||||
diff = num_to - num_from
|
|
||||||
if method == cls.LINEAR:
|
|
||||||
weights = torch.linspace(num_from, num_to, length)
|
|
||||||
elif method == cls.EASE_IN:
|
|
||||||
index = torch.linspace(0, 1, length)
|
|
||||||
weights = diff * np.power(index, 2) + num_from
|
|
||||||
elif method == cls.EASE_OUT:
|
|
||||||
index = torch.linspace(0, 1, length)
|
|
||||||
weights = diff * (1 - np.power(1 - index, 2)) + num_from
|
|
||||||
elif method == cls.EASE_IN_OUT:
|
|
||||||
index = torch.linspace(0, 1, length)
|
|
||||||
weights = diff * ((1 - np.cos(index * np.pi)) / 2) + num_from
|
|
||||||
else:
|
|
||||||
raise ValueError(f"Unrecognized interpolation method '{method}'.")
|
|
||||||
if reverse:
|
|
||||||
weights = weights.flip(dims=(0,))
|
|
||||||
return weights
|
|
||||||
|
|
||||||
class HookRef:
|
|
||||||
pass
|
pass
|
||||||
|
|
||||||
class Hook:
|
class Hook:
|
||||||
def __init__(self):
|
def __init__(self, hook_type: EnumHookType=None, hook_ref: _HookRef=None,
|
||||||
self.hook_ref = HookRef()
|
hook_keyframe: 'HookKeyframeGroup'=None):
|
||||||
self.hook_keyframe = HookKeyframeGroup()
|
self.hook_type = hook_type
|
||||||
|
self.hook_ref = hook_ref if hook_ref else _HookRef()
|
||||||
|
self.hook_keyframe = hook_keyframe if hook_keyframe else HookKeyframeGroup()
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def strength(self):
|
def strength(self):
|
||||||
@ -62,8 +45,11 @@ class Hook:
|
|||||||
def reset(self):
|
def reset(self):
|
||||||
self.hook_keyframe.reset()
|
self.hook_keyframe.reset()
|
||||||
|
|
||||||
def clone(self):
|
def clone(self, subtype: Callable=None):
|
||||||
c = Hook()
|
if subtype is None:
|
||||||
|
subtype = type(self)
|
||||||
|
c: Hook = subtype()
|
||||||
|
c.hook_type = self.hook_type
|
||||||
c.hook_ref = self.hook_ref
|
c.hook_ref = self.hook_ref
|
||||||
c.hook_keyframe = self.hook_keyframe
|
c.hook_keyframe = self.hook_keyframe
|
||||||
return c
|
return c
|
||||||
@ -74,6 +60,69 @@ class Hook:
|
|||||||
def __hash__(self):
|
def __hash__(self):
|
||||||
return hash(self.hook_ref)
|
return hash(self.hook_ref)
|
||||||
|
|
||||||
|
|
||||||
|
class WeightHook(Hook):
|
||||||
|
def __init__(self, strength_model=1.0, strength_clip=1.0):
|
||||||
|
super().__init__(hook_type=EnumHookType.Weight)
|
||||||
|
self.weights: Dict = None
|
||||||
|
self.weights_clip: Dict = None
|
||||||
|
self.need_weight_init = True
|
||||||
|
self._strength_model = strength_model
|
||||||
|
self._strength_clip = strength_clip
|
||||||
|
|
||||||
|
@property
|
||||||
|
def strength_model(self):
|
||||||
|
return self._strength_model * self.strength
|
||||||
|
|
||||||
|
@property
|
||||||
|
def strength_clip(self):
|
||||||
|
return self._strength_clip * self.strength
|
||||||
|
|
||||||
|
def add_hook_patches(self, model: 'ModelPatcher', target: EnumWeightTarget):
|
||||||
|
weights = None
|
||||||
|
if target == EnumWeightTarget.Model:
|
||||||
|
strength = self._strength_model
|
||||||
|
else:
|
||||||
|
strength = self._strength_clip
|
||||||
|
|
||||||
|
if self.need_weight_init:
|
||||||
|
key_map = {}
|
||||||
|
if target == EnumWeightTarget.Model:
|
||||||
|
key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)
|
||||||
|
else:
|
||||||
|
key_map = comfy.lora.model_lora_keys_clip(model.model, key_map)
|
||||||
|
weights = comfy.lora.load_lora(self.weights, key_map)
|
||||||
|
else:
|
||||||
|
if target == EnumWeightTarget.Model:
|
||||||
|
weights = self.weights
|
||||||
|
else:
|
||||||
|
weights = self.weights_clip
|
||||||
|
k = model.add_hook_patches(hook=self, patches=weights, strength_patch=strength)
|
||||||
|
# TODO: add logs about any keys that were not applied
|
||||||
|
|
||||||
|
def clone(self, subtype: Callable=None):
|
||||||
|
if subtype is None:
|
||||||
|
subtype = type(self)
|
||||||
|
c: WeightHook = super().clone(subtype)
|
||||||
|
c.weights = self.weights
|
||||||
|
c.weights_clip = self.weights_clip
|
||||||
|
c.need_weight_init = self.need_weight_init
|
||||||
|
c._strength_model = self._strength_model
|
||||||
|
c._strength_clip = self._strength_clip
|
||||||
|
return c
|
||||||
|
|
||||||
|
class PatchHook(Hook):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__(hook_type=EnumHookType.Patch)
|
||||||
|
self.patches: Dict = None
|
||||||
|
|
||||||
|
def clone(self, subtype: Callable=None):
|
||||||
|
if subtype is None:
|
||||||
|
subtype = type(self)
|
||||||
|
c: PatchHook = super().clone(type(self))
|
||||||
|
c.patches = self.patches
|
||||||
|
return c
|
||||||
|
|
||||||
class HookGroup:
|
class HookGroup:
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.hooks: List[Hook] = []
|
self.hooks: List[Hook] = []
|
||||||
@ -121,6 +170,7 @@ class HookGroup:
|
|||||||
final_hook = final_hook.clone_and_combine(hook)
|
final_hook = final_hook.clone_and_combine(hook)
|
||||||
return final_hook
|
return final_hook
|
||||||
|
|
||||||
|
|
||||||
class HookKeyframe:
|
class HookKeyframe:
|
||||||
def __init__(self, strength: float, start_percent=0.0, guarantee_steps=1):
|
def __init__(self, strength: float, start_percent=0.0, guarantee_steps=1):
|
||||||
self.strength = strength
|
self.strength = strength
|
||||||
@ -216,6 +266,35 @@ class HookKeyframeGroup:
|
|||||||
# return True if keyframe changed, False if no change
|
# return True if keyframe changed, False if no change
|
||||||
return prev_index != self._current_index
|
return prev_index != self._current_index
|
||||||
|
|
||||||
|
|
||||||
|
class InterpolationMethod:
|
||||||
|
LINEAR = "linear"
|
||||||
|
EASE_IN = "ease_in"
|
||||||
|
EASE_OUT = "ease_out"
|
||||||
|
EASE_IN_OUT = "ease_in_out"
|
||||||
|
|
||||||
|
_LIST = [LINEAR, EASE_IN, EASE_OUT, EASE_IN_OUT]
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def get_weights(cls, num_from: float, num_to: float, length: int, method: str, reverse=False):
|
||||||
|
diff = num_to - num_from
|
||||||
|
if method == cls.LINEAR:
|
||||||
|
weights = torch.linspace(num_from, num_to, length)
|
||||||
|
elif method == cls.EASE_IN:
|
||||||
|
index = torch.linspace(0, 1, length)
|
||||||
|
weights = diff * np.power(index, 2) + num_from
|
||||||
|
elif method == cls.EASE_OUT:
|
||||||
|
index = torch.linspace(0, 1, length)
|
||||||
|
weights = diff * (1 - np.power(1 - index, 2)) + num_from
|
||||||
|
elif method == cls.EASE_IN_OUT:
|
||||||
|
index = torch.linspace(0, 1, length)
|
||||||
|
weights = diff * ((1 - np.cos(index * np.pi)) / 2) + num_from
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unrecognized interpolation method '{method}'.")
|
||||||
|
if reverse:
|
||||||
|
weights = weights.flip(dims=(0,))
|
||||||
|
return weights
|
||||||
|
|
||||||
def get_sorted_list_via_attr(objects: List, attr: str) -> List:
|
def get_sorted_list_via_attr(objects: List, attr: str) -> List:
|
||||||
if not objects:
|
if not objects:
|
||||||
return objects
|
return objects
|
||||||
@ -239,15 +318,68 @@ def get_sorted_list_via_attr(objects: List, attr: str) -> List:
|
|||||||
sorted_list.extend(object_list)
|
sorted_list.extend(object_list)
|
||||||
return sorted_list
|
return sorted_list
|
||||||
|
|
||||||
|
def create_hook_lora(lora: Dict[str, torch.Tensor], strength_model: float, strength_clip: float):
|
||||||
|
hook_group = HookGroup()
|
||||||
|
hook = WeightHook(strength_model=strength_model, strength_clip=strength_clip)
|
||||||
|
hook_group.add(hook)
|
||||||
|
hook.weights = lora
|
||||||
|
hook.need_weight_init = True
|
||||||
|
return hook_group
|
||||||
|
|
||||||
|
def create_hook_model_as_lora(model: 'ModelPatcher', clip: 'CLIP',
|
||||||
|
model_loaded: 'ModelPatcher', clip_loaded: 'CLIP',
|
||||||
|
strength_model: float, strength_clip: float):
|
||||||
|
hook_group = HookGroup()
|
||||||
|
hook = WeightHook(strength_model=strength_model, strength_clip=strength_clip)
|
||||||
|
hook_group.add(hook)
|
||||||
|
if model is not None and model_loaded is not None:
|
||||||
|
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)
|
||||||
|
weights_model, k = model.get_weight_diffs(patches_model)
|
||||||
|
else:
|
||||||
|
weights_model = {}
|
||||||
|
k = ()
|
||||||
|
|
||||||
|
if clip is not None and clip_loaded is not None:
|
||||||
|
expected_clip_keys = clip_loaded.patcher.model.state_dict().copy()
|
||||||
|
patches_clip: Dict[str, torch.Tensor] = clip_loaded.cond_stage_model.state_dict()
|
||||||
|
weights_clip, k1 = clip.patcher.get_weight_diffs(patches_clip)
|
||||||
|
else:
|
||||||
|
weights_clip = {}
|
||||||
|
k1 = ()
|
||||||
|
|
||||||
|
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}")
|
||||||
|
|
||||||
|
hook.weights = weights_model
|
||||||
|
hook.weights_clip = weights_clip
|
||||||
|
hook.need_weight_init = False
|
||||||
|
return hook_group
|
||||||
|
|
||||||
def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: Dict[str, torch.Tensor],
|
def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: Dict[str, torch.Tensor],
|
||||||
hook: Hook, strength_model: float, strength_clip: float):
|
strength_model: float, strength_clip: float):
|
||||||
key_map = {}
|
key_map = {}
|
||||||
if model is not None:
|
if model is not None:
|
||||||
key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)
|
key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)
|
||||||
if clip is not None:
|
if clip is not None:
|
||||||
key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)
|
key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)
|
||||||
|
|
||||||
|
hook_group = HookGroup()
|
||||||
|
hook = WeightHook()
|
||||||
|
hook_group.add(hook)
|
||||||
loaded: Dict[str] = comfy.lora.load_lora(lora, key_map)
|
loaded: Dict[str] = comfy.lora.load_lora(lora, key_map)
|
||||||
if model is not None:
|
if model is not None:
|
||||||
new_modelpatcher = model.clone()
|
new_modelpatcher = model.clone()
|
||||||
@ -267,11 +399,14 @@ def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: Dict[st
|
|||||||
for x in loaded:
|
for x in loaded:
|
||||||
if (x not in k) and (x not in k1):
|
if (x not in k) and (x not in k1):
|
||||||
print(f"NOT LOADED {x}")
|
print(f"NOT LOADED {x}")
|
||||||
return (new_modelpatcher, new_clip)
|
return (new_modelpatcher, new_clip, hook_group)
|
||||||
|
|
||||||
def load_hook_model_as_lora_for_models(model: 'ModelPatcher', clip: 'CLIP',
|
def load_hook_model_as_lora_for_models(model: 'ModelPatcher', clip: 'CLIP',
|
||||||
model_loaded: 'ModelPatcher', clip_loaded: 'CLIP',
|
model_loaded: 'ModelPatcher', clip_loaded: 'CLIP',
|
||||||
hook: Hook, strength_model: float, strength_clip: float):
|
strength_model: float, strength_clip: float):
|
||||||
|
hook_group = HookGroup()
|
||||||
|
hook = WeightHook()
|
||||||
|
hook_group.add(hook)
|
||||||
if model is not None and model_loaded is not None:
|
if model is not None and model_loaded is not None:
|
||||||
new_modelpatcher = model.clone()
|
new_modelpatcher = model.clone()
|
||||||
expected_model_keys = set(model_loaded.model.state_dict().keys())
|
expected_model_keys = set(model_loaded.model.state_dict().keys())
|
||||||
@ -307,7 +442,7 @@ def load_hook_model_as_lora_for_models(model: 'ModelPatcher', clip: 'CLIP',
|
|||||||
if key not in k1:
|
if key not in k1:
|
||||||
print(f"CLIP-AS-LORA NOT LOADED {key}")
|
print(f"CLIP-AS-LORA NOT LOADED {key}")
|
||||||
|
|
||||||
return (new_modelpatcher, new_clip)
|
return (new_modelpatcher, new_clip, hook_group)
|
||||||
|
|
||||||
def set_hooks_for_conditioning(cond, hooks: HookGroup):
|
def set_hooks_for_conditioning(cond, hooks: HookGroup):
|
||||||
if hooks is None:
|
if hooks is None:
|
||||||
|
|||||||
@ -16,7 +16,7 @@
|
|||||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import Dict, List, Tuple, Optional
|
from typing import Dict, List, Tuple, Optional, Callable
|
||||||
import torch
|
import torch
|
||||||
import copy
|
import copy
|
||||||
import inspect
|
import inspect
|
||||||
@ -113,7 +113,8 @@ class ModelPatcher:
|
|||||||
self.weight_inplace_update = weight_inplace_update
|
self.weight_inplace_update = weight_inplace_update
|
||||||
self.patches_uuid = uuid.uuid4()
|
self.patches_uuid = uuid.uuid4()
|
||||||
|
|
||||||
self.hook_patches: Dict[comfy.hooks.HookRef] = {}
|
self.hook_patches: Dict[comfy.hooks._HookRef] = {}
|
||||||
|
self.hook_patches_backup: Dict[comfy.hooks._HookRef] = {}
|
||||||
self.hook_backup: Dict[str, Tuple[torch.Tensor, torch.device]] = {}
|
self.hook_backup: Dict[str, Tuple[torch.Tensor, torch.device]] = {}
|
||||||
self.cached_hook_patches: Dict[comfy.hooks.HookGroup, Dict[str, torch.Tensor]] = {}
|
self.cached_hook_patches: Dict[comfy.hooks.HookGroup, Dict[str, torch.Tensor]] = {}
|
||||||
self.current_hooks: Optional[comfy.hooks.HookGroup] = None
|
self.current_hooks: Optional[comfy.hooks.HookGroup] = None
|
||||||
@ -155,10 +156,8 @@ class ModelPatcher:
|
|||||||
n.object_patches_backup = self.object_patches_backup
|
n.object_patches_backup = self.object_patches_backup
|
||||||
|
|
||||||
# hooks
|
# hooks
|
||||||
for hook_ref in self.hook_patches:
|
n.hook_patches = self.create_hook_patches_clone(self.hook_patches)
|
||||||
n.hook_patches[hook_ref] = {}
|
n.hook_patches_backup = self.create_hook_patches_clone(self.hook_patches_backup)
|
||||||
for k in self.hook_patches[hook_ref]:
|
|
||||||
n.hook_patches[hook_ref][k] = self.hook_patches[hook_ref][k][:]
|
|
||||||
# TODO: do we really need to clone cached_hook_patches/current_hooks?
|
# TODO: do we really need to clone cached_hook_patches/current_hooks?
|
||||||
for group in self.cached_hook_patches:
|
for group in self.cached_hook_patches:
|
||||||
n.cached_hook_patches[group] = {}
|
n.cached_hook_patches[group] = {}
|
||||||
@ -170,6 +169,15 @@ class ModelPatcher:
|
|||||||
n.hook_mode = self.hook_mode
|
n.hook_mode = self.hook_mode
|
||||||
return n
|
return n
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def create_hook_patches_clone(orig_hook_patches):
|
||||||
|
new_hook_patches = {}
|
||||||
|
for hook_ref in orig_hook_patches:
|
||||||
|
new_hook_patches[hook_ref] = {}
|
||||||
|
for k in orig_hook_patches[hook_ref]:
|
||||||
|
new_hook_patches[hook_ref][k] = orig_hook_patches[hook_ref][k][:]
|
||||||
|
return new_hook_patches
|
||||||
|
|
||||||
def is_clone(self, other):
|
def is_clone(self, other):
|
||||||
if hasattr(other, 'model') and self.model is other.model:
|
if hasattr(other, 'model') and self.model is other.model:
|
||||||
return True
|
return True
|
||||||
@ -570,6 +578,12 @@ class ModelPatcher:
|
|||||||
|
|
||||||
def clean(self):
|
def clean(self):
|
||||||
self.clean_hooks()
|
self.clean_hooks()
|
||||||
|
self.restore_hook_patches()
|
||||||
|
|
||||||
|
def restore_hook_patches(self):
|
||||||
|
if len(self.hook_patches_backup) > 0:
|
||||||
|
self.hook_patches = self.hook_patches_backup
|
||||||
|
self.hook_patches_backup = {}
|
||||||
|
|
||||||
def set_hook_mode(self, hook_mode: comfy.hooks.EnumHookMode):
|
def set_hook_mode(self, hook_mode: comfy.hooks.EnumHookMode):
|
||||||
self.hook_mode = hook_mode
|
self.hook_mode = hook_mode
|
||||||
@ -591,8 +605,22 @@ class ModelPatcher:
|
|||||||
if cached_group.contains(hook):
|
if cached_group.contains(hook):
|
||||||
self.cached_hook_patches.pop(cached_group)
|
self.cached_hook_patches.pop(cached_group)
|
||||||
|
|
||||||
def add_hook_patches(self, hook: comfy.hooks.Hook, patches, strength_patch=1.0, strength_model=1.0, is_diff=False):
|
def register_all_hook_patches(self, hooks_dict: Dict[comfy.hooks.Hook, None], target: comfy.hooks.EnumWeightTarget):
|
||||||
|
self.restore_hook_patches()
|
||||||
|
weight_hooks_to_register: List[comfy.hooks.WeightHook] = []
|
||||||
|
for hook in hooks_dict:
|
||||||
|
if hook.hook_type == comfy.hooks.EnumHookType.Weight:
|
||||||
|
if hook.hook_ref not in self.hook_patches:
|
||||||
|
weight_hooks_to_register.append(hook)
|
||||||
|
if len(weight_hooks_to_register) > 0:
|
||||||
|
self.hook_patches_backup = self.create_hook_patches_clone(self.hook_patches)
|
||||||
|
for hook in weight_hooks_to_register:
|
||||||
|
hook.add_hook_patches(self, target)
|
||||||
|
|
||||||
|
def add_hook_patches(self, hook: comfy.hooks.WeightHook, patches, strength_patch=1.0, strength_model=1.0, is_diff=False):
|
||||||
# NOTE: this mirrors behavior of add_patches func
|
# NOTE: this mirrors behavior of add_patches func
|
||||||
|
if is_diff:
|
||||||
|
comfy.model_management.unload_model_clones(self)
|
||||||
current_hook_patches: Dict[str,List] = self.hook_patches.get(hook.hook_ref, {})
|
current_hook_patches: Dict[str,List] = self.hook_patches.get(hook.hook_ref, {})
|
||||||
p = set()
|
p = set()
|
||||||
model_sd = self.model.state_dict()
|
model_sd = self.model.state_dict()
|
||||||
@ -613,7 +641,6 @@ class ModelPatcher:
|
|||||||
if is_diff:
|
if is_diff:
|
||||||
# take difference between desired weight and existing weight to get diff
|
# take difference between desired weight and existing weight to get diff
|
||||||
# TODO: try to implement diff via strength_path/strength_model 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
|
model_dtype = comfy.utils.get_attr(self.model, key).dtype
|
||||||
if model_dtype in [torch.float8_e5m2, torch.float8_e4m3fn]:
|
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)
|
diff_weight = (patches[k].to(torch.float32)-comfy.utils.get_attr(self.model, key).to(torch.float32)).to(model_dtype)
|
||||||
@ -628,6 +655,22 @@ class ModelPatcher:
|
|||||||
self.patches_uuid = uuid.uuid4()
|
self.patches_uuid = uuid.uuid4()
|
||||||
return list(p)
|
return list(p)
|
||||||
|
|
||||||
|
def get_weight_diffs(self, patches):
|
||||||
|
comfy.model_management.unload_model_clones(self)
|
||||||
|
weights: Dict[str, Tuple] = {}
|
||||||
|
p = set()
|
||||||
|
model_sd = self.model.state_dict()
|
||||||
|
for k in patches:
|
||||||
|
if k in model_sd:
|
||||||
|
p.add(k)
|
||||||
|
model_dtype = comfy.utils.get_attr(self.model, k).dtype
|
||||||
|
if model_dtype in [torch.float8_e5m2, torch.float8_e4m3fn]:
|
||||||
|
diff_weight = (patches[k].to(torch.float32)-comfy.utils.get_attr(self.model, k).to(torch.float32)).to(model_dtype)
|
||||||
|
else:
|
||||||
|
diff_weight = patches[k]-comfy.utils.get_attr(self.model, k)
|
||||||
|
weights[k] = (diff_weight,)
|
||||||
|
return weights, p
|
||||||
|
|
||||||
def get_combined_hook_patches(self, hooks: comfy.hooks.HookGroup):
|
def get_combined_hook_patches(self, hooks: comfy.hooks.HookGroup):
|
||||||
# combined_patches will contain weights of all relevant hooks, per key
|
# combined_patches will contain weights of all relevant hooks, per key
|
||||||
combined_patches = {}
|
combined_patches = {}
|
||||||
|
|||||||
@ -1,6 +1,7 @@
|
|||||||
import torch
|
import torch
|
||||||
import comfy.model_management
|
import comfy.model_management
|
||||||
import comfy.conds
|
import comfy.conds
|
||||||
|
import comfy.hooks
|
||||||
|
|
||||||
def prepare_mask(noise_mask, shape, device):
|
def prepare_mask(noise_mask, shape, device):
|
||||||
"""ensures noise mask is of proper dimensions"""
|
"""ensures noise mask is of proper dimensions"""
|
||||||
@ -77,3 +78,13 @@ def cleanup_models(conds, models):
|
|||||||
control_cleanup += get_models_from_cond(conds[k], "control")
|
control_cleanup += get_models_from_cond(conds[k], "control")
|
||||||
|
|
||||||
cleanup_additional_models(set(control_cleanup))
|
cleanup_additional_models(set(control_cleanup))
|
||||||
|
|
||||||
|
def prepare_model_patcher(model, conds):
|
||||||
|
# check for hooks in conds - if not registered, see if can be applied
|
||||||
|
hooks = {}
|
||||||
|
for k in conds:
|
||||||
|
for cond in conds[k]:
|
||||||
|
if 'hooks' in cond:
|
||||||
|
for hook in cond['hooks'].hooks:
|
||||||
|
hooks[hook] = None
|
||||||
|
model.register_all_hook_patches(hooks, comfy.hooks.EnumWeightTarget.Model)
|
||||||
|
|||||||
@ -779,9 +779,12 @@ class CFGGuider:
|
|||||||
latent_image = latent_image.to(device)
|
latent_image = latent_image.to(device)
|
||||||
sigmas = sigmas.to(device)
|
sigmas = sigmas.to(device)
|
||||||
|
|
||||||
output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
|
try:
|
||||||
|
comfy.sampler_helpers.prepare_model_patcher(self.model_patcher, self.conds)
|
||||||
|
output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
|
||||||
|
finally:
|
||||||
|
self.model_patcher.clean()
|
||||||
|
|
||||||
self.model_patcher.clean()
|
|
||||||
comfy.sampler_helpers.cleanup_models(self.conds, self.loaded_models)
|
comfy.sampler_helpers.cleanup_models(self.conds, self.loaded_models)
|
||||||
del self.inner_model
|
del self.inner_model
|
||||||
del self.conds
|
del self.conds
|
||||||
|
|||||||
@ -8,6 +8,7 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
import comfy.hooks
|
import comfy.hooks
|
||||||
import comfy.sd
|
import comfy.sd
|
||||||
|
import comfy.utils
|
||||||
import folder_paths
|
import folder_paths
|
||||||
|
|
||||||
###########################################
|
###########################################
|
||||||
@ -39,7 +40,7 @@ class PairConditioningSetProperties:
|
|||||||
|
|
||||||
def set_properties(self, positive_NEW, negative_NEW,
|
def set_properties(self, positive_NEW, negative_NEW,
|
||||||
strength: float, set_cond_area: str,
|
strength: float, set_cond_area: str,
|
||||||
opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.Hook=None, opt_timesteps: Tuple=None):
|
opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.HookGroup=None, opt_timesteps: Tuple=None):
|
||||||
final_positive, final_negative = comfy.hooks.set_mask_conds(conds=[positive_NEW, negative_NEW],
|
final_positive, final_negative = comfy.hooks.set_mask_conds(conds=[positive_NEW, negative_NEW],
|
||||||
strength=strength, set_cond_area=set_cond_area,
|
strength=strength, set_cond_area=set_cond_area,
|
||||||
opt_mask=opt_mask, opt_hooks=opt_hooks, opt_timestep_range=opt_timesteps)
|
opt_mask=opt_mask, opt_hooks=opt_hooks, opt_timestep_range=opt_timesteps)
|
||||||
@ -70,10 +71,10 @@ class ConditioningSetProperties:
|
|||||||
|
|
||||||
def set_properties(self, cond_NEW,
|
def set_properties(self, cond_NEW,
|
||||||
strength: float, set_cond_area: str,
|
strength: float, set_cond_area: str,
|
||||||
opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.Hook=None, opt_timesteps: Tuple=None):
|
opt_mask: torch.Tensor=None, opt_hooks: comfy.hooks.HookGroup=None, opt_timesteps: Tuple=None):
|
||||||
(final_cond,) = comfy.hooks.set_mask_conds(conds=[cond_NEW],
|
(final_cond,) = comfy.hooks.set_mask_conds(conds=[cond_NEW],
|
||||||
strength=strength, set_cond_area=set_cond_area,
|
strength=strength, set_cond_area=set_cond_area,
|
||||||
opt_mask=opt_mask, opt_hooks=opt_hooks, opt_timestep_range=opt_timesteps)
|
opt_mask=opt_mask, opt_hooks=opt_hooks, opt_timestep_range=opt_timesteps)
|
||||||
return (final_cond,)
|
return (final_cond,)
|
||||||
|
|
||||||
class PairConditioningCombine:
|
class PairConditioningCombine:
|
||||||
@ -198,6 +199,124 @@ class ConditioningTimestepsRange:
|
|||||||
###########################################
|
###########################################
|
||||||
|
|
||||||
|
|
||||||
|
###########################################
|
||||||
|
# Create Hooks
|
||||||
|
#------------------------------------------
|
||||||
|
class CreateHookLora:
|
||||||
|
NodeId = 'CreateHookLora'
|
||||||
|
NodeName = 'Create Hook LoRA'
|
||||||
|
def __init__(self):
|
||||||
|
self.loaded_lora = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(s):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"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 = ("HOOKS",)
|
||||||
|
CATEGORY = "advanced/hooks/create"
|
||||||
|
FUNCTION = "create_hook"
|
||||||
|
|
||||||
|
def create_hook(self, lora_name: str, strength_model: float, strength_clip: float):
|
||||||
|
if strength_model == 0 and strength_clip == 0:
|
||||||
|
return (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)
|
||||||
|
|
||||||
|
hooks = comfy.hooks.create_hook_lora(lora=lora, strength_model=strength_model, strength_clip=strength_clip)
|
||||||
|
return (hooks,)
|
||||||
|
|
||||||
|
class CreateHookLoraModelOnly(CreateHookLora):
|
||||||
|
NodeId = 'CreateHookLoraModelOnly'
|
||||||
|
NodeName = 'Create Hook LoRA (MO)'
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(s):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"lora_name": (folder_paths.get_filename_list("loras"), ),
|
||||||
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("HOOKS",)
|
||||||
|
CATEGORY = "advanced/hooks/create"
|
||||||
|
FUNCTION = "create_hook_model_only"
|
||||||
|
|
||||||
|
def create_hook_model_only(self, lora_name: str, strength_model: float):
|
||||||
|
return self.create_hook(lora_name=lora_name, strength_model=strength_model, strength_clip=0)
|
||||||
|
|
||||||
|
class CreateHookModelAsLora:
|
||||||
|
NodeId = 'CreateHookModelAsLora'
|
||||||
|
NodeName = 'Create 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 = ("HOOKS",)
|
||||||
|
CATEGORY = "advanced/hooks/create"
|
||||||
|
FUNCTION = "create_hook"
|
||||||
|
|
||||||
|
def create_hook(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]
|
||||||
|
|
||||||
|
hooks = comfy.hooks.create_hook_model_as_lora(model=model, clip=clip,
|
||||||
|
model_loaded=model_loaded, clip_loaded=clip_loaded,
|
||||||
|
strength_model=strength_model, strength_clip=strength_clip)
|
||||||
|
return (hooks,)
|
||||||
|
|
||||||
|
class CreateHookModelAsLoraModelOnly:
|
||||||
|
NodeId = 'CreateHookModelAsLoraModelOnly'
|
||||||
|
NodeName = 'Create 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 = ("HOOKS",)
|
||||||
|
CATEGORY = "advanced/hooks/create"
|
||||||
|
FUNCTION = "create_hook_model_only"
|
||||||
|
|
||||||
|
def create_hook_model_only(self, model: 'ModelPatcher', ckpt_name: str, strength_model: float):
|
||||||
|
return CreateHookModelAsLora.create_hook(self, model=model, clip=None, ckpt_name=ckpt_name,
|
||||||
|
strength_model=strength_model, strength_clip=0)
|
||||||
|
#------------------------------------------
|
||||||
|
###########################################
|
||||||
|
|
||||||
|
|
||||||
###########################################
|
###########################################
|
||||||
# Register Hooks
|
# Register Hooks
|
||||||
#------------------------------------------
|
#------------------------------------------
|
||||||
@ -242,12 +361,9 @@ class RegisterHookLora:
|
|||||||
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||||
self.loaded_lora = (lora_path, lora)
|
self.loaded_lora = (lora_path, lora)
|
||||||
|
|
||||||
hook = comfy.hooks.Hook()
|
model_lora, clip_lora, hooks = comfy.hooks.load_hook_lora_for_models(model=model, clip=clip, lora=lora,
|
||||||
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)
|
strength_model=strength_model, strength_clip=strength_clip)
|
||||||
return (model_lora, clip_lora, hook_group)
|
return (model_lora, clip_lora, hooks)
|
||||||
|
|
||||||
class RegisterHookLoraModelOnly(RegisterHookLora):
|
class RegisterHookLoraModelOnly(RegisterHookLora):
|
||||||
NodeId = 'RegisterHookLoraModelOnly'
|
NodeId = 'RegisterHookLoraModelOnly'
|
||||||
@ -257,10 +373,8 @@ class RegisterHookLoraModelOnly(RegisterHookLora):
|
|||||||
return {
|
return {
|
||||||
"required": {
|
"required": {
|
||||||
"model": ("MODEL",),
|
"model": ("MODEL",),
|
||||||
"clip": ("CLIP",),
|
|
||||||
"lora_name": (folder_paths.get_filename_list("loras"), ),
|
"lora_name": (folder_paths.get_filename_list("loras"), ),
|
||||||
"strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
|
"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}),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -299,14 +413,10 @@ class RegisterHookModelAsLora:
|
|||||||
model_loaded = out[0]
|
model_loaded = out[0]
|
||||||
clip_loaded = out[1]
|
clip_loaded = out[1]
|
||||||
|
|
||||||
hook = comfy.hooks.Hook()
|
model_lora, clip_lora, hooks = comfy.hooks.load_hook_model_as_lora_for_models(model=model, clip=clip,
|
||||||
hook_group = comfy.hooks.HookGroup()
|
model_loaded=model_loaded, clip_loaded=clip_loaded,
|
||||||
hook_group.add(hook)
|
strength_model=strength_model, strength_clip=strength_clip)
|
||||||
model_lora, clip_lora = comfy.hooks.load_hook_model_as_lora_for_models(model=model, clip=clip,
|
return (model_lora, clip_lora, hooks)
|
||||||
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:
|
class RegisterHookModelAsLoraModelOnly:
|
||||||
NodeId = 'RegisterHookModelAsLoraModelOnly'
|
NodeId = 'RegisterHookModelAsLoraModelOnly'
|
||||||
@ -551,6 +661,11 @@ class CombineHooksEight:
|
|||||||
###########################################
|
###########################################
|
||||||
|
|
||||||
node_list = [
|
node_list = [
|
||||||
|
# Create
|
||||||
|
CreateHookLora,
|
||||||
|
CreateHookLoraModelOnly,
|
||||||
|
CreateHookModelAsLora,
|
||||||
|
CreateHookModelAsLoraModelOnly,
|
||||||
# Register
|
# Register
|
||||||
RegisterHookLora,
|
RegisterHookLora,
|
||||||
RegisterHookLoraModelOnly,
|
RegisterHookLoraModelOnly,
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user