"""Enums, endpoint maps, and response normalization utilities.""" from __future__ import annotations from enum import Enum from typing import Any # --------------------------------------------------------------------------- # Model enums # --------------------------------------------------------------------------- class ImageModel(str, Enum): MYSTIC = "mystic" FLUX_KONTEXT_PRO = "flux-kontext-pro" FLUX_2_PRO = "flux-2-pro" FLUX_2_TURBO = "flux-2-turbo" FLUX_PRO_1_1 = "flux-pro-1.1" SEEDREAM_V4 = "seedream-v4" SEEDREAM_V4_5 = "seedream-v4-5" IDEOGRAM_V2 = "ideogram-v2" class VideoModel(str, Enum): KLING_O1_PRO = "kling-o1-pro" KLING_O1_STD = "kling-o1-std" KLING_ELEMENTS_PRO = "kling-elements-pro" KLING_ELEMENTS_STD = "kling-elements-std" MINIMAX_HAILUO = "minimax-hailuo" class UpscaleMode(str, Enum): CREATIVE = "creative" PRECISION = "precision" PRECISION_V2 = "precision-v2" class VideoUpscaleMode(str, Enum): STANDARD = "standard" TURBO = "turbo" class AspectRatio(str, Enum): LANDSCAPE = "16:9" PORTRAIT = "9:16" SQUARE = "1:1" CLASSIC = "4:3" WIDE = "21:9" class IconStyle(str, Enum): SOLID = "solid" OUTLINE = "outline" COLOR = "color" FLAT = "flat" STICKER = "sticker" # --------------------------------------------------------------------------- # Endpoint maps # --------------------------------------------------------------------------- IMAGE_POST_ENDPOINTS: dict[ImageModel, str] = { ImageModel.MYSTIC: "/v1/ai/mystic", ImageModel.FLUX_KONTEXT_PRO: "/v1/ai/text-to-image/flux-kontext-pro", ImageModel.FLUX_2_PRO: "/v1/ai/text-to-image/flux-2-pro", ImageModel.FLUX_2_TURBO: "/v1/ai/text-to-image/flux-2-turbo", ImageModel.FLUX_PRO_1_1: "/v1/ai/text-to-image/flux-pro-v1-1", ImageModel.SEEDREAM_V4: "/v1/ai/text-to-image/seedream-v4", ImageModel.SEEDREAM_V4_5: "/v1/ai/text-to-image/seedream-v4-5", ImageModel.IDEOGRAM_V2: "/v1/ai/text-to-image/ideogram-v2", } IMAGE_STATUS_ENDPOINTS: dict[ImageModel, str] = { ImageModel.MYSTIC: "/v1/ai/mystic/{task_id}", ImageModel.FLUX_KONTEXT_PRO: "/v1/ai/text-to-image/flux-kontext-pro/{task_id}", ImageModel.FLUX_2_PRO: "/v1/ai/text-to-image/flux-2-pro/{task_id}", ImageModel.FLUX_2_TURBO: "/v1/ai/text-to-image/flux-2-turbo/{task_id}", ImageModel.FLUX_PRO_1_1: "/v1/ai/text-to-image/flux-pro-v1-1/{task_id}", ImageModel.SEEDREAM_V4: "/v1/ai/text-to-image/seedream-v4/{task_id}", ImageModel.SEEDREAM_V4_5: "/v1/ai/text-to-image/seedream-v4-5/{task_id}", ImageModel.IDEOGRAM_V2: "/v1/ai/text-to-image/ideogram-v2/{task_id}", } VIDEO_POST_ENDPOINTS: dict[VideoModel, str] = { VideoModel.KLING_O1_PRO: "/v1/ai/image-to-video/kling-o1-pro", VideoModel.KLING_O1_STD: "/v1/ai/image-to-video/kling-o1-std", VideoModel.KLING_ELEMENTS_PRO: "/v1/ai/image-to-video/kling-elements-pro", VideoModel.KLING_ELEMENTS_STD: "/v1/ai/image-to-video/kling-elements-std", VideoModel.MINIMAX_HAILUO: "/v1/ai/image-to-video/minimax-hailuo-02-1080p", } VIDEO_STATUS_ENDPOINTS: dict[VideoModel, str] = { VideoModel.KLING_O1_PRO: "/v1/ai/image-to-video/kling-o1/{task_id}", VideoModel.KLING_O1_STD: "/v1/ai/image-to-video/kling-o1/{task_id}", VideoModel.KLING_ELEMENTS_PRO: "/v1/ai/image-to-video/kling-elements-pro/{task_id}", VideoModel.KLING_ELEMENTS_STD: "/v1/ai/image-to-video/kling-elements-std/{task_id}", VideoModel.MINIMAX_HAILUO: "/v1/ai/image-to-video/minimax-hailuo-02-1080p/{task_id}", } # Per-model image input field name VIDEO_IMAGE_FIELDS: dict[VideoModel, str] = { VideoModel.KLING_O1_PRO: "first_frame", VideoModel.KLING_O1_STD: "first_frame", VideoModel.KLING_ELEMENTS_PRO: "images", # expects an array VideoModel.KLING_ELEMENTS_STD: "images", # expects an array VideoModel.MINIMAX_HAILUO: "image", } # Kling Elements requires named aspect ratio slugs VIDEO_SLUG_ASPECT_RATIO_MODELS: set[VideoModel] = { VideoModel.KLING_ELEMENTS_PRO, VideoModel.KLING_ELEMENTS_STD, } _VIDEO_RATIO_TO_SLUG: dict[str, str] = { "1:1": "square_1_1", "16:9": "widescreen_16_9", "9:16": "social_story_9_16", } def normalize_aspect_ratio_video(ratio: str, model: VideoModel) -> str: """Convert a user-facing aspect ratio to the format required by the video model.""" if model not in VIDEO_SLUG_ASPECT_RATIO_MODELS: return ratio return _VIDEO_RATIO_TO_SLUG.get(ratio, ratio) UPSCALE_POST_ENDPOINTS: dict[UpscaleMode, str] = { UpscaleMode.CREATIVE: "/v1/ai/image-upscaler", UpscaleMode.PRECISION: "/v1/ai/image-upscaler-precision", UpscaleMode.PRECISION_V2: "/v1/ai/image-upscaler-precision-v2", } UPSCALE_STATUS_ENDPOINTS: dict[UpscaleMode, str] = { UpscaleMode.CREATIVE: "/v1/ai/image-upscaler/{task_id}", UpscaleMode.PRECISION: "/v1/ai/image-upscaler-precision/{task_id}", UpscaleMode.PRECISION_V2: "/v1/ai/image-upscaler-precision-v2/{task_id}", } VIDEO_UPSCALE_POST_ENDPOINTS: dict[VideoUpscaleMode, str] = { VideoUpscaleMode.STANDARD: "/v1/ai/video-upscaler", VideoUpscaleMode.TURBO: "/v1/ai/video-upscaler/turbo", } VIDEO_UPSCALE_STATUS_ENDPOINT = "/v1/ai/video-upscaler/{task_id}" # --------------------------------------------------------------------------- # Aspect ratio normalization # --------------------------------------------------------------------------- # Models that require named slug aspect ratios instead of "W:H" strings SLUG_ASPECT_RATIO_MODELS: set[ImageModel] = { ImageModel.MYSTIC, ImageModel.FLUX_PRO_1_1, ImageModel.SEEDREAM_V4, ImageModel.SEEDREAM_V4_5, } # User-friendly "W:H" → API slug mapping _RATIO_TO_SLUG: dict[str, str] = { "1:1": "square_1_1", "16:9": "widescreen_16_9", "9:16": "social_story_9_16", "4:3": "classic_4_3", "3:4": "traditional_3_4", "3:2": "standard_3_2", "2:3": "portrait_2_3", "2:1": "horizontal_2_1", "1:2": "vertical_1_2", "4:5": "social_post_4_5", "21:9": "widescreen_16_9", # closest match } def normalize_aspect_ratio(ratio: str, model: ImageModel) -> str: """Convert a user-facing aspect ratio to the format required by the model.""" if model not in SLUG_ASPECT_RATIO_MODELS: return ratio # free-form models accept "1:1" directly slug = _RATIO_TO_SLUG.get(ratio) if slug: return slug # Already a slug (user passed "square_1_1" directly) — pass through return ratio # --------------------------------------------------------------------------- # Response normalization helpers # --------------------------------------------------------------------------- def get_task_id(raw: dict[str, Any]) -> str: """Extract task_id from any response shape.""" data = raw.get("data", raw) task_id = data.get("task_id") or data.get("id") or raw.get("task_id") or raw.get("id") if not task_id: raise ValueError(f"No task_id found in response: {raw}") return str(task_id) def get_status(raw: dict[str, Any]) -> str: """Extract normalized status string from any response shape.""" data = raw.get("data", raw) status = data.get("status") or raw.get("status") or "PENDING" return status.upper().replace(" ", "_") def get_output_urls(raw: dict[str, Any]) -> list[str]: """Extract all output file URLs from a completed task response.""" data = raw.get("data", raw) # Try common key names in order of likelihood for key in ("generated", "output", "outputs", "result", "results", "images", "videos"): items = data.get(key) if items is None: continue if isinstance(items, list) and items: urls = [] for item in items: if isinstance(item, dict): url = item.get("url") or item.get("download_url") or item.get("src") if url: urls.append(url) elif isinstance(item, str): urls.append(item) if urls: return urls elif isinstance(items, dict): url = items.get("url") or items.get("download_url") or items.get("src") if url: return [url] elif isinstance(items, str): return [items] # Fallback: top-level url field if "url" in data: return [data["url"]] return []