app.py

import torch
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.utils import export_to_video
import gradio as gr
import tempfile
import os
import spaces
from huggingface_hub import hf_hub_download

# --- Global Model Loading (runs once when the script starts) ---
MODEL_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers"


print(f"Loading VAE for {MODEL_ID}...")
# Use float32 for VAE for stability, as recommended in some cases
vae = AutoencoderKLWan.from_pretrained(
    MODEL_ID,
    subfolder="vae",
    torch_dtype=torch.float32
)
print(f"Loading Pipeline {MODEL_ID}...")
# Use bfloat16 for the main pipeline for memory efficiency and speed
pipe = WanPipeline.from_pretrained(
    MODEL_ID,
    vae=vae,
    torch_dtype=torch.bfloat16
)
print("Moving pipeline to CUDA...")
pipe.to("cuda")
causvid_path = hf_hub_download(repo_id="Kijai/WanVideo_comfy", filename="Wan21_CausVid_14B_T2V_lora_rank32.safetensors")
pipe.load_lora_weights(causvid_path)

# --- Gradio Interface Function ---
@spaces.GPU
def generate_video(prompt, negative_prompt, height, width, num_frames, guidance_scale, fps):

    print("Starting video generation...")
    print(f"  Prompt: {prompt}")
    print(f"  Negative Prompt: {negative_prompt if negative_prompt else 'None'}")
    print(f"  Height: {height}, Width: {width}")
    print(f"  Num Frames: {num_frames}, FPS: {fps}")
    print(f"  Guidance Scale: {guidance_scale}")

    # Ensure height and width are multiples of 8 (common requirement for VAEs)
    height = (int(height) // 8) * 8
    width = (int(width) // 8) * 8
    num_frames = int(num_frames)
    fps = int(fps)

    with torch.inference_mode(): # Conserve memory
        output_frames_list = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_frames=num_frames,
            guidance_scale=float(guidance_scale),
            # num_inference_steps=25 # Default is 25, can be exposed if needed
        ).frames

    if not output_frames_list or not output_frames_list[0]:
        raise gr.Error("Model returned empty frames. Check parameters or try a different prompt.")

    output_frames = output_frames_list[0] # The actual list of PIL Image frames

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

    export_to_video(output_frames, video_path, fps=fps)
    print(f"Video successfully generated and saved to {video_path}")

    return video_path


# --- Gradio UI Definition ---
default_prompt = "A cat walks on the grass, realistic"
default_negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"

with gr.Blocks() as demo:
    gr.Markdown(f"""
    # Text-to-Video with Wan 2.1 (14B)
    Powered by `diffusers` and `Wan-AI/{MODEL_ID}`.
    Model is loaded into memory when the app starts. This might take a few minutes.
    Ensure you have a GPU with sufficient VRAM (e.g., ~24GB+ for these default settings).
    """)

    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(label="Prompt", value=default_prompt, lines=3)
            negative_prompt_input = gr.Textbox(
                label="Negative Prompt (Optional)",
                value=default_negative_prompt,
                lines=3
            )
            with gr.Row():
                height_input = gr.Slider(minimum=256, maximum=768, step=64, value=480, label="Height (multiple of 8)")
                width_input = gr.Slider(minimum=256, maximum=1024, step=64, value=832, label="Width (multiple of 8)")
            with gr.Row():
                num_frames_input = gr.Slider(minimum=16, maximum=100, step=1, value=25, label="Number of Frames") # Reduced default for faster demo
                fps_input = gr.Slider(minimum=5, maximum=30, step=1, value=15, label="Output FPS")
            guidance_scale_input = gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=5.0, label="Guidance Scale")

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

        with gr.Column(scale=3):
            video_output = gr.Video(label="Generated Video")

    generate_button.click(
        fn=generate_video,
        inputs=[
            prompt_input,
            negative_prompt_input,
            height_input,
            width_input,
            num_frames_input,
            guidance_scale_input,
            fps_input
        ],
        outputs=video_output
    )

    gr.Examples(
        examples=[
            ["A panda eating bamboo in a lush forest, cinematic lighting", default_negative_prompt, 480, 832, 25, 5.0, 15],
            ["A majestic eagle soaring over snowy mountains", default_negative_prompt, 512, 768, 30, 7.0, 12],
            ["Timelapse of a flower blooming, vibrant colors", "static, ugly", 384, 640, 40, 6.0, 20],
            ["Astronaut walking on the moon, Earth in the background, highly detailed", default_negative_prompt, 480, 832, 20, 5.5, 10],
        ],
        inputs=[prompt_input, negative_prompt_input, height_input, width_input, num_frames_input, guidance_scale_input, fps_input],
        outputs=video_output,
        fn=generate_video,
        cache_examples=False # Caching videos can take a lot of space quickly
    )

if __name__ == "__main__":
    # The share=True option will create a public temporary link if you run this on Colab or similar
    demo.queue().launch(share=True, debug=True)