"""
HF ZeroGPU Space: LTX-Video text-to-video — a free-GPU video backend for the
faceless-reels pipeline (docs §9 hybrid).

Exposes a Gradio API function `generate` that returns the MP4 as base64 (so the
pipeline can decode it in src/lib/video.ts without downloading a hosted file).
The @spaces.GPU decorator gives the function on-demand ZeroGPU access.

Deploy: create a Gradio Space, set Hardware = ZeroGPU, add app.py + requirements.txt.
See hf-space/README.md.
"""
import base64
import os
import tempfile

import gradio as gr
import spaces
import torch
from diffusers import LTXPipeline
from diffusers.utils import export_to_video

MODEL_ID = os.environ.get("MODEL_ID", "Lightricks/LTX-Video")

# Loaded on CPU at startup (ZeroGPU only exposes the GPU inside @spaces.GPU).
# The weights download once when the Space boots.
pipe = LTXPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
# The LTX VAE decode in bf16 can emit NaN -> uniform black/gray/white frames.
# Upcast the VAE to float32 to keep the decode numerically stable (the known fix
# for LTX black-screen output). Transformer + text encoder stay bf16.
pipe.vae.to(torch.float32)


def _round32(x) -> int:
    return max(32, (int(x) // 32) * 32)


def _num_frames(duration_sec: float, fps: int) -> int:
    n = int(round(duration_sec * fps))
    return max(9, ((n - 1) // 8) * 8 + 1)  # LTX requires num_frames % 8 == 1


@spaces.GPU(duration=120)
def generate(prompt: str, width: float = 704, height: float = 1280, duration: float = 5) -> str:
    pipe.to("cuda")
    pipe.vae.to(torch.float32)  # re-assert float32 VAE on-device (avoids black frames)
    fps = 24
    result = pipe(
        prompt=prompt,
        negative_prompt="worst quality, inconsistent motion, blurry, jittery, distorted",
        width=_round32(width),
        height=_round32(height),
        num_frames=_num_frames(float(duration), fps),
        num_inference_steps=40,
        guidance_scale=3.0,
        # Timestep-aware VAE decode (LTX 0.9.1+): the documented values that
        # avoid black/NaN frames.
        decode_timestep=0.03,
        decode_noise_scale=0.025,
    )
    frames = result.frames[0]

    path = tempfile.mktemp(suffix=".mp4")
    export_to_video(frames, path, fps=fps)
    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("utf-8")
    os.remove(path)
    return b64


with gr.Blocks() as demo:
    gr.Markdown("# Faceless LTX-Video\nReturns the generated clip as base64.")
    prompt = gr.Textbox(label="prompt")
    width = gr.Number(value=704, label="width")
    height = gr.Number(value=1280, label="height")
    duration = gr.Number(value=5, label="duration (s)")
    out = gr.Textbox(label="video_base64")
    btn = gr.Button("Generate")
    # api_name="generate" fixes the API route the pipeline calls.
    btn.click(generate, [prompt, width, height, duration], out, api_name="generate")

demo.queue().launch()