valknar
0de3f7d6bc
Each video model uses a different image input field: - kling-o1-pro/std: first_frame (not image) - kling-elements-pro/std: images (array) - minimax-hailuo: image, duration fixed at "6" Also: - kling-elements requires slug aspect ratios (square_1_1, etc.) - Remove wan-2.5 and runway-gen4 which return 404 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
240 lines
8.3 KiB
Python
240 lines
8.3 KiB
Python
"""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 []
|