add key renaming and key pruning nodes

2025-10-12 05:25:27 +02:00
parent 3553f17fa2
commit 4f08fe69b2
6 changed files with 286 additions and 39 deletions
@@ -5,11 +5,27 @@ NODE_CLASS_MAPPINGS = {
    "CLIPMetaKeys": CLIPMetaKeys,
    "LoRAMetaKeys": LoRAMetaKeys,
    "CheckpointMetaKeys": CheckpointMetaKeys,
+    "UNetRenameKeys": UNetRenameKeys,
+    "CLIPRenameKeys": CLIPRenameKeys,
+    "LoRARenameKeys": LoRARenameKeys,
+    "CheckpointRenameKeys": CheckpointRenameKeys,
+    "UNetPruneKeys": UNetPruneKeys,
+    "CLIPPruneKeys": CLIPPruneKeys,
+    "LoRAPruneKeys": LoRAPruneKeys,
+    "CheckpointPruneKeys": CheckpointPruneKeys,
 }

 NODE_DISPLAY_NAME_MAPPINGS = {
-    "UNetMetaKeys": "Get UNet Metadata and Keys",
-    "CLIPMetaKeys": "Get CLIP Metadata and Keys",
-    "LoRAMetaKeys": "Get LoRA Metadata and Keys",
-    "CheckpointMetaKeys": "Get Checkpoint Metadata and Keys",
+    "UNetMetaKeys": "Get UNet Meta Keys",
+    "CLIPMetaKeys": "Get CLIP Meta Keys",
+    "LoRAMetaKeys": "Get LoRA Meta Keys",
+    "CheckpointMetaKeys": "Get Checkpoint Meta Keys",
+    "UNetRenameKeys": "Rename UNet Keys",
+    "CLIPRenameKeys": "Rename CLIP Keys",
+    "LoRARenameKeys": "Rename LoRA Keys",
+    "CheckpointRenameKeys": "Rename Checkpoint Keys",
+    "UNetPruneKeys": "Prune UNet Keys",
+    "CLIPPruneKeys": "Prune CLIP Keys",
+    "LoRAPruneKeys": "Prune LoRA Keys",
+    "CheckpointPruneKeys": "Prune Checkpoint Keys",
 }
@@ -1,8 +1,18 @@
 from .metakeys import *
+from .renamekeys import *
+from .prunekeys import *

-__all__ = (
+__all__ = [
    "UNetMetaKeys",
    "CLIPMetaKeys",
    "LoRAMetaKeys",
    "CheckpointMetaKeys",
-)
+    "UNetRenameKeys",
+    "CLIPRenameKeys",
+    "LoRARenameKeys",
+    "CheckpointRenameKeys",
+    "UNetPruneKeys",
+    "CLIPPruneKeys",
+    "LoRAPruneKeys",
+    "CheckpointPruneKeys",
+]
@@ -1,40 +1,8 @@
-import torch
-from safetensors import safe_open
-from safetensors.torch import save_file
 import os
 import folder_paths
-from comfy.model_management import get_torch_device
+from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors


-def convert_pt_to_safetensors(pt_path, safe_path):
-    try:
-        model = torch.load(pt_path, map_location="cpu")
-        if "state_dict" in model:
-            state_dict = model["state_dict"]
-        elif "model" in model:
-            state_dict = model["model"]
-        else:
-            state_dict = model
-        state_dict = {k.replace("state_dict.", ""): v for k, v in state_dict.items()}
-        metadata = {"format": "pt"}
-        save_file(state_dict, safe_path, metadata)
-        return True, ""
-    except Exception as e:
-        return False, str(e)
-
-def load_metadata_from_safetensors(safetensors_file: str) -> dict:
-    try:
-        with safe_open(safetensors_file, framework="pt", device="cpu") as f:
-            model_weights = {}
-            for key in f.keys():
-                model_weights[key] = f.get_tensor(key)
-            metadata = f.metadata()
-        if metadata is None:
-            metadata = {}
-        return model_weights, metadata
-    except Exception as e:
-        print(f"Error loading safetensors file: {e}")
-        return {}, {}

 class BaseMetaKeys:
    RETURN_TYPES = ("STRING", "STRING")
