Document download + run commands; ship live_infer.py

README.md

@@ -1,49 +1,65 @@
----
-license: apache-2.0
-library_name: pytorch
-pipeline_tag: other
-tags:
-  - world-model
-  - diffusion
-  - pong
-  - video
-  - causal
-  - game
----
-
-# Diffusion Pong
-
-A small action-conditioned world model that learned to keep a Pong match going. You hold **W** / **S**, it invents the next frame.
-
-![demo](demo.gif)
-
-128×128 at 6 FPS with a 12-frame latent history. Codec is a frozen SDXL VAE (`madebyollin/sdxl-vae-fp16-fix`).
-
-## Play it
-
-From the livediffusion source tree (after `pip install -e .`):
-
-```bash
-python scripts/live_infer.py
-```
-
-First run pulls this repo into the HF cache. Needs a CUDA GPU with BF16.
-
-```bash
-# optional knobs
-python scripts/live_infer.py --steps 2 --window-scale 7 --seed 42
-```
-
-Controls: **W** up, **S** down. Click the window once so it has focus. Clicking the canvas drops a yellow ball into the latent history for a few frames if you want to poke it.
-
-## What's in the files
-
-| file | what |
-|---|---|
-| `ema.safetensors` | playable weights |
-| `config.json` | video / model / codec settings used at train time |
-| `demo.gif` | live rollout clip |
-
-Action vector is `[W, S]` as floats in `{0,1}`. Hold both and it treats that as neutral.
-
-If something breaks, it's probably the VAE download or a GPU that can't do BF16.
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: other
+tags:
+  - world-model
+  - diffusion
+  - pong
+  - video
+  - causal
+  - game
+---
+
+# Diffusion Pong
+
+A small action-conditioned world model that learned to keep a Pong match going. You hold **W** / **S**, it invents the next frame.
+
+![demo](demo.gif)
+
+128×128 at 6 FPS with a 12-frame latent history. Codec is a frozen SDXL VAE (`madebyollin/sdxl-vae-fp16-fix`). Needs a CUDA GPU with BF16.
+
+## Play it
+
+Install the livediffusion package from this project's source tree, then grab the play script + weights from the Hub.
+
+```bash
+# from the livediffusion repo root
+pip install -e .
+pip install -U "huggingface_hub[cli]"
+```
+
+Download the script (and optionally the weights into a local folder):
+
+```bash
+# play script
+hf download kerzgrr/diffusionpong live_infer.py --local-dir scripts
+
+# weights + config (optional — the script can also pull these on its own)
+hf download kerzgrr/diffusionpong --local-dir checkpoints/diffusionpong --include "ema.safetensors" --include "config.json"
+```
+
+Run:
+
+```bash
+# downloads into the HF cache automatically if you skip --local-dir
+python scripts/live_infer.py --steps 2
+
+# or point at a folder you already downloaded
+python scripts/live_infer.py --repo kerzgrr/diffusionpong --local-dir checkpoints/diffusionpong --steps 2 --window-scale 7 --seed 42
+```
+
+Controls: **W** up, **S** down. Click the window once so it has focus. Clicking the canvas drops a yellow ball into the latent history for a few frames if you want to poke it.
+
+## What's in the files
+
+| file | what |
+|---|---|
+| `ema.safetensors` | playable weights |
+| `config.json` | video / model / codec settings used at train time |
+| `live_infer.py` | live play script (same as `scripts/live_infer.py`) |
+| `demo.gif` | live rollout clip |
+
+Action vector is `[W, S]` as floats in `{0,1}`. Hold both and it treats that as neutral.
+
+If something breaks, it's probably the VAE download or a GPU that can't do BF16.

live_infer.py

