live_infer.py

#!/usr/bin/env python3
"""Run Diffusion Breakout.

Deps:
  pip install torch numpy pillow safetensors huggingface_hub diffusers

Run:
  python live_infer.py
  python live_infer.py --steps 1 --window-scale 7 --seed 42

Controls: A/D or arrow keys. Click window for focus.
"""
from __future__ import annotations

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

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

DEFAULT_REPO = "kerzgrr/diffusionbreakout"


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


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


@dataclass(frozen=True, slots=True)
class CodecSettings:
    model_id: str = "madebyollin/sdxl-vae-fp16-fix"
    subfolder: str = ""
    revision: str = "main"
    latent_channels: int = 4
    spatial_compression: int = 8
    frame_batch_size: int = 256
    sample_posterior: bool = False
    enable_slicing: bool = True
    enable_tiling: bool = False
_BREAKOUT_BRICK_COLORS = (
    (220, 72, 88),
    (230, 140, 54),
    (236, 196, 72),
    (72, 186, 120),
    (72, 148, 220),
    (156, 102, 220),
)


class BreakoutSimulation:
    """Classic Breakout with continuous physics and a scripted bottom paddle."""

    def __init__(self, video_config: VideoConfig, seed: int) -> None:
        self.video_config = video_config
        self.rng = np.random.default_rng(seed)
        self.dt = 1.0 / video_config.fps
        self.paddle_width = max(18.0, video_config.width * 0.18)
        self.paddle_height = max(4.0, video_config.height * 0.03)
        self.ball_radius = max(3.0, min(video_config.width, video_config.height) * 0.025)
        self.paddle_y = video_config.height * 0.90
        self.brick_rows = 5
        self.brick_cols = 8
        self.brick_top = video_config.height * 0.14
        self.brick_gap = 1.5
        usable_width = video_config.width - 10.0
        self.brick_width = (usable_width - self.brick_gap * (self.brick_cols - 1)) / self.brick_cols
        self.brick_height = max(5.0, video_config.height * 0.045)
        self.paddle_x = video_config.width / 2
        self.score = 0
        self.elapsed = 0.0
        self.aim_bias = float(self.rng.uniform(-8.0, 8.0))
        self.reaction_delay = int(self.rng.integers(0, 4))
        self.randomness = float(self.rng.uniform(0.10, 0.28))
        self._observed_ball_x: deque[float] = deque(
            [video_config.width / 2] * (self.reaction_delay + 1),
            maxlen=self.reaction_delay + 1,
        )
        self._held_action = 0
        self._action_hold_frames = 0
        self.ball_position = np.zeros(2, dtype=np.float32)
        self.ball_velocity = np.zeros(2, dtype=np.float32)
        self.bricks = np.ones((self.brick_rows, self.brick_cols), dtype=bool)
        # Punch a few random holes so early frames aren't always a full wall.
        missing = int(self.rng.integers(0, 6))
        for _ in range(missing):
            row = int(self.rng.integers(0, self.brick_rows))
            col = int(self.rng.integers(0, self.brick_cols))
            self.bricks[row, col] = False
        self._serve()
        for _ in range(10):
            self.step()

    @classmethod
    def from_seed(
        cls,
        video_config: VideoConfig,
        seed: int,
        *,
        at_start: bool = False,
    ) -> BreakoutSimulation:
        simulation = cls(video_config, seed)
        if at_start:
            simulation.reset_to_start()
        return simulation

    def reset_to_start(self) -> None:
        """Reset to a fresh serve against a full brick wall (for inference)."""
        self.paddle_x = self.video_config.width / 2
        self.score = 0
        self.elapsed = 0.0
        self._held_action = 0
        self._action_hold_frames = 0
        self._observed_ball_x.clear()
        self._observed_ball_x.extend(
            [self.video_config.width / 2] * (self.reaction_delay + 1)
        )
        self.bricks[:] = True
        self._serve()

    def _brick_rect(self, row: int, col: int) -> tuple[float, float, float, float]:
        x0 = 5.0 + col * (self.brick_width + self.brick_gap)
        y0 = self.brick_top + row * (self.brick_height + self.brick_gap)
        return x0, y0, x0 + self.brick_width, y0 + self.brick_height

    def _serve(self) -> None:
        angle = float(self.rng.uniform(-0.65, 0.65))
        speed = float(self.rng.uniform(48.0, 64.0))
        self.ball_position[:] = (
            self.paddle_x + float(self.rng.uniform(-6.0, 6.0)),
            self.paddle_y - self.paddle_height - self.ball_radius - 2.0,
        )
        self.ball_velocity[:] = (
            speed * math.sin(angle),
            -speed * math.cos(angle),
        )

    def _reset_bricks(self) -> None:
        self.bricks[:] = True

    def _move_paddle(self, dt: float, action: int | None) -> None:
        width = float(self.video_config.width)
        paddle_speed = width * 0.85
        if action is None:
            target = float(self.ball_position[0]) + self.aim_bias
            delta = max(-paddle_speed * dt, min(paddle_speed * dt, target - self.paddle_x))
        else:
            delta = float(max(-1, min(1, action))) * paddle_speed * dt
        self.paddle_x += delta
        half = self.paddle_width / 2
        self.paddle_x = max(half + 2.0, min(width - half - 2.0, self.paddle_x))

    def _collide_bricks(self) -> None:
        x, y = map(float, self.ball_position)
        r = self.ball_radius
        for row in range(self.brick_rows):
            for col in range(self.brick_cols):
                if not self.bricks[row, col]:
                    continue
                x0, y0, x1, y1 = self._brick_rect(row, col)
                closest_x = max(x0, min(x, x1))
                closest_y = max(y0, min(y, y1))
                dx = x - closest_x
                dy = y - closest_y
                if dx * dx + dy * dy > r * r:
                    continue
                self.bricks[row, col] = False
                self.score += 1
                # Reflect on the dominant penetration axis.
                if abs(dx) > abs(dy):
                    self.ball_velocity[0] = -self.ball_velocity[0]
                    self.ball_position[0] = x0 - r if dx < 0 else x1 + r
                else:
                    self.ball_velocity[1] = -self.ball_velocity[1]
                    self.ball_position[1] = y0 - r if dy < 0 else y1 + r
                if not self.bricks.any():
                    self._reset_bricks()
                return

    def step(self, action: int | None = None) -> None:
        width = float(self.video_config.width)
        height = float(self.video_config.height)
        substeps = 4
        sub_dt = self.dt / substeps
        for _ in range(substeps):
            self._move_paddle(sub_dt, action)
            self.ball_position += self.ball_velocity * sub_dt
            x, y = map(float, self.ball_position)
            if x - self.ball_radius < 2:
                self.ball_position[0] = 2 + self.ball_radius
                self.ball_velocity[0] = abs(self.ball_velocity[0])
            elif x + self.ball_radius > width - 2:
                self.ball_position[0] = width - 2 - self.ball_radius
                self.ball_velocity[0] = -abs(self.ball_velocity[0])
            if y - self.ball_radius < 2:
                self.ball_position[1] = 2 + self.ball_radius
                self.ball_velocity[1] = abs(self.ball_velocity[1])
            paddle_top = self.paddle_y - self.paddle_height / 2
            hitting_paddle = (
                self.ball_velocity[1] > 0
                and y + self.ball_radius >= paddle_top
                and y - self.ball_radius <= self.paddle_y + self.paddle_height / 2
                and abs(x - self.paddle_x) <= self.paddle_width / 2 + self.ball_radius
            )
            if hitting_paddle:
                self.ball_position[1] = paddle_top - self.ball_radius
                self.ball_velocity[1] = -abs(self.ball_velocity[1]) * 1.01
                self.ball_velocity[0] += (x - self.paddle_x) * 1.35
                speed = float(np.linalg.norm(self.ball_velocity))
                target = float(np.clip(speed, 46.0, 78.0))
                if speed > 1e-5:
                    self.ball_velocity *= target / speed
            self._collide_bricks()
            if self.ball_position[1] > height + self.ball_radius * 2:
                self._serve()
            self.elapsed += sub_dt

    def render(self) -> np.ndarray:
        scale = 2
        width = self.video_config.width
        height = self.video_config.height
        image = Image.new("RGB", (width * scale, height * scale), (6, 8, 18))
        draw = ImageDraw.Draw(image)
        draw.rectangle((0, 0, width * scale, 3 * scale), fill=(36, 48, 72))
        for row in range(self.brick_rows):
            color = _BREAKOUT_BRICK_COLORS[row % len(_BREAKOUT_BRICK_COLORS)]
            for col in range(self.brick_cols):
                if not self.bricks[row, col]:
                    continue
                x0, y0, x1, y1 = self._brick_rect(row, col)
                draw.rounded_rectangle(
                    (x0 * scale, y0 * scale, x1 * scale, y1 * scale),
                    radius=max(1, round(scale)),
                    fill=color,
                )
        draw.rounded_rectangle(
            (
                (self.paddle_x - self.paddle_width / 2) * scale,
                (self.paddle_y - self.paddle_height / 2) * scale,
                (self.paddle_x + self.paddle_width / 2) * scale,
                (self.paddle_y + self.paddle_height / 2) * scale,
            ),
            radius=2 * scale,
            fill=(230, 238, 245),
        )
        ball_x, ball_y = self.ball_position * scale
        radius = self.ball_radius * scale
        draw.ellipse(
            (ball_x - radius, ball_y - radius, ball_x + radius, ball_y + radius),
            fill=(255, 209, 72),
        )
        image = image.resize((width, height), Image.Resampling.LANCZOS)
        return np.asarray(image, dtype=np.uint8)

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

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

    def _scripted_action(self) -> int:
        self._observed_ball_x.append(float(self.ball_position[0]))
        if self._action_hold_frames > 0:
            self._action_hold_frames -= 1
            return self._held_action
        self._action_hold_frames = int(self.rng.integers(1, 5))
        random_value = float(self.rng.random())
        delayed_target = self._observed_ball_x[0] + self.aim_bias
        error = delayed_target - self.paddle_x
        ideal_action = -1 if error < -4 else (1 if error > 4 else 0)
        if random_value < self.randomness * 0.35:
            action = 0
        elif random_value < self.randomness * 0.75:
            action = int(self.rng.choice((-1, 0, 1)))
        elif random_value < self.randomness:
            action = -ideal_action if ideal_action != 0 else int(self.rng.choice((-1, 1)))
        else:
            action = ideal_action
        self._held_action = action
        return action

    def sample_gameplay(self, count: int) -> tuple[np.ndarray, np.ndarray]:
        frames: list[np.ndarray] = []
        actions: list[tuple[float, float]] = []
        for _ in range(count):
            frames.append(self.render())
            action = self._scripted_action()
            # [A/left, D/right] — same layout we'll use for action fine-tuning later.
            actions.append((float(action < 0), float(action > 0)))
            self.step(action)
        return np.stack(frames), np.asarray(actions, dtype=np.float32)


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

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

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

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


