"""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)