179 lines
6.4 KiB
Python
179 lines
6.4 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"
|
||
|
|
WAN_2_5 = "wan-2.5"
|
||
|
|
RUNWAY_GEN4 = "runway-gen4"
|
||
|
|
|
||
|
|
|
||
|
|
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",
|
||
|
|
VideoModel.WAN_2_5: "/v1/ai/image-to-video/wan-2-5",
|
||
|
|
VideoModel.RUNWAY_GEN4: "/v1/ai/image-to-video/runway-gen4",
|
||
|
|
}
|
||
|
|
|
||
|
|
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}",
|
||
|
|
VideoModel.WAN_2_5: "/v1/ai/image-to-video/wan-2-5/{task_id}",
|
||
|
|
VideoModel.RUNWAY_GEN4: "/v1/ai/image-to-video/runway-gen4/{task_id}",
|
||
|
|
}
|
||
|
|
|
||
|
|
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}"
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# 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 []
|