studio/backends/hf_space.py

"""HF Spaces backend adapter — extended with LTX-Video I2V endpoint."""
from __future__ import annotations
from dataclasses import dataclass
import os
import tempfile
from pathlib import Path
from typing import Optional, Sequence

import numpy as np
from PIL import Image as PILImage

from pixel_cursor import PixelCursor, Image, IdentityLock, FrameStack, ops
from pixel_cursor.artifact import _new_framestack
from .animatediff import MOTION_LORA_MAP


INSTALL_HINT = (
    "HF Space backend requires gradio_client. Install with:\n"
    "  pip install -e '.[hf]'   (or: pip install gradio_client imageio imageio-ffmpeg)"
)


def _resolve_hf_token(explicit: Optional[str] = None) -> Optional[str]:
    if explicit:
        return explicit
    if os.environ.get("HF_TOKEN"):
        return os.environ["HF_TOKEN"]
    try:
        from huggingface_hub import HfFolder
        return HfFolder.get_token()
    except ImportError:
        return None


@dataclass
class HFSpaceAdapter:
    """Adapter that routes write_motion through a deployed HF Space."""

    space_id: str
    api_name: str = "/infer"
    hf_token: Optional[str] = None
    SOURCE_TAG: str = "hf_space"

    def register(self) -> None:
        ops.register_backend("bind_motion_exemplar", "hf_space", self.bind_motion_exemplar)
        ops.register_backend("write_motion", "hf_space", self.write_motion)

    def bind_motion_exemplar(
        self, cursor: PixelCursor, exemplar: Sequence[str]
    ) -> PixelCursor:
        if len(exemplar) != 1:
            raise ValueError(
                "hf_space.bind_motion_exemplar expects exactly one preset name."
            )
        preset = exemplar[0]
        if preset not in MOTION_LORA_MAP:
            valid = sorted(MOTION_LORA_MAP)
            raise ValueError(f"Unknown preset {preset!r}. Valid: {valid}")
        lock = IdentityLock(
            embedding=np.zeros(0, dtype=np.float32),
            source=f"{self.SOURCE_TAG}:{preset}:{MOTION_LORA_MAP[preset]}",
        )
        return cursor.lock_identity(lock)

    def dry_run(self, cursor: PixelCursor, motion_spec: dict | None = None) -> dict:
        if cursor.identity_lock is None or not cursor.identity_lock.source.startswith(
            f"{self.SOURCE_TAG}:"
        ):
            raise ValueError("dry_run requires bind_motion_exemplar(..., backend='hf_space') first")
        if not isinstance(cursor.artifact, Image):
            raise TypeError("HF Space dry_run expects an Image artifact")

        _, preset, lora_id = cursor.identity_lock.source.split(":", 2)
        ms = motion_spec or {}
        token_present = _resolve_hf_token() is not None
        return {
            "backend": "hf_space",
            "space_id": self.space_id,
            "api_name": self.api_name,
            "preset": preset,
            "motion_lora_id": lora_id,
            "input_image_shape": tuple(cursor.artifact.pixels.shape),
            "num_frames": ms.get("num_frames", 16),
            "expected_output_shape": (ms.get("num_frames", 16), *cursor.artifact.pixels.shape),
            "hf_token_resolved": token_present,
            "token_source": (
                "explicit" if self.hf_token
                else "env:HF_TOKEN" if os.environ.get("HF_TOKEN")
                else "~/.cache/huggingface/token" if token_present
                else "NONE"
            ),
            "gradio_client_installed": _gradio_client_available(),
        }

    def write_motion(self, cursor: PixelCursor, motion_spec: dict | None = None) -> FrameStack:
        if cursor.identity_lock is None or not cursor.identity_lock.source.startswith(
            f"{self.SOURCE_TAG}:"
        ):
            raise ValueError(
                "write_motion(hf_space) requires bind_motion_exemplar(..., backend='hf_space') first"
            )
        if not isinstance(cursor.artifact, Image):
            raise TypeError("HF Space write_motion expects an Image artifact")

        try:
            from gradio_client import Client, handle_file
        except ImportError as e:
            raise ImportError(INSTALL_HINT) from e

        token = _resolve_hf_token(self.hf_token)
        client = Client(self.space_id, token=token)

        _, preset, _ = cursor.identity_lock.source.split(":", 2)
        ms = motion_spec or {}

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            PILImage.fromarray(cursor.artifact.pixels).save(tmp.name)
            video_path = client.predict(
                handle_file(tmp.name),
                preset,
                ms.get("num_frames", 16),
                ms.get("num_inference_steps", 25),
                ms.get("guidance_scale", 7.5),
                ms.get("prompt", "high quality, detailed"),
                ms.get("negative_prompt", "bad quality, blurry"),
                ms.get("seed", 42),
                api_name=self.api_name,
            )

        return _video_path_to_framestack(video_path, preset)


