app.py

import gc
import os
import random
import threading
import uuid
from pathlib import Path

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import AutoencoderKLWan, WanImageToVideoPipeline
from diffusers.utils import export_to_video
from PIL import Image, ImageOps
from transformers import CLIPVisionModel


MODEL_ID = "Wan-AI/Wan2.1-FLF2V-14B-720P-diffusers"
VIDEO_FPS = 16
OUTPUT_DIR = Path("outputs")
OUTPUT_DIR.mkdir(exist_ok=True)

NEGATIVE_PROMPT = (
    "low quality, worst quality, blurry, overexposed, static, distorted, "
    "deformed, disfigured, duplicate, watermark, text, logo, artifacts"
)

RESOLUTIONS = {
    "480p (faster)": 480 * 832,
    "720p (best quality)": 720 * 1280,
}

pipe = None
model_lock = threading.Lock()


def load_pipeline():
    """Load once, on the first request, to keep Space startup responsive."""
    global pipe
    if pipe is not None:
        return pipe

    with model_lock:
        if pipe is not None:
            return pipe

        if not torch.cuda.is_available():
            raise gr.Error(
                "A CUDA GPU is required. In the Space settings, select an A100 80GB "
                "or another GPU with enough memory."
            )

        image_encoder = CLIPVisionModel.from_pretrained(
            MODEL_ID,
            subfolder="image_encoder",
            torch_dtype=torch.float32,
        )
        vae = AutoencoderKLWan.from_pretrained(
            MODEL_ID,
            subfolder="vae",
            torch_dtype=torch.float32,
        )
        loaded_pipe = WanImageToVideoPipeline.from_pretrained(
            MODEL_ID,
            image_encoder=image_encoder,
            vae=vae,
            torch_dtype=torch.bfloat16,
        )
        loaded_pipe.vae.enable_tiling()
        loaded_pipe.to("cuda")
        pipe = loaded_pipe

    return pipe


def prepare_frame(image: Image.Image, max_area: int, size=None):
    if image is None:
        return None, None

    image = ImageOps.exif_transpose(image).convert("RGB")

    if size is None:
        aspect = image.height / image.width
        height = max(128, round(np.sqrt(max_area * aspect) / 16) * 16)
        width = max(128, round(np.sqrt(max_area / aspect) / 16) * 16)
        size = (width, height)

    # Crop only the overflow so both conditioning frames have identical geometry.
    return ImageOps.fit(image, size, method=Image.Resampling.LANCZOS), size


def seconds_to_frames(duration_seconds):
    """Wan accepts 4k+1 frame counts; whole seconds at 16 fps fit exactly."""
    seconds = max(1, min(3, int(duration_seconds)))
    return seconds * VIDEO_FPS + 1


def estimate_gpu_duration(
    _start_image,
    _end_image,
    _prompt,
    _negative_prompt,
    resolution,
    duration_seconds,
    steps,
    _guidance,
    _seed,
):
    """Reserve only the free ZeroGPU time appropriate for this request."""
    seconds = max(1, min(3, int(duration_seconds)))
    resolution_factor = 1.6 if resolution == "720p (best quality)" else 1.0
    estimate = (8 + 7 * seconds) * (int(steps) / 8) * resolution_factor
    return max(12, min(60, int(round(estimate))))


@spaces.GPU(size="xlarge", duration=estimate_gpu_duration)
def generate_video(
    start_image,
    end_image,
    prompt,
    negative_prompt,
    resolution,
    duration_seconds,
    steps,
    guidance,
    seed,
    progress=gr.Progress(track_tqdm=False),
):
    if start_image is None:
        raise gr.Error("Please upload a start image.")
    if not prompt or not prompt.strip():
        raise gr.Error("Please describe the motion or scene in the prompt.")

    progress(0, desc="Loading the video model…")
    pipeline = load_pipeline()

    max_area = RESOLUTIONS[resolution]
    first_frame, target_size = prepare_frame(start_image, max_area)
    has_end_frame = end_image is not None
    if has_end_frame:
        last_frame, _ = prepare_frame(end_image, max_area, target_size)
    else:
        # The FLF2V checkpoint always expects two CLIP frame embeddings. Reusing
        # the first frame keeps the end-image input optional and creates a loop.
        last_frame = first_frame.copy()

    width, height = target_size
    duration_seconds = max(1, min(3, int(duration_seconds)))
    num_frames = seconds_to_frames(duration_seconds)
    actual_seed = random.randint(0, 2**31 - 1) if int(seed) < 0 else int(seed)
    generator = torch.Generator(device="cpu").manual_seed(actual_seed)

    def update_progress(_pipeline, step_index, _timestep, callback_kwargs):
        progress((step_index + 1) / int(steps), desc=f"Generating frame sequence · step {step_index + 1}/{steps}")
        return callback_kwargs

    output_path = OUTPUT_DIR / f"wan_{actual_seed}_{uuid.uuid4().hex[:8]}_{width}x{height}.mp4"

    try:
        with model_lock, torch.inference_mode():
            frames = pipeline(
                image=first_frame,
                last_image=last_frame,
                prompt=prompt.strip(),
                negative_prompt=(negative_prompt or "").strip(),
                height=height,
                width=width,
                num_frames=int(num_frames),
                num_inference_steps=int(steps),
                guidance_scale=float(guidance),
                generator=generator,
                callback_on_step_end=update_progress,
            ).frames[0]
        export_to_video(frames, str(output_path), fps=VIDEO_FPS)
    except torch.cuda.OutOfMemoryError as exc:
        gc.collect()
        torch.cuda.empty_cache()
        raise gr.Error("The GPU ran out of memory. Try 480p, fewer frames, or an A100 80GB GPU.") from exc

    progress(1, desc="Video ready")
    mode = "start → end" if has_end_frame else "loop"
    info = (
        f"Seed **{actual_seed}** · {width}×{height} · "
        f"{duration_seconds}s ({num_frames} frames at {VIDEO_FPS} fps) · {mode} mode"
    )
    return str(output_path), info, actual_seed


