# app.py — Wan 2.2 Lite Edition (8GB RAM Safe)

import os
import spaces
import torch
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
from diffusers.utils.export_utils import export_to_video
import gradio as gr
import tempfile
import numpy as np
from PIL import Image
import random
import gc
import aoti

from torchao.quantization import quantize_
from torchao.quantization import (
    Float8DynamicActivationFloat8WeightConfig,
    Int8WeightOnlyConfig
)

# =====================================================================
# BASIC CONFIG FOR LOW-RAM SYSTEMS
# =====================================================================
MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
from dotenv import load_dotenv
load_dotenv()

HF_TOKEN = os.environ.get("HF_TOKEN")

MAX_DIM = 480
MIN_DIM = 320
MULTIPLE_OF = 16
MAX_SEED = np.iinfo(np.int32).max

FPS = 12
MIN_FRAMES = 8
MAX_FRAMES = 200     # around 2 seconds max

# =====================================================================
# LOAD PIPELINE — STRIPPED DOWN (NO TRANSFORMER_2)
# =====================================================================
pipe = WanImageToVideoPipeline.from_pretrained(
    MODEL_ID,
    transformer=WanTransformer3DModel.from_pretrained(
        MODEL_ID,
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
        token=HF_TOKEN,
    ),
    torch_dtype=torch.bfloat16,
).to("cuda")

# =====================================================================
# QUANTIZATION FOR 8GB VRAM
# =====================================================================
quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())

# AOTI compile small blocks
aoti.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/Wan2", variant="fp8da")

# =====================================================================
# IMAGE RESIZE — 480p MAX
# =====================================================================
def resize_image(image):
    w, h = image.size
    ratio = w / h

    if w > h:
        new_w = MAX_DIM
        new_h = int(new_w / ratio)
    else:
        new_h = MAX_DIM
        new_w = int(new_h * ratio)

    new_w = max(MIN_DIM, round(new_w / MULTIPLE_OF) * MULTIPLE_OF)
    new_h = max(MIN_DIM, round(new_h / MULTIPLE_OF) * MULTIPLE_OF)

    return image.resize((new_w, new_h), Image.LANCZOS)


# =====================================================================
# FRAME COUNT
# =====================================================================
def get_num_frames(seconds):
    return np.clip(int(seconds * FPS), MIN_FRAMES, MAX_FRAMES)


# =====================================================================
# MAIN GENERATE FUNCTION
# =====================================================================
@spaces.GPU()
def generate_video(
    input_image,
    prompt,
    steps=1,
    negative_prompt="",
    duration_seconds=1.5,
    seed=42,
    randomize_seed=False,
):

    if input_image is None:
        raise gr.Error("Upload an input image.")

    gc.collect()
    torch.cuda.empty_cache()

    num_frames = get_num_frames(duration_seconds)
    seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    img = resize_image(input_image)

    out = pipe(
        image=img,
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_frames=num_frames,
        height=img.height,
        width=img.width,
        guidance_scale=0.5,
        num_inference_steps=int(steps),
        generator=torch.Generator("cuda").manual_seed(seed),
    )

    frames = out.frames[0]

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
        path = tmp.name

    export_to_video(frames, path, fps=FPS)

    gc.collect()
    torch.cuda.empty_cache()

    return path, seed


# =====================================================================
# GRADIO UI — LIGHT VERSION
# =====================================================================
with gr.Blocks() as demo:

    gr.Markdown("## 🚀 Wan 2.2 — Lite Edition (8GB RAM Optimized)")
    gr.Markdown("Runs at max 480p, 2-second videos, FP8 + INT8 optimized.")

    with gr.Row():
        with gr.Column():
            img = gr.Image(label="Input Image", type="pil")
            prompt_box = gr.Textbox("make this image move smoothly")
            duration = gr.Slider(0.8, 2.0, 1.5, label="Duration (seconds)")
            steps = gr.Slider(1, 3, 1, step=1, label="Steps (1 = fastest)")

            seed_box = gr.Slider(0, MAX_SEED, 42, label="Seed")
            rand_seed = gr.Checkbox(True, label="Randomize Seed")

            btn = gr.Button("Generate Video", variant="primary")

        with gr.Column():
            out_video = gr.Video(label="Result")

    btn.click(
        generate_video,
        inputs=[img, prompt_box, steps, None, duration, seed_box, rand_seed],
        outputs=[out_video, seed_box],
    )

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