live_infer.py

#!/usr/bin/env python3
"""Roll out Diffusion Pong Base (unconditional).

Deps:
  pip install torch numpy pillow safetensors huggingface_hub diffusers

Run:
  python live_infer_base.py
  python live_infer_base.py --steps 2 --window-scale 7 --seed 42

Unconditional next-frame model — no paddle controls. Click canvas to inject a ball.
"""
from __future__ import annotations

import argparse
import json
import math
import random
import time
import tkinter as tk
from collections import deque
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from diffusers import AutoencoderKL
from huggingface_hub import hf_hub_download
from PIL import Image, ImageDraw, ImageTk
from safetensors.torch import load_file
from torch import nn
from torch.utils.checkpoint import checkpoint

DEFAULT_REPO = "kerzgrr/diffusionpong-base"


@dataclass(frozen=True, slots=True)
class VideoConfig:
    height: int = 128
    width: int = 128
    fps: float = 6.0
    history_frames: int = 12
    future_frames: int = 5


@dataclass(frozen=True, slots=True)
class ModelConfig:
    target_patch_size: int = 2
    memory_patch_size: int = 4
    fine_memory_patch_size: int = 2
    fine_history_frames: int = 2
    hidden_size: int = 384
    depth: int = 12
    memory_depth: int = 3
    num_heads: int = 8
    mlp_ratio: float = 4.0
    qk_norm: bool = True
    generation_frames: int = 1
    action_dim: int = 0
    context_dim: int = 0


@dataclass(frozen=True, slots=True)
class CodecSettings:
    model_id: str = "madebyollin/sdxl-vae-fp16-fix"
    subfolder: str = ""
    revision: str = "main"
    latent_channels: int = 4
    spatial_compression: int = 8
    frame_batch_size: int = 256
    sample_posterior: bool = False
    enable_slicing: bool = True
    enable_tiling: bool = False
