Add Diffusion Miner teacher checkpoint, Tile VAE, live infer, and gameplay GIF

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+demo.gif filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,77 @@
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: other
+tags:
+  - world-model
+  - diffusion
+  - mining
+  - video
+  - causal
+  - game
+---
+
+# Diffusion Miner
+
+An action-conditioned world model that learned a falling-sand mining sandbox. You hold **A / D** to walk, **W** or Space to jump, and the arrows to dig; it invents the next frame.
+
+![demo](demo.gif)
+
+256×256 at 12 FPS. The attention prefix is 12 latent frames; a ConvGRU world memory scans 48. About 220M parameters. The codec is a frozen Tile VAE trained on this game, not SDXL. Needs a CUDA GPU with BF16.
+
+This release is the teacher-stage instantaneous flow checkpoint. One step is the right setting — there is no distilled student.
+
+## Play it
+
+```bash
+pip install torch numpy pillow safetensors huggingface_hub
+
+hf download kerzgrr/diffusionminer --local-dir checkpoints/diffusionminer
+python checkpoints/diffusionminer/live_infer.py \
+  --local-dir checkpoints/diffusionminer \
+  --steps 1
+```
+
+If the weights are already in the current directory:
+
+```bash
+python live_infer.py --steps 1 --window-scale 3 --seed 2026
+```
+
+Click the window once so it has focus. Esc quits. Clicking the canvas drops a miner into the latent history if you want to poke it.
+
+Write a short gameplay GIF with the scripted training miner:
+
+```bash
+python live_infer.py \
+  --local-dir checkpoints/diffusionminer \
+  --steps 1 \
+  --seed-play \
+  --autoplay \
+  --frames 144 \
+  --output diffusionminer.gif
+```
+
+## Controls
+
+| key | action |
+|---|---|
+| `A` | move left |
+| `D` | move right |
+| `W` / Space | jump |
+| ← ↑ → ↓ | dig that way |
+
+Action vector is `[left, right, jump, dig_left, dig_right, dig_up, dig_down]` as floats in `{0,1}`. Opposite move keys cancel. Only one dig direction is active at a time.
+
+## What's in the files
+
+| file | what |
+|---|---|
+| `ema.safetensors` | playable teacher weights |
+| `codec.pt` | frozen Tile VAE |
+| `config.json` | video / model / codec settings used at train time |
+| `livediffusion/` | inference package used by `live_infer.py` |
+| `live_infer.py` | live play launcher |
+| `demo.gif` | scripted gameplay clip |
+
+If something breaks, it is probably a GPU that cannot do BF16, or running `live_infer.py` without the rest of this snapshot on `sys.path`.

codec.pt

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

config.json

@@ -0,0 +1,117 @@
+{
+  "codec": {
+    "checkpoint_path": "codec.pt",
+    "enable_slicing": false,
+    "enable_tiling": false,
+    "frame_batch_size": 64,
+    "kind": "tile_vae",
+    "latent_channels": 4,
+    "model_id": "",
+    "revision": "main",
+    "sample_posterior": false,
+    "spatial_compression": 8,
+    "subfolder": "",
+    "temporal_compression": 1
+  },
+  "data": {
+    "batch_size": 16,
+    "data_path": "",
+    "dataset_kind": "mining",
+    "frame_stride": 1,
+    "hf_repo_id": "",
+    "hf_revision": "main",
+    "latent_cache_path": "",
+    "latent_prefetch_batches": 6,
+    "live_monitor": 1,
+    "live_mouse_scale": 80.0,
+    "network_clip_prefetch": 2,
+    "num_workers": 12,
+    "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": 240,
+    "steps": 1
+  },
+  "model": {
+    "action_dim": 7,
+    "context_dim": 0,
+    "depth": 13,
+    "fine_history_frames": 4,
+    "fine_memory_patch_size": 2,
+    "generation_frames": 1,
+    "hidden_size": 768,
+    "kv_heads": 3,
+    "memory_channels": 128,
+    "memory_depth": 4,
+    "memory_patch_size": 4,
+    "mlp_ratio": 4.0,
+    "num_heads": 12,
+    "prefix_encoder_depth": 4,
+    "qk_norm": true,
+    "register_tokens": 16,
+    "shared_modulation": true,
+    "target_patch_size": 2
+  },
+  "schema_version": 4,
+  "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": 2,
+    "history_noise_std": 0.01,
+    "keep_checkpoints": 5,
+    "learning_rate": 0.00018,
+    "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/mining-220m-memory",
+    "reward_loss_weight": 0.5,
+    "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.2,
+    "validate_every": 50,
+    "validation_batches": 1,
+    "warmup_steps": 500,
+    "weight_decay": 0.0
+  },
+  "video": {
+    "fps": 12.0,
+    "future_frames": 5,
+    "height": 256,
+    "history_frames": 12,
+    "memory_frames": 48,
+    "width": 256
+  }
+}

ema.safetensors

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

live_infer.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+"""Launch Diffusion Miner from this repo or a Hugging Face snapshot."""
+
+from __future__ import annotations
+
+import runpy
+import sys
+from pathlib import Path
+
+
+def _prepare_path() -> None:
+    here = Path(__file__).resolve().parent
+    for candidate in (here, here.parent / "src"):
+        if (candidate / "livediffusion" / "infer_mining.py").exists():
+            sys.path.insert(0, str(candidate))
+            return
+    raise FileNotFoundError("Could not find the livediffusion package next to this script.")
+
+
+if __name__ == "__main__":
+    _prepare_path()
+    runpy.run_module("livediffusion.infer_mining", run_name="__main__")

livediffusion/__init__.py

@@ -0,0 +1,3 @@
+"""Diffusion Miner inference package."""
+
+__version__ = "0.3.0"

livediffusion/config.py