MemoryCache = PrefixCache


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


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

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

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


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

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


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

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


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

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


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

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

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

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


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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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


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

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


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


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


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


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


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


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


def _inject_ball_into_history(
    history: deque[torch.Tensor],
    codec: FrameAutoencoderCodec,
    position: tuple[float, float],
) -> Image.Image:
    history_tensor = torch.stack(tuple(history), dim=1)
    device = history_tensor.device
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        decoded = codec.decode(history_tensor)
    painted_images: list[Image.Image] = []
    painted_frames: list[torch.Tensor] = []
    for index in range(decoded.shape[1]):
        image = _inject_ball(_tensor_to_image(decoded[0, index]), position)
        painted_images.append(image)
        painted_frames.append(_rgb_image_to_frame_tensor(image, device))
    rgb = torch.stack(painted_frames, dim=0).unsqueeze(0)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        latents = codec.encode(rgb)
    history.clear()
    history.extend(latents[:, index] for index in range(latents.shape[1]))
    return painted_images[-1]


def _encode_ball_frame(
    image: Image.Image,
    position: tuple[float, float],
    codec: FrameAutoencoderCodec,
    device: torch.device,
) -> tuple[Image.Image, torch.Tensor]:
    painted = _inject_ball(image, position)
    rgb = _rgb_image_to_frame_tensor(painted, device).reshape(
        1, 1, 3, painted.height, painted.width
    )
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        latent = codec.encode(rgb)[:, 0]
    return painted, latent


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

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

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

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

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

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

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

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


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

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

    device = torch.device("cuda")
    torch.set_float32_matmul_precision("high")
    seed = args.seed if args.seed is not None else random.SystemRandom().randrange(2**31)
    simulation = BreakoutSimulation.from_seed(video_cfg, seed, at_start=True)
    history_rgb = _frames_to_tensor(
        simulation.sample_frames(video_cfg.history_frames),
        device,
    )

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

    steps = max(1, args.steps)
    frame_limit = args.frames if args.frames > 0 else 10_000_000
    window = LiveWindow(args.window_scale)
    generator = torch.Generator(device=device).manual_seed(seed + 7)
    injected_position: tuple[float, float] | None = None
    force_ball_frames = 0
    print("Ready. Click the window, then hold A/D or arrow keys.")

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


if __name__ == "__main__":
    main()