@@ -0,0 +1,112 @@
+import os
+import re
+import folder_paths
+from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors
+from safetensors.torch import save_file
+
+class BasePruneKeys:
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("output_path",)
+    FUNCTION = "prune_keys"
+    CATEGORY = "ModelUtils"
+    DESCRIPTION = "Loads a model, removes specified keys, and saves it as a new safetensors file."
+
+    def _prune_keys(self, model_name, model_type, keys_to_prune_str, use_regex, output_filename):
+        model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
+
+        if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
+            temp_safe_path = model_path + ".safetensors"
+            if not os.path.exists(temp_safe_path):
+                conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
+                if not conversion_successful:
+                    raise Exception(f"Conversion failed: {error_message}")
+            model_path_to_load = temp_safe_path
+        else:
+            model_path_to_load = model_path
+
+        model_weights, metadata = load_metadata_from_safetensors(model_path_to_load)
+
+        patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
+
+        if not patterns:
+            raise ValueError("No keys/patterns provided to prune.")
+
+        pruned_tensors = {}
+        for key, tensor in model_weights.items():
+            is_match = False
+            if use_regex:
+                if any(re.search(pattern, key) for pattern in patterns):
+                    is_match = True
+            else:
+                if any(pattern in key for pattern in patterns):
+                    is_match = True
+
+            if not is_match:
+                pruned_tensors[key] = tensor
+
+        model_dir = folder_paths.get_folder_paths(model_type)[0]
+        output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
+
+        save_file(pruned_tensors, output_path, metadata)
+
+        return (output_path,)
+
+class UNetPruneKeys(BasePruneKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "unet_name": (folder_paths.get_filename_list("diffusion_models"), ),
+                "keys_to_prune": ("STRING", {"multiline": True, "default": ""}),
+                "use_regex": ("BOOLEAN", {"default": False}),
+                "output_filename": ("STRING", {"default": "pruned_unet"}),
+            }
+        }
+
+    def prune_keys(self, unet_name, keys_to_prune, use_regex, output_filename):
+        return self._prune_keys(unet_name, "diffusion_models", keys_to_prune, use_regex, output_filename)
+
+class CLIPPruneKeys(BasePruneKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
+                "keys_to_prune": ("STRING", {"multiline": True, "default": ""}),
+                "use_regex": ("BOOLEAN", {"default": False}),
+                "output_filename": ("STRING", {"default": "pruned_clip"}),
+            }
+        }
+
+    def prune_keys(self, clip_name, keys_to_prune, use_regex, output_filename):
+        return self._prune_keys(clip_name, "text_encoders", keys_to_prune, use_regex, output_filename)
+
+class LoRAPruneKeys(BasePruneKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "lora_name": (folder_paths.get_filename_list("loras"), ),
+                "keys_to_prune": ("STRING", {"multiline": True, "default": ""}),
+                "use_regex": ("BOOLEAN", {"default": False}),
+                "output_filename": ("STRING", {"default": "pruned_lora"}),
+            }
+        }
+
+    def prune_keys(self, lora_name, keys_to_prune, use_regex, output_filename):
+        return self._prune_keys(lora_name, "loras", keys_to_prune, use_regex, output_filename)
+
+class CheckpointPruneKeys(BasePruneKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
+                "keys_to_prune": ("STRING", {"multiline": True, "default": ""}),
+                "use_regex": ("BOOLEAN", {"default": False}),
+                "output_filename": ("STRING", {"default": "pruned_checkpoint"}),
+            }
+        }
+
+    def prune_keys(self, ckpt_name, keys_to_prune, use_regex, output_filename):
+        return self._prune_keys(ckpt_name, "checkpoints", keys_to_prune, use_regex, output_filename)
@@ -0,0 +1,107 @@
+import os
+import folder_paths
+from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors
+from safetensors.torch import save_file
+
+class BaseRenameKeys:
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("output_path",)
+    FUNCTION = "rename_keys"
+    CATEGORY = "ModelUtils"
+    DESCRIPTION = "Loads a model, renames specified keys, and saves it as a new safetensors file."
+
+    def _rename_keys(self, model_name, model_type, old_keys_str, new_keys_str, output_filename):
+        model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
+
+        if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
+            temp_safe_path = model_path + ".safetensors"
+            if not os.path.exists(temp_safe_path):
+                conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
+                if not conversion_successful:
+                    raise Exception(f"Conversion failed: {error_message}")
+            model_path_to_load = temp_safe_path
+        else:
+            model_path_to_load = model_path
+
+        model_weights, metadata = load_metadata_from_safetensors(model_path_to_load)
+
+        old_keys = [key.strip() for key in old_keys_str.strip().split('\n') if key.strip()]
+        new_keys = [key.strip() for key in new_keys_str.strip().split('\n') if key.strip()]
+
+        if len(old_keys) != len(new_keys):
+            raise ValueError("The number of old keys must match the number of new keys.")
+
+        key_map = dict(zip(old_keys, new_keys))
+
+        original_keys = list(model_weights.keys())
+        renamed_tensors = {}
+        for key in original_keys:
+            new_key = key_map.get(key, key)
+            renamed_tensors[new_key] = model_weights[key]
+
+        model_dir = folder_paths.get_folder_paths(model_type)[0]
+        output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
+
+        save_file(renamed_tensors, output_path, metadata)
+
+        return (output_path,)
+
+class UNetRenameKeys(BaseRenameKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "unet_name": (folder_paths.get_filename_list("diffusion_models"), ),
+                "old_keys": ("STRING", {"multiline": True, "default": ""}),
+                "new_keys": ("STRING", {"multiline": True, "default": ""}),
+                "output_filename": ("STRING", {"default": "renamed_unet"}),
+            }
+        }
+
+    def rename_keys(self, unet_name, old_keys, new_keys, output_filename):
+        return self._rename_keys(unet_name, "diffusion_models", old_keys, new_keys, output_filename)
+
+class CLIPRenameKeys(BaseRenameKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
+                "old_keys": ("STRING", {"multiline": True, "default": ""}),
+                "new_keys": ("STRING", {"multiline": True, "default": ""}),
+                "output_filename": ("STRING", {"default": "renamed_clip"}),
+            }
+        }
+
+    def rename_keys(self, clip_name, old_keys, new_keys, output_filename):
+        return self._rename_keys(clip_name, "text_encoders", old_keys, new_keys, output_filename)
+
+class LoRARenameKeys(BaseRenameKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "lora_name": (folder_paths.get_filename_list("loras"), ),
+                "old_keys": ("STRING", {"multiline": True, "default": ""}),
+                "new_keys": ("STRING", {"multiline": True, "default": ""}),
+                "output_filename": ("STRING", {"default": "renamed_lora"}),
+            }
+        }
+
+    def rename_keys(self, lora_name, old_keys, new_keys, output_filename):
+        return self._rename_keys(lora_name, "loras", old_keys, new_keys, output_filename)
+
+class CheckpointRenameKeys(BaseRenameKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
+                "old_keys": ("STRING", {"multiline": True, "default": ""}),
+                "new_keys": ("STRING", {"multiline": True, "default": ""}),
+                "output_filename": ("STRING", {"default": "renamed_checkpoint"}),
+            }
+        }
+
+    def rename_keys(self, ckpt_name, old_keys, new_keys, output_filename):
+        return self._rename_keys(ckpt_name, "checkpoints", old_keys, new_keys, output_filename)
@@ -0,0 +1,34 @@
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+from comfy.model_management import get_torch_device
+
+def convert_pt_to_safetensors(pt_path, safe_path):
+    try:
+        model = torch.load(pt_path, map_location="cpu")
+        if "state_dict" in model:
+            state_dict = model["state_dict"]
+        elif "model" in model:
+            state_dict = model["model"]
+        else:
+            state_dict = model
+        state_dict = {k.replace("state_dict.", ""): v for k, v in state_dict.items()}
+        metadata = {"format": "pt"}
+        save_file(state_dict, safe_path, metadata)
+        return True, ""
+    except Exception as e:
+        return False, str(e)
+
+def load_metadata_from_safetensors(safetensors_file: str) -> dict:
+    try:
+        with safe_open(safetensors_file, framework="pt", device="cpu") as f:
+            model_weights = {}
+            for key in f.keys():
+                model_weights[key] = f.get_tensor(key)
+            metadata = f.metadata()
+        if metadata is None:
+            metadata = {}
+        return model_weights, metadata
+    except Exception as e:
+        print(f"Error loading safetensors file: {e}")
+        return {}, {}