From 4f08fe69b2c3792c06859669fbc04c7e8a74b661 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sun, 12 Oct 2025 05:25:27 +0200 Subject: [PATCH] add key renaming and key pruning nodes --- __init__.py | 24 ++++++++-- nodes/__init__.py | 14 +++++- nodes/metakeys.py | 34 +------------- nodes/prunekeys.py | 112 ++++++++++++++++++++++++++++++++++++++++++++ nodes/renamekeys.py | 107 ++++++++++++++++++++++++++++++++++++++++++ nodes/utils.py | 34 ++++++++++++++ 6 files changed, 286 insertions(+), 39 deletions(-) create mode 100644 nodes/prunekeys.py create mode 100644 nodes/renamekeys.py create mode 100644 nodes/utils.py diff --git a/__init__.py b/__init__.py index bf2efef..7c1e909 100644 --- a/__init__.py +++ b/__init__.py @@ -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", } \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py index c59fb29..deae698 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -1,8 +1,18 @@ from .metakeys import * +from .renamekeys import * +from .prunekeys import * -__all__ = ( +__all__ = [ "UNetMetaKeys", "CLIPMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys", -) \ No newline at end of file + "UNetRenameKeys", + "CLIPRenameKeys", + "LoRARenameKeys", + "CheckpointRenameKeys", + "UNetPruneKeys", + "CLIPPruneKeys", + "LoRAPruneKeys", + "CheckpointPruneKeys", +] \ No newline at end of file diff --git a/nodes/metakeys.py b/nodes/metakeys.py index e4e7359..f7f5773 100644 --- a/nodes/metakeys.py +++ b/nodes/metakeys.py @@ -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") diff --git a/nodes/prunekeys.py b/nodes/prunekeys.py new file mode 100644 index 0000000..5316122 --- /dev/null +++ b/nodes/prunekeys.py @@ -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) \ No newline at end of file diff --git a/nodes/renamekeys.py b/nodes/renamekeys.py new file mode 100644 index 0000000..cf235cd --- /dev/null +++ b/nodes/renamekeys.py @@ -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) \ No newline at end of file diff --git a/nodes/utils.py b/nodes/utils.py new file mode 100644 index 0000000..23dd765 --- /dev/null +++ b/nodes/utils.py @@ -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 {}, {}