# ZeroGPU emulates CUDA during startup, allowing weights to be prepared before
# a real GPU is assigned to a generation request.
if os.getenv("SPACE_ID"):
    load_pipeline()


CSS = """
:root { --ink: #171512; --paper: #f6f2e9; --accent: #ee5b35; }
.gradio-container { max-width: 1180px !important; margin: 0 auto !important; background: var(--paper); }
.hero { padding: 2.25rem 0 1rem; }
.hero h1 { font-size: clamp(2.3rem, 6vw, 5rem); line-height: .92; letter-spacing: -.055em; color: var(--ink); margin: 0; }
.hero p { max-width: 650px; font-size: 1.05rem; color: #5b554c; margin-top: 1.1rem; }
.eyebrow { color: var(--accent); font-weight: 750; letter-spacing: .15em; text-transform: uppercase; font-size: .76rem; }
.frame-card { border: 1px solid #d9d1c3 !important; border-radius: 18px !important; background: rgba(255,255,255,.48) !important; }
.generate-btn { background: var(--accent) !important; color: white !important; border: none !important; font-weight: 750 !important; }
.output-video { border-radius: 18px; overflow: hidden; }
.footer-note { color: #766e62; font-size: .83rem; text-align: center; padding: 1rem; }
"""


with gr.Blocks(css=CSS, title="Between Frames · Image to Video") as demo:
    gr.HTML(
        """
        <section class="hero">
          <div class="eyebrow">Wan 2.1 · First / Last Frame to Video</div>
          <h1>Turn two stills<br>into one moving moment.</h1>
          <p>Choose where the shot begins, optionally choose where it ends, and describe what happens between them. Without an end image, the shot loops back to its first frame.</p>
        </section>
        """
    )

    with gr.Row(equal_height=False):
        with gr.Column(scale=6):
            with gr.Row():
                start_image = gr.Image(
                    type="pil",
                    image_mode="RGB",
                    label="01 · Start image",
                    sources=["upload", "clipboard"],
                    elem_classes="frame-card",
                    height=330,
                )
                end_image = gr.Image(
                    type="pil",
                    image_mode="RGB",
                    label="02 · End image (optional)",
                    sources=["upload", "clipboard"],
                    elem_classes="frame-card",
                    height=330,
                )

            prompt = gr.Textbox(
                label="03 · Describe the movement",
                placeholder="The camera slowly pushes in as wind moves through her hair; cinematic light shifts from dusk to night…",
                lines=4,
            )

            with gr.Accordion("Generation controls", open=False):
                negative_prompt = gr.Textbox(label="Negative prompt", value=NEGATIVE_PROMPT, lines=2)
                with gr.Row():
                    resolution = gr.Radio(list(RESOLUTIONS), value="480p (faster)", label="Resolution")
                    duration_seconds = gr.Slider(
                        1,
                        3,
                        value=1,
                        step=1,
                        label="Video duration (seconds)",
                        info="Longer videos use more of the free daily GPU quota.",
                    )
                with gr.Row():
                    steps = gr.Slider(8, 16, value=8, step=1, label="Inference steps")
                    guidance = gr.Slider(1, 6, value=1.0, step=0.1, label="Prompt guidance")
                    seed = gr.Number(value=-1, precision=0, label="Seed (−1 = random)")

            generate_btn = gr.Button("Generate video", variant="primary", size="lg", elem_classes="generate-btn")

        with gr.Column(scale=5):
            video = gr.Video(label="Generated video", autoplay=True, elem_classes="output-video")
            generation_info = gr.Markdown("Your generation details will appear here.")

    gr.HTML('<div class="footer-note">Large video models need a GPU. 720p generation can take several minutes.</div>')

    inputs = [
        start_image,
        end_image,
        prompt,
        negative_prompt,
        resolution,
        duration_seconds,
        steps,
        guidance,
        seed,
    ]
    generate_btn.click(
        fn=generate_video,
        inputs=inputs,
        outputs=[video, generation_info, seed],
        api_name="generate",
        concurrency_id="gpu_queue",
        concurrency_limit=1,
    )
    prompt.submit(
        fn=generate_video,
        inputs=inputs,
        outputs=[video, generation_info, seed],
        concurrency_id="gpu_queue",
        concurrency_limit=1,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=8).launch()