Action-conditioned Breakout EMA step 500 + standalone inference

README.md

@@ -0,0 +1,59 @@
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: other
+tags:
+  - world-model
+  - diffusion
+  - breakout
+  - video
+  - causal
+  - game
+---
+
+# Diffusion Breakout
+
+An action-conditioned Breakout world model. Hold **A / D** (or the arrow keys) and the model generates the next frame.
+
+![validation preview](preview.png)
+
+128x128 at 6 FPS with 12 frames of latent history. The model has about 70M parameters and uses the frozen SDXL VAE from `madebyollin/sdxl-vae-fp16-fix`.
+
+## Run it
+
+Requires a CUDA GPU with BF16 support.
+
+```bash
+pip install torch numpy pillow safetensors huggingface_hub diffusers
+
+hf download kerzgrr/diffusionbreakout live_infer.py --local-dir .
+python live_infer.py --steps 1
+```
+
+To keep the weights in a local folder:
+
+```bash
+hf download kerzgrr/diffusionbreakout \
+  --local-dir checkpoints/diffusionbreakout \
+  --include "ema.safetensors" \
+  --include "config.json" \
+  --include "live_infer.py"
+
+python checkpoints/diffusionbreakout/live_infer.py \
+  --local-dir checkpoints/diffusionbreakout \
+  --steps 1 \
+  --window-scale 6
+```
+
+Click the window for focus. **A / Left** moves left; **D / Right** moves right.
+
+## Files
+
+| file | contents |
+|---|---|
+| `ema.safetensors` | action-conditioned EMA weights |
+| `config.json` | model, video, codec, and training settings |
+| `live_infer.py` | single-file live inference |
+| `preview.png` | validation strip |
+
+Actions are encoded as `[left, right]`.

config.json

@@ -0,0 +1,108 @@
+{
+  "codec": {
+    "enable_slicing": true,
+    "enable_tiling": false,
+    "frame_batch_size": 256,
+    "kind": "frame_vae",
+    "latent_channels": 4,
+    "model_id": "madebyollin/sdxl-vae-fp16-fix",
+    "revision": "main",
+    "sample_posterior": false,
+    "spatial_compression": 8,
+    "subfolder": "",
+    "temporal_compression": 1
+  },
+  "data": {
+    "batch_size": 128,
+    "data_path": "",
+    "dataset_kind": "breakout",
+    "frame_stride": 1,
+    "hf_repo_id": "",
+    "hf_revision": "main",
+    "latent_cache_path": "",
+    "latent_prefetch_batches": 8,
+    "network_clip_prefetch": 2,
+    "num_workers": 8,
+    "persistent_workers": true,
+    "pin_memory": true,
+    "prefetch_batches": 2,
+    "sample_shuffle": 64,
+    "shard_shuffle": 256,
+    "train_samples": 500000,
+    "validation_samples": 1024,
+    "windows_per_video": 4
+  },
+  "inference": {
+    "frames": 120,
+    "steps": 1
+  },
+  "model": {
+    "action_dim": 2,
+    "context_dim": 0,
+    "depth": 12,
+    "fine_history_frames": 2,
+    "fine_memory_patch_size": 2,
+    "generation_frames": 1,
+    "hidden_size": 416,
+    "memory_depth": 3,
+    "memory_patch_size": 4,
+    "mlp_ratio": 4.0,
+    "num_heads": 8,
+    "qk_norm": true,
+    "target_patch_size": 2
+  },
+  "schema_version": 3,
+  "simulation": {
+    "max_shapes": 2,
+    "max_size": 6.0,
+    "max_speed": 20.0,
+    "min_shapes": 1,
+    "min_size": 3.0,
+    "min_speed": 8.0,
+    "supersample": 2,
+    "warmup_frames": 8
+  },
+  "training": {
+    "activation_checkpointing": true,
+    "critic_learning_rate": 0.0001,
+    "distribution_matching_start_step": 30000,
+    "distribution_matching_weight": 0.05,
+    "ema_decay": 0.9999,
+    "epochs": 100,
+    "flow_matching_ratio": 0.25,
+    "full_interval_ratio": 0.25,
+    "gradient_accumulation_steps": 1,
+    "history_noise_std": 0.01,
+    "keep_checkpoints": 5,
+    "learning_rate": 0.0003,
+    "log_every": 20,
+    "max_grad_norm": 1.0,
+    "min_learning_rate": 1e-05,
+    "on_policy_max_horizon": 4,
+    "on_policy_ratio": 0.15,
+    "on_policy_sampling_steps": 4,
+    "on_policy_start_step": 15000,
+    "output_dir": "runs/breakout-action-world-model",
+    "reward_loss_weight": 0.0,
+    "save_every": 50,
+    "seed": 2026,
+    "teacher_endpoint_ratio": 0.1,
+    "teacher_integration_steps": 8,
+    "teacher_logit_mean": 1.0,
+    "teacher_logit_std": 1.0,
+    "teacher_steps": 10000,
+    "teacher_x0_weight": 0.25,
+    "terminal_loss_weight": 0.0,
+    "validate_every": 50,
+    "validation_batches": 1,
+    "warmup_steps": 500,
+    "weight_decay": 0.0
+  },
+  "video": {
+    "fps": 6.0,
+    "future_frames": 5,
+    "height": 128,
+    "history_frames": 12,
+    "width": 128
+  }
+}

ema.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9f2a9ab8aa958c532fc5f010d0f429e49762803a4203269296e7ae48fdf8191
+size 281086680

live_infer.py

