app.py

"""LDR — Extrapolative Video World Models via Latent Dynamics Reasoning.

Upload three conditioning frames, pick a task type, and the model rolls out
~29 future frames of physical motion (uniform, parabola, collision, looming,
or bouncing) and returns them as a video.
"""
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST be before torch
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gradio as gr
import tempfile
import imageio
import imageio.v3 as iio
from pathlib import Path
from huggingface_hub import hf_hub_download

from ldr import build_ldr

# ---------------------------------------------------------------------------
# Model loading (module scope — ZeroGPU intercepts .to("cuda"))
# ---------------------------------------------------------------------------

MODEL_ID = "haodongli/LDR"

# The joint-5task checkpoint handles all five tasks in one model.
# Single-task checkpoints are smaller but only work for one task.
# We load both 256x256 single-task and joint so users can pick.
CKPT_CONFIGS = {
    "uniform":   ("256x256/single_task/uniform.pt",   256),
    "parabola":  ("256x256/single_task/parabola.pt",  256),
    "collision": ("256x256/single_task/collision.pt", 256),
    "looming":   ("256x256/single_task/looming.pt",   256),
    "bouncing":  ("256x256/single_task/bouncing.pt",  256),
    "joint":     ("256x256/joint_task/joint_5task.pt", 256),
}

OURS_ARCH = dict(num_pred=29, width=256, accel_scale=0.5, n_kp=16,
                 warp_flow_res=64, kappa_init=0.15)


def _load_checkpoint(ckpt_path, img_size):
    """Load a .pt checkpoint into an LDR model (mirrors eval.build_model)."""
    ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    arch = dict(OURS_ARCH)
    cargs = ck.get("args") if isinstance(ck, dict) else None
    if cargs:
        arch.update(width=cargs.get("accel_width", 256),
                    n_kp=cargs.get("n_kp", 16),
                    accel_scale=cargs.get("accel_scale", 0.5),
                    kappa_init=cargs.get("kappa_init", cargs.get("accel_init_damp", 0.15)))
    model = build_ldr(**arch).to("cuda").eval()
    sd = ck["model"] if (isinstance(ck, dict) and "model" in ck) else ck
    sd = {k[7:] if k.startswith("module.") else k: v for k, v in sd.items()}
    model.load_state_dict(sd)
    return model


# Pre-download all checkpoints at module scope so they are on disk for the
# worker to stream.
_ckpt_cache = {}  # task_name -> (model, img_size)
for _task, (_rel, _img) in CKPT_CONFIGS.items():
    _path = hf_hub_download(MODEL_ID, _rel, repo_type="model")
    _model = _load_checkpoint(_path, _img)
    _ckpt_cache[_task] = (_model, _img)


def _get_model(task):
    return _ckpt_cache[task]


# ---------------------------------------------------------------------------
# Frame utilities (ported from infer.py)
# ---------------------------------------------------------------------------

def _resize(ft, fr, img_size):
    """Resize a frame tensor (N, 3, H, W) and its numpy companion to img_size."""
    if ft.shape[-1] != img_size:
        ft = nn.functional.interpolate(ft, size=(img_size, img_size),
                                       mode="bilinear", align_corners=False)
        fr = (((ft + 1) * 127.5).clamp(0, 255).byte().permute(0, 2, 3, 1).numpy())
    return ft, fr


def _load_frames(images, img_size):
    """Load a list of PIL images → (float_tensor [-1,1], numpy uint8 frames)."""
    fr = np.stack([np.asarray(img.convert("RGB")) for img in images], 0)
    ft = (torch.from_numpy(fr.astype(np.float32)).permute(0, 3, 1, 2) / 127.5 - 1.0)
    return _resize(ft, fr, img_size)


def _gen(model, frames_t, frames_np, nc):
    """Run the model and concatenate conditioning frames with predictions."""
    cond_t = frames_t[:nc].to("cuda")
    with torch.no_grad():
        pred = model(cond_t.unsqueeze(0), nc, full=False,
                      cond_img=cond_t[nc - 1:nc])
        pf = ((pred.squeeze(0).clamp(-1, 1) + 1) * 127.5).byte().cpu() \
              .permute(0, 2, 3, 1).numpy()
    return np.concatenate([frames_np[:nc], pf], axis=0)


