#!/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("", self._on_click) self.root.bind("", self._on_key_press) self.root.bind("", 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()