diff --git a/__init__.py b/__init__.py index 1ed6dad..82c8b43 100644 --- a/__init__.py +++ b/__init__.py @@ -1,57 +1,52 @@ -from .nodes import * +from typing_extensions import override +from comfy_api.latest import ComfyExtension, io -NODE_CLASS_MAPPINGS = { - "ModelMetaKeys": ModelMetaKeys, - "TextEncoderMetaKeys": TextEncoderMetaKeys, - "LoRAMetaKeys": LoRAMetaKeys, - "CheckpointMetaKeys": CheckpointMetaKeys, - "EmbeddingMetaKeys": EmbeddingMetaKeys, - "ModelRenameKeys": ModelRenameKeys, - "TextEncoderRenameKeys": TextEncoderRenameKeys, - "LoRARenameKeys": LoRARenameKeys, - "CheckpointRenameKeys": CheckpointRenameKeys, - "EmbeddingRenameKeys": EmbeddingRenameKeys, - "ModelPruneKeys": ModelPruneKeys, - "TextEncoderPruneKeys": TextEncoderPruneKeys, - "LoRAPruneKeys": LoRAPruneKeys, - "CheckpointPruneKeys": CheckpointPruneKeys, - "EmbeddingPruneKeys": EmbeddingPruneKeys, - "ModelTwoMerger": ModelTwoMerger, - "TextEncoderTwoMerger": TextEncoderTwoMerger, - "LoRATwoMerger": LoRATwoMerger, - "CheckpointTwoMerger": CheckpointTwoMerger, - "EmbeddingTwoMerger": EmbeddingTwoMerger, - "ModelThreeMerger": ModelThreeMerger, - "TextEncoderThreeMerger": TextEncoderThreeMerger, - "LoRAThreeMerger": LoRAThreeMerger, - "CheckpointThreeMerger": CheckpointThreeMerger, - "EmbeddingThreeMerger": EmbeddingThreeMerger, -} +from .nodes.metakeys import ( + ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys, + CheckpointMetaKeys, EmbeddingMetaKeys +) +from .nodes.renamekeys import ( + ModelRenameKeys, TextEncoderRenameKeys, LoRARenameKeys, + CheckpointRenameKeys, EmbeddingRenameKeys +) +from .nodes.prunekeys import ( + ModelPruneKeys, TextEncoderPruneKeys, LoRAPruneKeys, + CheckpointPruneKeys, EmbeddingPruneKeys +) +from .nodes.merger import ( + ModelTwoMerger, TextEncoderTwoMerger, LoRATwoMerger, + CheckpointTwoMerger, EmbeddingTwoMerger, + ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger, + CheckpointThreeMerger, EmbeddingThreeMerger +) +from .nodes.lora_extract import ( + LoRASaveFromFile, LoRACheckpointSaveFromFile +) -NODE_DISPLAY_NAME_MAPPINGS = { - "ModelMetaKeys": "Get Diffusion Model Metadata & Keys", - "TextEncoderMetaKeys": "Get Text Encoder Metadata & Keys", - "LoRAMetaKeys": "Get LoRA Metadata & Keys", - "CheckpointMetaKeys": "Get Checkpoint Metadata & Keys", - "EmbeddingMetaKeys": "Get Embedding Metadata & Keys", - "ModelRenameKeys": "ModelRenameKeys", - "TextEncoderRenameKeys": "TextEncoderRenameKeys", - "LoRARenameKeys": "LoRARenameKeys", - "CheckpointRenameKeys": "CheckpointRenameKeys", - "EmbeddingRenameKeys": "EmbeddingRenameKeys", - "ModelPruneKeys": "ModelPruneKeys", - "TextEncoderPruneKeys": "TextEncoderPruneKeys", - "LoRAPruneKeys": "LoRAPruneKeys", - "CheckpointPruneKeys": "Prune Checkpoint Keys", - "EmbeddingPruneKeys": "Prune Embedding Keys", - "ModelTwoMerger": "Merge Models (2 Models)", - "TextEncoderTwoMerger": "Merge Text Encoders (2 Models)", - "LoRATwoMerger": "Merge LoRAs (2 Models)", - "CheckpointTwoMerger": "Merge Checkpoints (2 Models)", - "EmbeddingTwoMerger": "Merge Embeddings (2 Models)", - "ModelThreeMerger": "Merge Models (3 Models)", - "TextEncoderThreeMerger": "Merge Text Encoders (3 Models)", - "LoRAThreeMerger": "Merge LoRAs (3 Models)", - "CheckpointThreeMerger": "Merge Checkpoints (3 Models)", - "EmbeddingThreeMerger": "Merge Embeddings (3 Models)", -} \ No newline at end of file + +class ModelUtilsExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + # MetaKeys + ModelMetaKeys, TextEncoderMetaKeys, LoRAMetaKeys, + CheckpointMetaKeys, EmbeddingMetaKeys, + # RenameKeys + ModelRenameKeys, TextEncoderRenameKeys, LoRARenameKeys, + CheckpointRenameKeys, EmbeddingRenameKeys, + # PruneKeys + ModelPruneKeys, TextEncoderPruneKeys, LoRAPruneKeys, + CheckpointPruneKeys, EmbeddingPruneKeys, + # Two-Model Mergers + ModelTwoMerger, TextEncoderTwoMerger, LoRATwoMerger, + CheckpointTwoMerger, EmbeddingTwoMerger, + # Three-Model Mergers + ModelThreeMerger, TextEncoderThreeMerger, LoRAThreeMerger, + CheckpointThreeMerger, EmbeddingThreeMerger, + # LoRA Extraction + LoRASaveFromFile, LoRACheckpointSaveFromFile, + ] + + +async def comfy_entrypoint() -> ModelUtilsExtension: + return ModelUtilsExtension() \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py deleted file mode 100644 index 565aad8..0000000 --- a/nodes/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from .metakeys import * -from .renamekeys import * -from .prunekeys import * -from .merger import * - -__all__ = [ - "ModelMetaKeys", "TextEncoderMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys", "EmbeddingMetaKeys", - "ModelRenameKeys", "TextEncoderRenameKeys", "LoRARenameKeys", "CheckpointRenameKeys", "EmbeddingRenameKeys", - "ModelPruneKeys", "TextEncoderPruneKeys", "LoRAPruneKeys", "CheckpointPruneKeys", "EmbeddingPruneKeys", - "CheckpointTwoMerger", "ModelTwoMerger", "TextEncoderTwoMerger", "LoRATwoMerger", "EmbeddingTwoMerger", - "CheckpointThreeMerger", "ModelThreeMerger", "TextEncoderThreeMerger", "LoRAThreeMerger", "EmbeddingThreeMerger", -] \ No newline at end of file diff --git a/nodes/lora_extract.py b/nodes/lora_extract.py new file mode 100644 index 0000000..fe56431 --- /dev/null +++ b/nodes/lora_extract.py @@ -0,0 +1,350 @@ +""" +Memory-efficient LoRA extraction node. + +This node loads two models from files and extracts a LoRA by computing the difference, +without requiring the full models to be loaded into GPU memory via ComfyUI's model management. +""" +import os +import torch +import folder_paths +import comfy.utils +from comfy_api.latest import io +from safetensors.torch import save_file +from .merger_utils import MemoryEfficientSafeOpen + + +CLAMP_QUANTILE = 0.99 + + +def extract_lora(diff, rank): + """Extract LoRA weights from a weight difference tensor using SVD.""" + conv2d = (len(diff.shape) == 4) + kernel_size = None if not conv2d else diff.size()[2:4] + conv2d_3x3 = conv2d and kernel_size != (1, 1) + out_dim, in_dim = diff.size()[0:2] + rank = min(rank, in_dim, out_dim) + + if conv2d: + if conv2d_3x3: + diff = diff.flatten(start_dim=1) + else: + diff = diff.squeeze() + + U, S, Vh = torch.linalg.svd(diff.float()) + U = U[:, :rank] + S = S[:rank] + U = U @ torch.diag(S) + Vh = Vh[:rank, :] + + dist = torch.cat([U.flatten(), Vh.flatten()]) + hi_val = torch.quantile(dist, CLAMP_QUANTILE) + low_val = -hi_val + + U = U.clamp(low_val, hi_val) + Vh = Vh.clamp(low_val, hi_val) + if conv2d: + U = U.reshape(out_dim, rank, 1, 1) + Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1]) + return (U, Vh) + + +def extract_lora_from_files( + model_a_path: str, + model_b_path: str, + rank: int, + lora_type: str, + bias_diff: bool, + prefix_lora: str, + process_device: str, + save_dtype: str, +) -> dict[str, torch.Tensor]: + """ + Memory-efficient LoRA extraction by streaming tensors from safetensors files. + + Args: + model_a_path: Path to the finetuned/modified model + model_b_path: Path to the base/original model + rank: LoRA rank for SVD decomposition + lora_type: "standard" for LoRA, "full_diff" for full difference + bias_diff: Whether to include bias differences + prefix_lora: Prefix for LoRA keys (e.g., "diffusion_model.") + process_device: Device for processing ("cuda" or "cpu") + save_dtype: Output dtype ("fp16", "bf16", "fp32") + + Returns: + Dictionary of LoRA tensors + """ + output_sd = {} + + # Get dtype for saving + save_torch_dtype = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16 + }.get(save_dtype, torch.float16) + + # Open both files with memory-efficient handler + handler_a = MemoryEfficientSafeOpen(model_a_path) + handler_b = MemoryEfficientSafeOpen(model_b_path) + + try: + keys_a = set(handler_a.keys()) + keys_b = set(handler_b.keys()) + + # Only process keys that exist in both models + common_keys = keys_a & keys_b + weight_keys = [k for k in common_keys if k.endswith(".weight")] + bias_keys = [k for k in common_keys if k.endswith(".bias")] if bias_diff else [] + + # Log key differences + missing_in_b = keys_a - keys_b + extra_in_b = keys_b - keys_a + if missing_in_b: + print(f"[LoRA Extract] {len(missing_in_b)} keys in model A not in model B (skipped)") + if extra_in_b: + print(f"[LoRA Extract] {len(extra_in_b)} keys in model B not in model A (ignored)") + + pbar = comfy.utils.ProgressBar(len(weight_keys) + len(bias_keys)) + + for key in weight_keys: + # Load tensors one at a time + tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32) + tensor_b = handler_b.get_tensor(key).to(device=process_device, dtype=torch.float32) + + # Compute difference: finetuned - base + weight_diff = tensor_a - tensor_b + + # Clean up immediately + del tensor_a, tensor_b + + # Extract LoRA key name + lora_key = key[:-7] # Remove ".weight" + + if lora_type == "standard": + if weight_diff.ndim < 2: + # 1D weights (e.g., LayerNorm) - store as diff if bias_diff enabled + if bias_diff: + output_sd[f"{prefix_lora}{lora_key}.diff"] = weight_diff.to(save_torch_dtype).cpu() + else: + try: + lora_up, lora_down = extract_lora(weight_diff, rank) + output_sd[f"{prefix_lora}{lora_key}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu() + output_sd[f"{prefix_lora}{lora_key}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu() + except Exception as e: + print(f"[LoRA Extract] Could not extract LoRA for {key}: {e}") + + elif lora_type == "full_diff": + output_sd[f"{prefix_lora}{lora_key}.diff"] = weight_diff.to(save_torch_dtype).cpu() + + del weight_diff + pbar.update(1) + + # Process bias keys + for key in bias_keys: + tensor_a = handler_a.get_tensor(key).to(device=process_device, dtype=torch.float32) + tensor_b = handler_b.get_tensor(key).to(device=process_device, dtype=torch.float32) + + bias_diff_tensor = tensor_a - tensor_b + del tensor_a, tensor_b + + lora_key = key[:-5] # Remove ".bias" + output_sd[f"{prefix_lora}{lora_key}.diff_b"] = bias_diff_tensor.to(save_torch_dtype).cpu() + + del bias_diff_tensor + pbar.update(1) + + finally: + handler_a.__exit__(None, None, None) + handler_b.__exit__(None, None, None) + + return output_sd + + +class LoRASaveFromFile(io.ComfyNode): + """ + Memory-efficient LoRA extraction from safetensors files. + + Unlike the standard LoraSave node which requires pre-loaded and subtracted models, + this node loads models directly from files one tensor at a time, making it much + more memory efficient for extracting LoRAs from large models. + """ + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LoRASaveFromFile", + display_name="Extract LoRA from Files (Memory Efficient)", + category="ModelUtils/LoRA", + description="Extracts a LoRA by computing the difference between two model files. " + "Much more memory efficient than using ModelSubtract + LoraSave.", + inputs=[ + io.Combo.Input( + "model_a", + options=folder_paths.get_filename_list("diffusion_models"), + tooltip="The finetuned/modified model (A - B = LoRA)" + ), + io.Combo.Input( + "model_b", + options=folder_paths.get_filename_list("diffusion_models"), + tooltip="The base/original model (A - B = LoRA)" + ), + io.Int.Input("rank", default=64, min=1, max=4096, step=1, + tooltip="LoRA rank - higher preserves more detail but increases size"), + io.Combo.Input("lora_type", options=["standard", "full_diff"], default="standard", + tooltip="standard: SVD decomposition to lora_up/lora_down, full_diff: raw weight differences"), + io.Boolean.Input("bias_diff", default=True, + tooltip="Include bias differences in the LoRA"), + io.String.Input("output_filename", default="extracted_lora", + tooltip="Output filename (without extension)"), + io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"), + io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda", + tooltip="Device for SVD computation"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + is_output_node=True, + ) + + @classmethod + def execute( + cls, + model_a: str, + model_b: str, + rank: int, + lora_type: str, + bias_diff: bool, + output_filename: str, + save_dtype: str, + process_device: str, + ) -> io.NodeOutput: + # Get full paths + model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a) + model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b) + + # Extract LoRA + output_sd = extract_lora_from_files( + model_a_path=model_a_path, + model_b_path=model_b_path, + rank=rank, + lora_type=lora_type, + bias_diff=bias_diff, + prefix_lora="diffusion_model.", + process_device=process_device, + save_dtype=save_dtype, + ) + + if not output_sd: + raise ValueError("No LoRA weights extracted - models may be identical or incompatible") + + # Save to loras folder + output_dir = folder_paths.get_folder_paths("loras")[0] + os.makedirs(output_dir, exist_ok=True) + output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors") + + save_file(output_sd, output_path) + print(f"[LoRA Extract] Saved {len(output_sd)} tensors to {output_path}") + + return io.NodeOutput(output_path) + + +class LoRACheckpointSaveFromFile(io.ComfyNode): + """ + Memory-efficient LoRA extraction from checkpoint files (which may contain both model and text encoder). + """ + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LoRACheckpointSaveFromFile", + display_name="Extract LoRA from Checkpoints (Memory Efficient)", + category="ModelUtils/LoRA", + description="Extracts a LoRA by computing the difference between two checkpoint files. " + "Handles both diffusion model and text encoder weights.", + inputs=[ + io.Combo.Input( + "checkpoint_a", + options=folder_paths.get_filename_list("checkpoints"), + tooltip="The finetuned/modified checkpoint (A - B = LoRA)" + ), + io.Combo.Input( + "checkpoint_b", + options=folder_paths.get_filename_list("checkpoints"), + tooltip="The base/original checkpoint (A - B = LoRA)" + ), + io.Int.Input("rank", default=64, min=1, max=4096, step=1, + tooltip="LoRA rank - higher preserves more detail but increases size"), + io.Combo.Input("lora_type", options=["standard", "full_diff"], default="standard", + tooltip="standard: SVD decomposition, full_diff: raw weight differences"), + io.Boolean.Input("bias_diff", default=True, + tooltip="Include bias differences in the LoRA"), + io.String.Input("output_filename", default="extracted_lora", + tooltip="Output filename (without extension)"), + io.Combo.Input("save_dtype", options=["fp16", "bf16", "fp32"], default="fp16"), + io.Combo.Input("process_device", options=["cuda", "cpu"], default="cuda", + tooltip="Device for SVD computation"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + is_output_node=True, + ) + + @classmethod + def execute( + cls, + checkpoint_a: str, + checkpoint_b: str, + rank: int, + lora_type: str, + bias_diff: bool, + output_filename: str, + save_dtype: str, + process_device: str, + ) -> io.NodeOutput: + # Get full paths + checkpoint_a_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint_a) + checkpoint_b_path = folder_paths.get_full_path_or_raise("checkpoints", checkpoint_b) + + # Open handlers to detect prefixes + handler_a = MemoryEfficientSafeOpen(checkpoint_a_path) + keys = handler_a.keys() + handler_a.__exit__(None, None, None) + + # Auto-detect key prefixes + has_model = any(k.startswith("model.") for k in keys) + has_diffusion = any(k.startswith("diffusion_model.") for k in keys) + has_text_encoders = any(k.startswith("text_encoders.") or k.startswith("cond_stage_model.") for k in keys) + + output_sd = {} + + # Extract diffusion model LoRA + if has_model or has_diffusion: + prefix = "model.diffusion_model." if has_model else "diffusion_model." + model_lora = extract_lora_from_files( + model_a_path=checkpoint_a_path, + model_b_path=checkpoint_b_path, + rank=rank, + lora_type=lora_type, + bias_diff=bias_diff, + prefix_lora="diffusion_model.", + process_device=process_device, + save_dtype=save_dtype, + ) + output_sd.update(model_lora) + + # Note: Text encoder extraction would need additional prefix handling + # For full checkpoint support, this would need to be expanded + + if not output_sd: + raise ValueError("No LoRA weights extracted - checkpoints may be identical or incompatible") + + # Save to loras folder + output_dir = folder_paths.get_folder_paths("loras")[0] + os.makedirs(output_dir, exist_ok=True) + output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors") + + save_file(output_sd, output_path) + print(f"[LoRA Extract] Saved {len(output_sd)} tensors to {output_path}") + + return io.NodeOutput(output_path) diff --git a/nodes/merger.py b/nodes/merger.py index 5888107..3679565 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -3,6 +3,7 @@ import re import torch import folder_paths import comfy.utils +from comfy_api.latest import io from safetensors.torch import save_file from .merger_ops import TWO_MODEL_MODES, THREE_MODEL_MODES, MissingTensorBehavior, MissingTensorError from .merger_utils import MemoryEfficientSafeOpen @@ -45,60 +46,24 @@ def _matches_any_pattern(key, patterns): return False -# --- UI Dictionary Helper Functions --- -def _create_two_model_inputs(model_folder_name): - return {"required": { - "execution_mode": (["MERGE", "DOCUMENTATION ONLY"],), - "model_a": (["None"] + folder_paths.get_filename_list(model_folder_name),), - "model_b": (["None"] + folder_paths.get_filename_list(model_folder_name),), - "calc_mode": ([m.name for m in TWO_MODEL_MODES],), - "mismatch_mode": (["skip", "zeros", "error"], {"default": "skip"}), - "alpha": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), - "beta": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), - "gamma": ("FLOAT", {"default": 0.99, "min": 0.0, "max": 1.0, "step": 0.001}), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "output_filename": ("STRING", {"default": "merged_2_model"}), - "save_dtype": (["fp32", "fp16", "bf16"],), - "process_device": (["cuda", "cpu"],), - "exclude_patterns": ("STRING", {"default": "", "multiline": True, "placeholder": "Regex patterns (space-separated) - matching layers keep Model A only"}), - "discard_patterns": ("STRING", {"default": "", "multiline": True, "placeholder": "Regex patterns (space-separated) - matching layers removed from output"}), - }} - - -def _create_three_model_inputs(model_folder_name): - return {"required": { - "execution_mode": (["MERGE", "DOCUMENTATION ONLY"],), - "model_a": (["None"] + folder_paths.get_filename_list(model_folder_name),), - "model_b": (["None"] + folder_paths.get_filename_list(model_folder_name),), - "model_c": (["None"] + folder_paths.get_filename_list(model_folder_name),), - "calc_mode": ([m.name for m in THREE_MODEL_MODES],), - "mismatch_mode": (["skip", "zeros", "error"], {"default": "skip"}), - "alpha": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), - "beta": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), - "gamma": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), - "delta": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 3.0, "step": 0.01}), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "output_filename": ("STRING", {"default": "merged_3_model"}), - "save_dtype": (["fp32", "fp16", "bf16"],), - "process_device": (["cuda", "cpu"],), - "exclude_patterns": ("STRING", {"default": "", "multiline": True, "placeholder": "Regex patterns (space-separated) - matching layers keep Model A only"}), - "discard_patterns": ("STRING", {"default": "", "multiline": True, "placeholder": "Regex patterns (space-separated) - matching layers removed from output"}), - }} - - -# --- Base Classes for Backend Logic ONLY --- class MergerLogic: - def _execute_merge(self, model_names, calc_mode, all_modes, recipe_params, model_type): + """Shared merge execution logic for all merger nodes.""" + + @staticmethod + def execute_merge(model_names, calc_mode, all_modes, recipe_params, model_type): calc_mode_class = next((m for m in all_modes if m.name == calc_mode), None) - if not calc_mode_class: raise ValueError(f"Calc mode '{calc_mode}' not found.") + if not calc_mode_class: + raise ValueError(f"Calc mode '{calc_mode}' not found.") primary_model_name = model_names.get('model_a') - if not primary_model_name or primary_model_name == "None": raise ValueError("Model A is required to run the merge.") + if not primary_model_name or primary_model_name == "None": + raise ValueError("Model A is required to run the merge.") handlers = {} for name in model_names.values(): if name and name != "None": path = folder_paths.get_full_path(model_type, name) - if not path: raise FileNotFoundError(f"Model '{name}' not found.") + if not path: + raise FileNotFoundError(f"Model '{name}' not found.") handlers[name] = MemoryEfficientSafeOpen(path) primary_handler = handlers[primary_model_name] @@ -195,24 +160,58 @@ class MergerLogic: if error_keys: print(f"[Merger] Unexpected errors on {len(error_keys)} keys: {error_keys[:5]}{'...' if len(error_keys) > 5 else ''}") - for handler in handlers.values(): handler.__exit__(None, None, None) + for handler in handlers.values(): + handler.__exit__(None, None, None) output_folder = "loras" if calc_mode == "SVD LoRA Extraction" else model_type output_dir = folder_paths.get_folder_paths(output_folder)[0] os.makedirs(output_dir, exist_ok=True) output_filename = recipe_params.get("output_filename") output_path = os.path.join(output_dir, f"{output_filename}.safetensors") save_file(merged_state_dict, output_path, metadata=metadata) - return (f"{output_filename}.safetensors",) + return f"{output_filename}.safetensors" -class BaseTwoMerger(MergerLogic): - FUNCTION = "merge"; CATEGORY = "ModelUtils/Merging" - RETURN_TYPES = ("STRING", "STRING",); RETURN_NAMES = ("output_filename", "documentation",) +# --- Two-Model Merger Nodes --- - def merge(self, execution_mode, model_a, model_b, calc_mode, mismatch_mode, alpha, beta, gamma, seed, output_filename, save_dtype, process_device, exclude_patterns, discard_patterns): +class CheckpointTwoMerger(io.ComfyNode): + MODEL_TYPE = "checkpoints" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="CheckpointTwoMerger", + display_name="Merge Checkpoints (2 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("checkpoints")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("checkpoints")), + io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_2_checkpoint"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) + + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: doc = load_documentation_from_file('merger_2_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": - return ("Documentation mode active. No merge performed.", doc) + return io.NodeOutput("Documentation mode active. No merge performed.", doc) recipe_params = { "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, @@ -223,18 +222,265 @@ class BaseTwoMerger(MergerLogic): "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, } model_names = {"model_a": model_a, "model_b": model_b} - filename, = self._execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, self.MODEL_TYPE) - return (filename, doc) + filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) -class BaseThreeMerger(MergerLogic): - FUNCTION = "merge"; CATEGORY = "ModelUtils/Merging" - RETURN_TYPES = ("STRING", "STRING",); RETURN_NAMES = ("output_filename", "documentation",) +class ModelTwoMerger(io.ComfyNode): + MODEL_TYPE = "diffusion_models" - def merge(self, execution_mode, model_a, model_b, model_c, calc_mode, mismatch_mode, alpha, beta, gamma, delta, seed, output_filename, save_dtype, process_device, exclude_patterns, discard_patterns): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="ModelTwoMerger", + display_name="Merge Models (2 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("diffusion_models")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("diffusion_models")), + io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_2_model"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) + + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: + doc = load_documentation_from_file('merger_2_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, + "mismatch_mode": mismatch_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32, + "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, + } + model_names = {"model_a": model_a, "model_b": model_b} + filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) + + +class TextEncoderTwoMerger(io.ComfyNode): + MODEL_TYPE = "text_encoders" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderTwoMerger", + display_name="Merge Text Encoders (2 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("text_encoders")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("text_encoders")), + io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_2_textencoder"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) + + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: + doc = load_documentation_from_file('merger_2_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, + "mismatch_mode": mismatch_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32, + "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, + } + model_names = {"model_a": model_a, "model_b": model_b} + filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) + + +class LoRATwoMerger(io.ComfyNode): + MODEL_TYPE = "loras" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LoRATwoMerger", + display_name="Merge LoRAs (2 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("loras")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("loras")), + io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_2_lora"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) + + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: + doc = load_documentation_from_file('merger_2_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, + "mismatch_mode": mismatch_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32, + "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, + } + model_names = {"model_a": model_a, "model_b": model_b} + filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) + + +class EmbeddingTwoMerger(io.ComfyNode): + MODEL_TYPE = "embeddings" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="EmbeddingTwoMerger", + display_name="Merge Embeddings (2 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("embeddings")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("embeddings")), + io.Combo.Input("calc_mode", options=[m.name for m in TWO_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_2_embedding"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) + + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: + doc = load_documentation_from_file('merger_2_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "calc_mode": calc_mode, + "mismatch_mode": mismatch_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32, + "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, + } + model_names = {"model_a": model_a, "model_b": model_b} + filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) + + +# --- Three-Model Merger Nodes --- + +class CheckpointThreeMerger(io.ComfyNode): + MODEL_TYPE = "checkpoints" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="CheckpointThreeMerger", + display_name="Merge Checkpoints (3 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("checkpoints")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("checkpoints")), + io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("checkpoints")), + io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_3_checkpoint"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) + + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: doc = load_documentation_from_file('merger_3_model_modes.md') if execution_mode == "DOCUMENTATION ONLY": - return ("Documentation mode active. No merge performed.", doc) + return io.NodeOutput("Documentation mode active. No merge performed.", doc) recipe_params = { "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, @@ -245,55 +491,225 @@ class BaseThreeMerger(MergerLogic): "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, } model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} - filename, = self._execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, self.MODEL_TYPE) - return (filename, doc) + filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) -# --- Concrete 2-Model Nodes --- -class CheckpointTwoMerger(BaseTwoMerger): - MODEL_TYPE = "checkpoints" - @classmethod - def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE) -class ModelTwoMerger(BaseTwoMerger): +class ModelThreeMerger(io.ComfyNode): MODEL_TYPE = "diffusion_models" - @classmethod - def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE) -class TextEncoderTwoMerger(BaseTwoMerger): - MODEL_TYPE = "text_encoders" - @classmethod - def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE) -class LoRATwoMerger(BaseTwoMerger): - MODEL_TYPE = "loras" - @classmethod - def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE) -class EmbeddingTwoMerger(BaseTwoMerger): - MODEL_TYPE = "embeddings" - @classmethod - def INPUT_TYPES(s): return _create_two_model_inputs(s.MODEL_TYPE) -# --- Concrete 3-Model Nodes --- -class CheckpointThreeMerger(BaseThreeMerger): - MODEL_TYPE = "checkpoints" @classmethod - def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE) -class ModelThreeMerger(BaseThreeMerger): - MODEL_TYPE = "diffusion_models" - @classmethod - def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE) -class TextEncoderThreeMerger(BaseThreeMerger): - MODEL_TYPE = "text_encoders" - @classmethod - def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE) -class LoRAThreeMerger(BaseThreeMerger): - MODEL_TYPE = "loras" - @classmethod - def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE) -class EmbeddingThreeMerger(BaseThreeMerger): - MODEL_TYPE = "embeddings" - @classmethod - def INPUT_TYPES(s): return _create_three_model_inputs(s.MODEL_TYPE) + def define_schema(cls): + return io.Schema( + node_id="ModelThreeMerger", + display_name="Merge Models (3 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("diffusion_models")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("diffusion_models")), + io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("diffusion_models")), + io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_3_model"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) -__all__ = [ - "CheckpointTwoMerger", "ModelTwoMerger", "TextEncoderTwoMerger", "LoRATwoMerger", "EmbeddingTwoMerger", - "CheckpointThreeMerger", "ModelThreeMerger", "TextEncoderThreeMerger", "LoRAThreeMerger", "EmbeddingThreeMerger" -] \ No newline at end of file + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: + doc = load_documentation_from_file('merger_3_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, + "mismatch_mode": mismatch_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32, + "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, + } + model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} + filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) + + +class TextEncoderThreeMerger(io.ComfyNode): + MODEL_TYPE = "text_encoders" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncoderThreeMerger", + display_name="Merge Text Encoders (3 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("text_encoders")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("text_encoders")), + io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("text_encoders")), + io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_3_textencoder"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) + + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: + doc = load_documentation_from_file('merger_3_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, + "mismatch_mode": mismatch_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32, + "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, + } + model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} + filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) + + +class LoRAThreeMerger(io.ComfyNode): + MODEL_TYPE = "loras" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LoRAThreeMerger", + display_name="Merge LoRAs (3 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("loras")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("loras")), + io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("loras")), + io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_3_lora"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) + + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: + doc = load_documentation_from_file('merger_3_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, + "mismatch_mode": mismatch_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32, + "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, + } + model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} + filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) + + +class EmbeddingThreeMerger(io.ComfyNode): + MODEL_TYPE = "embeddings" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="EmbeddingThreeMerger", + display_name="Merge Embeddings (3 Models)", + category="ModelUtils/Merging", + inputs=[ + io.Combo.Input("execution_mode", options=["MERGE", "DOCUMENTATION ONLY"]), + io.Combo.Input("model_a", options=["None"] + folder_paths.get_filename_list("embeddings")), + io.Combo.Input("model_b", options=["None"] + folder_paths.get_filename_list("embeddings")), + io.Combo.Input("model_c", options=["None"] + folder_paths.get_filename_list("embeddings")), + io.Combo.Input("calc_mode", options=[m.name for m in THREE_MODEL_MODES]), + io.Combo.Input("mismatch_mode", options=["skip", "zeros", "error"], default="skip"), + io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff), + io.String.Input("output_filename", default="merged_3_embedding"), + io.Combo.Input("save_dtype", options=["fp32", "fp16", "bf16"]), + io.Combo.Input("process_device", options=["cuda", "cpu"]), + io.String.Input("exclude_patterns", default="", multiline=True), + io.String.Input("discard_patterns", default="", multiline=True), + ], + outputs=[ + io.String.Output(display_name="output_filename"), + io.String.Output(display_name="documentation"), + ], + ) + + @classmethod + def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str, + calc_mode: str, mismatch_mode: str, alpha: float, beta: float, + gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str, + process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput: + doc = load_documentation_from_file('merger_3_model_modes.md') + if execution_mode == "DOCUMENTATION ONLY": + return io.NodeOutput("Documentation mode active. No merge performed.", doc) + + recipe_params = { + "model_a": model_a, "model_b": model_b, "model_c": model_c, "calc_mode": calc_mode, + "mismatch_mode": mismatch_mode, + "alpha": alpha, "beta": beta, "gamma": gamma, "delta": delta, "seed": seed, + "output_filename": output_filename, "save_dtype": save_dtype, + "device": process_device, "dtype": torch.float32, + "exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns, + } + model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c} + filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE) + return io.NodeOutput(filename, doc) \ No newline at end of file diff --git a/nodes/metakeys.py b/nodes/metakeys.py index decd91e..626a383 100644 --- a/nodes/metakeys.py +++ b/nodes/metakeys.py @@ -1,98 +1,159 @@ import os import folder_paths +from comfy_api.latest import io from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors +def _get_metakeys(model_name: str, model_type: str) -> tuple[str, str]: + """Shared logic for all MetaKeys nodes.""" + model_path = folder_paths.get_full_path_or_raise(model_type, model_name) + model_weights = {} + metadata = {} + layer_shapes = "" -class BaseMetaKeys: - RETURN_TYPES = ("STRING", "STRING") - RETURN_NAMES = ("metadata", "keys") - FUNCTION = "get_metakeys" - CATEGORY = "ModelUtils/Keys" - DESCRIPTION = "Loads a model and returns its metadata and layer names." + 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: + return f"Conversion failed: {error_message}", "" + model_path = temp_safe_path - def _get_metakeys(self, model_name, model_type): - model_path = folder_paths.get_full_path_or_raise(model_type, model_name) - model_weights = {} - metadata = {} - layer_shapes = "" - - 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: - return f"Conversion failed: {error_message}", "" - model_path = temp_safe_path - - model_weights, metadata = load_metadata_from_safetensors(model_path) - metadata_str = str(metadata) - for layer_name, tensor in model_weights.items(): - shape = str(tensor.shape) - shape = shape.replace("torch.Size(", "").replace(")", "") - name_shape = f"{layer_name}, {shape}\n" - layer_shapes += name_shape - return metadata_str, layer_shapes + model_weights, metadata = load_metadata_from_safetensors(model_path) + metadata_str = str(metadata) + for layer_name, tensor in model_weights.items(): + shape = str(tensor.shape) + shape = shape.replace("torch.Size(", "").replace(")", "") + name_shape = f"{layer_name}, {shape}\n" + layer_shapes += name_shape + return metadata_str, layer_shapes - -class ModelMetaKeys(BaseMetaKeys): +class ModelMetaKeys(io.ComfyNode): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "diffusionmodel_name": (folder_paths.get_filename_list("diffusion_models"), ), - } - } + def define_schema(cls): + return io.Schema( + node_id="ModelMetaKeys", + display_name="Get Diffusion Model Metadata & Keys", + category="ModelUtils/Keys", + description="Loads a diffusion model and returns its metadata and layer names.", + inputs=[ + io.Combo.Input( + "diffusionmodel_name", + options=folder_paths.get_filename_list("diffusion_models"), + ), + ], + outputs=[ + io.String.Output(display_name="metadata"), + io.String.Output(display_name="keys"), + ], + ) - def get_metakeys(self, diffusionmodel_name): - return self._get_metakeys(diffusionmodel_name, "diffusion_models") - -class TextEncoderMetaKeys(BaseMetaKeys): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "textencoder_name": (folder_paths.get_filename_list("text_encoders"), ), - } - } + def execute(cls, diffusionmodel_name: str) -> io.NodeOutput: + metadata, keys = _get_metakeys(diffusionmodel_name, "diffusion_models") + return io.NodeOutput(metadata, keys) - def get_metakeys(self, textencoder_name): - return self._get_metakeys(textencoder_name, "text_encoders") -class LoRAMetaKeys(BaseMetaKeys): +class TextEncoderMetaKeys(io.ComfyNode): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "lora_name": (folder_paths.get_filename_list("loras"), ), - } - } - def get_metakeys(self, lora_name): - return self._get_metakeys(lora_name, "loras") + def define_schema(cls): + return io.Schema( + node_id="TextEncoderMetaKeys", + display_name="Get Text Encoder Metadata & Keys", + category="ModelUtils/Keys", + description="Loads a text encoder and returns its metadata and layer names.", + inputs=[ + io.Combo.Input( + "textencoder_name", + options=folder_paths.get_filename_list("text_encoders"), + ), + ], + outputs=[ + io.String.Output(display_name="metadata"), + io.String.Output(display_name="keys"), + ], + ) -class CheckpointMetaKeys(BaseMetaKeys): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), - } - } + def execute(cls, textencoder_name: str) -> io.NodeOutput: + metadata, keys = _get_metakeys(textencoder_name, "text_encoders") + return io.NodeOutput(metadata, keys) - def get_metakeys(self, ckpt_name): - return self._get_metakeys(ckpt_name, "checkpoints") -class EmbeddingMetaKeys(BaseMetaKeys): +class LoRAMetaKeys(io.ComfyNode): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "embedding": (folder_paths.get_filename_list("embeddings"), ), - } - } + def define_schema(cls): + return io.Schema( + node_id="LoRAMetaKeys", + display_name="Get LoRA Metadata & Keys", + category="ModelUtils/Keys", + description="Loads a LoRA and returns its metadata and layer names.", + inputs=[ + io.Combo.Input( + "lora_name", + options=folder_paths.get_filename_list("loras"), + ), + ], + outputs=[ + io.String.Output(display_name="metadata"), + io.String.Output(display_name="keys"), + ], + ) - def get_metakeys(self, embedding): - return self._get_metakeys(embedding, "embeddings") + @classmethod + def execute(cls, lora_name: str) -> io.NodeOutput: + metadata, keys = _get_metakeys(lora_name, "loras") + return io.NodeOutput(metadata, keys) -__all__ = ("ModelMetaKeys", "TextEncoderMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys", "EmbeddingMetaKeys") + +class CheckpointMetaKeys(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="CheckpointMetaKeys", + display_name="Get Checkpoint Metadata & Keys", + category="ModelUtils/Keys", + description="Loads a checkpoint and returns its metadata and layer names.", + inputs=[ + io.Combo.Input( + "ckpt_name", + options=folder_paths.get_filename_list("checkpoints"), + ), + ], + outputs=[ + io.String.Output(display_name="metadata"), + io.String.Output(display_name="keys"), + ], + ) + + @classmethod + def execute(cls, ckpt_name: str) -> io.NodeOutput: + metadata, keys = _get_metakeys(ckpt_name, "checkpoints") + return io.NodeOutput(metadata, keys) + + +class EmbeddingMetaKeys(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="EmbeddingMetaKeys", + display_name="Get Embedding Metadata & Keys", + category="ModelUtils/Keys", + description="Loads an embedding and returns its metadata and layer names.", + inputs=[ + io.Combo.Input( + "embedding", + options=folder_paths.get_filename_list("embeddings"), + ), + ], + outputs=[ + io.String.Output(display_name="metadata"), + io.String.Output(display_name="keys"), + ], + ) + + @classmethod + def execute(cls, embedding: str) -> io.NodeOutput: + metadata, keys = _get_metakeys(embedding, "embeddings") + return io.NodeOutput(metadata, keys) diff --git a/nodes/prunekeys.py b/nodes/prunekeys.py index fb505cf..883fa80 100644 --- a/nodes/prunekeys.py +++ b/nodes/prunekeys.py @@ -1,127 +1,199 @@ import os import re import folder_paths -from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors +from comfy_api.latest import io from safetensors.torch import save_file +from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors -class BasePruneKeys: - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("output_path",) - FUNCTION = "prune_keys" - CATEGORY = "ModelUtils/Keys" - 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) +def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str, + use_regex: bool, output_filename: str) -> str: + """Shared logic for all PruneKeys nodes.""" + 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 + 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: - model_path_to_load = model_path + if any(pattern in key for pattern in patterns): + is_match = True - model_weights, metadata = load_metadata_from_safetensors(model_path_to_load) + if not is_match: + pruned_tensors[key] = tensor - patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()] + model_dir = folder_paths.get_folder_paths(model_type)[0] + output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors") - if not patterns: - raise ValueError("No keys/patterns provided to prune.") + save_file(pruned_tensors, output_path, metadata) - 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 + return output_path - 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 ModelPruneKeys(BasePruneKeys): +class ModelPruneKeys(io.ComfyNode): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "diffusionmodel_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_model"}), - } - } + def define_schema(cls): + return io.Schema( + node_id="ModelPruneKeys", + display_name="Prune Diffusion Model Keys", + category="ModelUtils/Keys", + description="Loads a diffusion model, removes specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "diffusionmodel_name", + options=folder_paths.get_filename_list("diffusion_models"), + ), + io.String.Input("keys_to_prune", multiline=True, default=""), + io.Boolean.Input("use_regex", default=False), + io.String.Input("output_filename", default="pruned_model"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) - def prune_keys(self, diffusionmodel_name, keys_to_prune, use_regex, output_filename): - return self._prune_keys(diffusionmodel_name, "diffusion_models", keys_to_prune, use_regex, output_filename) - -class TextEncoderPruneKeys(BasePruneKeys): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "textencoder_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_textencoder"}), - } - } + def execute(cls, diffusionmodel_name: str, keys_to_prune: str, + use_regex: bool, output_filename: str) -> io.NodeOutput: + path = _prune_keys(diffusionmodel_name, "diffusion_models", + keys_to_prune, use_regex, output_filename) + return io.NodeOutput(path) - def prune_keys(self, textencoder_name, keys_to_prune, use_regex, output_filename): - return self._prune_keys(textencoder_name, "text_encoders", keys_to_prune, use_regex, output_filename) -class LoRAPruneKeys(BasePruneKeys): +class TextEncoderPruneKeys(io.ComfyNode): @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 define_schema(cls): + return io.Schema( + node_id="TextEncoderPruneKeys", + display_name="Prune Text Encoder Keys", + category="ModelUtils/Keys", + description="Loads a text encoder, removes specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "textencoder_name", + options=folder_paths.get_filename_list("text_encoders"), + ), + io.String.Input("keys_to_prune", multiline=True, default=""), + io.Boolean.Input("use_regex", default=False), + io.String.Input("output_filename", default="pruned_textencoder"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) - 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 execute(cls, textencoder_name: str, keys_to_prune: str, + use_regex: bool, output_filename: str) -> io.NodeOutput: + path = _prune_keys(textencoder_name, "text_encoders", + keys_to_prune, use_regex, output_filename) + return io.NodeOutput(path) - 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) -class EmbeddingPruneKeys(BasePruneKeys): +class LoRAPruneKeys(io.ComfyNode): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "embedding": (folder_paths.get_filename_list("embeddings"), ), - "keys_to_prune": ("STRING", {"multiline": True, "default": ""}), - "use_regex": ("BOOLEAN", {"default": False}), - "output_filename": ("STRING", {"default": "pruned_embedding"}), - } - } + def define_schema(cls): + return io.Schema( + node_id="LoRAPruneKeys", + display_name="Prune LoRA Keys", + category="ModelUtils/Keys", + description="Loads a LoRA, removes specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "lora_name", + options=folder_paths.get_filename_list("loras"), + ), + io.String.Input("keys_to_prune", multiline=True, default=""), + io.Boolean.Input("use_regex", default=False), + io.String.Input("output_filename", default="pruned_lora"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) - def prune_keys(self, embedding, keys_to_prune, use_regex, output_filename): - return self._prune_keys(embedding, "embeddings", keys_to_prune, use_regex, output_filename) \ No newline at end of file + @classmethod + def execute(cls, lora_name: str, keys_to_prune: str, + use_regex: bool, output_filename: str) -> io.NodeOutput: + path = _prune_keys(lora_name, "loras", + keys_to_prune, use_regex, output_filename) + return io.NodeOutput(path) + + +class CheckpointPruneKeys(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="CheckpointPruneKeys", + display_name="Prune Checkpoint Keys", + category="ModelUtils/Keys", + description="Loads a checkpoint, removes specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "ckpt_name", + options=folder_paths.get_filename_list("checkpoints"), + ), + io.String.Input("keys_to_prune", multiline=True, default=""), + io.Boolean.Input("use_regex", default=False), + io.String.Input("output_filename", default="pruned_checkpoint"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) + + @classmethod + def execute(cls, ckpt_name: str, keys_to_prune: str, + use_regex: bool, output_filename: str) -> io.NodeOutput: + path = _prune_keys(ckpt_name, "checkpoints", + keys_to_prune, use_regex, output_filename) + return io.NodeOutput(path) + + +class EmbeddingPruneKeys(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="EmbeddingPruneKeys", + display_name="Prune Embedding Keys", + category="ModelUtils/Keys", + description="Loads an embedding, removes specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "embedding", + options=folder_paths.get_filename_list("embeddings"), + ), + io.String.Input("keys_to_prune", multiline=True, default=""), + io.Boolean.Input("use_regex", default=False), + io.String.Input("output_filename", default="pruned_embedding"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) + + @classmethod + def execute(cls, embedding: str, keys_to_prune: str, + use_regex: bool, output_filename: str) -> io.NodeOutput: + path = _prune_keys(embedding, "embeddings", + keys_to_prune, use_regex, output_filename) + return io.NodeOutput(path) \ No newline at end of file diff --git a/nodes/renamekeys.py b/nodes/renamekeys.py index 269bfcc..d856a70 100644 --- a/nodes/renamekeys.py +++ b/nodes/renamekeys.py @@ -1,122 +1,194 @@ import os import folder_paths -from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors +from comfy_api.latest import io from safetensors.torch import save_file +from .utils import convert_pt_to_safetensors, load_metadata_from_safetensors -class BaseRenameKeys: - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("output_path",) - FUNCTION = "rename_keys" - CATEGORY = "ModelUtils/Keys" - 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) +def _rename_keys(model_name: str, model_type: str, old_keys_str: str, + new_keys_str: str, output_filename: str) -> str: + """Shared logic for all RenameKeys nodes.""" + 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 + 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) + 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()] + 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.") + 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)) + 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] + 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") + 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) + save_file(renamed_tensors, output_path, metadata) - return (output_path,) + return output_path -class ModelRenameKeys(BaseRenameKeys): + +class ModelRenameKeys(io.ComfyNode): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "diffusionmodel_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_model"}), - } - } + def define_schema(cls): + return io.Schema( + node_id="ModelRenameKeys", + display_name="Rename Diffusion Model Keys", + category="ModelUtils/Keys", + description="Loads a diffusion model, renames specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "diffusionmodel_name", + options=folder_paths.get_filename_list("diffusion_models"), + ), + io.String.Input("old_keys", multiline=True, default=""), + io.String.Input("new_keys", multiline=True, default=""), + io.String.Input("output_filename", default="renamed_model"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) - def rename_keys(self, diffusionmodel_name, old_keys, new_keys, output_filename): - return self._rename_keys(diffusionmodel_name, "diffusion_models", old_keys, new_keys, output_filename) - -class TextEncoderRenameKeys(BaseRenameKeys): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "textencoder_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_textencoder"}), - } - } + def execute(cls, diffusionmodel_name: str, old_keys: str, + new_keys: str, output_filename: str) -> io.NodeOutput: + path = _rename_keys(diffusionmodel_name, "diffusion_models", + old_keys, new_keys, output_filename) + return io.NodeOutput(path) - def rename_keys(self, textencoder_name, old_keys, new_keys, output_filename): - return self._rename_keys(textencoder_name, "text_encoders", old_keys, new_keys, output_filename) -class LoRARenameKeys(BaseRenameKeys): +class TextEncoderRenameKeys(io.ComfyNode): @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 define_schema(cls): + return io.Schema( + node_id="TextEncoderRenameKeys", + display_name="Rename Text Encoder Keys", + category="ModelUtils/Keys", + description="Loads a text encoder, renames specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "textencoder_name", + options=folder_paths.get_filename_list("text_encoders"), + ), + io.String.Input("old_keys", multiline=True, default=""), + io.String.Input("new_keys", multiline=True, default=""), + io.String.Input("output_filename", default="renamed_textencoder"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) - 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 execute(cls, textencoder_name: str, old_keys: str, + new_keys: str, output_filename: str) -> io.NodeOutput: + path = _rename_keys(textencoder_name, "text_encoders", + old_keys, new_keys, output_filename) + return io.NodeOutput(path) - 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) -class EmbeddingRenameKeys(BaseRenameKeys): +class LoRARenameKeys(io.ComfyNode): @classmethod - def INPUT_TYPES(s): - return { - "required": { - "embedding": (folder_paths.get_filename_list("embeddings"), ), - "old_keys": ("STRING", {"multiline": True, "default": ""}), - "new_keys": ("STRING", {"multiline": True, "default": ""}), - "output_filename": ("STRING", {"default": "renamed_embedding"}), - } - } + def define_schema(cls): + return io.Schema( + node_id="LoRARenameKeys", + display_name="Rename LoRA Keys", + category="ModelUtils/Keys", + description="Loads a LoRA, renames specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "lora_name", + options=folder_paths.get_filename_list("loras"), + ), + io.String.Input("old_keys", multiline=True, default=""), + io.String.Input("new_keys", multiline=True, default=""), + io.String.Input("output_filename", default="renamed_lora"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) - def rename_keys(self, embedding, old_keys, new_keys, output_filename): - return self._rename_keys(embedding, "embeddings", old_keys, new_keys, output_filename) \ No newline at end of file + @classmethod + def execute(cls, lora_name: str, old_keys: str, + new_keys: str, output_filename: str) -> io.NodeOutput: + path = _rename_keys(lora_name, "loras", + old_keys, new_keys, output_filename) + return io.NodeOutput(path) + + +class CheckpointRenameKeys(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="CheckpointRenameKeys", + display_name="Rename Checkpoint Keys", + category="ModelUtils/Keys", + description="Loads a checkpoint, renames specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "ckpt_name", + options=folder_paths.get_filename_list("checkpoints"), + ), + io.String.Input("old_keys", multiline=True, default=""), + io.String.Input("new_keys", multiline=True, default=""), + io.String.Input("output_filename", default="renamed_checkpoint"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) + + @classmethod + def execute(cls, ckpt_name: str, old_keys: str, + new_keys: str, output_filename: str) -> io.NodeOutput: + path = _rename_keys(ckpt_name, "checkpoints", + old_keys, new_keys, output_filename) + return io.NodeOutput(path) + + +class EmbeddingRenameKeys(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="EmbeddingRenameKeys", + display_name="Rename Embedding Keys", + category="ModelUtils/Keys", + description="Loads an embedding, renames specified keys, and saves it as a new safetensors file.", + inputs=[ + io.Combo.Input( + "embedding", + options=folder_paths.get_filename_list("embeddings"), + ), + io.String.Input("old_keys", multiline=True, default=""), + io.String.Input("new_keys", multiline=True, default=""), + io.String.Input("output_filename", default="renamed_embedding"), + ], + outputs=[ + io.String.Output(display_name="output_path"), + ], + ) + + @classmethod + def execute(cls, embedding: str, old_keys: str, + new_keys: str, output_filename: str) -> io.NodeOutput: + path = _rename_keys(embedding, "embeddings", + old_keys, new_keys, output_filename) + return io.NodeOutput(path) \ No newline at end of file