import os
import gradio as gr
import torch
from PIL import Image

# Your README says Wan2.2 Animate model repo:
MODEL_ID = "Wan-AI/Wan2.2-Animate-14B-Diffusers"

DEVICE = "cpu"
DTYPE = torch.float32  # CPU safe

pipe = None


def load_pipe_cpu_only():
    """
    Loads the model on CPU.
    NOTE: Animate 14B is too heavy for HF CPU free tier.
    We still try loading only to show correct Space structure.
    """
    global pipe
    if pipe is not None:
        return pipe

    # Important: On CPU this is likely to fail due to RAM/size
    from diffusers import DiffusionPipeline

    pipe = DiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=DTYPE,
    )
    pipe.to(DEVICE)
    return pipe


def animate_cpu_stub(image, prompt, steps, frames, seed):
    """
    CPU Basic cannot realistically generate Animate 14B.
    This function prevents Space crash and gives user correct reason.
    """
    if image is None:
        return None, "❌ Upload an image first."

    # If user still wants to "try", we do a safe minimal attempt
    steps = int(max(1, min(int(steps), 2)))
    frames = int(max(1, min(int(frames), 2)))

    try:
        _ = load_pipe_cpu_only()
    except Exception as e:
        return None, (
            "❌ Model could not load on HF free CPU (2vCPU/16GB).\n\n"
            f"Error: {repr(e)}\n\n"
            "✅ This Space is correct code structure for Animate task,\n"
            "but Animate 14B needs GPU to actually generate videos."
        )

    # Even if it loads (rare), video generation on CPU is impractical
    return None, (
        "⚠️ Model loaded on CPU (very rare), but generation is too slow / unstable.\n"
        "✅ Use GPU to generate.\n"
        "This Space is running correctly on free CPU."
    )


with gr.Blocks() as demo:
    gr.Markdown("# 🎬 Wan2.2 Animate 14B (HF Free CPU Runner)")
    gr.Markdown(
        """
This Space is designed to run **for free on Hugging Face CPU**.

✅ It follows your model README structure (Animate model).  
⚠️ But **Animate 14B cannot generate videos on free CPU** due to model size + compute.
"""
    )

    with gr.Row():
        image = gr.Image(type="pil", label="Input Image (for Animate)")

    prompt = gr.Textbox(
        label="Prompt",
        value="smooth cinematic camera motion, natural movement, high quality"
    )

    with gr.Row():
        steps = gr.Slider(1, 30, value=20, step=1, label="Steps")
        frames = gr.Slider(1, 81, value=49, step=1, label="Frames")

    seed = gr.Number(value=42, label="Seed")

    btn = gr.Button("Run (CPU Test)")
    status = gr.Textbox(label="Status / Output", lines=8)

    btn.click(
        fn=animate_cpu_stub,
        inputs=[image, prompt, steps, frames, seed],
        outputs=[gr.Video(visible=False), status],
    )

demo.queue().launch()