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

@@ -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',
+# NOTE: this function shows how to register weight hooks directly on the ModelPatchers
-                              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],
                               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

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:
+                    current_patches.append((strength_patch, patches[k], strength_model, offset, function))
-                        # 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_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 = {}

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,