import spaces
import logging
import secrets
import tempfile
import time

import gradio as gr
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

logging.basicConfig(level=logging.INFO)
MODEL = "hunyuanvideo-community/HunyuanVideo"
SIZES = {"Landscape": (512, 320), "Portrait": (320, 512), "Square": (384, 384)}
EXAMPLES = [
    ["A cat walks on the grass, realistic style."],
    ["A fluffy teddy bear sits on a bed of soft pillows surrounded by children's toys."],
    ["A cinematic tracking shot of a sailboat crossing a calm lake at sunrise, soft golden light and gentle ripples."],
]

transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    MODEL, subfolder="transformer", torch_dtype=torch.bfloat16
)
pipe = HunyuanVideoPipeline.from_pretrained(
    MODEL, transformer=transformer, torch_dtype=torch.float16
).to("cuda")
pipe.vae.enable_tiling()
logging.info("HunyuanVideo model loaded successfully")


@spaces.GPU(duration=180, size="xlarge")
def generate(
    prompt: str,
    aspect: str = "Landscape",
    frames: int = 33,
    steps: int = 30,
    seed: int = -1,
    guidance: float = 6.0,
    progress=gr.Progress(track_tqdm=True),
) -> tuple[str, int, str]:
    """Generate a silent HunyuanVideo MP4 from a scene description and return its seed and runtime."""
    prompt = prompt.strip()
    if not prompt:
        raise gr.Error("Describe the video you want to create first.")
    if len(prompt) > 1500:
        raise gr.Error("Keep your scene description under 1,500 characters.")
    if aspect not in SIZES or int(frames) not in (17, 33, 49):
        raise gr.Error("Choose one of the supported aspect ratios and clip lengths.")
    if not 10 <= int(steps) <= 40 or not 1 <= float(guidance) <= 10:
        raise gr.Error("Steps must be 10–40 and guidance must be 1–10.")
    if int(seed) < -1 or int(seed) > 2147483647:
        raise gr.Error("Seed must be -1 (random) or between 0 and 2147483647.")
    actual_seed = secrets.randbelow(2147483648) if int(seed) == -1 else int(seed)
    width, height = SIZES[aspect]
    started = time.perf_counter()
    with torch.inference_mode():
        video = pipe(
            prompt=prompt, width=width, height=height,
            num_frames=int(frames), num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            generator=torch.Generator(device="cuda").manual_seed(actual_seed),
        ).frames[0]
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as handle:
        output = handle.name
    export_to_video(video, output, fps=16)
    elapsed = time.perf_counter() - started
    logging.info("Generated %sx%s/%s frames in %.1fs", width, height, frames, elapsed)
    return output, actual_seed, f"Complete · {width} × {height} · {int(frames) / 16:.2f}s · generated in {elapsed:.1f}s"


with gr.Blocks(title="HunyuanVideo Studio", delete_cache=(3600, 3600)) as demo:
    gr.Markdown("# HunyuanVideo Studio\nTurn a scene in your imagination into a short moving picture.")
    gr.Markdown("Powered by Tencent's original **HunyuanVideo** model. Describe the subject, movement, lighting, and camera angle for the best results.")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Your scene", lines=5, max_length=1500,
                placeholder="A cat walks through sunlit grass. The camera follows at ground level…")
            with gr.Row():
                aspect = gr.Dropdown(list(SIZES), value="Landscape", label="Aspect ratio")
                frames = gr.Dropdown([("Quick · 1.06s", 17), ("Standard · 2.06s", 33), ("Longer · 3.06s", 49)], value=33, label="Clip length")
            with gr.Accordion("Generation settings", open=False):
                steps = gr.Slider(10, 40, value=30, step=1, label="Denoising steps")
                seed = gr.Number(value=-1, precision=0, minimum=-1, maximum=2147483647, label="Seed (-1 for a new result)")
                guidance = gr.Slider(1, 10, value=6, step=0.5, label="Prompt guidance")
            button = gr.Button("Generate video", variant="primary")
        with gr.Column():
            output = gr.Video(label="Your video", interactive=False, buttons=["download"])
            used_seed = gr.Number(label="Seed used", interactive=False)
            status = gr.Markdown("Your generated video will appear here.")
    button.click(generate, [prompt, aspect, frames, steps, seed, guidance],
                 [output, used_seed, status], api_name="generate", concurrency_limit=1)
    gr.Examples(EXAMPLES, inputs=[prompt], outputs=[output, used_seed, status],
                fn=generate, cache_examples=True, cache_mode="lazy", label="Try a scene")
    gr.Markdown("Short, silent previews at 16 fps. Generation uses your Hugging Face ZeroGPU quota; queues and available quota can vary.\n\n[Original model](https://huggingface.co/tencent/HunyuanVideo) · [Diffusers weights](https://huggingface.co/hunyuanvideo-community/HunyuanVideo) · [Model license](https://huggingface.co/tencent/HunyuanVideo/blob/main/LICENSE)")

if __name__ == "__main__":
    demo.queue(max_size=8).launch(mcp_server=True, theme=gr.themes.Soft(primary_hue="indigo"))