from .nodes_t2i import PrunaTextToImage
from .nodes_lora import PrunaTextToImageLoRA
from .nodes_edit import PrunaImageEdit, PrunaImageEditLoRA
from .nodes_upscale import PrunaImageUpscale
from .nodes_video import PrunaVideo

NODE_CLASS_MAPPINGS = {
    "PrunaTextToImage": PrunaTextToImage,
    "PrunaTextToImageLoRA": PrunaTextToImageLoRA,
    "PrunaImageEdit": PrunaImageEdit,
    "PrunaImageEditLoRA": PrunaImageEditLoRA,
    "PrunaImageUpscale": PrunaImageUpscale,
    "PrunaVideo": PrunaVideo,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PrunaTextToImage": "p-image",
    "PrunaTextToImageLoRA": "p-image-lora",
    "PrunaImageEdit": "p-image-edit",
    "PrunaImageEditLoRA": "p-image-edit-lora",
    "PrunaImageUpscale": "p-image-upscale",
    "PrunaVideo": "p-video",
}

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]