generate_video_client.py

import base64
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, Optional, Tuple

import requests
from PIL import Image


RUNPOD_API_BASE = os.environ.get("RUNPOD_API_BASE", "https://api.runpod.ai/v2")


def _strip_data_url_prefix(data_url: str) -> str:
    if "," in data_url and data_url.strip().lower().startswith("data:"):
        return data_url.split(",", 1)[1]
    return data_url


def _encode_image_to_base64(path: Path) -> str:
    img = Image.open(path).convert("RGB")
    from io import BytesIO

    buf = BytesIO()
    img.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{b64}"


@dataclass
class GenerateVideoClient:
    """
    Simple HTTP client for the LTX-2 RunPod Serverless endpoint.

    It uses the RunPod v2 REST API directly instead of the Python SDK, to keep
    the template dependency-light.
    """

    runpod_endpoint_id: str
    runpod_api_key: str
    api_base: str = RUNPOD_API_BASE

    def _runsync(self, input_payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Call the endpoint using the synchronous runsync API.
        """
        url = f"{self.api_base.rstrip('/')}/{self.runpod_endpoint_id}/runsync"
        headers = {
            "Authorization": f"Bearer {self.runpod_api_key}",
            "Content-Type": "application/json",
        }
        body = {"input": input_payload}
        resp = requests.post(url, json=body, headers=headers, timeout=600)
        resp.raise_for_status()
        data = resp.json()
        # RunPod wraps the handler's return value under "output".
        return data.get("output") or data

    # --------------------------------------------------------------------- #
    # Public methods
    # --------------------------------------------------------------------- #

    def create_video_t2v(
        self,
        prompt: str,
        negative_prompt: Optional[str] = None,
        width: Optional[int] = None,
        height: Optional[int] = None,
        num_frames: Optional[int] = None,
        fps: Optional[float] = None,
        steps: Optional[int] = None,
        seed: Optional[int] = None,
        cfg: Optional[float] = None,
        ckpt_name: Optional[str] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Text-to-video generation (mode='t2v').
        """
        if not prompt:
            raise ValueError("prompt is required")

        payload: Dict[str, Any] = {
            "mode": "t2v",
            "prompt": prompt,
        }
        if negative_prompt:
            payload["negative_prompt"] = negative_prompt
        if width is not None:
            payload["width"] = width
        if height is not None:
            payload["height"] = height
        if num_frames is not None:
            payload["num_frames"] = num_frames
        if fps is not None:
            payload["fps"] = fps
        if steps is not None:
            payload["steps"] = steps
        if seed is not None:
            payload["seed"] = seed
        if cfg is not None:
            payload["cfg"] = cfg
        if ckpt_name is not None:
            payload["ckpt_name"] = ckpt_name
        if extra:
            payload.update(extra)

        return self._runsync(payload)

    def create_video_i2v(
        self,
        image_path: Optional[str] = None,
        image_url: Optional[str] = None,
        image_base64: Optional[str] = None,
        prompt: str = "",
        negative_prompt: Optional[str] = None,
        width: Optional[int] = None,
        height: Optional[int] = None,
        num_frames: Optional[int] = None,
        fps: Optional[float] = None,
        steps: Optional[int] = None,
        seed: Optional[int] = None,
        cfg: Optional[float] = None,
        ckpt_name: Optional[str] = None,
        extra: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Image-to-video generation (mode='i2v').

        Provide exactly one of:
        - image_path
        - image_url
        - image_base64
        """
        if not prompt:
            raise ValueError("prompt is required")

        provided = [v for v in (image_path, image_url, image_base64) if v]
        if len(provided) != 1:
            raise ValueError("Provide exactly one of image_path, image_url, image_base64")

        payload: Dict[str, Any] = {
            "mode": "i2v",
            "prompt": prompt,
        }

        if image_path:
            payload["image_base64"] = _encode_image_to_base64(Path(image_path))
        elif image_url:
            payload["image_url"] = image_url
        elif image_base64:
            payload["image_base64"] = image_base64

        if negative_prompt:
            payload["negative_prompt"] = negative_prompt
        if width is not None:
            payload["width"] = width
        if height is not None:
            payload["height"] = height
        if num_frames is not None:
            payload["num_frames"] = num_frames
        if fps is not None:
            payload["fps"] = fps
        if steps is not None:
            payload["steps"] = steps
        if seed is not None:
            payload["seed"] = seed
        if cfg is not None:
            payload["cfg"] = cfg
        if ckpt_name is not None:
            payload["ckpt_name"] = ckpt_name
        if extra:
            payload.update(extra)

        return self._runsync(payload)

    def batch_process_images(
        self,
        image_folder_path: str,
        output_folder_path: str,
        prompt: str,
        negative_prompt: Optional[str] = None,
        valid_extensions: Tuple[str, ...] = (".jpg", ".jpeg", ".png", ".bmp", ".tiff"),
        **kwargs: Any,
    ) -> Dict[str, Any]:
        """
        Convenience helper to process a folder of images with mode='i2v'.
        """
        src_dir = Path(image_folder_path)
        out_dir = Path(output_folder_path)
        out_dir.mkdir(parents=True, exist_ok=True)

        files: Iterable[Path] = sorted(
            p for p in src_dir.iterdir() if p.is_file() and p.suffix.lower() in valid_extensions
        )

        total = 0
        successful = 0
        errors: Dict[str, str] = {}

        for img_path in files:
            total += 1
            try:
                result = self.create_video_i2v(
                    image_path=str(img_path),
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    **kwargs,
                )
                if "video" not in result:
                    errors[img_path.name] = f"no video field in result: {result}"
                    continue

                out_file = out_dir / f"{img_path.stem}.mp4"
                self.save_video_result(result, str(out_file))
                successful += 1
            except Exception as exc:  # noqa: BLE001
                errors[img_path.name] = str(exc)

        return {
            "total_files": total,
            "successful": successful,
            "errors": errors,
        }

    @staticmethod
    def save_video_result(result: Dict[str, Any], output_path: str) -> None:
        """
        Save the video Base64 from a handler result to disk.

        Accepts either:
        - {"video": "data:video/mp4;base64,..."}
        - {"output": {"video": "data:video/mp4;base64,..."}}
        """
        video_str = result.get("video")
        if video_str is None and isinstance(result.get("output"), dict):
            video_str = result["output"].get("video")
        if not video_str:
            raise ValueError("Result does not contain a 'video' field.")

        raw_b64 = _strip_data_url_prefix(video_str)
        data = base64.b64decode(raw_b64)
        out_path = Path(output_path)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_bytes(data)