From 5909b06b2e90ae29a8e4d89d825bd477e0579e59 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 12 Nov 2024 09:05:19 -0600 Subject: [PATCH] Removed Register/Direct hook nodes since they were present only for testing, removed diff-related weight hook calculation as improved_memory removes unload_model_clones and using sample time registered hooks is less hacky --- comfy/hooks.py | 91 +---------------- comfy/model_patcher.py | 32 +----- comfy_extras/nodes_hooks.py | 197 ------------------------------------ 3 files changed, 4 insertions(+), 316 deletions(-) diff --git a/comfy/hooks.py b/comfy/hooks.py index fd92a3f85..7530523ba 100644 --- a/comfy/hooks.py +++ b/comfy/hooks.py @@ -97,7 +97,6 @@ class WeightHook(Hook): super().__init__(hook_type=EnumHookType.Weight) self.weights: dict = None self.weights_clip: dict = None - self.is_diff = False self.need_weight_init = True self._strength_model = strength_model self._strength_clip = strength_clip @@ -131,7 +130,7 @@ class WeightHook(Hook): weights = self.weights else: weights = self.weights_clip - k = model.add_hook_patches(hook=self, patches=weights, strength_patch=strength, is_diff=self.is_diff) + k = model.add_hook_patches(hook=self, patches=weights, strength_patch=strength) return True # TODO: add logs about any keys that were not applied @@ -142,7 +141,6 @@ class WeightHook(Hook): c.weights = self.weights c.weights_clip = self.weights_clip c.need_weight_init = self.need_weight_init - c.is_diff = self.is_diff c._strength_model = self._strength_model c._strength_clip = self._strength_clip return c @@ -543,49 +541,7 @@ def get_patch_weights_from_model(model: 'ModelPatcher', discard_model_sampling=F patches_model.pop(key, None) return patches_model -def create_hook_model_as_lora_precalc(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 - +# NOTE: this function shows how to register weight hooks directly on the ModelPatchers def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: dict[str, torch.Tensor], strength_model: float, strength_clip: float): key_map = {} @@ -618,49 +574,6 @@ def load_hook_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', lora: dict[st print(f"NOT LOADED {x}") return (new_modelpatcher, new_clip, hook_group) -def load_hook_model_as_lora_for_models(model: 'ModelPatcher', clip: 'CLIP', - model_loaded: 'ModelPatcher', clip_loaded: 'CLIP', - 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: - new_modelpatcher = model.clone() - 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 - - 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 = new_clip.patcher.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, hook_group) - def set_hooks_for_conditioning(cond, hooks: HookGroup): if hooks is None: return cond diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 99c520cef..50fa977f7 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -920,11 +920,9 @@ class ModelPatcher: for callback in self.get_all_callbacks(CallbacksMP.ON_REGISTER_ALL_HOOK_PATCHES): callback(self, hooks_dict, target) - def add_hook_patches(self, hook: comfy.hooks.WeightHook, patches, strength_patch=1.0, strength_model=1.0, is_diff=False): + def add_hook_patches(self, hook: comfy.hooks.WeightHook, patches, strength_patch=1.0, strength_model=1.0): with self.use_ejected(): # 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, {}) p = set() model_sd = self.model.state_dict() @@ -942,39 +940,13 @@ class ModelPatcher: if key in model_sd: p.add(k) current_patches: list[tuple] = current_hook_patches.get(key, []) - if is_diff: - # take difference between desired weight and existing weight to get diff - 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) - else: - diff_weight = patches[k]-comfy.utils.get_attr(self.model, key) - current_patches.append((strength_patch, (diff_weight,), strength_model, offset, function)) - else: - current_patches.append((strength_patch, patches[k], strength_model, offset, function)) + current_patches.append((strength_patch, patches[k], strength_model, offset, function)) current_hook_patches[key] = current_patches self.hook_patches[hook.hook_ref] = current_hook_patches # since should care about these patches too to determine if same model, reroll patches_uuid self.patches_uuid = uuid.uuid4() return list(p) - def get_weight_diffs(self, patches): - with self.use_ejected(): - 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): # combined_patches will contain weights of all relevant hooks, per key combined_patches = {} diff --git a/comfy_extras/nodes_hooks.py b/comfy_extras/nodes_hooks.py index 157294c12..7fea2ea24 100644 --- a/comfy_extras/nodes_hooks.py +++ b/comfy_extras/nodes_hooks.py @@ -343,70 +343,6 @@ class CreateHookLoraModelOnly(CreateHookLora): def create_hook_model_only(self, lora_name: str, strength_model: float, prev_hooks: comfy.hooks.HookGroup=None): return self.create_hook(lora_name=lora_name, strength_model=strength_model, strength_clip=0, prev_hooks=prev_hooks) -class CreateHookModelAsLoraDirect: - NodeId = 'CreateHookModelAsLoraDirect' - NodeName = 'Create Hook Model as LoRA Direct' - @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}), - }, - "optional": { - "prev_hooks": ("HOOKS",) - } - } - - 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, - prev_hooks: comfy.hooks.HookGroup=None): - if prev_hooks is None: - prev_hooks = comfy.hooks.HookGroup() - prev_hooks.clone() - - 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_precalc(model=model, clip=clip, - model_loaded=model_loaded, clip_loaded=clip_loaded, - strength_model=strength_model, strength_clip=strength_clip) - return (prev_hooks.clone_and_combine(hooks),) - -class CreateHookModelAsLoraDirectModelOnly: - NodeId = 'CreateHookModelAsLoraDirectModelOnly' - NodeName = 'Create Hook Model as LoRA Direct (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}), - }, - "optional": { - "prev_hooks": ("HOOKS",) - } - } - - 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, - prev_hooks: comfy.hooks.HookGroup=None): - return CreateHookModelAsLoraDirect.create_hook(self, model=model, clip=None, ckpt_name=ckpt_name, - strength_model=strength_model, strength_clip=0, prev_hooks=prev_hooks) - class CreateHookModelAsLora: NodeId = 'CreateHookModelAsLora' NodeName = 'Create Hook Model as LoRA' @@ -487,132 +423,6 @@ class CreateHookModelAsLoraModelOnly(CreateHookModelAsLora): ########################################### -########################################### -# 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) - - model_lora, clip_lora, hooks = comfy.hooks.load_hook_lora_for_models(model=model, clip=clip, lora=lora, - strength_model=strength_model, strength_clip=strength_clip) - return (model_lora, clip_lora, hooks) - -class RegisterHookLoraModelOnly(RegisterHookLora): - NodeId = 'RegisterHookLoraModelOnly' - NodeName = 'Register Hook LoRA (MO)' - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "model": ("MODEL",), - "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 = ("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] - - model_lora, clip_lora, hooks = comfy.hooks.load_hook_model_as_lora_for_models(model=model, clip=clip, - model_loaded=model_loaded, clip_loaded=clip_loaded, - strength_model=strength_model, strength_clip=strength_clip) - return (model_lora, clip_lora, hooks) - -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 #------------------------------------------ @@ -836,13 +646,6 @@ node_list = [ CreateHookLoraModelOnly, CreateHookModelAsLora, CreateHookModelAsLoraModelOnly, - CreateHookModelAsLoraDirect, # TODO: remove before merge - CreateHookModelAsLoraDirectModelOnly, # TODO: remove before merge - # Register - RegisterHookLora, - RegisterHookLoraModelOnly, - RegisterHookModelAsLora, - RegisterHookModelAsLoraModelOnly, # Scheduling SetHookKeyframes, CreateHookKeyframe,