@@ -0,0 +1,356 @@
+from __future__ import annotations
+
+import json
+import tomllib
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+SCHEMA_VERSION = 4
+
+
+@dataclass(frozen=True, slots=True)
+class VideoConfig:
+    height: int = 64
+    width: int = 64
+    fps: float = 2.0
+    history_frames: int = 5
+    future_frames: int = 5
+    # Frames scanned into the recurrent world memory before the attention prefix.
+    # Zero keeps memory and attention on the same window.
+    memory_frames: int = 0
+
+    @property
+    def scanned_frames(self) -> int:
+        return max(self.memory_frames, self.history_frames)
+
+
+@dataclass(frozen=True, slots=True)
+class SimulationConfig:
+    min_shapes: int = 1
+    max_shapes: int = 4
+    min_size: float = 5.0
+    max_size: float = 11.0
+    min_speed: float = 12.0
+    max_speed: float = 32.0
+    supersample: int = 2
+    warmup_frames: int = 8
+
+
+@dataclass(frozen=True, slots=True)
+class CodecConfig:
+    kind: str = "frame_vae"
+    model_id: str = "madebyollin/sdxl-vae-fp16-fix"
+    checkpoint_path: str = ""
+    subfolder: str = ""
+    revision: str = "main"
+    latent_channels: int = 4
+    spatial_compression: int = 8
+    temporal_compression: int = 1
+    frame_batch_size: int = 32
+    sample_posterior: bool = False
+    enable_slicing: bool = True
+    enable_tiling: bool = True
+
+
+@dataclass(frozen=True, slots=True)
+class DataConfig:
+    dataset_kind: str = "simulation"
+    data_path: str = ""
+    latent_cache_path: str = ""
+    hf_repo_id: str = ""
+    hf_revision: str = "main"
+    shard_shuffle: int = 256
+    sample_shuffle: int = 64
+    windows_per_video: int = 4
+    network_clip_prefetch: int = 2
+    train_samples: int = 50_000
+    validation_samples: int = 1_024
+    batch_size: int = 16
+    num_workers: int = 8
+    prefetch_batches: int = 2
+    latent_prefetch_batches: int = 6
+    pin_memory: bool = True
+    persistent_workers: bool = True
+    frame_stride: int = 1
+    live_monitor: int = 1
+    live_mouse_scale: float = 80.0
+
+
+@dataclass(frozen=True, slots=True)
+class ModelConfig:
+    target_patch_size: int = 1
+    memory_patch_size: int = 1
+    fine_memory_patch_size: int = 1
+    fine_history_frames: int = 2
+    hidden_size: int = 384
+    depth: int = 12
+    memory_depth: int = 3
+    prefix_encoder_depth: int = 0
+    num_heads: int = 8
+    kv_heads: int = 0
+    mlp_ratio: float = 4.0
+    qk_norm: bool = True
+    register_tokens: int = 0
+    shared_modulation: bool = False
+    memory_channels: int = 0
+    generation_frames: int = 1
+    action_dim: int = 0
+    context_dim: int = 0
+
+
+@dataclass(frozen=True, slots=True)
+class TrainingConfig:
+    output_dir: str = "runs/latent-world-model"
+    seed: int = 1_337
+    epochs: int = 100
+    learning_rate: float = 2e-4
+    min_learning_rate: float = 1e-5
+    weight_decay: float = 0.0
+    warmup_steps: int = 1_000
+    teacher_steps: int = 20_000
+    teacher_integration_steps: int = 8
+    teacher_endpoint_ratio: float = 0.1
+    teacher_x0_weight: float = 0.25
+    teacher_logit_mean: float = 1.0
+    teacher_logit_std: float = 1.0
+    gradient_accumulation_steps: int = 1
+    max_grad_norm: float = 1.0
+    flow_matching_ratio: float = 0.25
+    full_interval_ratio: float = 0.25
+    on_policy_start_step: int = 30_000
+    on_policy_ratio: float = 0.15
+    on_policy_max_horizon: int = 4
+    on_policy_sampling_steps: int = 4
+    history_noise_std: float = 0.01
+    reward_loss_weight: float = 1.0
+    terminal_loss_weight: float = 0.1
+    distribution_matching_start_step: int = 60_000
+    distribution_matching_weight: float = 0.05
+    critic_learning_rate: float = 1e-4
+    ema_decay: float = 0.9999
+    log_every: int = 20
+    validate_every: int = 1_000
+    validation_batches: int = 8
+    save_every: int = 5_000
+    keep_checkpoints: int = 5
+    activation_checkpointing: bool = True
+
+
+@dataclass(frozen=True, slots=True)
+class InferenceConfig:
+    steps: int = 2
+    frames: int = 120
+
+
+@dataclass(frozen=True, slots=True)
+class ProjectConfig:
+    schema_version: int
+    video: VideoConfig
+    simulation: SimulationConfig
+    codec: CodecConfig
+    data: DataConfig
+    model: ModelConfig
+    training: TrainingConfig
+    inference: InferenceConfig
+
+    def as_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    def validate(self) -> None:
+        video = self.video
+        simulation = self.simulation
+        codec = self.codec
+        data = self.data
+        model = self.model
+        training = self.training
+        if self.schema_version != SCHEMA_VERSION:
+            raise ValueError(
+                f"Config schema {self.schema_version} is unsupported; expected {SCHEMA_VERSION}."
+            )
+        if video.history_frames < 1:
+            raise ValueError("At least one history frame is required.")
+        if video.height <= 0 or video.width <= 0 or video.fps <= 0:
+            raise ValueError("Video dimensions and FPS must be positive.")
+        latent_future_frames = (
+            (video.future_frames - 1) // codec.temporal_compression + 1
+        )
+        if (
+            latent_future_frames
+            < training.on_policy_max_horizon + model.generation_frames
+        ):
+            raise ValueError(
+                "Future latent frames must cover on-policy horizon and generation window."
+            )
+        if simulation.min_shapes < 1 or simulation.max_shapes < simulation.min_shapes:
+            raise ValueError("Invalid shape-count range.")
+        if not 0 < simulation.min_size <= simulation.max_size < min(video.height, video.width) / 2:
+            raise ValueError("Shape sizes must fit inside the frame.")
+        if simulation.min_speed <= 0 or simulation.max_speed < simulation.min_speed:
+            raise ValueError("Invalid speed range.")
+        if codec.kind not in {"frame_vae", "ltx2_video", "tile_vae", "identity"}:
+            raise ValueError(
+                "codec.kind must be 'frame_vae', 'ltx2_video', 'tile_vae', "
+                "or 'identity'."
+            )
+        if codec.kind == "tile_vae" and not codec.checkpoint_path:
+            raise ValueError("codec.checkpoint_path is required for kind='tile_vae'.")
+        if codec.spatial_compression < 1 or codec.temporal_compression < 1:
+            raise ValueError("Codec compression ratios must be positive.")
+        if video.height % codec.spatial_compression or video.width % codec.spatial_compression:
+            raise ValueError("Video dimensions must be divisible by codec spatial compression.")
+        latent_height = video.height // codec.spatial_compression
+        latent_width = video.width // codec.spatial_compression
+        for patch_size in (model.target_patch_size, model.memory_patch_size):
+            if latent_height % patch_size or latent_width % patch_size:
+                raise ValueError(f"Latent dimensions must be divisible by patch size {patch_size}.")
+        if (
+            latent_height % model.fine_memory_patch_size
+            or latent_width % model.fine_memory_patch_size
+        ):
+            raise ValueError("Latent dimensions must be divisible by fine memory patch size.")
+        if not 1 <= model.fine_history_frames <= video.history_frames:
+            raise ValueError("fine_history_frames must fit inside the history window.")
+        if model.hidden_size % model.num_heads:
+            raise ValueError("Hidden size must be divisible by the number of heads.")
+        if model.kv_heads < 0:
+            raise ValueError("kv_heads cannot be negative.")
+        if model.kv_heads and model.num_heads % model.kv_heads:
+            raise ValueError("num_heads must be divisible by kv_heads.")
+        if model.memory_channels < 0:
+            raise ValueError("memory_channels cannot be negative.")
+        if video.memory_frames and video.memory_frames < video.history_frames:
+            raise ValueError("memory_frames must be at least history_frames.")
+        if video.memory_frames and not model.memory_channels:
+            raise ValueError("memory_frames requires memory_channels to be positive.")
+        if model.generation_frames < 1:
+            raise ValueError("generation_frames must be positive.")
+        if model.action_dim > 0 and model.generation_frames != 1:
+            raise ValueError(
+                "Interactive action-conditioned models must generate one frame per step."
+            )
+        if (model.hidden_size // model.num_heads) < 12:
+            raise ValueError("Attention head size is too small for 3D rotary embeddings.")
+        if model.memory_depth < 0:
+            raise ValueError("memory_depth cannot be negative.")
+        if model.prefix_encoder_depth < 0:
+            raise ValueError("prefix_encoder_depth cannot be negative.")
+        if model.register_tokens < 0:
+            raise ValueError("register_tokens cannot be negative.")
+        if data.dataset_kind not in {
+            "simulation",
+            "pong",
+            "breakout",
+            "asteroids",
+            "racing",
+            "tank_duel",
+            "mining",
+            "frame_folder",
+            "video_folder",
+            "latent_cache",
+            "hf_webdataset",
+            "live_screen",
+            "minecraft_vpt",
+        }:
+            raise ValueError(f"Unsupported dataset kind: {data.dataset_kind}.")
+        if (
+            data.dataset_kind
+            not in {
+                "simulation",
+                "pong",
+                "breakout",
+                "asteroids",
+                "racing",
+                "tank_duel",
+                "mining",
+                "hf_webdataset",
+                "live_screen",
+            }
+            and not data.data_path
+            and not data.latent_cache_path
+        ):
+            raise ValueError("A data or latent-cache path is required for non-simulation datasets.")
+        if data.dataset_kind == "minecraft_vpt" and not data.data_path:
+            raise ValueError("data_path is required for minecraft_vpt ArrayRecord shards.")
+        if data.dataset_kind == "hf_webdataset" and not data.hf_repo_id:
+            raise ValueError("hf_repo_id is required for a Hugging Face WebDataset.")
+        if data.dataset_kind == "live_screen":
+            if data.num_workers != 0:
+                raise ValueError("live_screen requires num_workers = 0.")
+            if data.persistent_workers:
+                raise ValueError("live_screen requires persistent_workers = false.")
+            if data.live_monitor < 0:
+                raise ValueError("live_monitor cannot be negative.")
+            if data.live_mouse_scale <= 0:
+                raise ValueError("live_mouse_scale must be positive.")
+        if data.windows_per_video < 0:
+            raise ValueError("windows_per_video cannot be negative.")
+        if data.network_clip_prefetch < 1:
+            raise ValueError("network_clip_prefetch must be positive.")
+        if data.batch_size < 1 or data.prefetch_batches != 2:
+            raise ValueError(
+                "Batch size must be positive and prefetch_batches must be exactly two."
+            )
+        if data.latent_prefetch_batches < 1:
+            raise ValueError("latent_prefetch_batches must be positive.")
+        if data.frame_stride < 1:
+            raise ValueError("frame_stride must be positive.")
+        if data.num_workers == 0 and data.persistent_workers:
+            raise ValueError("persistent_workers requires at least one worker.")
+        if training.teacher_steps < training.warmup_steps:
+            raise ValueError("teacher_steps must be at least warmup_steps.")
+        if training.teacher_integration_steps < 2:
+            raise ValueError("Teacher integration requires at least two steps.")
+        if not 0 <= training.teacher_endpoint_ratio <= 1:
+            raise ValueError("teacher_endpoint_ratio must be in [0, 1].")
+        if training.teacher_logit_std <= 0:
+            raise ValueError("teacher_logit_std must be positive.")
+        if training.flow_matching_ratio + training.full_interval_ratio > 1:
+            raise ValueError("Flow-matching and full-interval ratios cannot exceed one.")
+        if not 0 <= training.on_policy_ratio <= 1:
+            raise ValueError("on_policy_ratio must be in [0, 1].")
+        if training.gradient_accumulation_steps < 1:
+            raise ValueError("gradient_accumulation_steps must be positive.")
+        if self.inference.steps < 1:
+            raise ValueError("Inference steps must be positive.")
+
+
+def _config_from_raw(raw: dict[str, Any]) -> ProjectConfig:
+    sections = {
+        "schema_version",
+        "video",
+        "simulation",
+        "codec",
+        "data",
+        "model",
+        "training",
+        "inference",
+    }
+    unknown = set(raw) - sections
+    missing = sections - set(raw)
+    if unknown or missing:
+        raise ValueError(
+            f"Invalid config sections; missing={sorted(missing)}, unknown={sorted(unknown)}"
+        )
+    config = ProjectConfig(
+        schema_version=int(raw["schema_version"]),
+        video=VideoConfig(**raw["video"]),
+        simulation=SimulationConfig(**raw["simulation"]),
+        codec=CodecConfig(**raw["codec"]),
+        data=DataConfig(**raw["data"]),
+        model=ModelConfig(**raw["model"]),
+        training=TrainingConfig(**raw["training"]),
+        inference=InferenceConfig(**raw["inference"]),
+    )
+    config.validate()
+    return config
+
+
+def load_config(path: str | Path) -> ProjectConfig:
+    config_path = Path(path)
+    if config_path.suffix.lower() == ".json":
+        raw = json.loads(config_path.read_text(encoding="utf-8"))
+        return _config_from_raw(raw)
+    with config_path.open("rb") as file:
+        raw = tomllib.load(file)
+    return _config_from_raw(raw)

livediffusion/infer_mining.py

@@ -0,0 +1,526 @@
+#!/usr/bin/env python3
+"""Play Diffusion Miner — an action-conditioned mining world model.
+
+From this repo:
+
+  python3 -m livediffusion.infer_mining --checkpoint runs/mining-220m-memory --live
+
+From a Hugging Face snapshot:
+
+  python live_infer.py
+"""
+
+from __future__ import annotations
+
+import argparse
+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
+from PIL import Image, ImageOps, ImageTk
+from safetensors.torch import load_file
+
+from livediffusion.config import ProjectConfig, load_config
+from livediffusion.meanflow import sample_next_latent
+from livediffusion.mining import MiningSimulation, paint_player
+from livediffusion.model import CausalLatentVideoDiT
+from livediffusion.tile_vae import TileVAE
+
+DEFAULT_REPO = "kerzgrr/diffusionminer"
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Diffusion Miner live world model.")
+    parser.add_argument("--repo", default=DEFAULT_REPO)
+    parser.add_argument("--revision", default="main")
+    parser.add_argument("--local-dir", default="")
+    parser.add_argument("--checkpoint", default="")
+    parser.add_argument("--config", default="")
+    parser.add_argument("--codec", default="")
+    parser.add_argument("--steps", type=int, default=1)
+    parser.add_argument(
+        "--frames",
+        type=int,
+        default=0,
+        help="Frames to generate. 0 with a live window runs until quit.",
+    )
+    parser.add_argument("--seed", type=int, default=2026)
+    parser.add_argument("--window-scale", type=int, default=3)
+    parser.add_argument("--gif-scale", type=int, default=2)
+    parser.add_argument("--init", default="", help="Optional photo to seed history.")
+    parser.add_argument("--output", default="", help="Write a GIF instead of a window.")
+    parser.add_argument(
+        "--autoplay",
+        action="store_true",
+        help="Drive the model with the scripted miner used in training.",
+    )
+    parser.add_argument(
+        "--seed-play",
+        action="store_true",
+        help="Fill history with scripted gameplay instead of a still spawn.",
+    )
+    parser.add_argument(
+        "--live",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+    )
+    return parser.parse_args()
+
+
+def _resolve_checkpoint_dir(path: Path) -> Path:
+    if path.is_file():
+        return path.parent
+    latest = path / "latest.txt"
+    if (path / "ema.safetensors").exists():
+        return path
+    if latest.exists():
+        checkpoint = path / latest.read_text(encoding="utf-8").strip()
+        if (checkpoint / "ema.safetensors").exists():
+            return checkpoint
+    raise FileNotFoundError(f"No ema.safetensors under {path}.")
+
+
+def _resolve_codec(config_path: Path, config: ProjectConfig, codec_arg: str) -> Path:
+    if codec_arg:
+        return Path(codec_arg)
+    path = Path(config.codec.checkpoint_path)
+    if path.exists():
+        return path
+    sibling = config_path.parent / path.name
+    if sibling.exists():
+        return sibling
+    raise FileNotFoundError(
+        f"Tile VAE checkpoint {path} not found. Pass --codec or place codec.pt "
+        f"next to {config_path}."
+    )
+
+
+def _discover_local_bundle() -> Path | None:
+    here = Path(__file__).resolve()
+    for candidate in (Path.cwd(), here.parent, here.parents[1]):
+        if (candidate / "ema.safetensors").exists() and (candidate / "codec.pt").exists():
+            return candidate
+    return None
+
+
+def _resolve_assets(args: argparse.Namespace) -> tuple[Path, Path, Path]:
+    if not args.checkpoint and not args.local_dir:
+        discovered = _discover_local_bundle()
+        if discovered is not None:
+            args.local_dir = str(discovered)
+    if args.checkpoint:
+        checkpoint = _resolve_checkpoint_dir(Path(args.checkpoint))
+        weights = checkpoint / "ema.safetensors"
+        config_path = Path(args.config) if args.config else checkpoint / "config.json"
+        codec_path = _resolve_codec(config_path, load_config(config_path), args.codec)
+        return weights, config_path, codec_path
+    root = Path(args.local_dir) if args.local_dir else None
+    names = ("ema.safetensors", "config.json", "codec.pt")
+    if root is not None and all((root / name).exists() for name in names):
+        return root / "ema.safetensors", root / "config.json", root / "codec.pt"
+    downloaded = [
+        Path(
+            hf_hub_download(
+                args.repo,
+                filename=name,
+                revision=args.revision,
+                local_dir=root,
+            )
+        )
+        for name in names
+    ]
+    return downloaded[0], downloaded[1], downloaded[2]
+
+
+def _load_photo(path: Path, width: int, height: int) -> np.ndarray:
+    with Image.open(path) as raw:
+        image = ImageOps.exif_transpose(raw.convert("RGB"))
+    fitted = ImageOps.fit(image, (width, height), method=Image.Resampling.LANCZOS)
+    return np.asarray(fitted, dtype=np.uint8)
+
+
+def _frames_to_tensor(frames: np.ndarray, device: torch.device) -> torch.Tensor:
+    tensor = torch.from_numpy(np.ascontiguousarray(frames)).permute(0, 3, 1, 2).float()
+    return tensor.div_(127.5).sub_(1.0).unsqueeze(0).to(device)
+
+
+def _tensor_to_image(frame: torch.Tensor) -> Image.Image:
+    array = frame.float().clamp(-1.0, 1.0).add(1.0).mul(127.5).byte()
+    return Image.fromarray(array.permute(1, 2, 0).cpu().numpy(), mode="RGB")
+
+
+def _image_to_frame(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)
+    )
+
+
+class LiveWindow:
+    def __init__(self, scale: int) -> None:
+        self.scale = max(1, scale)
+        self.is_open = True
+        self.pressed: set[str] = set()
+        self.pending_click: tuple[float, float] | None = None
+        self.generated_width = 0
+        self.generated_height = 0
+        self.root = tk.Tk()
+        self.root.title("Diffusion Miner")
+        self.label = tk.Label(self.root, background="black")
+        self.label.pack()
+        self.photo: ImageTk.PhotoImage | None = None
+        self.root.bind("<KeyPress>", self._on_press)
+        self.root.bind("<KeyRelease>", self._on_release)
+        self.root.bind("<Button-1>", self._on_click)
+        self.root.protocol("WM_DELETE_WINDOW", self.close)
+        self.root.focus_force()
+
+    def _on_press(self, event: tk.Event) -> None:
+        key = str(event.keysym).lower()
+        if key == "escape":
+            self.close()
+            return
+        if key in {"w", "s", "a", "d", "left", "right", "up", "down", "space"}:
+            self.pressed.add(key)
+
+    def _on_release(self, event: tk.Event) -> None:
+        self.pressed.discard(str(event.keysym).lower())
+
+    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 mining_action(self) -> tuple[float, ...]:
+        left = "a" in self.pressed
+        right = "d" in self.pressed
+        if left and right:
+            left = right = False
+        dig_left = "left" in self.pressed
+        dig_right = "right" in self.pressed and not dig_left
+        dig_up = "up" in self.pressed and not (dig_left or dig_right)
+        dig_down = "down" in self.pressed and not (dig_left or dig_right or dig_up)
+        jump = bool(self.pressed & {"w", "space"})
+        return (
+            float(left),
+            float(right),
+            float(jump),
+            float(dig_left),
+            float(dig_right),
+            float(dig_up),
+            float(dig_down),
+        )
+
+    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()
+        extra = f"  held {held}" if held else ""
+        self.root.title(
+            f"Diffusion Miner — frame {frame_index + 1}  {latency_ms:.0f} ms{extra}"
+        )
+        try:
+            self.root.update_idletasks()
+            self.root.update()
+        except tk.TclError:
+            self.is_open = False
+        return self.is_open
+
+
+def _encode_frames(vae: TileVAE, frames: torch.Tensor, batch_size: int) -> torch.Tensor:
+    batch, count, channels, height, width = frames.shape
+    flat = frames.reshape(batch * count, channels, height, width)
+    encoded = [
+        vae.encode(chunk.contiguous(memory_format=torch.channels_last), sample=False)
+        for chunk in flat.split(batch_size)
+    ]
+    latents = torch.cat(encoded)
+    return latents.reshape(
+        batch,
+        count,
+        vae.latent_channels,
+        height // vae.spatial_compression,
+        width // vae.spatial_compression,
+    )
+
+
+def _decode_latents(vae: TileVAE, latents: torch.Tensor, batch_size: int) -> torch.Tensor:
+    batch, count, channels, height, width = latents.shape
+    flat = latents.reshape(batch * count, channels, height, width)
+    decoded = [
+        vae.decode(chunk.contiguous(memory_format=torch.channels_last))
+        for chunk in flat.split(batch_size)
+    ]
+    return torch.cat(decoded).reshape(
+        batch,
+        count,
+        3,
+        height * vae.spatial_compression,
+        width * vae.spatial_compression,
+    ).clamp(-1.0, 1.0)
+
+
+def _paint_player(image: Image.Image, position: tuple[float, float]) -> Image.Image:
+    frame = np.asarray(image, dtype=np.uint8).copy()
+    tile = max(4, min(frame.shape[0], frame.shape[1]) // 32)
+    paint_player(
+        frame,
+        round(position[0] - tile / 2),
+        round(position[1] - tile),
+        tile,
+        tile * 2,
+    )
+    return Image.fromarray(frame, mode="RGB")
+
+
+def _inject_player(
+    history: deque[torch.Tensor],
+    vae: TileVAE,
+    batch_size: int,
+    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 = _decode_latents(vae, history_tensor, batch_size)
+    painted: list[torch.Tensor] = []
+    preview = None
+    for index in range(decoded.shape[1]):
+        image = _paint_player(_tensor_to_image(decoded[0, index]), position)
+        preview = image
+        painted.append(_image_to_frame(image, device))
+    rgb = torch.stack(painted, dim=0).unsqueeze(0)
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        latents = _encode_frames(vae, rgb, batch_size)
+    history.clear()
+    history.extend(latents[:, index] for index in range(latents.shape[1]))
+    if preview is None:
+        raise RuntimeError("History was empty.")
+    return preview
+
+
+def _seed_history(
+    args: argparse.Namespace,
+    config: ProjectConfig,
+    vae: TileVAE,
+    device: torch.device,
+) -> tuple[deque[torch.Tensor], deque[torch.Tensor], MiningSimulation]:
+    scanned = config.video.scanned_frames
+    simulation = MiningSimulation.from_seed(config.video, args.seed, at_start=True)
+    actions = np.zeros((scanned, MiningSimulation.ACTION_DIM), dtype=np.float32)
+    if args.init:
+        frame = _load_photo(Path(args.init), config.video.width, config.video.height)
+        rgb = np.repeat(frame[None], scanned, axis=0)
+    elif args.seed_play:
+        rgb, actions, _, _ = simulation.sample_gameplay(scanned)
+    else:
+        rgb = np.repeat(simulation.render()[None], scanned, axis=0)
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        history_tensor = _encode_frames(
+            vae,
+            _frames_to_tensor(rgb, device),
+            config.codec.frame_batch_size,
+        )
+    history = deque(
+        (history_tensor[:, index] for index in range(history_tensor.shape[1])),
+        maxlen=history_tensor.shape[1],
+    )
+    action_history = deque(
+        (
+            torch.as_tensor(actions[index], device=device, dtype=torch.float32).unsqueeze(0)
+            for index in range(scanned)
+        ),
+        maxlen=scanned,
+    )
+    return history, action_history, simulation
+
+
+def _next_action(
+    args: argparse.Namespace,
+    simulation: MiningSimulation,
+    window: LiveWindow | None,
+    device: torch.device,
+) -> torch.Tensor:
+    if args.autoplay:
+        action = simulation.scripted_action()
+        simulation.step(action)
+        values = action
+    elif window is not None:
+        values = window.mining_action()
+    else:
+        values = (0.0,) * MiningSimulation.ACTION_DIM
+    return torch.tensor(values, device=device, dtype=torch.float32).reshape(
+        1,
+        1,
+        MiningSimulation.ACTION_DIM,
+    )
+
+
+@torch.inference_mode()
+def main() -> None:
+    args = _parse_args()
+    live = not bool(args.output) if args.live is None else args.live
+    if args.frames == 0:
+        if not live:
+            raise ValueError("--frames 0 requires a live window.")
+        frame_count = 10**9
+    else:
+        frame_count = args.frames
+    if args.steps < 1:
+        raise ValueError("steps must be positive.")
+    if args.gif_scale < 1:
+        raise ValueError("gif-scale must be positive.")
+    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
+        raise RuntimeError("Diffusion Miner needs a BF16 CUDA GPU.")
+
+    weights_path, config_path, codec_path = _resolve_assets(args)
+    config = load_config(config_path)
+    device = torch.device("cuda")
+    torch.set_float32_matmul_precision("high")
+
+    vae = TileVAE.load(codec_path, device=device, dtype=torch.bfloat16)
+    vae.requires_grad_(False)
+    vae.to(memory_format=torch.channels_last)
+    model = CausalLatentVideoDiT(
+        config.model,
+        latent_channels=vae.latent_channels,
+        history_frames=config.video.history_frames,
+    ).to(device).eval()
+    model.load_state_dict(load_file(weights_path, device=str(device)), strict=True)
+
+    history, action_history, simulation = _seed_history(args, config, vae, device)
+    window = LiveWindow(args.window_scale) if live else None
+    images: list[Image.Image] = []
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+    print(
+        "Diffusion Miner ready. A/D move, W/Space jump, arrows dig, click to "
+        "drop a miner, Esc quit."
+        if live
+        else f"Writing {frame_count} frames to {args.output}."
+    )
+    frame_index = 0
+    force_sprite_frames = 0
+    injected: tuple[float, float] | None = None
+    while frame_index < frame_count:
+        if window is not None and not window.is_open:
+            break
+        click = window.consume_click() if window is not None else None
+        if click is not None:
+            preview = _inject_player(
+                history,
+                vae,
+                config.codec.frame_batch_size,
+                click,
+            )
+            injected = click
+            force_sprite_frames = len(history)
+            if not window.show(preview, max(frame_index - 1, 0), 0.0):
+                break
+        action = _next_action(args, simulation, window, device)
+        history_tensor = torch.stack(tuple(history), dim=1)
+        history_actions = torch.stack(tuple(action_history), dim=1)
+        torch.cuda.synchronize()
+        started = time.perf_counter()
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            cache = model.build_memory_cache(
+                history_tensor,
+                history_actions=history_actions,
+            )
+            latent = sample_next_latent(
+                model,
+                history_tensor,
+                steps=args.steps,
+                actions=action,
+                generator=generator,
+                memory_cache=cache,
+                instantaneous=True,
+                history_actions=history_actions,
+            )
+            frame = _decode_latents(
+                vae,
+                latent,
+                config.codec.frame_batch_size,
+            )[0, 0]
+        torch.cuda.synchronize()
+        latency_ms = (time.perf_counter() - started) * 1_000
+        image = _tensor_to_image(frame)
+        frame_latent = latent[:, 0]
+        if injected is not None and force_sprite_frames > 0:
+            painted = _paint_player(image, injected)
+            rgb = _image_to_frame(painted, device).reshape(
+                1,
+                1,
+                3,
+                painted.height,
+                painted.width,
+            )
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                frame_latent = _encode_frames(
+                    vae,
+                    rgb,
+                    config.codec.frame_batch_size,
+                )[:, 0]
+            image = painted
+            force_sprite_frames -= 1
+            if force_sprite_frames == 0:
+                injected = None
+        history.append(frame_latent)
+        action_history.append(action[:, 0])
+        if window is not None and not window.show(image, frame_index, latency_ms):
+            break
+        if args.output:
+            if args.gif_scale > 1:
+                image = image.resize(
+                    (image.width * args.gif_scale, image.height * args.gif_scale),
+                    Image.Resampling.NEAREST,
+                )
+            images.append(image)
+        frame_index += 1
+    if window is not None:
+        window.close()
+    if args.output:
+        if not images:
+            raise RuntimeError("No frames generated.")
+        output = Path(args.output)
+        output.parent.mkdir(parents=True, exist_ok=True)
+        images[0].save(
+            output,
+            save_all=True,
+            append_images=images[1:],
+            duration=round(1_000 / config.video.fps),
+            loop=0,
+            optimize=False,
+        )
+        print(f"Saved {len(images)} frames to {output}")
+
+
+if __name__ == "__main__":
+    main()

livediffusion/meanflow.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+
+@dataclass(slots=True)
+class FlowLoss:
+    loss: torch.Tensor
+    raw_mse: torch.Tensor
+    x0_mse: torch.Tensor
+    prediction_rms: torch.Tensor
+    target_rms: torch.Tensor
+    mean_interval: torch.Tensor
+
+
+def _adaptive_loss(
+    prediction: torch.Tensor,
+    target: torch.Tensor,
+    interval: torch.Tensor,
+    gamma: float = 0.5,
+    epsilon: float = 1e-3,
+) -> FlowLoss:
+    squared_error = (prediction.float() - target.float()).square()
+    per_sample_mse = squared_error.flatten(1).mean(dim=1)
+    weights = (per_sample_mse + epsilon).pow(gamma - 1.0).detach()
+    return FlowLoss(
+        loss=(weights * per_sample_mse).mean(),
+        raw_mse=per_sample_mse.mean().detach(),
+        x0_mse=torch.zeros((), device=prediction.device),
+        prediction_rms=prediction.float().square().mean().sqrt().detach(),
+        target_rms=target.float().square().mean().sqrt().detach(),
+        mean_interval=interval.float().mean().detach(),
+    )
+
+
+def sample_time_intervals(
+    batch_size: int,
+    device: torch.device,
+    flow_matching_ratio: float,
+    full_interval_ratio: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    t = torch.rand(batch_size, device=device, dtype=torch.float32)
+    r = torch.rand(batch_size, device=device, dtype=torch.float32) * t
+    selector = torch.rand(batch_size, device=device)
+    full_mask = selector < full_interval_ratio
+    flow_mask = (selector >= full_interval_ratio) & (
+        selector < full_interval_ratio + flow_matching_ratio
+    )
+    t = torch.where(full_mask, torch.ones_like(t), t)
+    r = torch.where(full_mask, torch.zeros_like(r), r)
+    r = torch.where(flow_mask, t, r)
+    return r, t
+
+
+def flow_matching_loss(
+    model: nn.Module,
+    clean_target: torch.Tensor,
+    history: torch.Tensor,
+    actions: torch.Tensor | None = None,
+    history_noise_std: float = 0.0,
+    endpoint_ratio: float = 0.0,
+    history_actions: torch.Tensor | None = None,
+    x0_weight: float = 0.0,
+    logit_mean: float = 0.0,
+    logit_std: float = 1.0,
+) -> FlowLoss:
+    batch_size = clean_target.shape[0]
+    logits = (
+        torch.randn(batch_size, device=clean_target.device, dtype=torch.float32)
+        * logit_std
+        + logit_mean
+    )
+    t = logits.sigmoid()
+    endpoint_mask = torch.rand(batch_size, device=clean_target.device) < endpoint_ratio
+    t = torch.where(endpoint_mask, torch.ones_like(t), t)
+    noise = torch.randn_like(clean_target)
+    velocity = noise - clean_target
+    broadcast_t = t.reshape(batch_size, *([1] * (clean_target.ndim - 1)))
+    noisy_target = clean_target + broadcast_t.to(clean_target.dtype) * velocity
+    if history_noise_std > 0:
+        history = history + torch.randn_like(history) * history_noise_std
+    prediction = model(
+        noisy_target,
+        history,
+        t,
+        t,
+        actions=actions,
+        history_actions=history_actions,
+    )
+    result = _adaptive_loss(
+        prediction,
+        velocity,
+        torch.zeros_like(t),
+        gamma=1.0,
+    )
+    predicted_x0 = noisy_target - broadcast_t.to(prediction.dtype) * prediction
+    x0_mse = (
+        predicted_x0.float() - clean_target.float()
+    ).square().flatten(1).mean(dim=1).mean()
+    result.loss = result.loss + x0_weight * x0_mse
+    result.x0_mse = x0_mse.detach()
+    return result
+
+
+@torch.no_grad()
+def _teacher_average_velocity(
+    teacher: nn.Module,
+    noisy_target: torch.Tensor,
+    history: torch.Tensor,
+    r: torch.Tensor,
+    t: torch.Tensor,
+    integration_steps: int,
+    actions: torch.Tensor | None,
+    history_actions: torch.Tensor | None = None,
+) -> torch.Tensor:
+    state = noisy_target
+    interval = t - r
+    step_size = interval / integration_steps
+    broadcast_step = step_size.reshape(
+        step_size.shape[0],
+        *([1] * (noisy_target.ndim - 1)),
+    ).to(noisy_target.dtype)
+    for index in range(integration_steps):
+        current_t = t - step_size * index
+        velocity = teacher(
+            state,
+            history,
+            current_t,
+            current_t,
+            actions=actions,
+            history_actions=history_actions,
+        )
+        state = state - broadcast_step * velocity
+    broadcast_interval = interval.reshape(
+        interval.shape[0],
+        *([1] * (noisy_target.ndim - 1)),
+    ).to(noisy_target.dtype)
+    return (noisy_target - state) / broadcast_interval.clamp_min(1e-5)
+
+
+def flow_map_distillation_loss(
+    student: nn.Module,
+    teacher: nn.Module,
+    clean_target: torch.Tensor,
+    history: torch.Tensor,
+    flow_matching_ratio: float,
+    full_interval_ratio: float,
+    integration_steps: int,
+    actions: torch.Tensor | None = None,
+    history_noise_std: float = 0.0,
+    history_actions: torch.Tensor | None = None,
+) -> FlowLoss:
+    batch_size = clean_target.shape[0]
+    r, t = sample_time_intervals(
+        batch_size,
+        clean_target.device,
+        flow_matching_ratio,
+        full_interval_ratio,
+    )
+    noise = torch.randn_like(clean_target)
+    velocity = noise - clean_target
+    broadcast_t = t.reshape(batch_size, *([1] * (clean_target.ndim - 1)))
+    noisy_target = clean_target + broadcast_t.to(clean_target.dtype) * velocity
+    if history_noise_std > 0:
+        history = history + torch.randn_like(history) * history_noise_std
+    interval = t - r
+    teacher_target = _teacher_average_velocity(
+        teacher,
+        noisy_target,
+        history,
+        r,
+        t,
+        integration_steps,
+        actions,
+        history_actions,
+    )
+    flow_mask = interval <= 1e-6
+    broadcast_mask = flow_mask.reshape(
+        batch_size,
+        *([1] * (clean_target.ndim - 1)),
+    )
+    target = torch.where(broadcast_mask, velocity, teacher_target).detach()
+    prediction = student(
+        noisy_target,
+        history,
+        r,
+        t,
+        actions=actions,
+        history_actions=history_actions,
+    )
+    return _adaptive_loss(prediction, target, interval)
+
+
+@torch.no_grad()
+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,
+    history_actions: torch.Tensor | None = None,
+) -> 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,
+            history_actions=history_actions,
+        )
+        state = state - (t_value - r_value).to(state.dtype) * average_velocity
+    return state
+
+
+sample_next_frame = sample_next_latent

livediffusion/mining.py

@@ -0,0 +1,644 @@
+from __future__ import annotations
+
+import numpy as np
+
+from livediffusion.config import VideoConfig
+
+AIR = 0
+DIRT = 1
+STONE = 2
+SAND = 3
+GRAVEL = 4
+WATER = 5
+ORE = 6
+BEDROCK = 7
+
+TILE_COUNT = 8
+
+# Flat, high-contrast palette. No thin lines anywhere, so an 8x VAE keeps the edges.
+TILE_COLORS = np.array(
+    [
+        (16, 18, 30),
+        (126, 84, 48),
+        (108, 112, 122),
+        (214, 182, 96),
+        (52, 54, 62),
+        (44, 104, 216),
+        (72, 240, 236),
+        (74, 62, 102),
+    ],
+    dtype=np.int16,
+)
+# Every tile hue is separable from every other one and from the player, so a
+# reconstruction error is never ambiguous about which material it was.
+_DIG_FLASH_COLOR = (255, 84, 240)
+_PLAYER_HEAD_COLOR = (248, 252, 255)
+_PLAYER_BODY_COLOR = (232, 96, 72)
+
+# Water is passable; everything else with a body blocks the player.
+TILE_SOLID = np.array([0, 1, 1, 1, 1, 0, 1, 1], dtype=bool)
+# Sand and gravel fall; bedrock never moves; ore and stone stay until mined.
+TILE_FALLS = np.array([0, 0, 0, 1, 1, 0, 0, 0], dtype=bool)
+TILE_DIGGABLE = np.array([0, 1, 1, 1, 1, 0, 1, 0], dtype=bool)
+# Per-tile shade jitter is a pure function of position, so it is memorizable structure
+# rather than noise the model cannot predict.
+_SHADE_RANGE = 10
+_PLAYER_EYE_COLOR = (24, 28, 40)
+
+
+def paint_player(
+    frame: np.ndarray,
+    left: int,
+    top: int,
+    width: int,
+    height: int,
+    facing: int = 1,
+) -> None:
+    """Draw the miner onto an RGB frame. Used by the sim and by infer click-spawn."""
+    right = min(frame.shape[1], left + width)
+    bottom = min(frame.shape[0], top + height)
+    if right <= left or bottom <= top:
+        return
+    body_top = min(bottom, top + height // 2)
+    frame[max(0, top) : body_top, max(0, left) : right] = _PLAYER_HEAD_COLOR
+    frame[body_top:bottom, max(0, left) : right] = _PLAYER_BODY_COLOR
+    eye = left + (width - width // 3 if facing > 0 else 0)
+    eye_right = min(frame.shape[1], eye + max(1, width // 3))
+    eye_bottom = min(bottom, max(0, top) + max(1, height // 6))
+    frame[max(0, top) : eye_bottom, max(0, eye) : eye_right] = _PLAYER_EYE_COLOR
+
+
+def _shift_down(mask: np.ndarray) -> np.ndarray:
+    shifted = np.zeros_like(mask)
+    shifted[1:] = mask[:-1]
+    return shifted
+
+
+def _shift_horizontal(mask: np.ndarray, step: int) -> np.ndarray:
+    shifted = np.zeros_like(mask)
+    if step > 0:
+        shifted[:, step:] = mask[:, :-step]
+    else:
+        shifted[:, :step] = mask[:, -step:]
+    return shifted
+
+
+def _smooth(values: np.ndarray, passes: int) -> np.ndarray:
+    for _ in range(passes):
+        values = (values + np.roll(values, 1) + np.roll(values, -1)) / 3.0
+    return values
+
+
+class MiningSimulation:
+    """Single-screen tile sandbox with irreversible digging, falling sand, and water."""
+
+    # [left, right, jump, dig_left, dig_right, dig_up, dig_down]
+    ACTION_DIM = 7
+
+    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.tile = max(4, min(video_config.width, video_config.height) // 32)
+        self.grid_width = video_config.width // self.tile
+        self.grid_height = video_config.height // self.tile
+        self.gravity = self.tile * 42.0
+        self.move_speed = self.tile * 7.5
+        self.jump_speed = self.tile * 17.0
+        self.terminal_speed = self.tile * 30.0
+        self.swim_drag = 0.45
+        self.dig_cooldown = 0.16
+        self.player_width = float(self.tile)
+        self.player_height = float(self.tile * 2)
+        self.frame_index = 0
+        self.ore_collected = 0
+        self.deaths = 0
+        self.grid = np.zeros((self.grid_height, self.grid_width), dtype=np.uint8)
+        self.shade = self._make_shade()
+        self.position = np.zeros(2, dtype=np.float64)
+        self.velocity = np.zeros(2, dtype=np.float64)
+        self.facing = 1
+        self.dig_timer = 0.0
+        self.dig_target: tuple[tuple[int, int], ...] = ()
+        self.dig_flash = 0.0
+        self.drown_timer = 0.0
+        self.pending_reward = 0.0
+        self.pending_terminal = 0.0
+        self._action = (0,) * self.ACTION_DIM
+        self._hold_frames = 0
+        self._stuck_frames = 0
+        self._last_cell = (-1, -1)
+        self._reset_round()
+        for _ in range(6):
+            self.step()
+
+    @classmethod
+    def from_seed(
+        cls,
+        video_config: VideoConfig,
+        seed: int,
+        *,
+        at_start: bool = False,
+    ) -> MiningSimulation:
+        simulation = cls(video_config, seed)
+        if at_start:
+            simulation.reset_to_start()
+        return simulation
+
+    def reset_to_start(self) -> None:
+        self.ore_collected = 0
+        self.deaths = 0
+        self.frame_index = 0
+        self._reset_round()
+
+    def _make_shade(self) -> np.ndarray:
+        rows = np.arange(self.grid_height)[:, None]
+        columns = np.arange(self.grid_width)[None, :]
+        mixed = (rows * 73_856_093) ^ (columns * 19_349_663)
+        return ((mixed % (2 * _SHADE_RANGE + 1)) - _SHADE_RANGE).astype(np.int16)
+
+    def _generate_world(self) -> None:
+        height = self.grid_height
+        width = self.grid_width
+        grid = np.full((height, width), STONE, dtype=np.uint8)
+
+        walk = self.rng.normal(0.0, 1.0, size=width)
+        surface = _smooth(walk, 6)
+        surface = surface / max(float(np.abs(surface).max()), 1e-6)
+        base = height * 0.28
+        amplitude = height * 0.08
+        surface_rows = np.clip(
+            np.round(base + surface * amplitude).astype(int),
+            2,
+            height - 8,
+        )
+
+        rows = np.arange(height)[:, None]
+        grid[rows < surface_rows[None, :]] = AIR
+        dirt_depth = self.rng.integers(3, 6, size=width)
+        dirt_limit = surface_rows + dirt_depth
+        grid[(rows >= surface_rows[None, :]) & (rows < dirt_limit[None, :])] = DIRT
+
+        self._carve_caves(grid, surface_rows)
+        self._scatter(grid, ORE, count=self.rng.integers(9, 15), radius=1, in_stone=True)
+        self._scatter(grid, SAND, count=self.rng.integers(3, 7), radius=2, in_stone=False)
+        self._scatter(grid, GRAVEL, count=self.rng.integers(2, 5), radius=1, in_stone=True)
+        self._pour_water(grid, surface_rows)
+
+        grid[0, :] = BEDROCK
+        grid[-1, :] = BEDROCK
+        grid[:, 0] = BEDROCK
+        grid[:, -1] = BEDROCK
+        # Keep the spawn column open so the miner never starts entombed.
+        spawn_column = int(self.rng.integers(3, width - 3))
+        spawn_row = int(surface_rows[spawn_column])
+        grid[max(1, spawn_row - 3) : spawn_row, spawn_column] = AIR
+        self.grid = grid
+        self._spawn_column = spawn_column
+        self._spawn_row = spawn_row
+
+    def _carve_caves(self, grid: np.ndarray, surface_rows: np.ndarray) -> None:
+        cave_count = int(self.rng.integers(3, 7))
+        for _ in range(cave_count):
+            column = float(self.rng.integers(2, self.grid_width - 2))
+            row = float(self.rng.integers(int(surface_rows.max()) + 3, self.grid_height - 3))
+            length = int(self.rng.integers(6, 18))
+            floor = int(surface_rows.max()) + 2
+            for _ in range(length):
+                radius = float(self.rng.uniform(1.0, 2.4))
+                self._stamp(grid, row, column, radius, AIR, in_stone=False)
+                angle = float(self.rng.uniform(0.0, np.pi * 2.0))
+                column = float(np.clip(column + np.cos(angle) * 1.6, 2, self.grid_width - 3))
+                row = float(np.clip(row + np.sin(angle) * 1.2, floor, self.grid_height - 3))
+
+    def _stamp(
+        self,
+        grid: np.ndarray,
+        row: float,
+        column: float,
+        radius: float,
+        tile: int,
+        in_stone: bool,
+    ) -> None:
+        top = max(1, int(row - radius))
+        bottom = min(self.grid_height - 1, int(row + radius) + 1)
+        left = max(1, int(column - radius))
+        right = min(self.grid_width - 1, int(column + radius) + 1)
+        if top >= bottom or left >= right:
+            return
+        rows = np.arange(top, bottom)[:, None]
+        columns = np.arange(left, right)[None, :]
+        inside = (rows - row) ** 2 + (columns - column) ** 2 <= radius**2
+        region = grid[top:bottom, left:right]
+        if in_stone:
+            inside &= region == STONE
+        region[inside] = tile
+
+    def _scatter(
+        self,
+        grid: np.ndarray,
+        tile: int,
+        count: int,
+        radius: int,
+        in_stone: bool,
+    ) -> None:
+        for _ in range(int(count)):
+            row = float(self.rng.integers(2, self.grid_height - 2))
+            column = float(self.rng.integers(2, self.grid_width - 2))
+            self._stamp(grid, row, column, float(radius), tile, in_stone)
+
+    def _pour_water(self, grid: np.ndarray, surface_rows: np.ndarray) -> None:
+        pool_count = int(self.rng.integers(1, 4))
+        for _ in range(pool_count):
+            column = int(self.rng.integers(2, self.grid_width - 2))
+            row = int(self.rng.integers(int(surface_rows.max()) + 2, self.grid_height - 3))
+            span = int(self.rng.integers(2, 5))
+            left = max(1, column - span)
+            right = min(self.grid_width - 1, column + span)
+            region = grid[row : row + 2, left:right]
+            region[region == AIR] = WATER
+
+    def _reset_round(self) -> None:
+        self._generate_world()
+        self._respawn_player()
+
+    def _respawn_player(self) -> None:
+        # The world is never regenerated on death: tunnels the miner dug stay dug,
+        # which is the persistence signal the model is meant to learn.
+        column = self._spawn_column
+        row = self._spawn_row - 2
+        for candidate in range(1, self.grid_height - 2):
+            if not TILE_SOLID[self.grid[candidate, column]] and not TILE_SOLID[
+                self.grid[candidate + 1, column]
+            ]:
+                row = candidate
+                break
+        else:
+            self.grid[row : row + 2, column] = AIR
+        self.position[0] = float(column * self.tile)
+        self.position[1] = float(row * self.tile)
+        self.velocity[:] = 0.0
+        self.facing = 1
+        self.dig_timer = 0.0
+        self.dig_target = ()
+        self.dig_flash = 0.0
+        self.drown_timer = 0.0
+        self._action = (0,) * self.ACTION_DIM
+        self._hold_frames = 0
+
+    def _tile_at(self, row: int, column: int) -> int:
+        if row < 0 or row >= self.grid_height or column < 0 or column >= self.grid_width:
+            return BEDROCK
+        return int(self.grid[row, column])
+
+    def _solid_at(self, row: int, column: int) -> bool:
+        return bool(TILE_SOLID[self._tile_at(row, column)])
+
+    def _box_overlaps_solid(self, x: float, y: float) -> bool:
+        left = int(np.floor(x / self.tile))
+        right = int(np.floor((x + self.player_width - 1e-6) / self.tile))
+        top = int(np.floor(y / self.tile))
+        bottom = int(np.floor((y + self.player_height - 1e-6) / self.tile))
+        for row in range(top, bottom + 1):
+            for column in range(left, right + 1):
+                if self._solid_at(row, column):
+                    return True
+        return False
+
+    def _player_tiles(self) -> tuple[int, int, int]:
+        column = int((self.position[0] + self.player_width * 0.5) // self.tile)
+        head = int(self.position[1] // self.tile)
+        feet = int((self.position[1] + self.player_height - 1e-6) // self.tile)
+        return column, head, feet
+
+    def _footprint(self) -> tuple[int, int, int, int]:
+        left = int(self.position[0] // self.tile)
+        right = int((self.position[0] + self.player_width - 1e-6) // self.tile)
+        head = int(self.position[1] // self.tile)
+        feet = int((self.position[1] + self.player_height - 1e-6) // self.tile)
+        return left, right, head, feet
+
+    def _on_ground(self) -> bool:
+        return self._box_overlaps_solid(self.position[0], self.position[1] + 1.0)
+
+    def _in_water(self) -> bool:
+        column, head, feet = self._player_tiles()
+        return self._tile_at(head, column) == WATER or self._tile_at(feet, column) == WATER
+
+    def _step_physics(self, action: tuple[int, ...]) -> None:
+        left, right, jump = action[0], action[1], action[2]
+        submerged = self._in_water()
+        direction = float(right) - float(left)
+        if direction != 0.0:
+            self.facing = 1 if direction > 0 else -1
+        speed = self.move_speed * (0.6 if submerged else 1.0)
+        self.velocity[0] = direction * speed
+        grounded = self._on_ground()
+        if jump and (grounded or submerged):
+            self.velocity[1] = -self.jump_speed * (0.62 if submerged else 1.0)
+        gravity = self.gravity * (0.35 if submerged else 1.0)
+        self.velocity[1] = min(self.velocity[1] + gravity * self.dt, self.terminal_speed)
+        if submerged:
+            self.velocity[1] *= 1.0 - self.swim_drag * self.dt * self.video_config.fps * 0.1
+
+        step_x = float(self.velocity[0] * self.dt)
+        candidate = self.position[0] + step_x
+        if self._box_overlaps_solid(candidate, self.position[1]):
+            # Step up one tile so single-block ledges do not need a jump.
+            lifted = self.position[1] - self.tile
+            climbable = step_x != 0.0 and self._on_ground()
+            if climbable and not self._box_overlaps_solid(candidate, lifted):
+                self.position[0] = candidate
+                self.position[1] = lifted
+            else:
+                self.velocity[0] = 0.0
+        else:
+            self.position[0] = candidate
+
+        step_y = float(self.velocity[1] * self.dt)
+        candidate = self.position[1] + step_y
+        if self._box_overlaps_solid(self.position[0], candidate):
+            while abs(step_y) > 0.5:
+                step_y *= 0.5
+                probe = self.position[1] + step_y
+                if not self._box_overlaps_solid(self.position[0], probe):
+                    self.position[1] = probe
+            self.velocity[1] = 0.0
+        else:
+            self.position[1] = candidate
+
+        self.position[0] = float(
+            np.clip(self.position[0], self.tile, self.video_config.width - self.tile * 2)
+        )
+        self.position[1] = float(
+            np.clip(self.position[1], self.tile, self.video_config.height - self.tile * 3)
+        )
+
+    def _dig(self, action: tuple[int, ...]) -> None:
+        self.dig_flash = max(0.0, self.dig_flash - self.dt)
+        self.dig_timer = max(0.0, self.dig_timer - self.dt)
+        dig_left, dig_right, dig_up, dig_down = action[3], action[4], action[5], action[6]
+        if not (dig_left or dig_right or dig_up or dig_down):
+            self.dig_target = ()
+            return
+        left, right, head, feet = self._footprint()
+        # A dig clears the player's whole cross-section, so the shaft it opens is
+        # always wide enough to move through.
+        if dig_left:
+            self.facing = -1
+            target = tuple((row, left - 1) for row in range(head, feet + 1))
+        elif dig_right:
+            self.facing = 1
+            target = tuple((row, right + 1) for row in range(head, feet + 1))
+        elif dig_up:
+            target = tuple((head - 1, column) for column in range(left, right + 1))
+        else:
+            target = tuple((feet + 1, column) for column in range(left, right + 1))
+        self.dig_target = target
+        if self.dig_timer > 0.0:
+            return
+        removed = False
+        for row, column in target:
+            tile = self._tile_at(row, column)
+            if not TILE_DIGGABLE[tile]:
+                continue
+            self.grid[row, column] = AIR
+            removed = True
+            if tile == ORE:
+                self.ore_collected += 1
+                self.pending_reward += 1.0
+        if removed:
+            self.dig_timer = self.dig_cooldown
+            self.dig_flash = self.dig_cooldown
+
+    def _step_tiles(self) -> None:
+        grid = self.grid
+        falls = TILE_FALLS[grid]
+        below = np.zeros_like(grid)
+        below[:-1] = grid[1:]
+        below[-1] = BEDROCK
+        open_below = (below == AIR) | (below == WATER)
+        moving = falls & open_below
+        moving[-1] = False
+        if moving.any():
+            destination = _shift_down(moving)
+            carried = grid[moving]
+            displaced = grid[destination]
+            grid[destination] = carried
+            grid[moving] = displaced
+
+        # Sand slides off peaks so it forms cones instead of towers.
+        side = 1 if self.frame_index % 2 == 0 else -1
+        sand = grid == SAND
+        below[:-1] = grid[1:]
+        below[-1] = BEDROCK
+        blocked = sand & ~((below == AIR) | (below == WATER))
+        diagonal = np.zeros_like(grid)
+        if side > 0:
+            diagonal[:-1, :-1] = grid[1:, 1:]
+        else:
+            diagonal[:-1, 1:] = grid[1:, :-1]
+        slide = blocked & (diagonal == AIR)
+        slide[-1] = False
+        if slide.any():
+            destination = _shift_down(_shift_horizontal(slide, side))
+            grid[destination] = grid[slide]
+            grid[slide] = AIR
+
+        water = grid == WATER
+        below[:-1] = grid[1:]
+        below[-1] = BEDROCK
+        settled = water & (below != AIR)
+        neighbour = np.zeros_like(grid)
+        if side > 0:
+            neighbour[:, :-1] = grid[:, 1:]
+        else:
+            neighbour[:, 1:] = grid[:, :-1]
+        spread = settled & (neighbour == AIR)
+        if spread.any():
+            destination = _shift_horizontal(spread, side)
+            grid[destination] = WATER
+            grid[spread] = AIR
+
+    def _check_fatal(self) -> None:
+        buried = False
+        if self._box_overlaps_solid(self.position[0], self.position[1]):
+            # Material settling onto the miner shoves it upward; only a fully
+            # sealed pocket actually kills.
+            buried = True
+            for offset in range(1, self.tile * 2 + 1):
+                lifted = self.position[1] - offset
+                if lifted < self.tile:
+                    break
+                if not self._box_overlaps_solid(self.position[0], lifted):
+                    self.position[1] = float(lifted)
+                    self.velocity[1] = 0.0
+                    buried = False
+                    break
+        column, head, _ = self._player_tiles()
+        if self._tile_at(head, column) == WATER:
+            self.drown_timer += self.dt
+        else:
+            self.drown_timer = max(0.0, self.drown_timer - self.dt * 0.5)
+        if buried or self.drown_timer > 2.5:
+            self.deaths += 1
+            self.pending_terminal = 1.0
+            self.drown_timer = 0.0
+            self._respawn_player()
+
+    def scripted_action(self) -> tuple[int, ...]:
+        """Next training-policy action. Does not step the world."""
+        return self._scripted_action()
+
+    def _scripted_action(self) -> tuple[int, ...]:
+        column, _, feet = self._player_tiles()
+        cell = (feet, column)
+        if cell == self._last_cell:
+            self._stuck_frames += 1
+        else:
+            self._stuck_frames = 0
+            self._last_cell = cell
+        if self._hold_frames > 0:
+            self._hold_frames -= 1
+            return self._action
+        if self._stuck_frames > 18:
+            # Break out of dead ends instead of grinding at an unreachable target.
+            action = [0] * self.ACTION_DIM
+            action[int(self.rng.integers(0, self.ACTION_DIM))] = 1
+            action[2] = 1
+            self._action = tuple(action)
+            self._hold_frames = 4
+            self._stuck_frames = 0
+            return self._action
+        action = self._plan_action()
+        self._action = action
+        self._hold_frames = int(self.rng.integers(0, 3))
+        return action
+
+    def _plan_action(self) -> tuple[int, ...]:
+        if self.rng.random() < 0.06:
+            choice = int(self.rng.integers(0, self.ACTION_DIM))
+            action = [0] * self.ACTION_DIM
+            action[choice] = 1
+            return tuple(action)
+
+        column, head, feet = self._player_tiles()
+        action = [0] * self.ACTION_DIM
+
+        # Escape falling material before anything else.
+        if self._tile_at(head - 1, column) in {SAND, GRAVEL}:
+            action[1 if self.facing > 0 else 0] = 1
+            action[2] = 1
+            return tuple(action)
+        if self.drown_timer > 0.6:
+            action[2] = 1
+            action[5] = 1
+            return tuple(action)
+
+        ore = np.argwhere(self.grid == ORE)
+        if ore.size == 0:
+            action[6] = 1
+            return tuple(action)
+        distances = np.abs(ore[:, 0] - feet) + np.abs(ore[:, 1] - column)
+        target_row, target_column = ore[int(np.argmin(distances))]
+        delta_column = int(target_column) - column
+        delta_row = int(target_row) - feet
+
+        if abs(delta_column) <= 1 and abs(delta_row) <= 1:
+            if delta_column < 0:
+                action[3] = 1
+            elif delta_column > 0:
+                action[4] = 1
+            elif delta_row < 0:
+                action[5] = 1
+            else:
+                action[6] = 1
+            return tuple(action)
+
+        if delta_column != 0:
+            forward = 1 if delta_column > 0 else 0
+            action[forward] = 1
+            ahead = column + (1 if delta_column > 0 else -1)
+            if self._solid_at(feet, ahead):
+                # Bedrock never yields, so climb over it instead of digging forever.
+                if TILE_DIGGABLE[self._tile_at(feet, ahead)]:
+                    action[4 if delta_column > 0 else 3] = 1
+                else:
+                    action[2] = 1
+            elif self._solid_at(head, ahead):
+                action[2] = 1
+            return tuple(action)
+
+        if delta_row > 0:
+            action[6] = 1
+        else:
+            action[2] = 1
+            action[5] = 1
+        return tuple(action)
+
+    def step(self, action: tuple[int, ...] | None = None) -> None:
+        resolved = self._scripted_action() if action is None else tuple(action)
+        if len(resolved) != self.ACTION_DIM:
+            raise ValueError(f"Mining actions must have {self.ACTION_DIM} entries.")
+        self.pending_reward = 0.0
+        self.pending_terminal = 0.0
+        self._dig(resolved)
+        self._step_physics(resolved)
+        self._step_tiles()
+        self._check_fatal()
+        self.frame_index += 1
+
+    def render(self) -> np.ndarray:
+        colors = TILE_COLORS[self.grid] + self.shade[:, :, None]
+        if self.dig_flash > 0.0:
+            for row, column in self.dig_target:
+                if 0 <= row < self.grid_height and 0 <= column < self.grid_width:
+                    colors[row, column] = _DIG_FLASH_COLOR
+        frame = np.clip(colors, 0, 255).astype(np.uint8)
+        frame = frame.repeat(self.tile, axis=0).repeat(self.tile, axis=1)
+
+        paint_player(
+            frame,
+            round(float(self.position[0])),
+            round(float(self.position[1])),
+            int(self.player_width),
+            int(self.player_height),
+            self.facing,
+        )
+
+        # Ore pips give a fontless score readout that stays stable at training size.
+        for index in range(min(self.ore_collected, 8)):
+            x = self.tile + index * (self.tile + 2)
+            pip = max(2, self.tile // 2)
+            frame[2 : 2 + pip, x : x + pip] = TILE_COLORS[ORE]
+        return frame
+
+    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 sample_gameplay(
+        self,
+        count: int,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        frames: list[np.ndarray] = []
+        actions: list[tuple[int, ...]] = []
+        rewards: list[float] = []
+        terminals: list[float] = []
+        for _ in range(count):
+            frames.append(self.render())
+            action = self.scripted_action()
+            actions.append(action)
+            self.step(action)
+            rewards.append(self.pending_reward)
+            terminals.append(self.pending_terminal)
+        return (
+            np.stack(frames),
+            np.asarray(actions, dtype=np.float32),
+            np.asarray(rewards, dtype=np.float32),
+            np.asarray(terminals, dtype=np.float32),
+        )

livediffusion/model.py

@@ -0,0 +1,920 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from livediffusion.config import ModelConfig
+
+
+@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):
+    """3D RoPE with a shorter temporal base so time and space do not share a spectrum."""
+
+    def __init__(
+        self,
+        head_size: int,
+        time_base: float = 1_000.0,
+        space_base: float = 10_000.0,
+    ) -> None:
+        super().__init__()
+        self.rotary_size = (head_size // 6) * 6
+        self.axis_size = self.rotary_size // 3
+        axis_half = self.axis_size // 2
+        steps = torch.arange(axis_half, dtype=torch.float32) / max(axis_half, 1)
+        time_frequencies = torch.exp(-math.log(time_base) * steps)
+        space_frequencies = torch.exp(-math.log(space_base) * steps)
+        self.register_buffer(
+            "frequencies",
+            torch.stack((time_frequencies, space_frequencies, space_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[axis][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.
+
+    Keys and values may use fewer heads than queries. Prefix and target keys are
+    concatenated, so both sides share the same reduced head count.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        qk_norm: bool,
+        kv_heads: int = 0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_size = hidden_size // num_heads
+        self.kv_heads = kv_heads or num_heads
+        kv_size = self.kv_heads * self.head_size
+        self.target_qkv = nn.Linear(hidden_size, hidden_size + kv_size * 2, bias=False)
+        self.prefix_key_value = nn.Linear(hidden_size, kv_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.kv_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
+        kv_size = self.kv_heads * self.head_size
+        projected = self.target_qkv(target)
+        query, key, value = projected.split((channels, kv_size, kv_size), dim=-1)
+        query = query.reshape(batch, tokens, self.num_heads, self.head_size)
+        key = key.reshape(batch, tokens, self.kv_heads, self.head_size)
+        value = value.reshape(batch, tokens, self.kv_heads, self.head_size)
+        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,
+            enable_gqa=self.kv_heads != self.num_heads,
+        )
+        return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels))
+
+
+@dataclass(slots=True)
+class WorldMemoryState:
+    plane: torch.Tensor
+    registers: torch.Tensor
+
+
+class WorldMemory(nn.Module):
+    """Gated recurrent memory at latent resolution, plus global register state.
+
+    The plane is a convolutional GRU over one cell per latent position, so an
+    irreversible world edit stays represented after it scrolls out of the
+    attention prefix. Registers carry frame-global state the plane cannot hold.
+    """
+
+    def __init__(
+        self,
+        latent_channels: int,
+        hidden_size: int,
+        memory_channels: int,
+        read_patch: int,
+        register_tokens: int,
+    ) -> None:
+        super().__init__()
+        self.memory_channels = memory_channels
+        self.read_patch = read_patch
+        self.register_tokens = register_tokens
+        self.cell = nn.Conv2d(
+            latent_channels + memory_channels,
+            memory_channels * 3,
+            kernel_size=3,
+            padding=1,
+        )
+        self.read = nn.Conv2d(
+            memory_channels,
+            hidden_size,
+            kernel_size=read_patch,
+            stride=read_patch,
+        )
+        self.read_norm = nn.RMSNorm(hidden_size, eps=1e-6)
+        self.scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
+        self.register_update = (
+            nn.Linear(memory_channels + hidden_size, hidden_size * 2)
+            if register_tokens > 0
+            else None
+        )
+
+    def initial_state(
+        self,
+        batch: int,
+        height: int,
+        width: int,
+        registers: torch.Tensor,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> WorldMemoryState:
+        plane = torch.zeros(
+            batch,
+            self.memory_channels,
+            height,
+            width,
+            device=device,
+            dtype=dtype,
+        )
+        return WorldMemoryState(plane, registers)
+
+    def advance(
+        self,
+        state: WorldMemoryState,
+        frames: torch.Tensor,
+    ) -> WorldMemoryState:
+        plane = state.plane
+        for index in range(frames.shape[1]):
+            latent = frames[:, index].to(plane.dtype)
+            reset, update, candidate = self.cell(
+                torch.cat((latent, plane), dim=1)
+            ).chunk(3, dim=1)
+            gate = torch.sigmoid(update)
+            proposal = torch.tanh(candidate + torch.sigmoid(reset) * plane)
+            plane = plane * (1.0 - gate) + proposal * gate
+        registers = state.registers
+        if self.register_update is not None:
+            pooled = plane.mean(dim=(2, 3))
+            summary = torch.cat(
+                (pooled[:, None].expand(-1, registers.shape[1], -1), registers),
+                dim=-1,
+            )
+            gate, proposal = self.register_update(summary).chunk(2, dim=-1)
+            gate = torch.sigmoid(gate)
+            registers = registers * (1.0 - gate) + torch.tanh(proposal) * gate
+        return WorldMemoryState(plane, registers)
+
+    def tokens(self, state: WorldMemoryState) -> torch.Tensor:
+        read = self.read(state.plane).flatten(2).transpose(1, 2)
+        return self.read_norm(read) + self.scale_embedding
+
+
+class SharedModulation(nn.Module):
+    """One adaLN projection shared by every block, with per-block offsets.
+
+    Per-block modulation costs 6 * hidden^2 each; sharing recovers that budget
+    for depth while keeping each block independently steerable.
+    """
+
+    def __init__(self, hidden_size: int, blocks: int) -> None:
+        super().__init__()
+        self.projection = nn.Linear(hidden_size, hidden_size * 6)
+        self.offsets = nn.Parameter(torch.zeros(blocks, hidden_size * 6))
+
+    def forward(self, condition: torch.Tensor, block_index: int) -> torch.Tensor:
+        return self.projection(F.silu(condition)) + self.offsets[block_index]
+
+
+class PrefixEncoderBlock(nn.Module):
+    """Clean-context stack that runs before the joint prefix/target DiT."""
+
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        hidden_size = config.hidden_size
+        self.norm1 = nn.RMSNorm(hidden_size, eps=1e-6)
+        self.attention = SelfAttention(
+            hidden_size,
+            config.num_heads,
+            config.qk_norm,
+        )
+        self.norm2 = nn.RMSNorm(hidden_size, eps=1e-6)
+        self.mlp = SwiGLU(hidden_size, config.mlp_ratio)
+        self.attn_scale = nn.Parameter(torch.ones(hidden_size))
+        self.mlp_scale = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, prefix: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        prefix = prefix + self.attn_scale * self.attention(self.norm1(prefix), positions)
+        return prefix + self.mlp_scale * self.mlp(self.norm2(prefix))
+
+
+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,
+            config.kv_heads,
+        )
+        self.target_norm2 = nn.RMSNorm(
+            hidden_size,
+            eps=1e-6,
+            elementwise_affine=False,
+        )
+        self.target_mlp = SwiGLU(hidden_size, config.mlp_ratio)
+        self.modulation = (
+            None
+            if config.shared_modulation
+            else 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],
+        modulation: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if modulation is None:
+            if self.modulation is None:
+                raise RuntimeError("Shared modulation must be supplied by the trunk.")
+            modulation = self.modulation(F.silu(condition))
+        shift1, scale1, gate1, shift2, scale2, gate2 = modulation.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,
+        modulation: torch.Tensor | None = None,
+    ) -> 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,
+            modulation,
+        )
+        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.register_tokens = (
+            nn.Parameter(torch.randn(1, model_config.register_tokens, hidden_size) * 0.02)
+            if model_config.register_tokens > 0
+            else None
+        )
+        self.memory_blocks = nn.ModuleList(
+            PrefixEncoderBlock(model_config)
+            for _ in range(model_config.prefix_encoder_depth)
+        )
+        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.shared_modulation = (
+            SharedModulation(hidden_size, model_config.depth)
+            if model_config.shared_modulation
+            else None
+        )
+        self.world_memory = (
+            WorldMemory(
+                latent_channels,
+                hidden_size,
+                model_config.memory_channels,
+                coarse_patch,
+                model_config.register_tokens,
+            )
+            if model_config.memory_channels > 0
+            else None
+        )
+        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:
+            if block.modulation is not None:
+                nn.init.zeros_(block.modulation.weight)
+                nn.init.zeros_(block.modulation.bias)
+        if self.shared_modulation is not None:
+            nn.init.zeros_(self.shared_modulation.projection.weight)
+            nn.init.zeros_(self.shared_modulation.projection.bias)
+            nn.init.zeros_(self.shared_modulation.offsets)
+        if self.world_memory is not None:
+            nn.init.zeros_(self.world_memory.scale_embedding)
+        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,
+        history_actions: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if history.shape[1] < self.history_frames:
+            raise ValueError(
+                f"Expected at least {self.history_frames} history frames, "
+                f"got {history.shape[1]}."
+            )
+        batch = history.shape[0]
+        registers = (
+            self.register_tokens.expand(batch, -1, -1)
+            if self.register_tokens is not None
+            else None
+        )
+        memory_tokens: torch.Tensor | None = None
+        if self.world_memory is not None:
+            # Every observed frame is folded into the recurrent state, including
+            # frames older than the attention window.
+            state = self.world_memory.initial_state(
+                batch,
+                history.shape[-2],
+                history.shape[-1],
+                registers
+                if registers is not None
+                else history.new_zeros(batch, 0, self.model_config.hidden_size),
+                history.device,
+                history.dtype,
+            )
+            state = self.world_memory.advance(state, history)
+            memory_tokens = self.world_memory.tokens(state)
+            if registers is not None:
+                registers = state.registers
+        window = history[:, -self.history_frames :]
+        coarse, coarse_positions = self._patch(
+            self.coarse_prefix_embed,
+            window,
+            0,
+        )
+        if self.action_projection is not None and history_actions is not None:
+            window_actions = history_actions[:, -self.history_frames :]
+            tokens_per_frame = coarse.shape[1] // self.history_frames
+            coarse = coarse + self.action_projection(
+                window_actions
+            ).repeat_interleave(tokens_per_frame, dim=1)
+        history = window
+        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 memory_tokens is not None:
+            coarse_patch = self.model_config.memory_patch_size
+            memory_positions = _video_positions(
+                1,
+                history.shape[-2] // coarse_patch,
+                history.shape[-1] // coarse_patch,
+                -2,
+                history.device,
+            )
+            prefix = torch.cat((memory_tokens, prefix), dim=1)
+            positions = torch.cat((memory_positions, positions), dim=0)
+        if registers is not None:
+            register_count = registers.shape[1]
+            register_positions = torch.stack(
+                (
+                    torch.full((register_count,), -1.0, device=history.device),
+                    torch.zeros(register_count, device=history.device),
+                    torch.arange(register_count, device=history.device, dtype=torch.float32),
+                ),
+                dim=-1,
+            )
+            prefix = torch.cat((registers, prefix), dim=1)
+            positions = torch.cat((register_positions, 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,
+            )
+        for block in self.memory_blocks:
+            prefix = block(prefix, positions)
+        return prefix, positions
+
+    def _pool_prefix(self, prefix: torch.Tensor) -> torch.Tensor:
+        register_count = self.model_config.register_tokens
+        if register_count > 0:
+            return prefix[:, :register_count].mean(dim=1)
+        return prefix.mean(dim=1)
+
+    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(self._pool_prefix(prefix))
+
+    def build_memory_cache(
+        self,
+        history: torch.Tensor,
+        context: torch.Tensor | None = None,
+        history_actions: torch.Tensor | None = None,
+    ) -> PrefixCache:
+        prefix, positions = self._encode_prefix(history, context, history_actions)
+        pooled = self.prefix_condition(self._pool_prefix(prefix))
+        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,
+        history_actions: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if memory_cache is None:
+            prefix, _ = self._encode_prefix(history, None, history_actions)
+            pooled = self.prefix_condition(self._pool_prefix(prefix))
+        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,
+        history_actions: torch.Tensor | 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,
+                history_actions,
+            )
+            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):
+            modulation = (
+                self.shared_modulation(condition, block_index)
+                if self.shared_modulation is not None
+                else None
+            )
+            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,
+                        modulation,
+                        use_reentrant=False,
+                    )
+                else:
+                    prefix, target = block(
+                        prefix,
+                        target,
+                        condition,
+                        prefix_positions,
+                        target_positions,
+                        modulation,
+                    )
+            else:
+                target = block.advance_target(
+                    target,
+                    condition,
+                    target_positions,
+                    memory_cache.key_values[block_index],
+                    modulation,
+                )
+        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

livediffusion/tile_vae.py

@@ -0,0 +1,282 @@
+"""Compact frame autoencoder trained from scratch on tile-rendered gameplay.
+
+The frozen SDXL VAE spends 83M parameters and a mid-block attention layer on
+photographic detail that tile art does not contain. A purpose-built encoder
+reaches far lower reconstruction error on the same 8x grid at a fraction of the
+cost, which matters because every training step encodes the whole memory
+window rather than a single frame.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+@dataclass(frozen=True, slots=True)
+class TileVAEConfig:
+    latent_channels: int = 4
+    base_channels: int = 64
+    channel_multipliers: tuple[int, ...] = (1, 2, 4)
+    blocks_per_stage: int = 2
+
+    @property
+    def spatial_compression(self) -> int:
+        return 2 ** len(self.channel_multipliers)
+
+    def validate(self) -> None:
+        if self.latent_channels < 1:
+            raise ValueError("latent_channels must be positive.")
+        if self.base_channels < 8:
+            raise ValueError("base_channels must be at least 8.")
+        if not self.channel_multipliers:
+            raise ValueError("channel_multipliers cannot be empty.")
+        if any(multiplier < 1 for multiplier in self.channel_multipliers):
+            raise ValueError("channel_multipliers must be positive.")
+        if self.blocks_per_stage < 1:
+            raise ValueError("blocks_per_stage must be positive.")
+
+
+def _norm(channels: int) -> nn.GroupNorm:
+    groups = 32 if channels % 32 == 0 else 8 if channels % 8 == 0 else 1
+    return nn.GroupNorm(groups, channels, eps=1e-6)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.norm1 = _norm(in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = _norm(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+        self.skip = (
+            nn.Conv2d(in_channels, out_channels, 1)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        hidden = self.conv1(F.silu(self.norm1(features)))
+        hidden = self.conv2(F.silu(self.norm2(hidden)))
+        return self.skip(features) + hidden
+
+
+class Encoder(nn.Module):
+    def __init__(self, config: TileVAEConfig) -> None:
+        super().__init__()
+        channels = config.base_channels
+        self.stem = nn.Conv2d(3, channels, 3, padding=1)
+        stages: list[nn.Module] = []
+        for multiplier in config.channel_multipliers:
+            target = config.base_channels * multiplier
+            for _ in range(config.blocks_per_stage):
+                stages.append(ResidualBlock(channels, target))
+                channels = target
+            stages.append(nn.Conv2d(channels, channels, 3, stride=2, padding=1))
+        self.stages = nn.Sequential(*stages)
+        self.mid = ResidualBlock(channels, channels)
+        self.out_norm = _norm(channels)
+        self.out = nn.Conv2d(channels, config.latent_channels * 2, 3, padding=1)
+
+    def forward(self, frames: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        hidden = self.mid(self.stages(self.stem(frames)))
+        mean, log_variance = self.out(F.silu(self.out_norm(hidden))).chunk(2, dim=1)
+        return mean, log_variance.clamp(-30.0, 20.0)
+
+
+class Decoder(nn.Module):
+    def __init__(self, config: TileVAEConfig) -> None:
+        super().__init__()
+        multipliers = config.channel_multipliers
+        channels = config.base_channels * multipliers[-1]
+        self.stem = nn.Conv2d(config.latent_channels, channels, 3, padding=1)
+        self.mid = ResidualBlock(channels, channels)
+        stages: list[nn.Module] = []
+        for index in reversed(range(len(multipliers))):
+            target = config.base_channels * multipliers[index]
+            for _ in range(config.blocks_per_stage):
+                stages.append(ResidualBlock(channels, target))
+                channels = target
+            stages.append(Upsample(channels))
+        self.stages = nn.Sequential(*stages)
+        self.out_norm = _norm(channels)
+        self.out = nn.Conv2d(channels, 3, 3, padding=1)
+
+    def forward(self, latents: torch.Tensor) -> torch.Tensor:
+        hidden = self.stages(self.mid(self.stem(latents)))
+        return self.out(F.silu(self.out_norm(hidden)))
+
+
+class Upsample(nn.Module):
+    """Nearest-neighbour upsample then convolve, which avoids checkerboarding."""
+
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 3, padding=1)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        upsampled = F.interpolate(features.float(), scale_factor=2.0, mode="nearest")
+        return self.conv(upsampled.to(features.dtype))
+
+
+class TileVAE(nn.Module):
+    """Convolutional VAE over single frames, with unit-variance latent scaling."""
+
+    def __init__(self, config: TileVAEConfig | None = None) -> None:
+        super().__init__()
+        config = config or TileVAEConfig()
+        config.validate()
+        self.config = config
+        self.encoder = Encoder(config)
+        self.decoder = Decoder(config)
+        self.register_buffer("latent_shift", torch.zeros(config.latent_channels))
+        self.register_buffer("latent_scale", torch.ones(config.latent_channels))
+
+    @property
+    def latent_channels(self) -> int:
+        return self.config.latent_channels
+
+    @property
+    def spatial_compression(self) -> int:
+        return self.config.spatial_compression
+
+    def encode_moments(
+        self,
+        frames: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.encoder(frames)
+
+    def normalize(self, latents: torch.Tensor) -> torch.Tensor:
+        shift = self.latent_shift.reshape(1, -1, 1, 1).to(latents.dtype)
+        scale = self.latent_scale.reshape(1, -1, 1, 1).to(latents.dtype)
+        return (latents - shift) * scale
+
+    def denormalize(self, latents: torch.Tensor) -> torch.Tensor:
+        shift = self.latent_shift.reshape(1, -1, 1, 1).to(latents.dtype)
+        scale = self.latent_scale.reshape(1, -1, 1, 1).to(latents.dtype)
+        return latents / scale + shift
+
+    def encode(self, frames: torch.Tensor, sample: bool = False) -> torch.Tensor:
+        mean, log_variance = self.encode_moments(frames)
+        if sample:
+            mean = mean + torch.randn_like(mean) * (0.5 * log_variance).exp()
+        return self.normalize(mean)
+
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.denormalize(latents))
+
+    def forward(
+        self,
+        frames: torch.Tensor,
+        sample: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        mean, log_variance = self.encode_moments(frames)
+        latents = mean
+        if sample:
+            latents = mean + torch.randn_like(mean) * (0.5 * log_variance).exp()
+        return self.decoder(latents), mean, log_variance
+
+    def set_latent_statistics(self, mean: torch.Tensor, std: torch.Tensor) -> None:
+        self.latent_shift.copy_(mean.detach().flatten().to(self.latent_shift.dtype))
+        self.latent_scale.copy_(
+            (1.0 / std.detach().flatten().clamp_min(1e-5)).to(self.latent_scale.dtype)
+        )
+
+    def save(self, path: str | Path) -> Path:
+        destination = Path(path)
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        torch.save(
+            {
+                "format": "tile_vae/1",
+                "config": asdict(self.config),
+                "state_dict": self.state_dict(),
+            },
+            destination,
+        )
+        return destination
+
+    @classmethod
+    def load(
+        cls,
+        path: str | Path,
+        device: torch.device | str = "cpu",
+        dtype: torch.dtype = torch.float32,
+    ) -> TileVAE:
+        payload = torch.load(Path(path), map_location="cpu", weights_only=True)
+        if payload.get("format") != "tile_vae/1":
+            raise ValueError(f"{path} is not a tile_vae/1 checkpoint.")
+        raw = dict(payload["config"])
+        raw["channel_multipliers"] = tuple(raw["channel_multipliers"])
+        model = cls(TileVAEConfig(**raw))
+        model.load_state_dict(payload["state_dict"], strict=True)
+        return model.to(device=device, dtype=dtype).eval()
+
+
+@dataclass(slots=True)
+class ReconstructionLoss:
+    total: torch.Tensor
+    l1: torch.Tensor
+    mse: torch.Tensor
+    edge: torch.Tensor
+    kl: torch.Tensor
+    weights: dict[str, float] = field(default_factory=dict)
+
+
+def _spatial_gradients(frames: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        frames[..., :, 1:] - frames[..., :, :-1],
+        frames[..., 1:, :] - frames[..., :-1, :],
+    )
+
+
+def reconstruction_loss(
+    reconstruction: torch.Tensor,
+    target: torch.Tensor,
+    mean: torch.Tensor,
+    log_variance: torch.Tensor,
+    *,
+    l1_weight: float = 1.0,
+    mse_weight: float = 1.0,
+    edge_weight: float = 0.5,
+    kl_weight: float = 1e-6,
+) -> ReconstructionLoss:
+    """Pixel, gradient, and KL terms.
+
+    The gradient term is what keeps tile boundaries from bleeding: a one-pixel
+    smear costs almost nothing in L1 but is exactly the error that turns a
+    crisp 8px tile edge into an ambiguous one.
+    """
+    reconstruction = reconstruction.float()
+    target = target.float()
+    l1 = (reconstruction - target).abs().mean()
+    mse = (reconstruction - target).square().mean()
+    prediction_dx, prediction_dy = _spatial_gradients(reconstruction)
+    target_dx, target_dy = _spatial_gradients(target)
+    edge = (prediction_dx - target_dx).abs().mean() + (
+        prediction_dy - target_dy
+    ).abs().mean()
+    mean = mean.float()
+    log_variance = log_variance.float()
+    kl = 0.5 * (
+        mean.square() + log_variance.exp() - 1.0 - log_variance
+    ).flatten(1).sum(dim=1).mean()
+    total = (
+        l1_weight * l1 + mse_weight * mse + edge_weight * edge + kl_weight * kl
+    )
+    return ReconstructionLoss(
+        total=total,
+        l1=l1.detach(),
+        mse=mse.detach(),
+        edge=edge.detach(),
+        kl=kl.detach(),
+        weights={
+            "l1": l1_weight,
+            "mse": mse_weight,
+            "edge": edge_weight,
+            "kl": kl_weight,
+        },
+    )