@@ -0,0 +1,1368 @@
+#!/usr/bin/env python3
+"""Run Diffusion Breakout.
+
+Deps:
+  pip install torch numpy pillow safetensors huggingface_hub diffusers
+
+Run:
+  python live_infer.py
+  python live_infer.py --steps 1 --window-scale 7 --seed 42
+
+Controls: A/D or arrow keys. Click window for focus.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import time
+import tkinter as tk
+from collections import deque
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusers import AutoencoderKL
+from huggingface_hub import hf_hub_download
+from PIL import Image, ImageDraw, ImageTk
+from safetensors.torch import load_file
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+DEFAULT_REPO = "kerzgrr/diffusionbreakout"
+
+
+@dataclass(frozen=True, slots=True)
+class VideoConfig:
+    height: int = 128
+    width: int = 128
+    fps: float = 6.0
+    history_frames: int = 12
+    future_frames: int = 5
+
+
+@dataclass(frozen=True, slots=True)
+class ModelConfig:
+    target_patch_size: int = 2
+    memory_patch_size: int = 4
+    fine_memory_patch_size: int = 2
+    fine_history_frames: int = 2
+    hidden_size: int = 384
+    depth: int = 12
+    memory_depth: int = 3
+    num_heads: int = 8
+    mlp_ratio: float = 4.0
+    qk_norm: bool = True
+    generation_frames: int = 1
+    action_dim: int = 2
+    context_dim: int = 0
+
+
+@dataclass(frozen=True, slots=True)
+class CodecSettings:
+    model_id: str = "madebyollin/sdxl-vae-fp16-fix"
+    subfolder: str = ""
+    revision: str = "main"
+    latent_channels: int = 4
+    spatial_compression: int = 8
+    frame_batch_size: int = 256
+    sample_posterior: bool = False
+    enable_slicing: bool = True
+    enable_tiling: bool = False
+_BREAKOUT_BRICK_COLORS = (
+    (220, 72, 88),
+    (230, 140, 54),
+    (236, 196, 72),
+    (72, 186, 120),
+    (72, 148, 220),
+    (156, 102, 220),
+)
+
+
+class BreakoutSimulation:
+    """Classic Breakout with continuous physics and a scripted bottom paddle."""
+
+    def __init__(self, video_config: VideoConfig, seed: int) -> None:
+        self.video_config = video_config
+        self.rng = np.random.default_rng(seed)
+        self.dt = 1.0 / video_config.fps
+        self.paddle_width = max(18.0, video_config.width * 0.18)
+        self.paddle_height = max(4.0, video_config.height * 0.03)
+        self.ball_radius = max(3.0, min(video_config.width, video_config.height) * 0.025)
+        self.paddle_y = video_config.height * 0.90
+        self.brick_rows = 5
+        self.brick_cols = 8
+        self.brick_top = video_config.height * 0.14
+        self.brick_gap = 1.5
+        usable_width = video_config.width - 10.0
+        self.brick_width = (usable_width - self.brick_gap * (self.brick_cols - 1)) / self.brick_cols
+        self.brick_height = max(5.0, video_config.height * 0.045)
+        self.paddle_x = video_config.width / 2
+        self.score = 0
+        self.elapsed = 0.0
+        self.aim_bias = float(self.rng.uniform(-8.0, 8.0))
+        self.reaction_delay = int(self.rng.integers(0, 4))
+        self.randomness = float(self.rng.uniform(0.10, 0.28))
+        self._observed_ball_x: deque[float] = deque(
+            [video_config.width / 2] * (self.reaction_delay + 1),
+            maxlen=self.reaction_delay + 1,
+        )
+        self._held_action = 0
+        self._action_hold_frames = 0
+        self.ball_position = np.zeros(2, dtype=np.float32)
+        self.ball_velocity = np.zeros(2, dtype=np.float32)
+        self.bricks = np.ones((self.brick_rows, self.brick_cols), dtype=bool)
+        # Punch a few random holes so early frames aren't always a full wall.
+        missing = int(self.rng.integers(0, 6))
+        for _ in range(missing):
+            row = int(self.rng.integers(0, self.brick_rows))
+            col = int(self.rng.integers(0, self.brick_cols))
+            self.bricks[row, col] = False
+        self._serve()
+        for _ in range(10):
+            self.step()
+
+    @classmethod
+    def from_seed(
+        cls,
+        video_config: VideoConfig,
+        seed: int,
+        *,
+        at_start: bool = False,
+    ) -> BreakoutSimulation:
+        simulation = cls(video_config, seed)
+        if at_start:
+            simulation.reset_to_start()
+        return simulation
+
+    def reset_to_start(self) -> None:
+        """Reset to a fresh serve against a full brick wall (for inference)."""
+        self.paddle_x = self.video_config.width / 2
+        self.score = 0
+        self.elapsed = 0.0
+        self._held_action = 0
+        self._action_hold_frames = 0
+        self._observed_ball_x.clear()
+        self._observed_ball_x.extend(
+            [self.video_config.width / 2] * (self.reaction_delay + 1)
+        )
+        self.bricks[:] = True
+        self._serve()
+
+    def _brick_rect(self, row: int, col: int) -> tuple[float, float, float, float]:
+        x0 = 5.0 + col * (self.brick_width + self.brick_gap)
+        y0 = self.brick_top + row * (self.brick_height + self.brick_gap)
+        return x0, y0, x0 + self.brick_width, y0 + self.brick_height
+
+    def _serve(self) -> None:
+        angle = float(self.rng.uniform(-0.65, 0.65))
+        speed = float(self.rng.uniform(48.0, 64.0))
+        self.ball_position[:] = (
+            self.paddle_x + float(self.rng.uniform(-6.0, 6.0)),
+            self.paddle_y - self.paddle_height - self.ball_radius - 2.0,
+        )
+        self.ball_velocity[:] = (
+            speed * math.sin(angle),
+            -speed * math.cos(angle),
+        )
+
+    def _reset_bricks(self) -> None:
+        self.bricks[:] = True
+
+    def _move_paddle(self, dt: float, action: int | None) -> None:
+        width = float(self.video_config.width)
+        paddle_speed = width * 0.85
+        if action is None:
+            target = float(self.ball_position[0]) + self.aim_bias
+            delta = max(-paddle_speed * dt, min(paddle_speed * dt, target - self.paddle_x))
+        else:
+            delta = float(max(-1, min(1, action))) * paddle_speed * dt
+        self.paddle_x += delta
+        half = self.paddle_width / 2
+        self.paddle_x = max(half + 2.0, min(width - half - 2.0, self.paddle_x))
+
+    def _collide_bricks(self) -> None:
+        x, y = map(float, self.ball_position)
+        r = self.ball_radius
+        for row in range(self.brick_rows):
+            for col in range(self.brick_cols):
+                if not self.bricks[row, col]:
+                    continue
+                x0, y0, x1, y1 = self._brick_rect(row, col)
+                closest_x = max(x0, min(x, x1))
+                closest_y = max(y0, min(y, y1))
+                dx = x - closest_x
+                dy = y - closest_y
+                if dx * dx + dy * dy > r * r:
+                    continue
+                self.bricks[row, col] = False
+                self.score += 1
+                # Reflect on the dominant penetration axis.
+                if abs(dx) > abs(dy):
+                    self.ball_velocity[0] = -self.ball_velocity[0]
+                    self.ball_position[0] = x0 - r if dx < 0 else x1 + r
+                else:
+                    self.ball_velocity[1] = -self.ball_velocity[1]
+                    self.ball_position[1] = y0 - r if dy < 0 else y1 + r
+                if not self.bricks.any():
+                    self._reset_bricks()
+                return
+
+    def step(self, action: int | None = None) -> None:
+        width = float(self.video_config.width)
+        height = float(self.video_config.height)
+        substeps = 4
+        sub_dt = self.dt / substeps
+        for _ in range(substeps):
+            self._move_paddle(sub_dt, action)
+            self.ball_position += self.ball_velocity * sub_dt
+            x, y = map(float, self.ball_position)
+            if x - self.ball_radius < 2:
+                self.ball_position[0] = 2 + self.ball_radius
+                self.ball_velocity[0] = abs(self.ball_velocity[0])
+            elif x + self.ball_radius > width - 2:
+                self.ball_position[0] = width - 2 - self.ball_radius
+                self.ball_velocity[0] = -abs(self.ball_velocity[0])
+            if y - self.ball_radius < 2:
+                self.ball_position[1] = 2 + self.ball_radius
+                self.ball_velocity[1] = abs(self.ball_velocity[1])
+            paddle_top = self.paddle_y - self.paddle_height / 2
+            hitting_paddle = (
+                self.ball_velocity[1] > 0
+                and y + self.ball_radius >= paddle_top
+                and y - self.ball_radius <= self.paddle_y + self.paddle_height / 2
+                and abs(x - self.paddle_x) <= self.paddle_width / 2 + self.ball_radius
+            )
+            if hitting_paddle:
+                self.ball_position[1] = paddle_top - self.ball_radius
+                self.ball_velocity[1] = -abs(self.ball_velocity[1]) * 1.01
+                self.ball_velocity[0] += (x - self.paddle_x) * 1.35
+                speed = float(np.linalg.norm(self.ball_velocity))
+                target = float(np.clip(speed, 46.0, 78.0))
+                if speed > 1e-5:
+                    self.ball_velocity *= target / speed
+            self._collide_bricks()
+            if self.ball_position[1] > height + self.ball_radius * 2:
+                self._serve()
+            self.elapsed += sub_dt
+
+    def render(self) -> np.ndarray:
+        scale = 2
+        width = self.video_config.width
+        height = self.video_config.height
+        image = Image.new("RGB", (width * scale, height * scale), (6, 8, 18))
+        draw = ImageDraw.Draw(image)
+        draw.rectangle((0, 0, width * scale, 3 * scale), fill=(36, 48, 72))
+        for row in range(self.brick_rows):
+            color = _BREAKOUT_BRICK_COLORS[row % len(_BREAKOUT_BRICK_COLORS)]
+            for col in range(self.brick_cols):
+                if not self.bricks[row, col]:
+                    continue
+                x0, y0, x1, y1 = self._brick_rect(row, col)
+                draw.rounded_rectangle(
+                    (x0 * scale, y0 * scale, x1 * scale, y1 * scale),
+                    radius=max(1, round(scale)),
+                    fill=color,
+                )
+        draw.rounded_rectangle(
+            (
+                (self.paddle_x - self.paddle_width / 2) * scale,
+                (self.paddle_y - self.paddle_height / 2) * scale,
+                (self.paddle_x + self.paddle_width / 2) * scale,
+                (self.paddle_y + self.paddle_height / 2) * scale,
+            ),
+            radius=2 * scale,
+            fill=(230, 238, 245),
+        )
+        ball_x, ball_y = self.ball_position * scale
+        radius = self.ball_radius * scale
+        draw.ellipse(
+            (ball_x - radius, ball_y - radius, ball_x + radius, ball_y + radius),
+            fill=(255, 209, 72),
+        )
+        image = image.resize((width, height), Image.Resampling.LANCZOS)
+        return np.asarray(image, dtype=np.uint8)
+
+    def next_frame(self) -> np.ndarray:
+        frame = self.render()
+        self.step()
+        return frame
+
+    def sample_frames(self, count: int) -> np.ndarray:
+        return np.stack([self.next_frame() for _ in range(count)])
+
+    def _scripted_action(self) -> int:
+        self._observed_ball_x.append(float(self.ball_position[0]))
+        if self._action_hold_frames > 0:
+            self._action_hold_frames -= 1
+            return self._held_action
+        self._action_hold_frames = int(self.rng.integers(1, 5))
+        random_value = float(self.rng.random())
+        delayed_target = self._observed_ball_x[0] + self.aim_bias
+        error = delayed_target - self.paddle_x
+        ideal_action = -1 if error < -4 else (1 if error > 4 else 0)
+        if random_value < self.randomness * 0.35:
+            action = 0
+        elif random_value < self.randomness * 0.75:
+            action = int(self.rng.choice((-1, 0, 1)))
+        elif random_value < self.randomness:
+            action = -ideal_action if ideal_action != 0 else int(self.rng.choice((-1, 1)))
+        else:
+            action = ideal_action
+        self._held_action = action
+        return action
+
+    def sample_gameplay(self, count: int) -> tuple[np.ndarray, np.ndarray]:
+        frames: list[np.ndarray] = []
+        actions: list[tuple[float, float]] = []
+        for _ in range(count):
+            frames.append(self.render())
+            action = self._scripted_action()
+            # [A/left, D/right] — same layout we'll use for action fine-tuning later.
+            actions.append((float(action < 0), float(action > 0)))
+            self.step(action)
+        return np.stack(frames), np.asarray(actions, dtype=np.float32)
+
+
+class FrameAutoencoderCodec(nn.Module):
+    def __init__(self, config: CodecSettings, device: torch.device, dtype: torch.dtype) -> None:
+        super().__init__()
+        load_kwargs: dict[str, object] = {
+            "revision": config.revision,
+            "torch_dtype": dtype,
+            "low_cpu_mem_usage": True,
+        }
+        if config.subfolder:
+            load_kwargs["subfolder"] = config.subfolder
+        try:
+            self.vae = AutoencoderKL.from_pretrained(
+                config.model_id, local_files_only=True, **load_kwargs
+            )
+        except OSError:
+            self.vae = AutoencoderKL.from_pretrained(config.model_id, **load_kwargs)
+        self.vae.to(device=device).eval()
+        self.vae.to(memory_format=torch.channels_last)
+        self.vae.requires_grad_(False)
+        if config.enable_slicing:
+            self.vae.enable_slicing()
+        if config.enable_tiling:
+            self.vae.enable_tiling()
+        self.channels = int(self.vae.config.latent_channels)
+        self.spatial_compression = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.temporal_compression = 1
+        self.frame_batch_size = config.frame_batch_size
+        self.sample_posterior = config.sample_posterior
+        self.dtype = dtype
+        self.scaling_factor = float(self.vae.config.scaling_factor)
+        shift_factor = getattr(self.vae.config, "shift_factor", None)
+        self.shift_factor = float(shift_factor) if shift_factor is not None else 0.0
+
+    def latent_frame_count(self, frame_count: int) -> int:
+        return frame_count
+
+    @torch.no_grad()
+    def encode(self, frames: torch.Tensor) -> torch.Tensor:
+        batch, frame_count, channels, height, width = frames.shape
+        flat_frames = frames.reshape(batch * frame_count, channels, height, width)
+        encoded: list[torch.Tensor] = []
+        for frame_batch in flat_frames.split(self.frame_batch_size):
+            frame_batch = frame_batch.to(self.dtype).contiguous(
+                memory_format=torch.channels_last
+            )
+            posterior = self.vae.encode(frame_batch).latent_dist
+            latents = posterior.sample() if self.sample_posterior else posterior.mode()
+            encoded.append((latents - self.shift_factor) * self.scaling_factor)
+        merged = torch.cat(encoded)
+        return merged.reshape(
+            batch,
+            frame_count,
+            self.channels,
+            height // self.spatial_compression,
+            width // self.spatial_compression,
+        )
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        batch, frame_count, channels, height, width = latents.shape
+        flat_latents = latents.reshape(batch * frame_count, channels, height, width)
+        decoded: list[torch.Tensor] = []
+        for latent_batch in flat_latents.split(self.frame_batch_size):
+            denormalized = latent_batch / self.scaling_factor + self.shift_factor
+            denormalized = denormalized.to(self.dtype).contiguous(
+                memory_format=torch.channels_last
+            )
+            decoded.append(self.vae.decode(denormalized).sample)
+        merged = torch.cat(decoded)
+        return merged.reshape(
+            batch,
+            frame_count,
+            3,
+            height * self.spatial_compression,
+            width * self.spatial_compression,
+        ).clamp(-1.0, 1.0)
+@dataclass(slots=True)
+class PrefixCache:
+    key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...]
+    pooled: torch.Tensor
+
+
+MemoryCache = PrefixCache
+
+
+def _video_positions(
+    frames: int,
+    height: int,
+    width: int,
+    time_offset: int,
+    device: torch.device,
+) -> torch.Tensor:
+    time = torch.arange(time_offset, time_offset + frames, device=device)
+    y = torch.arange(height, device=device)
+    x = torch.arange(width, device=device)
+    grid = torch.meshgrid(time, y, x, indexing="ij")
+    return torch.stack(grid, dim=-1).reshape(-1, 3).float()
+
+
+class Rotary3D(nn.Module):
+    def __init__(self, head_size: int, base: float = 10_000.0) -> None:
+        super().__init__()
+        self.rotary_size = (head_size // 6) * 6
+        self.axis_size = self.rotary_size // 3
+        frequencies = torch.exp(
+            -math.log(base)
+            * torch.arange(self.axis_size // 2, dtype=torch.float32)
+            / (self.axis_size // 2)
+        )
+        self.register_buffer("frequencies", frequencies, persistent=False)
+
+    @staticmethod
+    def _rotate_half(inputs: torch.Tensor) -> torch.Tensor:
+        pairs = inputs.reshape(*inputs.shape[:-1], -1, 2)
+        first, second = pairs.unbind(dim=-1)
+        return torch.stack((-second, first), dim=-1).flatten(-2)
+
+    def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        angles = torch.cat(
+            [positions[:, axis, None] * self.frequencies[None] for axis in range(3)],
+            dim=-1,
+        )
+        cosine = angles.cos().repeat_interleave(2, dim=-1)[None, None]
+        sine = angles.sin().repeat_interleave(2, dim=-1)[None, None]
+        rotary, remainder = inputs.split(
+            (self.rotary_size, inputs.shape[-1] - self.rotary_size),
+            dim=-1,
+        )
+        rotated = rotary * cosine.to(rotary.dtype)
+        rotated = rotated + self._rotate_half(rotary) * sine.to(rotary.dtype)
+        return torch.cat((rotated, remainder), dim=-1)
+
+
+class ScalarEmbedding(nn.Module):
+    def __init__(self, hidden_size: int) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size),
+        )
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        value = value.reshape(-1).float()
+        half = self.hidden_size // 2
+        frequencies = torch.exp(
+            -math.log(10_000.0)
+            * torch.arange(half, device=value.device, dtype=torch.float32)
+            / half
+        )
+        angles = value[:, None] * frequencies[None] * 1_000.0
+        embedding = torch.cat((angles.cos(), angles.sin()), dim=-1)
+        return self.mlp(embedding.to(dtype=self.mlp[0].weight.dtype))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, hidden_size: int, ratio: float) -> None:
+        super().__init__()
+        inner_size = int(hidden_size * ratio * 2 / 3)
+        inner_size = ((inner_size + 63) // 64) * 64
+        self.input = nn.Linear(hidden_size, inner_size * 2, bias=False)
+        self.output = nn.Linear(inner_size, hidden_size, bias=False)
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        values, gates = self.input(inputs).chunk(2, dim=-1)
+        return self.output(values * F.silu(gates))
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, qk_norm: bool) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_size = hidden_size // num_heads
+        self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
+        self.output = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.q_norm = (
+            nn.RMSNorm(self.head_size, eps=1e-6)
+            if qk_norm
+            else nn.Identity()
+        )
+        self.k_norm = (
+            nn.RMSNorm(self.head_size, eps=1e-6)
+            if qk_norm
+            else nn.Identity()
+        )
+        self.rotary = Rotary3D(self.head_size)
+
+    def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        batch, tokens, channels = inputs.shape
+        query, key, value = self.qkv(inputs).reshape(
+            batch,
+            tokens,
+            3,
+            self.num_heads,
+            self.head_size,
+        ).unbind(dim=2)
+        query = self.rotary(self.q_norm(query).transpose(1, 2), positions)
+        key = self.rotary(self.k_norm(key).transpose(1, 2), positions)
+        value = value.transpose(1, 2)
+        attended = F.scaled_dot_product_attention(query, key, value)
+        return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels))
+
+
+class PrefixAttention(nn.Module):
+    """Target attention over the joint clean-prefix and noisy-target sequence."""
+
+    def __init__(self, hidden_size: int, num_heads: int, qk_norm: bool) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_size = hidden_size // num_heads
+        self.target_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
+        self.prefix_key_value = nn.Linear(hidden_size, hidden_size * 2, bias=False)
+        self.output = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.q_norm = (
+            nn.RMSNorm(self.head_size, eps=1e-6)
+            if qk_norm
+            else nn.Identity()
+        )
+        self.k_norm = (
+            nn.RMSNorm(self.head_size, eps=1e-6)
+            if qk_norm
+            else nn.Identity()
+        )
+        self.rotary = Rotary3D(self.head_size)
+
+    def project_prefix(
+        self,
+        prefix: torch.Tensor,
+        positions: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch, tokens, _ = prefix.shape
+        key, value = self.prefix_key_value(prefix).reshape(
+            batch,
+            tokens,
+            2,
+            self.num_heads,
+            self.head_size,
+        ).unbind(dim=2)
+        key = self.rotary(self.k_norm(key).transpose(1, 2), positions)
+        return key, value.transpose(1, 2)
+
+    def forward_cached(
+        self,
+        target: torch.Tensor,
+        target_positions: torch.Tensor,
+        prefix_key_value: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        batch, tokens, channels = target.shape
+        query, key, value = self.target_qkv(target).reshape(
+            batch,
+            tokens,
+            3,
+            self.num_heads,
+            self.head_size,
+        ).unbind(dim=2)
+        query = self.rotary(self.q_norm(query).transpose(1, 2), target_positions)
+        key = self.rotary(self.k_norm(key).transpose(1, 2), target_positions)
+        value = value.transpose(1, 2)
+        prefix_key, prefix_value = prefix_key_value
+        key = torch.cat((prefix_key, key), dim=2)
+        value = torch.cat((prefix_value, value), dim=2)
+        attended = F.scaled_dot_product_attention(query, key, value)
+        return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels))
+
+
+class PrefixDiTBlock(nn.Module):
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        hidden_size = config.hidden_size
+        self.prefix_norm1 = nn.RMSNorm(hidden_size, eps=1e-6)
+        self.prefix_attention = SelfAttention(
+            hidden_size,
+            config.num_heads,
+            config.qk_norm,
+        )
+        self.prefix_norm2 = nn.RMSNorm(hidden_size, eps=1e-6)
+        self.prefix_mlp = SwiGLU(hidden_size, config.mlp_ratio)
+        self.prefix_projection_norm = nn.RMSNorm(hidden_size, eps=1e-6)
+        self.target_norm1 = nn.RMSNorm(
+            hidden_size,
+            eps=1e-6,
+            elementwise_affine=False,
+        )
+        self.target_attention = PrefixAttention(
+            hidden_size,
+            config.num_heads,
+            config.qk_norm,
+        )
+        self.target_norm2 = nn.RMSNorm(
+            hidden_size,
+            eps=1e-6,
+            elementwise_affine=False,
+        )
+        self.target_mlp = SwiGLU(hidden_size, config.mlp_ratio)
+        self.modulation = nn.Linear(hidden_size, hidden_size * 6)
+
+    @staticmethod
+    def _modulate(
+        inputs: torch.Tensor,
+        shift: torch.Tensor,
+        scale: torch.Tensor,
+    ) -> torch.Tensor:
+        return inputs * (1.0 + scale[:, None]) + shift[:, None]
+
+    def advance_prefix(
+        self,
+        prefix: torch.Tensor,
+        positions: torch.Tensor,
+    ) -> torch.Tensor:
+        prefix = prefix + self.prefix_attention(self.prefix_norm1(prefix), positions)
+        return prefix + self.prefix_mlp(self.prefix_norm2(prefix))
+
+    def prefix_key_value(
+        self,
+        prefix: torch.Tensor,
+        positions: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.target_attention.project_prefix(
+            self.prefix_projection_norm(prefix),
+            positions,
+        )
+
+    def advance_target(
+        self,
+        target: torch.Tensor,
+        condition: torch.Tensor,
+        target_positions: torch.Tensor,
+        prefix_key_value: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        shift1, scale1, gate1, shift2, scale2, gate2 = self.modulation(
+            F.silu(condition)
+        ).chunk(6, dim=-1)
+        normalized = self._modulate(self.target_norm1(target), shift1, scale1)
+        target = target + gate1[:, None] * self.target_attention.forward_cached(
+            normalized,
+            target_positions,
+            prefix_key_value,
+        )
+        normalized = self._modulate(self.target_norm2(target), shift2, scale2)
+        return target + gate2[:, None] * self.target_mlp(normalized)
+
+    def forward(
+        self,
+        prefix: torch.Tensor,
+        target: torch.Tensor,
+        condition: torch.Tensor,
+        prefix_positions: torch.Tensor,
+        target_positions: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        prefix = self.advance_prefix(prefix, prefix_positions)
+        key_value = self.prefix_key_value(prefix, prefix_positions)
+        target = self.advance_target(
+            target,
+            condition,
+            target_positions,
+            key_value,
+        )
+        return prefix, target
+
+
+class CausalLatentVideoDiT(nn.Module):
+    """Multiscale clean-prefix/noisy-suffix latent video transformer."""
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        latent_channels: int,
+        history_frames: int = 12,
+    ) -> None:
+        super().__init__()
+        self.model_config = model_config
+        self.latent_channels = latent_channels
+        self.history_frames = history_frames
+        self.gradient_checkpointing = False
+        hidden_size = model_config.hidden_size
+        target_patch = model_config.target_patch_size
+        coarse_patch = model_config.memory_patch_size
+        fine_patch = model_config.fine_memory_patch_size
+        self.target_patch_embed = nn.Conv3d(
+            latent_channels,
+            hidden_size,
+            kernel_size=(1, target_patch, target_patch),
+            stride=(1, target_patch, target_patch),
+        )
+        self.coarse_prefix_embed = nn.Conv3d(
+            latent_channels,
+            hidden_size,
+            kernel_size=(1, coarse_patch, coarse_patch),
+            stride=(1, coarse_patch, coarse_patch),
+        )
+        self.fine_prefix_embed = nn.Conv3d(
+            latent_channels,
+            hidden_size,
+            kernel_size=(1, fine_patch, fine_patch),
+            stride=(1, fine_patch, fine_patch),
+        )
+        self.motion_prefix_embed = nn.Conv3d(
+            latent_channels,
+            hidden_size,
+            kernel_size=(1, fine_patch, fine_patch),
+            stride=(1, fine_patch, fine_patch),
+        )
+        self.coarse_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
+        self.fine_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
+        self.t_embedding = ScalarEmbedding(hidden_size)
+        self.r_embedding = ScalarEmbedding(hidden_size)
+        self.delta_embedding = ScalarEmbedding(hidden_size)
+        self.time_projection = nn.Sequential(
+            nn.Linear(hidden_size * 3, hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size),
+        )
+        self.prefix_condition = nn.Linear(hidden_size, hidden_size)
+        self.action_projection = (
+            nn.Sequential(
+                nn.Linear(model_config.action_dim, hidden_size),
+                nn.SiLU(),
+                nn.Linear(hidden_size, hidden_size),
+            )
+            if model_config.action_dim > 0
+            else None
+        )
+        self.context_projection = (
+            nn.Linear(model_config.context_dim, hidden_size, bias=False)
+            if model_config.context_dim > 0
+            else None
+        )
+        self.blocks = nn.ModuleList(
+            PrefixDiTBlock(model_config) for _ in range(model_config.depth)
+        )
+        self.output_norm = nn.RMSNorm(
+            hidden_size,
+            eps=1e-6,
+            elementwise_affine=False,
+        )
+        self.output_modulation = nn.Linear(hidden_size, hidden_size * 2)
+        self.output_projection = nn.Linear(
+            hidden_size,
+            target_patch * target_patch * latent_channels,
+        )
+        self.outcome_norm = nn.RMSNorm(hidden_size, eps=1e-6)
+        self.reward_head = nn.Linear(hidden_size, 1)
+        self.terminal_head = nn.Linear(hidden_size, 1)
+        self.reset_parameters()
+
+    def enable_activation_checkpointing(self, enabled: bool = True) -> None:
+        self.gradient_checkpointing = enabled
+
+    def reset_parameters(self) -> None:
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+        for embedding in (
+            self.target_patch_embed,
+            self.coarse_prefix_embed,
+            self.fine_prefix_embed,
+            self.motion_prefix_embed,
+        ):
+            nn.init.normal_(embedding.weight, std=0.02)
+        nn.init.normal_(self.coarse_scale_embedding, std=0.02)
+        nn.init.normal_(self.fine_scale_embedding, std=0.02)
+        for block in self.blocks:
+            nn.init.zeros_(block.modulation.weight)
+            nn.init.zeros_(block.modulation.bias)
+        nn.init.zeros_(self.output_modulation.weight)
+        nn.init.zeros_(self.output_modulation.bias)
+        nn.init.zeros_(self.output_projection.weight)
+        nn.init.zeros_(self.output_projection.bias)
+        nn.init.zeros_(self.reward_head.weight)
+        nn.init.zeros_(self.reward_head.bias)
+        nn.init.zeros_(self.terminal_head.weight)
+        nn.init.zeros_(self.terminal_head.bias)
+
+    @staticmethod
+    def _patch(
+        embedding: nn.Conv3d,
+        video: torch.Tensor,
+        time_offset: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        tokens = embedding(video.transpose(1, 2))
+        _, _, frames, height, width = tokens.shape
+        positions = _video_positions(
+            frames,
+            height,
+            width,
+            time_offset,
+            video.device,
+        )
+        return tokens.flatten(2).transpose(1, 2), positions
+
+    def _encode_prefix(
+        self,
+        history: torch.Tensor,
+        context: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if history.shape[1] != self.history_frames:
+            raise ValueError(
+                f"Expected {self.history_frames} history frames, got {history.shape[1]}."
+            )
+        coarse, coarse_positions = self._patch(
+            self.coarse_prefix_embed,
+            history,
+            0,
+        )
+        fine_frames = self.model_config.fine_history_frames
+        fine_history = history[:, -fine_frames:]
+        previous = history[:, -fine_frames - 1 : -1]
+        if previous.shape[1] != fine_frames:
+            previous = torch.cat((fine_history[:, :1], fine_history[:, :-1]), dim=1)
+        motion = fine_history - previous
+        fine, fine_positions = self._patch(
+            self.fine_prefix_embed,
+            fine_history,
+            self.history_frames - fine_frames,
+        )
+        motion_tokens, _ = self._patch(
+            self.motion_prefix_embed,
+            motion,
+            self.history_frames - fine_frames,
+        )
+        coarse = coarse + self.coarse_scale_embedding
+        fine = fine + motion_tokens + self.fine_scale_embedding
+        prefix = torch.cat((coarse, fine), dim=1)
+        positions = torch.cat((coarse_positions, fine_positions), dim=0)
+        if self.context_projection is not None and context is not None:
+            projected = self.context_projection(context)
+            prefix = torch.cat((prefix, projected), dim=1)
+            positions = torch.cat(
+                (
+                    positions,
+                    torch.zeros(projected.shape[1], 3, device=history.device),
+                ),
+                dim=0,
+            )
+        return prefix, positions
+
+    def _condition(
+        self,
+        prefix: torch.Tensor,
+        r: torch.Tensor,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        time = self.time_projection(
+            torch.cat(
+                (
+                    self.t_embedding(t),
+                    self.r_embedding(r),
+                    self.delta_embedding(t - r),
+                ),
+                dim=-1,
+            )
+        )
+        return time + self.prefix_condition(prefix.mean(dim=1))
+
+    def build_memory_cache(
+        self,
+        history: torch.Tensor,
+        context: torch.Tensor | None = None,
+    ) -> PrefixCache:
+        prefix, positions = self._encode_prefix(history, context)
+        pooled = self.prefix_condition(prefix.mean(dim=1))
+        key_values: list[tuple[torch.Tensor, torch.Tensor]] = []
+        for block in self.blocks:
+            prefix = block.advance_prefix(prefix, positions)
+            key_values.append(block.prefix_key_value(prefix, positions))
+        return PrefixCache(tuple(key_values), pooled)
+
+    def predict_outcomes(
+        self,
+        history: torch.Tensor,
+        actions: torch.Tensor | None = None,
+        frame_count: int = 1,
+        memory_cache: PrefixCache | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if memory_cache is None:
+            prefix, _ = self._encode_prefix(history)
+            pooled = self.prefix_condition(prefix.mean(dim=1))
+        else:
+            pooled = memory_cache.pooled
+        hidden = pooled[:, None].expand(-1, frame_count, -1)
+        if self.action_projection is not None:
+            if actions is None or actions.shape[1] != frame_count:
+                raise ValueError("Outcome actions must match frame_count.")
+            hidden = hidden + self.action_projection(actions)
+        hidden = self.outcome_norm(hidden)
+        return self.reward_head(hidden).squeeze(-1), self.terminal_head(hidden).squeeze(-1)
+
+    def _unpatchify(
+        self,
+        patches: torch.Tensor,
+        frames: int,
+        height: int,
+        width: int,
+    ) -> torch.Tensor:
+        batch = patches.shape[0]
+        patch = self.model_config.target_patch_size
+        patches = patches.reshape(
+            batch,
+            frames,
+            height,
+            width,
+            patch,
+            patch,
+            self.latent_channels,
+        )
+        return patches.permute(0, 1, 6, 2, 4, 3, 5).reshape(
+            batch,
+            frames,
+            self.latent_channels,
+            height * patch,
+            width * patch,
+        )
+
+    def forward(
+        self,
+        noisy_target: torch.Tensor,
+        history: torch.Tensor,
+        r: torch.Tensor,
+        t: torch.Tensor,
+        actions: torch.Tensor | None = None,
+        context: torch.Tensor | None = None,
+        memory_cache: PrefixCache | None = None,
+    ) -> torch.Tensor:
+        target_volume = self.target_patch_embed(noisy_target.transpose(1, 2))
+        batch, _, frames, height, width = target_volume.shape
+        target = target_volume.flatten(2).transpose(1, 2)
+        target_positions = _video_positions(
+            frames,
+            height,
+            width,
+            self.history_frames,
+            noisy_target.device,
+        )
+        if memory_cache is None:
+            prefix, prefix_positions = self._encode_prefix(history, context)
+            condition = self._condition(prefix, r, t)
+        else:
+            prefix = None
+            prefix_positions = None
+            condition = self.time_projection(
+                torch.cat(
+                    (
+                        self.t_embedding(t),
+                        self.r_embedding(r),
+                        self.delta_embedding(t - r),
+                    ),
+                    dim=-1,
+                )
+            ) + memory_cache.pooled
+        if self.action_projection is not None:
+            if actions is None or actions.shape[:2] != (batch, frames):
+                raise ValueError("Actions must match [batch, target_frames, action_dim].")
+            action_embedding = self.action_projection(actions)
+            target = target + action_embedding.repeat_interleave(height * width, dim=1)
+            condition = condition + action_embedding.mean(dim=1)
+        for block_index, block in enumerate(self.blocks):
+            if memory_cache is None:
+                if prefix is None or prefix_positions is None:
+                    raise RuntimeError("Prefix state is unavailable.")
+                if self.training and self.gradient_checkpointing:
+                    prefix, target = checkpoint(
+                        block,
+                        prefix,
+                        target,
+                        condition,
+                        prefix_positions,
+                        target_positions,
+                        use_reentrant=False,
+                    )
+                else:
+                    prefix, target = block(
+                        prefix,
+                        target,
+                        condition,
+                        prefix_positions,
+                        target_positions,
+                    )
+            else:
+                target = block.advance_target(
+                    target,
+                    condition,
+                    target_positions,
+                    memory_cache.key_values[block_index],
+                )
+        shift, scale = self.output_modulation(F.silu(condition)).chunk(2, dim=-1)
+        target = self.output_norm(target) * (1.0 + scale[:, None]) + shift[:, None]
+        return self._unpatchify(
+            self.output_projection(target),
+            frames,
+            height,
+            width,
+        )
+
+
+CausalVideoDiT = CausalLatentVideoDiT
+@torch.inference_mode()
+def sample_next_latent(
+    model: nn.Module,
+    history: torch.Tensor,
+    steps: int = 2,
+    actions: torch.Tensor | None = None,
+    generator: torch.Generator | None = None,
+    window_frames: int = 1,
+    memory_cache: object | None = None,
+    instantaneous: bool = False,
+) -> torch.Tensor:
+    if steps < 1:
+        raise ValueError("steps must be positive.")
+    batch, _, channels, height, width = history.shape
+    state = torch.randn(
+        (batch, window_frames, channels, height, width),
+        device=history.device,
+        dtype=history.dtype,
+        generator=generator,
+    )
+    boundaries = torch.linspace(1.0, 0.0, steps + 1, device=history.device)
+    for index in range(steps):
+        t_value = boundaries[index]
+        r_value = boundaries[index + 1]
+        t = t_value.expand(batch)
+        r = r_value.expand(batch)
+        model_r = t if instantaneous else r
+        average_velocity = model(
+            state,
+            history,
+            model_r,
+            t,
+            actions=actions,
+            memory_cache=memory_cache,
+        )
+        state = state - (t_value - r_value).to(state.dtype) * average_velocity
+    return state
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Diffusion Breakout.")
+    parser.add_argument("--repo", default=DEFAULT_REPO)
+    parser.add_argument("--revision", default="main")
+    parser.add_argument("--local-dir", default=None)
+    parser.add_argument("--steps", type=int, default=1)
+    parser.add_argument("--frames", type=int, default=0, help="0 = until window closed")
+    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)
+    return parser.parse_args()
+
+
+def _download(repo: str, revision: str, local_dir: str | None) -> Path:
+    if local_dir:
+        root = Path(local_dir)
+        root.mkdir(parents=True, exist_ok=True)
+        for name in ("config.json", "ema.safetensors"):
+            hf_hub_download(
+                repo_id=repo,
+                filename=name,
+                revision=revision,
+                local_dir=str(root),
+            )
+        return root
+    config_path = Path(
+        hf_hub_download(repo_id=repo, filename="config.json", revision=revision)
+    )
+    hf_hub_download(repo_id=repo, filename="ema.safetensors", revision=revision)
+    return config_path.parent
+
+
+def _load_settings(config_path: Path) -> tuple[VideoConfig, ModelConfig, CodecSettings]:
+    raw = json.loads(config_path.read_text(encoding="utf-8"))
+    video = VideoConfig(**{key: raw["video"][key] for key in VideoConfig.__dataclass_fields__})
+    model = ModelConfig(**{key: raw["model"][key] for key in ModelConfig.__dataclass_fields__})
+    codec_raw = raw["codec"]
+    codec = CodecSettings(
+        model_id=codec_raw.get("model_id", "madebyollin/sdxl-vae-fp16-fix"),
+        subfolder=codec_raw.get("subfolder", "") or "",
+        revision=codec_raw.get("revision", "main"),
+        latent_channels=int(codec_raw.get("latent_channels", 4)),
+        spatial_compression=int(codec_raw.get("spatial_compression", 8)),
+        frame_batch_size=int(codec_raw.get("frame_batch_size", 256)),
+        sample_posterior=bool(codec_raw.get("sample_posterior", False)),
+        enable_slicing=bool(codec_raw.get("enable_slicing", True)),
+        enable_tiling=bool(codec_raw.get("enable_tiling", False)),
+    )
+    return video, model, codec
+
+
+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)
+    return Image.fromarray(array.byte().permute(1, 2, 0).cpu().numpy(), 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 _rgb_image_to_frame_tensor(image: Image.Image, device: torch.device) -> torch.Tensor:
+    array = np.asarray(image, dtype=np.uint8).copy()
+    return (
+        torch.from_numpy(array)
+        .permute(2, 0, 1)
+        .float()
+        .div_(127.5)
+        .sub_(1.0)
+        .to(device, non_blocking=True)
+    )
+
+
+def _inject_ball_into_history(
+    history: deque[torch.Tensor],
+    codec: FrameAutoencoderCodec,
+    position: tuple[float, float],
+) -> Image.Image:
+    history_tensor = torch.stack(tuple(history), dim=1)
+    device = history_tensor.device
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        decoded = codec.decode(history_tensor)
+    painted_images: list[Image.Image] = []
+    painted_frames: list[torch.Tensor] = []
+    for index in range(decoded.shape[1]):
+        image = _inject_ball(_tensor_to_image(decoded[0, index]), position)
+        painted_images.append(image)
+        painted_frames.append(_rgb_image_to_frame_tensor(image, device))
+    rgb = torch.stack(painted_frames, dim=0).unsqueeze(0)
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        latents = codec.encode(rgb)
+    history.clear()
+    history.extend(latents[:, index] for index in range(latents.shape[1]))
+    return painted_images[-1]
+
+
+def _encode_ball_frame(
+    image: Image.Image,
+    position: tuple[float, float],
+    codec: FrameAutoencoderCodec,
+    device: torch.device,
+) -> tuple[Image.Image, torch.Tensor]:
+    painted = _inject_ball(image, position)
+    rgb = _rgb_image_to_frame_tensor(painted, device).reshape(
+        1, 1, 3, painted.height, painted.width
+    )
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        latent = codec.encode(rgb)[:, 0]
+    return painted, latent
+
+
+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 Breakout")
+        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", "a", "d", "left", "right"}:
+            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]:
+        negative = bool(self.pressed & {"w", "a", "left"})
+        positive = bool(self.pressed & {"s", "d", "right"})
+        if negative and positive:
+            return (0.0, 0.0)
+        return (float(negative), float(positive))
+
+    def _on_click(self, event: tk.Event) -> None:
+        if self.generated_width <= 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 Breakout — 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(args.repo, args.revision, args.local_dir)
+    video_cfg, model_cfg, codec_cfg = _load_settings(checkpoint / "config.json")
+    weights = checkpoint / "ema.safetensors"
+    if not weights.exists():
+        raise SystemExit(f"Missing {weights}")
+
+    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 = BreakoutSimulation.from_seed(video_cfg, seed, at_start=True)
+    history_rgb = _frames_to_tensor(
+        simulation.sample_frames(video_cfg.history_frames),
+        device,
+    )
+
+    print("Loading VAE + DiT …")
+    codec = FrameAutoencoderCodec(codec_cfg, device, torch.bfloat16)
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        latents = codec.encode(history_rgb)
+    history: deque[torch.Tensor] = deque(
+        (latents[:, index] for index in range(video_cfg.history_frames)),
+        maxlen=video_cfg.history_frames,
+    )
+    model = (
+        CausalLatentVideoDiT(
+            model_cfg,
+            latent_channels=codec.channels,
+            history_frames=codec.latent_frame_count(video_cfg.history_frames),
+        )
+        .to(device)
+        .eval()
+    )
+    model.load_state_dict(load_file(str(weights)), strict=True)
+
+    steps = max(1, args.steps)
+    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)
+    injected_position: tuple[float, float] | None = None
+    force_ball_frames = 0
+    print("Ready. Click the window, then hold A/D or arrow keys.")
+
+    for frame_index in range(frame_limit):
+        click = window.consume_click()
+        if click is not None:
+            preview = _inject_ball_into_history(history, codec, click)
+            injected_position = click
+            force_ball_frames = video_cfg.history_frames
+            if not window.show(preview, frame_index, 0.0):
+                break
+        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 if model_cfg.action_dim > 0 else None,
+                generator=generator,
+                window_frames=1,
+                memory_cache=memory_cache,
+                instantaneous=True,
+            )
+            decoded = codec.decode(prediction)
+        torch.cuda.synchronize()
+        latency_ms = (time.perf_counter() - started) * 1000
+        image = _tensor_to_image(decoded[0, 0])
+        frame_latent = prediction[:, 0]
+        if injected_position is not None and force_ball_frames > 0:
+            image, frame_latent = _encode_ball_frame(
+                image, injected_position, codec, device
+            )
+            force_ball_frames -= 1
+            if force_ball_frames == 0:
+                injected_position = 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()