class PongSimulation:
    """Deterministic Pong gameplay with continuous physics and scripted paddles."""

    def __init__(self, video_config: VideoConfig, seed: int) -> None:
        self.video_config = video_config
        self.rng = np.random.default_rng(seed)
        self.dt = 1.0 / video_config.fps
        self.paddle_height = max(18.0, video_config.height * 0.20)
        self.paddle_width = max(4.0, video_config.width * 0.025)
        self.ball_radius = max(3.0, min(video_config.width, video_config.height) * 0.025)
        self.left_x = video_config.width * 0.07
        self.right_x = video_config.width * 0.93
        self.left_y = video_config.height / 2
        self.right_y = video_config.height / 2
        self.left_score = 0
        self.right_score = 0
        self.elapsed = 0.0
        self.phase = float(self.rng.uniform(0.0, math.tau))
        self.left_reaction_delay = int(self.rng.integers(0, 5))
        self.left_aim_bias = float(self.rng.uniform(-10.0, 10.0))
        self.left_randomness = float(self.rng.uniform(0.12, 0.30))
        self.right_aim_bias = float(self.rng.uniform(-22.0, 22.0))
        self.right_speed_scale = float(self.rng.uniform(0.38, 0.58))
        self._observed_ball_y: deque[float] = deque(
            [video_config.height / 2] * (self.left_reaction_delay + 1),
            maxlen=self.left_reaction_delay + 1,
        )
        self._held_left_action = 0
        self._action_hold_frames = 0
        self.ball_position = np.zeros(2, dtype=np.float32)
        self.ball_velocity = np.zeros(2, dtype=np.float32)
        self._serve(direction=1 if self.rng.random() < 0.5 else -1)
        for _ in range(12):
            self.step()

    @classmethod
    def from_seed(cls, video_config: VideoConfig, seed: int) -> PongSimulation:
        return cls(video_config, seed)

    def _serve(self, direction: int) -> None:
        angle = float(self.rng.uniform(-0.55, 0.55))
        speed = float(self.rng.uniform(52.0, 68.0))
        self.ball_position[:] = (
            self.video_config.width / 2,
            self.video_config.height / 2,
        )
        self.ball_velocity[:] = (
            direction * speed * math.cos(angle),
            speed * math.sin(angle),
        )

    def _move_paddles(self, dt: float, left_action: int | None) -> None:
        height = float(self.video_config.height)
        right_target = float(
            self.ball_position[1]
            + self.right_aim_bias
            + math.sin(self.elapsed * 1.7 + self.phase) * 18.0
        )
        paddle_speed = height * 0.75
        right_paddle_speed = height * self.right_speed_scale
        if left_action is None:
            left_delta = max(
                -paddle_speed * dt,
                min(
                    paddle_speed * dt,
                    float(self.ball_position[1]) - self.left_y,
                ),
            )
        else:
            left_delta = float(max(-1, min(1, left_action))) * paddle_speed * dt
        self.left_y += left_delta
        right_delta = max(
            -right_paddle_speed * dt,
            min(right_paddle_speed * dt, right_target - self.right_y),
        )
        self.right_y += right_delta
        self.left_y = max(
            self.paddle_height / 2,
            min(height - self.paddle_height / 2, self.left_y),
        )
        self.right_y = max(
            self.paddle_height / 2,
            min(height - self.paddle_height / 2, self.right_y),
        )

    def _paddle_collision(self, x: float, y: float, left: bool) -> bool:
        paddle_x = self.left_x if left else self.right_x
        paddle_y = self.left_y if left else self.right_y
        moving_toward = self.ball_velocity[0] < 0 if left else self.ball_velocity[0] > 0
        touching_x = (
            x - self.ball_radius <= paddle_x + self.paddle_width / 2
            if left
            else x + self.ball_radius >= paddle_x - self.paddle_width / 2
        )
        touching_y = abs(y - paddle_y) <= self.paddle_height / 2 + self.ball_radius
        return moving_toward and touching_x and touching_y

    def step(self, left_action: int | None = None) -> None:
        width = float(self.video_config.width)
        height = float(self.video_config.height)
        substeps = 4
        sub_dt = self.dt / substeps
        for _ in range(substeps):
            self._move_paddles(sub_dt, left_action)
            self.ball_position += self.ball_velocity * sub_dt
            x, y = map(float, self.ball_position)
            if y - self.ball_radius < 2:
                self.ball_position[1] = 2 + self.ball_radius
                self.ball_velocity[1] = abs(self.ball_velocity[1])
            elif y + self.ball_radius > height - 2:
                self.ball_position[1] = height - 2 - self.ball_radius
                self.ball_velocity[1] = -abs(self.ball_velocity[1])
            if self._paddle_collision(x, y, left=True):
                self.ball_position[0] = self.left_x + self.paddle_width / 2 + self.ball_radius
                self.ball_velocity[0] = abs(self.ball_velocity[0]) * 1.015
                self.ball_velocity[1] += (y - self.left_y) * 1.2
            elif self._paddle_collision(x, y, left=False):
                self.ball_position[0] = (
                    self.right_x - self.paddle_width / 2 - self.ball_radius
                )
                self.ball_velocity[0] = -abs(self.ball_velocity[0]) * 1.015
                self.ball_velocity[1] += (y - self.right_y) * 1.2
            if self.ball_position[0] < -self.ball_radius:
                self.right_score += 1
                self._serve(direction=1)
            elif self.ball_position[0] > width + self.ball_radius:
                self.left_score += 1
                self._serve(direction=-1)
            self.elapsed += sub_dt

    def render(self) -> np.ndarray:
        scale = 2
        width = self.video_config.width
        height = self.video_config.height
        image = Image.new("RGB", (width * scale, height * scale), (3, 8, 16))
        draw = ImageDraw.Draw(image)
        white = (225, 239, 245)
        draw.rectangle((0, 0, width * scale, 3 * scale), fill=(42, 67, 78))
        draw.rectangle(
            (0, (height - 3) * scale, width * scale, height * scale),
            fill=(42, 67, 78),
        )
        for y in range(8, height - 8, 14):
            draw.rectangle(
                (
                    (width // 2 - 1) * scale,
                    y * scale,
                    (width // 2 + 1) * scale,
                    (y + 7) * scale,
                ),
                fill=(50, 75, 84),
            )
        for x, y in ((self.left_x, self.left_y), (self.right_x, self.right_y)):
            draw.rounded_rectangle(
                (
                    (x - self.paddle_width / 2) * scale,
                    (y - self.paddle_height / 2) * scale,
                    (x + self.paddle_width / 2) * scale,
                    (y + self.paddle_height / 2) * scale,
                ),
                radius=2 * scale,
                fill=white,
            )
        ball_x, ball_y = self.ball_position * scale
        radius = self.ball_radius * scale
        draw.ellipse(
            (ball_x - radius, ball_y - radius, ball_x + radius, ball_y + radius),
            fill=(255, 209, 72),
        )
        draw.text(
            (width * scale * 0.43, 7 * scale),
            f"{self.left_score % 10}   {self.right_score % 10}",
            fill=(115, 151, 160),
        )
        image = image.resize((width, height), Image.Resampling.LANCZOS)
        return np.asarray(image, dtype=np.uint8)

    def next_frame(self) -> np.ndarray:
        frame = self.render()
        self.step()
        return frame

    def sample_frames(self, count: int) -> np.ndarray:
        return np.stack([self.next_frame() for _ in range(count)])

    def _scripted_left_action(self) -> int:
        self._observed_ball_y.append(float(self.ball_position[1]))
        if self._action_hold_frames > 0:
            self._action_hold_frames -= 1
            return self._held_left_action
        self._action_hold_frames = int(self.rng.integers(1, 6))
        random_value = float(self.rng.random())
        delayed_target = self._observed_ball_y[0] + self.left_aim_bias
        error = delayed_target - self.left_y
        ideal_action = -1 if error < -5 else (1 if error > 5 else 0)
        if random_value < self.left_randomness * 0.35:
            action = 0
        elif random_value < self.left_randomness * 0.75:
            action = int(self.rng.choice((-1, 0, 1)))
        elif random_value < self.left_randomness:
            action = -ideal_action if ideal_action != 0 else int(
                self.rng.choice((-1, 1))
            )
        else:
            action = ideal_action
        self._held_left_action = action
        return action

    def sample_gameplay(self, count: int) -> tuple[np.ndarray, np.ndarray]:
        frames: list[np.ndarray] = []
        actions: list[tuple[float, float]] = []
        for _ in range(count):
            frames.append(self.render())
            action = self._scripted_left_action()
            actions.append((float(action < 0), float(action > 0)))
            self.step(action)
        return np.stack(frames), np.asarray(actions, dtype=np.float32)


_BREAKOUT_BRICK_COLORS = (
    (220, 72, 88),
    (230, 140, 54),
    (236, 196, 72),
    (72, 186, 120),
    (72, 148, 220),
    (156, 102, 220),
)

class FrameAutoencoderCodec(nn.Module):
    def __init__(self, config: CodecSettings, device: torch.device, dtype: torch.dtype) -> None:
        super().__init__()
        load_kwargs: dict[str, object] = {
            "revision": config.revision,
            "torch_dtype": dtype,
            "low_cpu_mem_usage": True,
        }
        if config.subfolder:
            load_kwargs["subfolder"] = config.subfolder
        try:
            self.vae = AutoencoderKL.from_pretrained(
                config.model_id, local_files_only=True, **load_kwargs
            )
        except OSError:
            self.vae = AutoencoderKL.from_pretrained(config.model_id, **load_kwargs)
        self.vae.to(device=device).eval()
        self.vae.to(memory_format=torch.channels_last)
        self.vae.requires_grad_(False)
        if config.enable_slicing:
            self.vae.enable_slicing()
        if config.enable_tiling:
            self.vae.enable_tiling()
        self.channels = int(self.vae.config.latent_channels)
        self.spatial_compression = 2 ** (len(self.vae.config.block_out_channels) - 1)
        self.temporal_compression = 1
        self.frame_batch_size = config.frame_batch_size
        self.sample_posterior = config.sample_posterior
        self.dtype = dtype
        self.scaling_factor = float(self.vae.config.scaling_factor)
        shift_factor = getattr(self.vae.config, "shift_factor", None)
        self.shift_factor = float(shift_factor) if shift_factor is not None else 0.0

    def latent_frame_count(self, frame_count: int) -> int:
        return frame_count

    @torch.no_grad()
    def encode(self, frames: torch.Tensor) -> torch.Tensor:
        batch, frame_count, channels, height, width = frames.shape
        flat_frames = frames.reshape(batch * frame_count, channels, height, width)
        encoded: list[torch.Tensor] = []
        for frame_batch in flat_frames.split(self.frame_batch_size):
            frame_batch = frame_batch.to(self.dtype).contiguous(
                memory_format=torch.channels_last
            )
            posterior = self.vae.encode(frame_batch).latent_dist
            latents = posterior.sample() if self.sample_posterior else posterior.mode()
            encoded.append((latents - self.shift_factor) * self.scaling_factor)
        merged = torch.cat(encoded)
        return merged.reshape(
            batch,
            frame_count,
            self.channels,
            height // self.spatial_compression,
            width // self.spatial_compression,
        )

    @torch.no_grad()
    def decode(self, latents: torch.Tensor) -> torch.Tensor:
        batch, frame_count, channels, height, width = latents.shape
        flat_latents = latents.reshape(batch * frame_count, channels, height, width)
        decoded: list[torch.Tensor] = []
        for latent_batch in flat_latents.split(self.frame_batch_size):
            denormalized = latent_batch / self.scaling_factor + self.shift_factor
            denormalized = denormalized.to(self.dtype).contiguous(
                memory_format=torch.channels_last
            )
            decoded.append(self.vae.decode(denormalized).sample)
        merged = torch.cat(decoded)
        return merged.reshape(
            batch,
            frame_count,
            3,
            height * self.spatial_compression,
            width * self.spatial_compression,
        ).clamp(-1.0, 1.0)
@dataclass(slots=True)
class PrefixCache:
    key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...]
    pooled: torch.Tensor


MemoryCache = PrefixCache


def _video_positions(
    frames: int,
    height: int,
    width: int,
    time_offset: int,
    device: torch.device,
) -> torch.Tensor:
    time = torch.arange(time_offset, time_offset + frames, device=device)
    y = torch.arange(height, device=device)
    x = torch.arange(width, device=device)
    grid = torch.meshgrid(time, y, x, indexing="ij")
    return torch.stack(grid, dim=-1).reshape(-1, 3).float()


class Rotary3D(nn.Module):
    def __init__(self, head_size: int, base: float = 10_000.0) -> None:
        super().__init__()
        self.rotary_size = (head_size // 6) * 6
        self.axis_size = self.rotary_size // 3
        frequencies = torch.exp(
            -math.log(base)
            * torch.arange(self.axis_size // 2, dtype=torch.float32)
            / (self.axis_size // 2)
        )
        self.register_buffer("frequencies", frequencies, persistent=False)

    @staticmethod
    def _rotate_half(inputs: torch.Tensor) -> torch.Tensor:
        pairs = inputs.reshape(*inputs.shape[:-1], -1, 2)
        first, second = pairs.unbind(dim=-1)
        return torch.stack((-second, first), dim=-1).flatten(-2)

    def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        angles = torch.cat(
            [positions[:, axis, None] * self.frequencies[None] for axis in range(3)],
            dim=-1,
        )
        cosine = angles.cos().repeat_interleave(2, dim=-1)[None, None]
        sine = angles.sin().repeat_interleave(2, dim=-1)[None, None]
        rotary, remainder = inputs.split(
            (self.rotary_size, inputs.shape[-1] - self.rotary_size),
            dim=-1,
        )
        rotated = rotary * cosine.to(rotary.dtype)
        rotated = rotated + self._rotate_half(rotary) * sine.to(rotary.dtype)
        return torch.cat((rotated, remainder), dim=-1)


class ScalarEmbedding(nn.Module):
    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
        )

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        value = value.reshape(-1).float()
        half = self.hidden_size // 2
        frequencies = torch.exp(
            -math.log(10_000.0)
            * torch.arange(half, device=value.device, dtype=torch.float32)
            / half
        )
        angles = value[:, None] * frequencies[None] * 1_000.0
        embedding = torch.cat((angles.cos(), angles.sin()), dim=-1)
        return self.mlp(embedding.to(dtype=self.mlp[0].weight.dtype))


class SwiGLU(nn.Module):
    def __init__(self, hidden_size: int, ratio: float) -> None:
        super().__init__()
        inner_size = int(hidden_size * ratio * 2 / 3)
        inner_size = ((inner_size + 63) // 64) * 64
        self.input = nn.Linear(hidden_size, inner_size * 2, bias=False)
        self.output = nn.Linear(inner_size, hidden_size, bias=False)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        values, gates = self.input(inputs).chunk(2, dim=-1)
        return self.output(values * F.silu(gates))


class SelfAttention(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int, qk_norm: bool) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_size = hidden_size // num_heads
        self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
        self.output = nn.Linear(hidden_size, hidden_size, bias=False)
        self.q_norm = (
            nn.RMSNorm(self.head_size, eps=1e-6)
            if qk_norm
            else nn.Identity()
        )
        self.k_norm = (
            nn.RMSNorm(self.head_size, eps=1e-6)
            if qk_norm
            else nn.Identity()
        )
        self.rotary = Rotary3D(self.head_size)

    def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        batch, tokens, channels = inputs.shape
        query, key, value = self.qkv(inputs).reshape(
            batch,
            tokens,
            3,
            self.num_heads,
            self.head_size,
        ).unbind(dim=2)
        query = self.rotary(self.q_norm(query).transpose(1, 2), positions)
        key = self.rotary(self.k_norm(key).transpose(1, 2), positions)
        value = value.transpose(1, 2)
        attended = F.scaled_dot_product_attention(query, key, value)
        return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels))


class PrefixAttention(nn.Module):
    """Target attention over the joint clean-prefix and noisy-target sequence."""

    def __init__(self, hidden_size: int, num_heads: int, qk_norm: bool) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_size = hidden_size // num_heads
        self.target_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
        self.prefix_key_value = nn.Linear(hidden_size, hidden_size * 2, bias=False)
        self.output = nn.Linear(hidden_size, hidden_size, bias=False)
        self.q_norm = (
            nn.RMSNorm(self.head_size, eps=1e-6)
            if qk_norm
            else nn.Identity()
        )
        self.k_norm = (
            nn.RMSNorm(self.head_size, eps=1e-6)
            if qk_norm
            else nn.Identity()
        )
        self.rotary = Rotary3D(self.head_size)

    def project_prefix(
        self,
        prefix: torch.Tensor,
        positions: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch, tokens, _ = prefix.shape
        key, value = self.prefix_key_value(prefix).reshape(
            batch,
            tokens,
            2,
            self.num_heads,
            self.head_size,
        ).unbind(dim=2)
        key = self.rotary(self.k_norm(key).transpose(1, 2), positions)
        return key, value.transpose(1, 2)

    def forward_cached(
        self,
        target: torch.Tensor,
        target_positions: torch.Tensor,
        prefix_key_value: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        batch, tokens, channels = target.shape
        query, key, value = self.target_qkv(target).reshape(
            batch,
            tokens,
            3,
            self.num_heads,
            self.head_size,
        ).unbind(dim=2)
        query = self.rotary(self.q_norm(query).transpose(1, 2), target_positions)
        key = self.rotary(self.k_norm(key).transpose(1, 2), target_positions)
        value = value.transpose(1, 2)
        prefix_key, prefix_value = prefix_key_value
        key = torch.cat((prefix_key, key), dim=2)
        value = torch.cat((prefix_value, value), dim=2)
        attended = F.scaled_dot_product_attention(query, key, value)
        return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels))


class PrefixDiTBlock(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        hidden_size = config.hidden_size
        self.prefix_norm1 = nn.RMSNorm(hidden_size, eps=1e-6)
        self.prefix_attention = SelfAttention(
            hidden_size,
            config.num_heads,
            config.qk_norm,
        )
        self.prefix_norm2 = nn.RMSNorm(hidden_size, eps=1e-6)
        self.prefix_mlp = SwiGLU(hidden_size, config.mlp_ratio)
        self.prefix_projection_norm = nn.RMSNorm(hidden_size, eps=1e-6)
        self.target_norm1 = nn.RMSNorm(
            hidden_size,
            eps=1e-6,
            elementwise_affine=False,
        )
        self.target_attention = PrefixAttention(
            hidden_size,
            config.num_heads,
            config.qk_norm,
        )
        self.target_norm2 = nn.RMSNorm(
            hidden_size,
            eps=1e-6,
            elementwise_affine=False,
        )
        self.target_mlp = SwiGLU(hidden_size, config.mlp_ratio)
        self.modulation = nn.Linear(hidden_size, hidden_size * 6)

    @staticmethod
    def _modulate(
        inputs: torch.Tensor,
        shift: torch.Tensor,
        scale: torch.Tensor,
    ) -> torch.Tensor:
        return inputs * (1.0 + scale[:, None]) + shift[:, None]

    def advance_prefix(
        self,
        prefix: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        prefix = prefix + self.prefix_attention(self.prefix_norm1(prefix), positions)
        return prefix + self.prefix_mlp(self.prefix_norm2(prefix))

    def prefix_key_value(
        self,
        prefix: torch.Tensor,
        positions: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        return self.target_attention.project_prefix(
            self.prefix_projection_norm(prefix),
            positions,
        )

    def advance_target(
        self,
        target: torch.Tensor,
        condition: torch.Tensor,
        target_positions: torch.Tensor,
        prefix_key_value: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        shift1, scale1, gate1, shift2, scale2, gate2 = self.modulation(
            F.silu(condition)
        ).chunk(6, dim=-1)
        normalized = self._modulate(self.target_norm1(target), shift1, scale1)
        target = target + gate1[:, None] * self.target_attention.forward_cached(
            normalized,
            target_positions,
            prefix_key_value,
        )
        normalized = self._modulate(self.target_norm2(target), shift2, scale2)
        return target + gate2[:, None] * self.target_mlp(normalized)

    def forward(
        self,
        prefix: torch.Tensor,
        target: torch.Tensor,
        condition: torch.Tensor,
        prefix_positions: torch.Tensor,
        target_positions: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        prefix = self.advance_prefix(prefix, prefix_positions)
        key_value = self.prefix_key_value(prefix, prefix_positions)
        target = self.advance_target(
            target,
            condition,
            target_positions,
            key_value,
        )
        return prefix, target


class CausalLatentVideoDiT(nn.Module):
    """Multiscale clean-prefix/noisy-suffix latent video transformer."""

    def __init__(
        self,
        model_config: ModelConfig,
        latent_channels: int,
        history_frames: int = 12,
    ) -> None:
        super().__init__()
        self.model_config = model_config
        self.latent_channels = latent_channels
        self.history_frames = history_frames
        self.gradient_checkpointing = False
        hidden_size = model_config.hidden_size
        target_patch = model_config.target_patch_size
        coarse_patch = model_config.memory_patch_size
        fine_patch = model_config.fine_memory_patch_size
        self.target_patch_embed = nn.Conv3d(
            latent_channels,
            hidden_size,
            kernel_size=(1, target_patch, target_patch),
            stride=(1, target_patch, target_patch),
        )
        self.coarse_prefix_embed = nn.Conv3d(
            latent_channels,
            hidden_size,
            kernel_size=(1, coarse_patch, coarse_patch),
            stride=(1, coarse_patch, coarse_patch),
        )
        self.fine_prefix_embed = nn.Conv3d(
            latent_channels,
            hidden_size,
            kernel_size=(1, fine_patch, fine_patch),
            stride=(1, fine_patch, fine_patch),
        )
        self.motion_prefix_embed = nn.Conv3d(
            latent_channels,
            hidden_size,
            kernel_size=(1, fine_patch, fine_patch),
            stride=(1, fine_patch, fine_patch),
        )
        self.coarse_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
        self.fine_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
        self.t_embedding = ScalarEmbedding(hidden_size)
        self.r_embedding = ScalarEmbedding(hidden_size)
        self.delta_embedding = ScalarEmbedding(hidden_size)
        self.time_projection = nn.Sequential(
            nn.Linear(hidden_size * 3, hidden_size),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
        )
        self.prefix_condition = nn.Linear(hidden_size, hidden_size)
        self.action_projection = (
            nn.Sequential(
                nn.Linear(model_config.action_dim, hidden_size),
                nn.SiLU(),
                nn.Linear(hidden_size, hidden_size),
            )
            if model_config.action_dim > 0
            else None
        )
        self.context_projection = (
            nn.Linear(model_config.context_dim, hidden_size, bias=False)
            if model_config.context_dim > 0
            else None
        )
        self.blocks = nn.ModuleList(
            PrefixDiTBlock(model_config) for _ in range(model_config.depth)
        )
        self.output_norm = nn.RMSNorm(
            hidden_size,
            eps=1e-6,
            elementwise_affine=False,
        )
        self.output_modulation = nn.Linear(hidden_size, hidden_size * 2)
        self.output_projection = nn.Linear(
            hidden_size,
            target_patch * target_patch * latent_channels,
        )
        self.outcome_norm = nn.RMSNorm(hidden_size, eps=1e-6)
        self.reward_head = nn.Linear(hidden_size, 1)
        self.terminal_head = nn.Linear(hidden_size, 1)
        self.reset_parameters()

    def enable_activation_checkpointing(self, enabled: bool = True) -> None:
        self.gradient_checkpointing = enabled

    def reset_parameters(self) -> None:
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
        for embedding in (
            self.target_patch_embed,
            self.coarse_prefix_embed,
            self.fine_prefix_embed,
            self.motion_prefix_embed,
        ):
            nn.init.normal_(embedding.weight, std=0.02)
        nn.init.normal_(self.coarse_scale_embedding, std=0.02)
        nn.init.normal_(self.fine_scale_embedding, std=0.02)
        for block in self.blocks:
            nn.init.zeros_(block.modulation.weight)
            nn.init.zeros_(block.modulation.bias)
        nn.init.zeros_(self.output_modulation.weight)
        nn.init.zeros_(self.output_modulation.bias)
        nn.init.zeros_(self.output_projection.weight)
        nn.init.zeros_(self.output_projection.bias)
        nn.init.zeros_(self.reward_head.weight)
        nn.init.zeros_(self.reward_head.bias)
        nn.init.zeros_(self.terminal_head.weight)
        nn.init.zeros_(self.terminal_head.bias)

    @staticmethod
    def _patch(
        embedding: nn.Conv3d,
        video: torch.Tensor,
        time_offset: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        tokens = embedding(video.transpose(1, 2))
        _, _, frames, height, width = tokens.shape
        positions = _video_positions(
            frames,
            height,
            width,
            time_offset,
            video.device,
        )
        return tokens.flatten(2).transpose(1, 2), positions

    def _encode_prefix(
        self,
        history: torch.Tensor,
        context: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if history.shape[1] != self.history_frames:
            raise ValueError(
                f"Expected {self.history_frames} history frames, got {history.shape[1]}."
            )
        coarse, coarse_positions = self._patch(
            self.coarse_prefix_embed,
            history,
            0,
        )
        fine_frames = self.model_config.fine_history_frames
        fine_history = history[:, -fine_frames:]
        previous = history[:, -fine_frames - 1 : -1]
        if previous.shape[1] != fine_frames:
            previous = torch.cat((fine_history[:, :1], fine_history[:, :-1]), dim=1)
        motion = fine_history - previous
        fine, fine_positions = self._patch(
            self.fine_prefix_embed,
            fine_history,
            self.history_frames - fine_frames,
        )
        motion_tokens, _ = self._patch(
            self.motion_prefix_embed,
            motion,
            self.history_frames - fine_frames,
        )
        coarse = coarse + self.coarse_scale_embedding
        fine = fine + motion_tokens + self.fine_scale_embedding
        prefix = torch.cat((coarse, fine), dim=1)
        positions = torch.cat((coarse_positions, fine_positions), dim=0)
        if self.context_projection is not None and context is not None:
            projected = self.context_projection(context)
            prefix = torch.cat((prefix, projected), dim=1)
            positions = torch.cat(
                (
                    positions,
                    torch.zeros(projected.shape[1], 3, device=history.device),
                ),
                dim=0,
            )
        return prefix, positions

    def _condition(
        self,
        prefix: torch.Tensor,
        r: torch.Tensor,
        t: torch.Tensor,
    ) -> torch.Tensor:
        time = self.time_projection(
            torch.cat(
                (
                    self.t_embedding(t),
                    self.r_embedding(r),
                    self.delta_embedding(t - r),
                ),
                dim=-1,
            )
        )
        return time + self.prefix_condition(prefix.mean(dim=1))

    def build_memory_cache(
        self,
        history: torch.Tensor,
        context: torch.Tensor | None = None,
    ) -> PrefixCache:
        prefix, positions = self._encode_prefix(history, context)
        pooled = self.prefix_condition(prefix.mean(dim=1))
        key_values: list[tuple[torch.Tensor, torch.Tensor]] = []
        for block in self.blocks:
            prefix = block.advance_prefix(prefix, positions)
            key_values.append(block.prefix_key_value(prefix, positions))
        return PrefixCache(tuple(key_values), pooled)

    def predict_outcomes(
        self,
        history: torch.Tensor,
        actions: torch.Tensor | None = None,
        frame_count: int = 1,
        memory_cache: PrefixCache | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if memory_cache is None:
            prefix, _ = self._encode_prefix(history)
            pooled = self.prefix_condition(prefix.mean(dim=1))
        else:
            pooled = memory_cache.pooled
        hidden = pooled[:, None].expand(-1, frame_count, -1)
        if self.action_projection is not None:
            if actions is None or actions.shape[1] != frame_count:
                raise ValueError("Outcome actions must match frame_count.")
            hidden = hidden + self.action_projection(actions)
        hidden = self.outcome_norm(hidden)
        return self.reward_head(hidden).squeeze(-1), self.terminal_head(hidden).squeeze(-1)

    def _unpatchify(
        self,
        patches: torch.Tensor,
        frames: int,
        height: int,
        width: int,
    ) -> torch.Tensor:
        batch = patches.shape[0]
        patch = self.model_config.target_patch_size
        patches = patches.reshape(
            batch,
            frames,
            height,
            width,
            patch,
            patch,
            self.latent_channels,
        )
        return patches.permute(0, 1, 6, 2, 4, 3, 5).reshape(
            batch,
            frames,
            self.latent_channels,
            height * patch,
            width * patch,
        )

    def forward(
        self,
        noisy_target: torch.Tensor,
        history: torch.Tensor,
        r: torch.Tensor,
        t: torch.Tensor,
        actions: torch.Tensor | None = None,
        context: torch.Tensor | None = None,
        memory_cache: PrefixCache | None = None,
    ) -> torch.Tensor:
        target_volume = self.target_patch_embed(noisy_target.transpose(1, 2))
        batch, _, frames, height, width = target_volume.shape
        target = target_volume.flatten(2).transpose(1, 2)
        target_positions = _video_positions(
            frames,
            height,
            width,
            self.history_frames,
            noisy_target.device,
        )
        if memory_cache is None:
            prefix, prefix_positions = self._encode_prefix(history, context)
            condition = self._condition(prefix, r, t)
        else:
            prefix = None
            prefix_positions = None
            condition = self.time_projection(
                torch.cat(
                    (
                        self.t_embedding(t),
                        self.r_embedding(r),
                        self.delta_embedding(t - r),
                    ),
                    dim=-1,
                )
            ) + memory_cache.pooled
        if self.action_projection is not None:
            if actions is None or actions.shape[:2] != (batch, frames):
                raise ValueError("Actions must match [batch, target_frames, action_dim].")
            action_embedding = self.action_projection(actions)
            target = target + action_embedding.repeat_interleave(height * width, dim=1)
            condition = condition + action_embedding.mean(dim=1)
        for block_index, block in enumerate(self.blocks):
            if memory_cache is None:
                if prefix is None or prefix_positions is None:
                    raise RuntimeError("Prefix state is unavailable.")
                if self.training and self.gradient_checkpointing:
                    prefix, target = checkpoint(
                        block,
                        prefix,
                        target,
                        condition,
                        prefix_positions,
                        target_positions,
                        use_reentrant=False,
                    )
                else:
                    prefix, target = block(
                        prefix,
                        target,
                        condition,
                        prefix_positions,
                        target_positions,
                    )
            else:
                target = block.advance_target(
                    target,
                    condition,
                    target_positions,
                    memory_cache.key_values[block_index],
                )
        shift, scale = self.output_modulation(F.silu(condition)).chunk(2, dim=-1)
        target = self.output_norm(target) * (1.0 + scale[:, None]) + shift[:, None]
        return self._unpatchify(
            self.output_projection(target),
            frames,
            height,
            width,
        )


CausalVideoDiT = CausalLatentVideoDiT
@torch.inference_mode()
def sample_next_latent(
    model: nn.Module,
    history: torch.Tensor,
    steps: int = 2,
    actions: torch.Tensor | None = None,
    generator: torch.Generator | None = None,
    window_frames: int = 1,
    memory_cache: object | None = None,
    instantaneous: bool = False,
) -> torch.Tensor:
    if steps < 1:
        raise ValueError("steps must be positive.")
    batch, _, channels, height, width = history.shape
    state = torch.randn(
        (batch, window_frames, channels, height, width),
        device=history.device,
        dtype=history.dtype,
        generator=generator,
    )
    boundaries = torch.linspace(1.0, 0.0, steps + 1, device=history.device)
    for index in range(steps):
        t_value = boundaries[index]
        r_value = boundaries[index + 1]
        t = t_value.expand(batch)
        r = r_value.expand(batch)
        model_r = t if instantaneous else r
        average_velocity = model(
            state,
            history,
            model_r,
            t,
            actions=actions,
            memory_cache=memory_cache,
        )
        state = state - (t_value - r_value).to(state.dtype) * average_velocity
    return state

def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Live Diffusion Pong Base (unconditional).")
    parser.add_argument("--repo", default=DEFAULT_REPO)
    parser.add_argument("--revision", default="main")
    parser.add_argument("--local-dir", default=None)
    parser.add_argument("--steps", type=int, default=2)
    parser.add_argument("--frames", type=int, default=0, help="0 = until window closed")
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--window-scale", type=int, default=6)
    parser.add_argument("--fps-cap", type=float, default=0.0)
    return parser.parse_args()


def _download(repo: str, revision: str, local_dir: str | None) -> Path:
    if local_dir:
        root = Path(local_dir)
        root.mkdir(parents=True, exist_ok=True)
        for name in ("config.json", "ema.safetensors"):
            hf_hub_download(
                repo_id=repo,
                filename=name,
                revision=revision,
                local_dir=str(root),
            )
        return root
    config_path = Path(
        hf_hub_download(repo_id=repo, filename="config.json", revision=revision)
    )
    hf_hub_download(repo_id=repo, filename="ema.safetensors", revision=revision)
    return config_path.parent


def _load_settings(config_path: Path) -> tuple[VideoConfig, ModelConfig, CodecSettings]:
    raw = json.loads(config_path.read_text(encoding="utf-8"))
    video = VideoConfig(**{key: raw["video"][key] for key in VideoConfig.__dataclass_fields__})
    model = ModelConfig(**{key: raw["model"][key] for key in ModelConfig.__dataclass_fields__})
    codec_raw = raw["codec"]
    codec = CodecSettings(
        model_id=codec_raw.get("model_id", "madebyollin/sdxl-vae-fp16-fix"),
        subfolder=codec_raw.get("subfolder", "") or "",
        revision=codec_raw.get("revision", "main"),
        latent_channels=int(codec_raw.get("latent_channels", 4)),
        spatial_compression=int(codec_raw.get("spatial_compression", 8)),
        frame_batch_size=int(codec_raw.get("frame_batch_size", 256)),
        sample_posterior=bool(codec_raw.get("sample_posterior", False)),
        enable_slicing=bool(codec_raw.get("enable_slicing", True)),
        enable_tiling=bool(codec_raw.get("enable_tiling", False)),
    )
    return video, model, codec


def _frames_to_tensor(frames: np.ndarray, device: torch.device) -> torch.Tensor:
    tensor = torch.from_numpy(frames.copy()).permute(0, 3, 1, 2).float()
    return tensor.div_(127.5).sub_(1.0).unsqueeze(0).to(device, non_blocking=True)


def _tensor_to_image(frame: torch.Tensor) -> Image.Image:
    array = frame.float().clamp(-1.0, 1.0).add(1.0).mul(127.5)
    return Image.fromarray(array.byte().permute(1, 2, 0).cpu().numpy(), mode="RGB")


def _inject_ball(image: Image.Image, position: tuple[float, float]) -> Image.Image:
    image = image.copy()
    draw = ImageDraw.Draw(image)
    radius = max(3.0, min(image.width, image.height) * 0.025)
    x, y = position
    draw.ellipse(
        (x - radius, y - radius, x + radius, y + radius),
        fill=(255, 209, 72),
        outline=(255, 247, 196),
        width=max(1, round(radius * 0.25)),
    )
    return image


def _image_to_video_tensor(image: Image.Image, device: torch.device) -> torch.Tensor:
    array = np.asarray(image, dtype=np.uint8).copy()
    tensor = torch.from_numpy(array).permute(2, 0, 1).float()
    return (
        tensor.div_(127.5)
        .sub_(1.0)
        .reshape(1, 1, 3, image.height, image.width)
        .to(device, non_blocking=True)
    )


class LiveWindow:
    def __init__(self, scale: int) -> None:
        self.scale = max(1, scale)
        self.is_open = True
        self.root = tk.Tk()
        self.root.title("Diffusion Pong Base")
        self.label = tk.Label(self.root, background="black")
        self.label.pack()
        self.photo: ImageTk.PhotoImage | None = None
        self.generated_width = 0
        self.generated_height = 0
        self.pending_click: tuple[float, float] | None = None
        self.pressed: set[str] = set()
        self.label.bind("<Button-1>", self._on_click)
        self.root.bind("<KeyPress>", self._on_key_press)
        self.root.bind("<KeyRelease>", self._on_key_release)
        self.root.protocol("WM_DELETE_WINDOW", self.close)
        self.root.focus_force()

    def _on_key_press(self, event: tk.Event) -> None:
        key = str(event.keysym).lower()
        if key in {"w", "s", "a", "d", "left", "right"}:
            self.pressed.add(key)

    def _on_key_release(self, event: tk.Event) -> None:
        self.pressed.discard(str(event.keysym).lower())

    def action(self) -> tuple[float, float]:
        negative = bool(self.pressed & {"w", "a", "left"})
        positive = bool(self.pressed & {"s", "d", "right"})
        if negative and positive:
            return (0.0, 0.0)
        return (float(negative), float(positive))

    def _on_click(self, event: tk.Event) -> None:
        if self.generated_width <= 0:
            return
        x = (float(event.x) / self.scale) % self.generated_width
        y = min(max(float(event.y) / self.scale, 0.0), self.generated_height - 1.0)
        self.pending_click = (x, y)

    def consume_click(self) -> tuple[float, float] | None:
        click = self.pending_click
        self.pending_click = None
        return click

    def close(self) -> None:
        if self.is_open:
            self.is_open = False
            self.root.destroy()

    def show(self, image: Image.Image, frame_index: int, latency_ms: float) -> bool:
        if not self.is_open:
            return False
        self.generated_width, self.generated_height = image.size
        display = image.resize(
            (image.width * self.scale, image.height * self.scale),
            Image.Resampling.NEAREST,
        )
        self.photo = ImageTk.PhotoImage(display)
        self.label.configure(image=self.photo)
        held = "+".join(sorted(self.pressed)).upper()
        held_text = f"  [{held}]" if held else ""
        self.root.title(
            f"Diffusion Pong Base — frame {frame_index + 1}  {latency_ms:.0f} ms{held_text}"
        )
        try:
            self.root.update_idletasks()
            self.root.update()
        except tk.TclError:
            self.is_open = False
        return self.is_open


@torch.inference_mode()
def main() -> None:
    args = _parse_args()
    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
        raise SystemExit("Need a CUDA GPU with BF16 support.")

    print(f"Downloading {args.repo} …")
    checkpoint = _download(args.repo, args.revision, args.local_dir)
    video_cfg, model_cfg, codec_cfg = _load_settings(checkpoint / "config.json")
    weights = checkpoint / "ema.safetensors"
    if not weights.exists():
        raise SystemExit(f"Missing {weights}")

    device = torch.device("cuda")
    torch.set_float32_matmul_precision("high")
    seed = args.seed if args.seed is not None else random.SystemRandom().randrange(2**31)
    simulation = PongSimulation.from_seed(video_cfg, seed)
    for _ in range(random.Random(seed).randint(0, 400)):
        simulation.step()
    history_rgb = _frames_to_tensor(
        simulation.sample_frames(video_cfg.history_frames),
        device,
    )

    print("Loading VAE + DiT …")
    codec = FrameAutoencoderCodec(codec_cfg, device, torch.bfloat16)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        latents = codec.encode(history_rgb)
    history: deque[torch.Tensor] = deque(
        (latents[:, index] for index in range(video_cfg.history_frames)),
        maxlen=video_cfg.history_frames,
    )
    model = (
        CausalLatentVideoDiT(
            model_cfg,
            latent_channels=codec.channels,
            history_frames=codec.latent_frame_count(video_cfg.history_frames),
        )
        .to(device)
        .eval()
    )
    model.load_state_dict(load_file(str(weights)), strict=True)

    steps = max(1, args.steps)
    frame_limit = args.frames if args.frames > 0 else 10_000_000
    window = LiveWindow(args.window_scale)
    generator = torch.Generator(device=device).manual_seed(seed + 7)
    inject_pos: tuple[float, float] | None = None
    inject_frames = 0
    print("Ready. Close the window to stop.")

    for frame_index in range(frame_limit):
        click = window.consume_click()
        if click is not None:
            inject_pos = click
            inject_frames = 3
        history_tensor = torch.stack(tuple(history), dim=1)
        action = torch.tensor(
            [window.action()],
            device=device,
            dtype=torch.float32,
        ).unsqueeze(0)
        torch.cuda.synchronize()
        started = time.perf_counter()
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            memory_cache = model.build_memory_cache(history_tensor)
            prediction = sample_next_latent(
                model,
                history_tensor,
                steps,
                actions=action if model_cfg.action_dim > 0 else None,
                generator=generator,
                window_frames=1,
                memory_cache=memory_cache,
                instantaneous=True,
            )
            decoded = codec.decode(prediction)
        torch.cuda.synchronize()
        latency_ms = (time.perf_counter() - started) * 1000
        image = _tensor_to_image(decoded[0, 0])
        frame_latent = prediction[:, 0]
        if inject_pos is not None and inject_frames > 0:
            image = _inject_ball(image, inject_pos)
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                frame_latent = codec.encode(_image_to_video_tensor(image, device))[:, 0]
            inject_frames -= 1
            if inject_frames == 0:
                inject_pos = None
        history.append(frame_latent)
        if not window.show(image, frame_index, latency_ms):
            break
        if args.fps_cap > 0:
            target = 1.0 / args.fps_cap
            elapsed = time.perf_counter() - started
            if elapsed < target:
                time.sleep(target - elapsed)
    if window.is_open:
        window.close()
    print("Done.")


if __name__ == "__main__":
    main()