def generate_headlocked_via_space(
    image_path: str,
    *,
    space_id: str,
    prompt: str,
    negative_prompt: str = (
        "head movement, swaying, bobbing, nodding, camera shake, "
        "zoom, pan, jump cut, cartoon, deformed, blurry"
    ),
    height: int = 576,
    width: int = 320,
    num_frames: int = 121,
    num_inference_steps: int = 25,
    guidance_scale: float = 3.0,
    seed: int = 42,
    hf_token: Optional[str] = None,
    api_name: str = "/infer_ltx_i2v",
) -> str:
    """Call the Space's LTX I2V endpoint and return the path to the generated mp4.

    Designed for hologram head-locked clip generation. Default params:
      - 576×320 portrait (9:16-ish, both div-32, confirmed working on A10G)
      - 121 frames = 5.04s @ 24fps (8*15+1)
      - guidance_scale=3.0 (LTX-Video optimal range is 2-4)

    The Space enforces div-32 and 8k+1 constraints internally — caller values
    are rounded up, not rejected.

    Example:
        path = generate_headlocked_via_space(
            "/path/to/chancellor-li-hq-smoothLIGHT-2144x3840.jpg",
            space_id="AlterProgramming/venture-studio",
            prompt="East-Asian man, 30s, dark navy suit ... head absolutely still ...",
        )
        # path is a local mp4 file → copy to v2-compatible/ and run motion_grammar
    """
    try:
        from gradio_client import Client, handle_file
    except ImportError as e:
        raise ImportError(INSTALL_HINT) from e

    token = _resolve_hf_token(hf_token)
    client = Client(space_id, token=token)

    result = client.predict(
        handle_file(image_path),
        prompt,
        negative_prompt,
        float(height),
        float(width),
        float(num_frames),
        float(num_inference_steps),
        float(guidance_scale),
        float(seed),
        api_name=api_name,
    )
    return result if isinstance(result, str) else result[0]


def generate_sprite_via_space(
    prompt: str,
    *,
    space_id: str,
    negative_prompt: str = "",
    num_inference_steps: int = 25,
    guidance_scale: float = 7.5,
    height: int = 512,
    width: int = 512,
    seed: int = 0,
    lora_weight: float = 0.9,
    hf_token: Optional[str] = None,
    api_name: str = "/infer_txt2img",
) -> PILImage.Image:
    """Call the Space's txt2img endpoint and return the generated sprite."""
    try:
        from gradio_client import Client
    except ImportError as e:
        raise ImportError(INSTALL_HINT) from e

    token = _resolve_hf_token(hf_token)
    client = Client(space_id, token=token)
    png_path = client.predict(
        prompt,
        negative_prompt,
        int(num_inference_steps),
        float(guidance_scale),
        int(height),
        int(width),
        int(seed),
        float(lora_weight),
        api_name=api_name,
    )
    return PILImage.open(png_path).convert("RGB")


def _video_path_to_framestack(video_path: str | Path, preset: str) -> FrameStack:
    try:
        import imageio.v3 as iio
    except ImportError as e:
        raise ImportError(
            "Loading the Space's video response requires imageio. Install with:\n"
            "  pip install -e '.[hf]'"
        ) from e
    frames = iio.imread(str(video_path))
    if frames.ndim != 4:
        raise ValueError(f"unexpected video shape from Space: {frames.shape}")
    return _new_framestack(frames.astype(np.uint8), fps=8, name=f"hf_space:{preset}")


def _gradio_client_available() -> bool:
    try:
        import gradio_client  # noqa: F401
        return True
    except ImportError:
        return False