import os
import uuid
import torch
import spaces

from diffusers import (
    WanPipeline,
    AutoencoderKLWan,
)

from diffusers.utils import export_to_video

MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"

pipe = None


def load_pipeline():

    global pipe

    if pipe is not None:
        return pipe

    vae = AutoencoderKLWan.from_pretrained(
        MODEL_ID,
        subfolder="vae",
        torch_dtype=torch.float32,
    )

    pipe = WanPipeline.from_pretrained(
        MODEL_ID,
        vae=vae,
        torch_dtype=torch.bfloat16,
    )

    pipe.enable_model_cpu_offload()

    return pipe


@spaces.GPU(duration=240)
def generate_video(
    prompt,
    negative_prompt,
    steps,
    guidance,
    frames,
    seed,
):

    if not prompt.strip():
        return None, "Please enter a prompt."

    pipe = load_pipeline()

    if seed == -1:
        seed = torch.seed()

    generator = torch.Generator("cpu").manual_seed(int(seed))

    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=480,
        width=832,
        num_frames=int(frames),
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        generator=generator,
    )

    video = result.frames[0]

    os.makedirs("outputs", exist_ok=True)

    filename = f"outputs/{uuid.uuid4().hex}.mp4"

    export_to_video(
        video,
        filename,
        fps=15,
    )

    return (
        filename,
        f"Finished ✓  Seed: {seed}",
    )