api_client.py

"""API client for daVinci-MagiHuman WebUI.

Sends generation requests to the remote Router.
No retry — if the router rejects (503 busy), the error is returned immediately.

Configure via environment variables:
    ROUTER_URL       Router endpoint (e.g. http://your-server:7860)
    ROUTER_TIMEOUT   Request timeout in seconds (default 660)
"""

import base64
import io
import os
import random
import uuid

import requests
from PIL import Image

ROUTER_URL = os.environ.get("ROUTER_URL", "http://localhost:7860").rstrip("/")
ROUTER_TIMEOUT = int(os.environ.get("ROUTER_TIMEOUT", "660"))


def _pil_to_base64(image: Image.Image) -> str:
    """Encode a PIL Image to a base64 string (PNG format)."""
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def generate(
    image: Image.Image,
    video_prompt: str,
    seed: int = -1,
    output_dir: str = "./outputs",
    seconds: int = 5,
) -> dict:
    """
    Send a generation request to the router and download the video.

    No retry. If the router returns 503 (busy), we immediately return the error.

    Returns:
        dict with keys: video_path (local), seed, error
    """
    if seed == -1:
        seed = random.randint(0, 2**31 - 1)

    os.makedirs(output_dir, exist_ok=True)
    image_base64 = _pil_to_base64(image)

    payload = {
        "task": "ti2av",
        "prompt": video_prompt,
        "image_base64": image_base64,
        "seed": seed,
        "output_dir": "/tmp/magihuman_outputs",
        "seconds": seconds,
        "sr_resolution": "540p",
    }

    result = {
        "video_path": "",
        "seed": seed,
        "error": None,
    }

    try:
        resp = requests.post(
            f"{ROUTER_URL}/generate/file",
            json=payload,
            timeout=ROUTER_TIMEOUT,
        )
        resp.raise_for_status()

        # Save mp4 bytes to local file
        local_path = os.path.join(output_dir, f"magihuman_{uuid.uuid4().hex[:8]}.mp4")
        with open(local_path, "wb") as f:
            f.write(resp.content)
        result["video_path"] = local_path

    except requests.HTTPError as e:
        detail = ""
        try:
            detail = e.response.json().get("detail", "")
        except Exception:
            detail = e.response.text[:200] if e.response else ""
        if e.response is not None and e.response.status_code == 503:
            result["error"] = "All servers are busy. Please try again later."
        else:
            result["error"] = f"HTTP {e.response.status_code}: {detail}" if e.response else str(e)
    except requests.ConnectionError:
        result["error"] = "Cannot reach the generation server. Please try again later."
    except requests.Timeout:
        result["error"] = "Request timed out. The server may be overloaded."
    except Exception as e:
        result["error"] = str(e)

    return result