diff --git a/__init__.py b/__init__.py index 7265f2f..0129cf1 100644 --- a/__init__.py +++ b/__init__.py @@ -3,26 +3,55 @@ from .nodes import * NODE_CLASS_MAPPINGS = { "UNetMetaKeys": UNetMetaKeys, "CLIPMetaKeys": CLIPMetaKeys, + "LoRAMetaKeys": LoRAMetaKeys, + "CheckpointMetaKeys": CheckpointMetaKeys, + "EmbeddingMetaKeys": EmbeddingMetaKeys, + "UNetRenameKeys": UNetRenameKeys, + "CLIPRenameKeys": CLIPRenameKeys, + "LoRARenameKeys": LoRARenameKeys, + "CheckpointRenameKeys": CheckpointRenameKeys, + "EmbeddingRenameKeys": EmbeddingRenameKeys, + "UNetPruneKeys": UNetPruneKeys, + "CLIPPruneKeys": CLIPPruneKeys, + "LoRAPruneKeys": LoRAPruneKeys, "CheckpointPruneKeys": CheckpointPruneKeys, - "CheckpointTwoMerger": CheckpointTwoMerger, + "EmbeddingPruneKeys": EmbeddingPruneKeys, "UNetTwoMerger": UNetTwoMerger, "CLIPTwoMerger": CLIPTwoMerger, "LoRATwoMerger": LoRATwoMerger, - "CheckpointThreeMerger": CheckpointThreeMerger, + "CheckpointTwoMerger": CheckpointTwoMerger, + "EmbeddingTwoMerger": EmbeddingTwoMerger, "UNetThreeMerger": UNetThreeMerger, "CLIPThreeMerger": CLIPThreeMerger, "LoRAThreeMerger": LoRAThreeMerger, + "CheckpointThreeMerger": CheckpointThreeMerger, + "EmbeddingThreeMerger": EmbeddingThreeMerger, } NODE_DISPLAY_NAME_MAPPINGS = { - "UNetMetaKeys": "Get UNet Meta Keys", + "UNetMetaKeys": "Get UNet Metadata & Keys", + "CLIPMetaKeys": "Get CLIP Metadata & Keys", + "LoRAMetaKeys": "Get LoRA Metadata & Keys", + "CheckpointMetaKeys": "Get Checkpoint Metadata & Keys", + "EmbeddingMetaKeys": "Get Embedding Metadata & Keys", + "UNetRenameKeys": "UNetRenameKeys", + "CLIPRenameKeys": "CLIPRenameKeys", + "LoRARenameKeys": "LoRARenameKeys", + "CheckpointRenameKeys": "CheckpointRenameKeys", + "EmbeddingRenameKeys": "EmbeddingRenameKeys", + "UNetPruneKeys": "UNetPruneKeys", + "CLIPPruneKeys": "CLIPPruneKeys", + "LoRAPruneKeys": "LoRAPruneKeys", "CheckpointPruneKeys": "Prune Checkpoint Keys", - "CheckpointTwoMerger": "Merge Checkpoints (2 Models)", + "EmbeddingPruneKeys": "Prune Embedding Keys", "UNetTwoMerger": "Merge UNets (2 Models)", "CLIPTwoMerger": "Merge CLIPs (2 Models)", "LoRATwoMerger": "Merge LoRAs (2 Models)", - "CheckpointThreeMerger": "Merge Checkpoints (3 Models)", + "CheckpointTwoMerger": "Merge Checkpoints (2 Models)", + "EmbeddingTwoMerger": "Merge Embeddings (2 Models)", "UNetThreeMerger": "Merge UNets (3 Models)", "CLIPThreeMerger": "Merge CLIPs (3 Models)", "LoRAThreeMerger": "Merge LoRAs (3 Models)", + "CheckpointThreeMerger": "Merge Checkpoints (3 Models)", + "EmbeddingThreeMerger": "Merge Embeddings (3 Models)", } \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py index c94a632..195e3a4 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -4,9 +4,9 @@ from .prunekeys import * from .merger import * __all__ = [ - "UNetMetaKeys", "CLIPMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys", - "UNetRenameKeys", "CLIPRenameKeys", "LoRARenameKeys", "CheckpointRenameKeys", - "UNetPruneKeys", "CLIPPruneKeys", "LoRAPruneKeys", "CheckpointPruneKeys", - "CheckpointTwoMerger", "UNetTwoMerger", "CLIPTwoMerger", "LoRATwoMerger", - "CheckpointThreeMerger", "UNetThreeMerger", "CLIPThreeMerger", "LoRAThreeMerger", + "UNetMetaKeys", "CLIPMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys", "EmbeddingMetaKeys", + "UNetRenameKeys", "CLIPRenameKeys", "LoRARenameKeys", "CheckpointRenameKeys", "EmbeddingRenameKeys", + "UNetPruneKeys", "CLIPPruneKeys", "LoRAPruneKeys", "CheckpointPruneKeys", "EmbeddingPruneKeys", + "CheckpointTwoMerger", "UNetTwoMerger", "CLIPTwoMerger", "LoRATwoMerger", "EmbeddingTwoMerger", + "CheckpointThreeMerger", "UNetThreeMerger", "CLIPThreeMerger", "LoRAThreeMerger", "EmbeddingThreeMerger", ] \ No newline at end of file diff --git a/nodes/merger.py b/nodes/merger.py index 61fb18c..f0f8de8 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -140,6 +140,10 @@ 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): @@ -158,8 +162,12 @@ 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) __all__ = [ - "CheckpointTwoMerger", "UNetTwoMerger", "CLIPTwoMerger", "LoRATwoMerger", - "CheckpointThreeMerger", "UNetThreeMerger", "CLIPThreeMerger", "LoRAThreeMerger" + "CheckpointTwoMerger", "UNetTwoMerger", "CLIPTwoMerger", "LoRATwoMerger", "EmbeddingTwoMerger", + "CheckpointThreeMerger", "UNetThreeMerger", "CLIPThreeMerger", "LoRAThreeMerger", "EmbeddingThreeMerger" ] \ No newline at end of file diff --git a/nodes/metakeys.py b/nodes/metakeys.py index f04fc78..f5d4af5 100644 --- a/nodes/metakeys.py +++ b/nodes/metakeys.py @@ -83,4 +83,16 @@ class CheckpointMetaKeys(BaseMetaKeys): def get_metakeys(self, ckpt_name): return self._get_metakeys(ckpt_name, "checkpoints") -__all__ = ("UNetMetaKeys", "CLIPMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys") +class EmbeddingMetaKeys(BaseMetaKeys): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "embedding": (folder_paths.get_filename_list("embeddings"), ), + } + } + + def get_metakeys(self, embedding): + return self._get_metakeys(embedding, "embeddings") + +__all__ = ("UNetMetaKeys", "CLIPMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys", "EmbeddingMetaKeys") diff --git a/nodes/prunekeys.py b/nodes/prunekeys.py index e333b6e..a864225 100644 --- a/nodes/prunekeys.py +++ b/nodes/prunekeys.py @@ -109,4 +109,19 @@ class CheckpointPruneKeys(BasePruneKeys): } 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 + return self._prune_keys(ckpt_name, "checkpoints", keys_to_prune, use_regex, output_filename) + +class EmbeddingPruneKeys(BasePruneKeys): + @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 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 diff --git a/nodes/renamekeys.py b/nodes/renamekeys.py index f82d5b3..726125c 100644 --- a/nodes/renamekeys.py +++ b/nodes/renamekeys.py @@ -104,4 +104,19 @@ class CheckpointRenameKeys(BaseRenameKeys): } 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 + return self._rename_keys(ckpt_name, "checkpoints", old_keys, new_keys, output_filename) + +class EmbeddingRenameKeys(BaseRenameKeys): + @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 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