"""
07 · ZeroGPU Animator   (a gr.Workflow `fn` node running real inference on ZeroGPU)
===================================================================================

    [image] ─┐
             ├─▶ (fn) animate  →  LTX-Video 0.9.7-distilled on ZeroGPU  ─▶  🎬 Video
    [prompt] ┘

The whole point of this demo: an `fn` node is just Python, so it can hold a diffusers
pipeline and burn a real GPU. The pipeline is built once at import (ZeroGPU lets you
place weights on `cuda` at module scope) and `animate` is decorated with `@spaces.GPU`,
which is what actually leases a GPU slice for the duration of the call.

Two ZeroGPU rules this demo exists to illustrate:

  1. A ZeroGPU Space MUST expose at least one `@spaces.GPU` function, or it never boots —
     it dies with "No @spaces.GPU function detected during startup".
  2. Only decorate functions that really touch the GPU. Wrapping a function that merely
     calls some *other* Space over the network spends the visitor's quota on a no-op, and
     the worker re-raises only the exception's class name, so real errors arrive as a bare
     `'AppError'`.

LTX-Video 0.9.7-distilled is the distilled checkpoint: 7 steps at guidance 1.0, no CFG,
no upsampler pass — fast enough to fit comfortably in one ZeroGPU lease.
"""

import base64
import os
import tempfile

import gradio as gr
import spaces
import torch
from diffusers import LTXConditionPipeline
from diffusers.utils import export_to_video, load_image

MODEL = "Lightricks/LTX-Video-0.9.7-distilled"

HEIGHT, WIDTH = 480, 832     # both must be divisible by 32 (the VAE's spatial ratio)
NUM_FRAMES = 97              # must be 8k + 1
FPS = 24
STEPS = 7                    # distilled: 7 steps, guidance 1.0, no CFG
NEGATIVE = "worst quality, inconsistent motion, blurry, jittery, distorted"
SEED = 42

# Built once at import. ZeroGPU allows CUDA placement at module scope; the *lease* is
# taken by @spaces.GPU below, not by this.
pipe = LTXConditionPipeline.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
pipe.vae.enable_tiling()
pipe.to("cuda")


def _img_src(image):
    """The canvas hands an fn node either a path/URL string or a {path,url} dict."""
    if isinstance(image, dict):
        return image.get("path") or image.get("url")
    return image


@spaces.GPU(duration=120)
def animate(image, prompt: str) -> dict:
    """Animate the conditioning image with LTX-Video, on this Space's own ZeroGPU slice."""
    src = _img_src(image)
    if not src:
        raise gr.Error("Connect an image to the animate node.")

    frames = pipe(
        image=load_image(src).convert("RGB"),
        frame_index=0,
        strength=1.0,
        prompt=(prompt or "gentle natural motion, cinematic"),
        negative_prompt=NEGATIVE,
        height=HEIGHT,
        width=WIDTH,
        num_frames=NUM_FRAMES,
        frame_rate=FPS,
        num_inference_steps=STEPS,
        guidance_scale=1.0,
        decode_timestep=0.05,
        decode_noise_scale=0.025,
        generator=torch.Generator("cuda").manual_seed(SEED),
    ).frames[0]

    path = os.path.join(tempfile.mkdtemp(), "animated.mp4")
    export_to_video(frames, path, fps=FPS)

    # {path} for the REST API, {url: data-uri} so the canvas video player can load it.
    data = open(path, "rb").read()
    return {"path": path, "url": "data:video/mp4;base64," + base64.b64encode(data).decode()}


BIND = {"animate": animate}

WORKFLOW = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflow.json")
demo = gr.Workflow(WORKFLOW, bind=BIND)

if __name__ == "__main__":
    demo.launch()