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,
    "EmbeddingPruneKeys": EmbeddingPruneKeys,
    "UNetTwoMerger": UNetTwoMerger,
    "CLIPTwoMerger": CLIPTwoMerger,
    "LoRATwoMerger": LoRATwoMerger,
    "CheckpointTwoMerger": CheckpointTwoMerger,
    "EmbeddingTwoMerger": EmbeddingTwoMerger,
    "UNetThreeMerger": UNetThreeMerger,
    "CLIPThreeMerger": CLIPThreeMerger,
    "LoRAThreeMerger": LoRAThreeMerger,
    "CheckpointThreeMerger": CheckpointThreeMerger,
    "EmbeddingThreeMerger": EmbeddingThreeMerger,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "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",
    "EmbeddingPruneKeys": "Prune Embedding Keys",
    "UNetTwoMerger": "Merge UNets (2 Models)",
    "CLIPTwoMerger": "Merge CLIPs (2 Models)",
    "LoRATwoMerger": "Merge LoRAs (2 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)",
}