# ---------------------------------------------------------------------------
# Inference function
# ---------------------------------------------------------------------------

@spaces.GPU(duration=30)
def generate(frame1, frame2, frame3, task="uniform",
             fps=8, progress=gr.Progress(track_tqdm=True)):
    """Roll out future video frames from three conditioning images.

    Args:
        frame1: First conditioning frame (earliest in the sequence).
        frame2: Second conditioning frame.
        frame3: Third conditioning frame (latest in the sequence).
        task: Physics task type — uniform, parabola, collision, looming, or bouncing.
             Use "joint" for the model trained on all five tasks.
        fps: Output video frames per second.
    Returns:
        Path to the generated MP4 video file.
    """
    model, img_size = _get_model(task)
    images = [frame1, frame2, frame3]
    ft, fr = _load_frames(images, img_size)
    all_frames = _gen(model, ft, fr, nc=3)

    # Write to a unique temp file (concurrency-safe)
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False, dir="/tmp")
    tmp.close()
    imageio.mimwrite(tmp.name, list(all_frames), fps=fps, codec="libx264",
                     quality=8, macro_block_size=1)
    return tmp.name


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown("""
    # Latent Dynamics Reasoning (LDR)
    Upload three conditioning frames of a physical scene and the model
    extrapolates ~29 future frames as video. Pick the physics task type that
    matches your input (uniform motion, parabola, collision, looming, or
    bouncing), or use the **joint** model trained on all five.

    [Model](https://huggingface.co/haodongli/LDR) ·
    [Code](https://github.com/Lat-Dyn-Reason/Lat-Dyn-Reason) ·
    [Project Page](https://lat-dyn-reason.github.io/)
    """)

    with gr.Row():
        with gr.Column(scale=3):
            frame1 = gr.Image(label="Frame 1 (t=0)", type="pil",
                              height=200)
            frame2 = gr.Image(label="Frame 2 (t=1)", type="pil",
                              height=200)
            frame3 = gr.Image(label="Frame 3 (t=2)", type="pil",
                              height=200)
        with gr.Column(scale=2):
            task = gr.Dropdown(
                choices=["uniform", "parabola", "collision", "looming",
                         "bouncing", "joint"],
                value="uniform",
                label="Task / Checkpoint",
                info="Match the physics type of your scene. "
                     "'joint' handles all five but slightly lower quality.")
            fps = gr.Slider(4, 16, value=8, step=1, label="Output FPS")
            run = gr.Button("Generate Video", variant="primary")

    output = gr.Video(label="Generated Video (conditioning + prediction)")

    with gr.Accordion("Examples", open=True):
        gr.Examples(
            examples=[
                # [frame1, frame2, frame3, task, fps]
                ["examples/uniform/00.png", "examples/uniform/01.png",
                 "examples/uniform/02.png", "uniform", 8],
                ["examples/parabola/00.png", "examples/parabola/01.png",
                 "examples/parabola/02.png", "parabola", 8],
                ["examples/collision/00.png", "examples/collision/01.png",
                 "examples/collision/02.png", "collision", 8],
                ["examples/looming/00.png", "examples/looming/01.png",
                 "examples/looming/02.png", "looming", 8],
                ["examples/bouncing/00.png", "examples/bouncing/01.png",
                 "examples/bouncing/02.png", "bouncing", 8],
                ["examples/uniform_pikachu/00.png",
                 "examples/uniform_pikachu/01.png",
                 "examples/uniform_pikachu/02.png", "uniform", 8],
                ["examples/uniform_soccer/00.png",
                 "examples/uniform_soccer/01.png",
                 "examples/uniform_soccer/02.png", "uniform", 8],
            ],
            inputs=[frame1, frame2, frame3, task, fps],
            outputs=output,
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        fn=generate,
        inputs=[frame1, frame2, frame3, task, fps],
        outputs=output,
        api_name="generate",
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)