freepik/freepik_cli/api/videos.py

"""Video generation API methods."""

from __future__ import annotations

from typing import Any, Optional, Tuple

from freepik_cli.api.client import FreepikClient
from freepik_cli.api.models import (
    VIDEO_POST_ENDPOINTS,
    VIDEO_STATUS_ENDPOINTS,
    VIDEO_IMAGE_FIELDS,
    VIDEO_ASPECT_RATIO_MODELS,
    VideoModel,
    get_output_urls,
    get_status,
    get_task_id,
    normalize_aspect_ratio_video,
)


class VideoAPI:
    def __init__(self, client: FreepikClient) -> None:
        self._client = client

    def generate(
        self,
        model: VideoModel,
        image_b64: str,
        prompt: Optional[str] = None,
        duration: int = 5,
        aspect_ratio: str = "16:9",
        seed: Optional[int] = None,
    ) -> str:
        """Submit an image-to-video task. Returns task_id."""
        # Video API expects raw base64, not a data URI
        if ";base64," in image_b64:
            image_b64 = image_b64.split(";base64,", 1)[1]

        image_field = VIDEO_IMAGE_FIELDS[model]

        # kling-elements uses an array; all others use a scalar
        if image_field == "images":
            payload: dict[str, Any] = {"images": [image_b64]}
        else:
            payload = {image_field: image_b64}

        if prompt:
            payload["prompt"] = prompt

        # minimax only supports duration=6; clamp silently
        effective_duration = duration
        if model == VideoModel.MINIMAX_HAILUO:
            effective_duration = 6
        payload["duration"] = str(effective_duration)

        if aspect_ratio and model in VIDEO_ASPECT_RATIO_MODELS:
            payload["aspect_ratio"] = normalize_aspect_ratio_video(aspect_ratio, model)

        if seed is not None:
            payload["seed"] = seed

        endpoint = VIDEO_POST_ENDPOINTS[model]
        raw = self._client.post(endpoint, json=payload)
        return get_task_id(raw)

    def get_status(self, model: VideoModel, task_id: str) -> Tuple[str, dict[str, Any]]:
        """Poll status. Returns (status_str, raw_response)."""
        endpoint = VIDEO_STATUS_ENDPOINTS[model].format(task_id=task_id)
        raw = self._client.get(endpoint)
        return get_status(raw), raw

    def get_output_urls(self, raw: dict[str, Any]) -> list[str]:
        return get_output_urls(raw)