@@ -0,0 +1,303 @@
+#!/usr/bin/env python3
+"""Play Diffusion Pong live. Downloads kerzgrr/diffusionpong on first run.
+
+Requires: a BF16 CUDA GPU, this repo installed (`pip install -e .`), and
+`huggingface_hub` (already a project dependency).
+
+Controls:
+  W / S  — move left paddle up / down
+  click  — inject a yellow ball for a few frames (debug)
+  close window or Ctrl+C to quit
+"""
+
+from __future__ import annotations
+
+import argparse
+import random
+import sys
+import time
+import tkinter as tk
+from collections import deque
+from pathlib import Path
+
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+from PIL import Image, ImageDraw, ImageTk
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+SRC = REPO_ROOT / "src"
+if SRC.is_dir() and str(SRC) not in sys.path:
+    sys.path.insert(0, str(SRC))
+
+from livediffusion.checkpoint import load_model_weights
+from livediffusion.codec import load_codec
+from livediffusion.config import load_config
+from livediffusion.meanflow import sample_next_latent
+from livediffusion.model import CausalLatentVideoDiT
+from livediffusion.simulator import PongSimulation
+
+DEFAULT_REPO = "kerzgrr/diffusionpong"
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Live play Diffusion Pong from the Hub.")
+    parser.add_argument("--repo", default=DEFAULT_REPO, help="Hugging Face model repo id")
+    parser.add_argument("--revision", default="main")
+    parser.add_argument(
+        "--local-dir",
+        default=None,
+        help="Optional folder to download weights into (default: HF cache).",
+    )
+    parser.add_argument("--steps", type=int, default=2)
+    parser.add_argument("--frames", type=int, default=0, help="0 = run until you close the window")
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--window-scale", type=int, default=6)
+    parser.add_argument("--fps-cap", type=float, default=0.0, help="Optional display cap; 0 = uncapped")
+    return parser.parse_args()
+
+
+def _download_checkpoint(repo: str, revision: str, local_dir: str | None) -> Path:
+    if local_dir:
+        root = Path(
+            snapshot_download(
+                repo_id=repo,
+                revision=revision,
+                local_dir=local_dir,
+                allow_patterns=["ema.safetensors", "model.safetensors", "config.json"],
+            )
+        )
+        return root
+    config_path = Path(
+        hf_hub_download(repo_id=repo, filename="config.json", revision=revision)
+    )
+    root = config_path.parent
+    for name in ("ema.safetensors", "model.safetensors"):
+        try:
+            hf_hub_download(repo_id=repo, filename=name, revision=revision)
+        except Exception:
+            if name == "ema.safetensors":
+                raise
+    return root
+
+
+def _frames_to_tensor(frames: np.ndarray, device: torch.device) -> torch.Tensor:
+    tensor = torch.from_numpy(frames.copy()).permute(0, 3, 1, 2).float()
+    return tensor.div_(127.5).sub_(1.0).unsqueeze(0).to(device, non_blocking=True)
+
+
+def _tensor_to_image(frame: torch.Tensor) -> Image.Image:
+    array = frame.float().clamp(-1.0, 1.0).add(1.0).mul(127.5)
+    array = array.byte().permute(1, 2, 0).cpu().numpy()
+    return Image.fromarray(array, mode="RGB")
+
+
+def _inject_ball(image: Image.Image, position: tuple[float, float]) -> Image.Image:
+    image = image.copy()
+    draw = ImageDraw.Draw(image)
+    radius = max(3.0, min(image.width, image.height) * 0.025)
+    x, y = position
+    draw.ellipse(
+        (x - radius, y - radius, x + radius, y + radius),
+        fill=(255, 209, 72),
+        outline=(255, 247, 196),
+        width=max(1, round(radius * 0.25)),
+    )
+    return image
+
+
+def _image_to_video_tensor(image: Image.Image, device: torch.device) -> torch.Tensor:
+    array = np.asarray(image, dtype=np.uint8).copy()
+    tensor = torch.from_numpy(array).permute(2, 0, 1).float()
+    return (
+        tensor.div_(127.5)
+        .sub_(1.0)
+        .reshape(1, 1, 3, image.height, image.width)
+        .to(device, non_blocking=True)
+    )
+
+
+class LiveWindow:
+    def __init__(self, scale: int) -> None:
+        self.scale = max(1, scale)
+        self.is_open = True
+        self.root = tk.Tk()
+        self.root.title("Diffusion Pong")
+        self.label = tk.Label(self.root, background="black")
+        self.label.pack()
+        self.photo: ImageTk.PhotoImage | None = None
+        self.generated_width = 0
+        self.generated_height = 0
+        self.pending_click: tuple[float, float] | None = None
+        self.pressed: set[str] = set()
+        self.label.bind("<Button-1>", self._on_click)
+        self.root.bind("<KeyPress>", self._on_key_press)
+        self.root.bind("<KeyRelease>", self._on_key_release)
+        self.root.protocol("WM_DELETE_WINDOW", self.close)
+        self.root.focus_force()
+
+    def _on_key_press(self, event: tk.Event) -> None:
+        key = str(event.keysym).lower()
+        if key in {"w", "s"}:
+            self.pressed.add(key)
+
+    def _on_key_release(self, event: tk.Event) -> None:
+        self.pressed.discard(str(event.keysym).lower())
+
+    def action(self) -> tuple[float, float]:
+        if {"w", "s"}.issubset(self.pressed):
+            return (0.0, 0.0)
+        return (float("w" in self.pressed), float("s" in self.pressed))
+
+    def _on_click(self, event: tk.Event) -> None:
+        if self.generated_width <= 0 or self.generated_height <= 0:
+            return
+        x = (float(event.x) / self.scale) % self.generated_width
+        y = min(max(float(event.y) / self.scale, 0.0), self.generated_height - 1.0)
+        self.pending_click = (x, y)
+
+    def consume_click(self) -> tuple[float, float] | None:
+        click = self.pending_click
+        self.pending_click = None
+        return click
+
+    def close(self) -> None:
+        if self.is_open:
+            self.is_open = False
+            self.root.destroy()
+
+    def show(self, image: Image.Image, frame_index: int, latency_ms: float) -> bool:
+        if not self.is_open:
+            return False
+        self.generated_width, self.generated_height = image.size
+        display = image.resize(
+            (image.width * self.scale, image.height * self.scale),
+            Image.Resampling.NEAREST,
+        )
+        self.photo = ImageTk.PhotoImage(display)
+        self.label.configure(image=self.photo)
+        held = "+".join(sorted(self.pressed)).upper()
+        held_text = f"  [{held}]" if held else ""
+        self.root.title(
+            f"Diffusion Pong — frame {frame_index + 1}  {latency_ms:.0f} ms{held_text}"
+        )
+        try:
+            self.root.update_idletasks()
+            self.root.update()
+        except tk.TclError:
+            self.is_open = False
+        return self.is_open
+
+
+@torch.inference_mode()
+def main() -> None:
+    args = _parse_args()
+    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
+        raise SystemExit("Need a CUDA GPU with BF16 support.")
+
+    print(f"Downloading {args.repo} …")
+    checkpoint = _download_checkpoint(args.repo, args.revision, args.local_dir)
+    config_path = checkpoint / "config.json"
+    if not config_path.exists():
+        raise SystemExit(f"Missing config.json in {checkpoint}")
+    config = load_config(config_path)
+    device = torch.device("cuda")
+    torch.set_float32_matmul_precision("high")
+
+    seed = args.seed if args.seed is not None else random.SystemRandom().randrange(2**31)
+    simulation = PongSimulation.from_seed(config.video, seed)
+    for _ in range(random.Random(seed).randint(0, 400)):
+        simulation.step()
+    history_rgb = _frames_to_tensor(
+        simulation.sample_frames(config.video.history_frames),
+        device,
+    )
+
+    print("Loading VAE + DiT …")
+    codec = load_codec(config.codec, device)
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        history = deque(
+            (
+                codec.encode(history_rgb)[:, index]
+                for index in range(config.video.history_frames)
+            ),
+            maxlen=config.video.history_frames,
+        )
+    model = (
+        CausalLatentVideoDiT(
+            config.model,
+            latent_channels=codec.channels,
+            history_frames=codec.latent_frame_count(config.video.history_frames),
+        )
+        .to(device)
+        .eval()
+    )
+    load_model_weights(checkpoint, model, use_ema=True)
+
+    steps = max(1, args.steps)
+    # Teacher-stage checkpoint: use instantaneous velocity sampling.
+    instantaneous = True
+    frame_limit = args.frames if args.frames > 0 else 10_000_000
+    window = LiveWindow(args.window_scale)
+    generator = torch.Generator(device=device).manual_seed(seed + 7)
+    inject_pos: tuple[float, float] | None = None
+    inject_frames = 0
+    print("Ready. Click the window, then hold W/S.")
+
+    for frame_index in range(frame_limit):
+        click = window.consume_click()
+        if click is not None:
+            inject_pos = click
+            inject_frames = 3
+
+        history_tensor = torch.stack(tuple(history), dim=1)
+        action = torch.tensor(
+            [window.action()],
+            device=device,
+            dtype=torch.float32,
+        ).unsqueeze(0)
+
+        torch.cuda.synchronize()
+        started = time.perf_counter()
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            memory_cache = model.build_memory_cache(history_tensor)
+            prediction = sample_next_latent(
+                model,
+                history_tensor,
+                steps,
+                actions=action,
+                generator=generator,
+                window_frames=1,
+                memory_cache=memory_cache,
+                instantaneous=instantaneous,
+            )
+            decoded = codec.decode(prediction)
+        torch.cuda.synchronize()
+        latency_ms = (time.perf_counter() - started) * 1_000
+
+        image = _tensor_to_image(decoded[0, 0])
+        frame_latent = prediction[:, 0]
+        if inject_pos is not None and inject_frames > 0:
+            image = _inject_ball(image, inject_pos)
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                frame_latent = codec.encode(_image_to_video_tensor(image, device))[:, 0]
+            inject_frames -= 1
+            if inject_frames == 0:
+                inject_pos = None
+
+        history.append(frame_latent)
+        if not window.show(image, frame_index, latency_ms):
+            break
+        if args.fps_cap > 0:
+            target = 1.0 / args.fps_cap
+            elapsed = time.perf_counter() - started
+            if elapsed < target:
+                time.sleep(target - elapsed)
+
+    if window.is_open:
+        window.close()
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()