from .aimlapi_image_editor import AIMLAPIImageEditor
from .aimlapi_image_generator import AIMLAPIImageGenerator
from .aimlapi_video_first_last import AIMLAPIVideoFirstLast
from .aimlapi_video_generator import AIMLAPIVideoGenerator
from .aimlapi_video_with_audio import AIMLAPIVideoWithAudio


NODE_CLASS_MAPPINGS = {
    "AIMLAPIImageGenerator": AIMLAPIImageGenerator,
    "AIMLAPIImageEditor": AIMLAPIImageEditor,
    "AIMLAPIVideoGenerator": AIMLAPIVideoGenerator,
    "AIMLAPIVideoWithAudio": AIMLAPIVideoWithAudio,
    "AIMLAPIVideoFirstLast": AIMLAPIVideoFirstLast,
}


NODE_DISPLAY_NAME_MAPPINGS = {
    "AIMLAPIImageGenerator": "AI/ML API Image Generator (Nano Banana Pro / Seedream 4.5)",
    "AIMLAPIImageEditor": "AI/ML API Image Editor (Nano Banana Pro Edit / Seedream 4.5)",
    "AIMLAPIVideoGenerator": "AI/ML API Video Generator (LTXV / Veo 3.1)",
    "AIMLAPIVideoWithAudio": "AI/ML API Video + Audio (Fabric 1.0)",
    "AIMLAPIVideoFirstLast": "AI/ML API Video Keyframes (Kling / Veo / Runway)",
}