#!/usr/bin/env python3 """Roll out Diffusion Pong Base (unconditional). Deps: pip install torch numpy pillow safetensors huggingface_hub diffusers Run: python live_infer_base.py python live_infer_base.py --steps 2 --window-scale 7 --seed 42 Unconditional next-frame model — no paddle controls. Click canvas to inject a ball. """ from __future__ import annotations import argparse import json import math import random import time import tkinter as tk from collections import deque from dataclasses import dataclass from pathlib import Path import numpy as np import torch import torch.nn.functional as F from diffusers import AutoencoderKL from huggingface_hub import hf_hub_download from PIL import Image, ImageDraw, ImageTk from safetensors.torch import load_file from torch import nn from torch.utils.checkpoint import checkpoint DEFAULT_REPO = "kerzgrr/diffusionpong-base" @dataclass(frozen=True, slots=True) class VideoConfig: height: int = 128 width: int = 128 fps: float = 6.0 history_frames: int = 12 future_frames: int = 5 @dataclass(frozen=True, slots=True) class ModelConfig: target_patch_size: int = 2 memory_patch_size: int = 4 fine_memory_patch_size: int = 2 fine_history_frames: int = 2 hidden_size: int = 384 depth: int = 12 memory_depth: int = 3 num_heads: int = 8 mlp_ratio: float = 4.0 qk_norm: bool = True generation_frames: int = 1 action_dim: int = 0 context_dim: int = 0 @dataclass(frozen=True, slots=True) class CodecSettings: model_id: str = "madebyollin/sdxl-vae-fp16-fix" subfolder: str = "" revision: str = "main" latent_channels: int = 4 spatial_compression: int = 8 frame_batch_size: int = 256 sample_posterior: bool = False enable_slicing: bool = True enable_tiling: bool = False class PongSimulation: """Deterministic Pong gameplay with continuous physics and scripted paddles.""" def __init__(self, video_config: VideoConfig, seed: int) -> None: self.video_config = video_config self.rng = np.random.default_rng(seed) self.dt = 1.0 / video_config.fps self.paddle_height = max(18.0, video_config.height * 0.20) self.paddle_width = max(4.0, video_config.width * 0.025) self.ball_radius = max(3.0, min(video_config.width, video_config.height) * 0.025) self.left_x = video_config.width * 0.07 self.right_x = video_config.width * 0.93 self.left_y = video_config.height / 2 self.right_y = video_config.height / 2 self.left_score = 0 self.right_score = 0 self.elapsed = 0.0 self.phase = float(self.rng.uniform(0.0, math.tau)) self.left_reaction_delay = int(self.rng.integers(0, 5)) self.left_aim_bias = float(self.rng.uniform(-10.0, 10.0)) self.left_randomness = float(self.rng.uniform(0.12, 0.30)) self.right_aim_bias = float(self.rng.uniform(-22.0, 22.0)) self.right_speed_scale = float(self.rng.uniform(0.38, 0.58)) self._observed_ball_y: deque[float] = deque( [video_config.height / 2] * (self.left_reaction_delay + 1), maxlen=self.left_reaction_delay + 1, ) self._held_left_action = 0 self._action_hold_frames = 0 self.ball_position = np.zeros(2, dtype=np.float32) self.ball_velocity = np.zeros(2, dtype=np.float32) self._serve(direction=1 if self.rng.random() < 0.5 else -1) for _ in range(12): self.step() @classmethod def from_seed(cls, video_config: VideoConfig, seed: int) -> PongSimulation: return cls(video_config, seed) def _serve(self, direction: int) -> None: angle = float(self.rng.uniform(-0.55, 0.55)) speed = float(self.rng.uniform(52.0, 68.0)) self.ball_position[:] = ( self.video_config.width / 2, self.video_config.height / 2, ) self.ball_velocity[:] = ( direction * speed * math.cos(angle), speed * math.sin(angle), ) def _move_paddles(self, dt: float, left_action: int | None) -> None: height = float(self.video_config.height) right_target = float( self.ball_position[1] + self.right_aim_bias + math.sin(self.elapsed * 1.7 + self.phase) * 18.0 ) paddle_speed = height * 0.75 right_paddle_speed = height * self.right_speed_scale if left_action is None: left_delta = max( -paddle_speed * dt, min( paddle_speed * dt, float(self.ball_position[1]) - self.left_y, ), ) else: left_delta = float(max(-1, min(1, left_action))) * paddle_speed * dt self.left_y += left_delta right_delta = max( -right_paddle_speed * dt, min(right_paddle_speed * dt, right_target - self.right_y), ) self.right_y += right_delta self.left_y = max( self.paddle_height / 2, min(height - self.paddle_height / 2, self.left_y), ) self.right_y = max( self.paddle_height / 2, min(height - self.paddle_height / 2, self.right_y), ) def _paddle_collision(self, x: float, y: float, left: bool) -> bool: paddle_x = self.left_x if left else self.right_x paddle_y = self.left_y if left else self.right_y moving_toward = self.ball_velocity[0] < 0 if left else self.ball_velocity[0] > 0 touching_x = ( x - self.ball_radius <= paddle_x + self.paddle_width / 2 if left else x + self.ball_radius >= paddle_x - self.paddle_width / 2 ) touching_y = abs(y - paddle_y) <= self.paddle_height / 2 + self.ball_radius return moving_toward and touching_x and touching_y def step(self, left_action: int | None = None) -> None: width = float(self.video_config.width) height = float(self.video_config.height) substeps = 4 sub_dt = self.dt / substeps for _ in range(substeps): self._move_paddles(sub_dt, left_action) self.ball_position += self.ball_velocity * sub_dt x, y = map(float, self.ball_position) if y - self.ball_radius < 2: self.ball_position[1] = 2 + self.ball_radius self.ball_velocity[1] = abs(self.ball_velocity[1]) elif y + self.ball_radius > height - 2: self.ball_position[1] = height - 2 - self.ball_radius self.ball_velocity[1] = -abs(self.ball_velocity[1]) if self._paddle_collision(x, y, left=True): self.ball_position[0] = self.left_x + self.paddle_width / 2 + self.ball_radius self.ball_velocity[0] = abs(self.ball_velocity[0]) * 1.015 self.ball_velocity[1] += (y - self.left_y) * 1.2 elif self._paddle_collision(x, y, left=False): self.ball_position[0] = ( self.right_x - self.paddle_width / 2 - self.ball_radius ) self.ball_velocity[0] = -abs(self.ball_velocity[0]) * 1.015 self.ball_velocity[1] += (y - self.right_y) * 1.2 if self.ball_position[0] < -self.ball_radius: self.right_score += 1 self._serve(direction=1) elif self.ball_position[0] > width + self.ball_radius: self.left_score += 1 self._serve(direction=-1) self.elapsed += sub_dt def render(self) -> np.ndarray: scale = 2 width = self.video_config.width height = self.video_config.height image = Image.new("RGB", (width * scale, height * scale), (3, 8, 16)) draw = ImageDraw.Draw(image) white = (225, 239, 245) draw.rectangle((0, 0, width * scale, 3 * scale), fill=(42, 67, 78)) draw.rectangle( (0, (height - 3) * scale, width * scale, height * scale), fill=(42, 67, 78), ) for y in range(8, height - 8, 14): draw.rectangle( ( (width // 2 - 1) * scale, y * scale, (width // 2 + 1) * scale, (y + 7) * scale, ), fill=(50, 75, 84), ) for x, y in ((self.left_x, self.left_y), (self.right_x, self.right_y)): draw.rounded_rectangle( ( (x - self.paddle_width / 2) * scale, (y - self.paddle_height / 2) * scale, (x + self.paddle_width / 2) * scale, (y + self.paddle_height / 2) * scale, ), radius=2 * scale, fill=white, ) ball_x, ball_y = self.ball_position * scale radius = self.ball_radius * scale draw.ellipse( (ball_x - radius, ball_y - radius, ball_x + radius, ball_y + radius), fill=(255, 209, 72), ) draw.text( (width * scale * 0.43, 7 * scale), f"{self.left_score % 10} {self.right_score % 10}", fill=(115, 151, 160), ) image = image.resize((width, height), Image.Resampling.LANCZOS) return np.asarray(image, dtype=np.uint8) def next_frame(self) -> np.ndarray: frame = self.render() self.step() return frame def sample_frames(self, count: int) -> np.ndarray: return np.stack([self.next_frame() for _ in range(count)]) def _scripted_left_action(self) -> int: self._observed_ball_y.append(float(self.ball_position[1])) if self._action_hold_frames > 0: self._action_hold_frames -= 1 return self._held_left_action self._action_hold_frames = int(self.rng.integers(1, 6)) random_value = float(self.rng.random()) delayed_target = self._observed_ball_y[0] + self.left_aim_bias error = delayed_target - self.left_y ideal_action = -1 if error < -5 else (1 if error > 5 else 0) if random_value < self.left_randomness * 0.35: action = 0 elif random_value < self.left_randomness * 0.75: action = int(self.rng.choice((-1, 0, 1))) elif random_value < self.left_randomness: action = -ideal_action if ideal_action != 0 else int( self.rng.choice((-1, 1)) ) else: action = ideal_action self._held_left_action = action return action def sample_gameplay(self, count: int) -> tuple[np.ndarray, np.ndarray]: frames: list[np.ndarray] = [] actions: list[tuple[float, float]] = [] for _ in range(count): frames.append(self.render()) action = self._scripted_left_action() actions.append((float(action < 0), float(action > 0))) self.step(action) return np.stack(frames), np.asarray(actions, dtype=np.float32) _BREAKOUT_BRICK_COLORS = ( (220, 72, 88), (230, 140, 54), (236, 196, 72), (72, 186, 120), (72, 148, 220), (156, 102, 220), ) class FrameAutoencoderCodec(nn.Module): def __init__(self, config: CodecSettings, device: torch.device, dtype: torch.dtype) -> None: super().__init__() load_kwargs: dict[str, object] = { "revision": config.revision, "torch_dtype": dtype, "low_cpu_mem_usage": True, } if config.subfolder: load_kwargs["subfolder"] = config.subfolder try: self.vae = AutoencoderKL.from_pretrained( config.model_id, local_files_only=True, **load_kwargs ) except OSError: self.vae = AutoencoderKL.from_pretrained(config.model_id, **load_kwargs) self.vae.to(device=device).eval() self.vae.to(memory_format=torch.channels_last) self.vae.requires_grad_(False) if config.enable_slicing: self.vae.enable_slicing() if config.enable_tiling: self.vae.enable_tiling() self.channels = int(self.vae.config.latent_channels) self.spatial_compression = 2 ** (len(self.vae.config.block_out_channels) - 1) self.temporal_compression = 1 self.frame_batch_size = config.frame_batch_size self.sample_posterior = config.sample_posterior self.dtype = dtype self.scaling_factor = float(self.vae.config.scaling_factor) shift_factor = getattr(self.vae.config, "shift_factor", None) self.shift_factor = float(shift_factor) if shift_factor is not None else 0.0 def latent_frame_count(self, frame_count: int) -> int: return frame_count @torch.no_grad() def encode(self, frames: torch.Tensor) -> torch.Tensor: batch, frame_count, channels, height, width = frames.shape flat_frames = frames.reshape(batch * frame_count, channels, height, width) encoded: list[torch.Tensor] = [] for frame_batch in flat_frames.split(self.frame_batch_size): frame_batch = frame_batch.to(self.dtype).contiguous( memory_format=torch.channels_last ) posterior = self.vae.encode(frame_batch).latent_dist latents = posterior.sample() if self.sample_posterior else posterior.mode() encoded.append((latents - self.shift_factor) * self.scaling_factor) merged = torch.cat(encoded) return merged.reshape( batch, frame_count, self.channels, height // self.spatial_compression, width // self.spatial_compression, ) @torch.no_grad() def decode(self, latents: torch.Tensor) -> torch.Tensor: batch, frame_count, channels, height, width = latents.shape flat_latents = latents.reshape(batch * frame_count, channels, height, width) decoded: list[torch.Tensor] = [] for latent_batch in flat_latents.split(self.frame_batch_size): denormalized = latent_batch / self.scaling_factor + self.shift_factor denormalized = denormalized.to(self.dtype).contiguous( memory_format=torch.channels_last ) decoded.append(self.vae.decode(denormalized).sample) merged = torch.cat(decoded) return merged.reshape( batch, frame_count, 3, height * self.spatial_compression, width * self.spatial_compression, ).clamp(-1.0, 1.0) @dataclass(slots=True) class PrefixCache: key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...] pooled: torch.Tensor MemoryCache = PrefixCache def _video_positions( frames: int, height: int, width: int, time_offset: int, device: torch.device, ) -> torch.Tensor: time = torch.arange(time_offset, time_offset + frames, device=device) y = torch.arange(height, device=device) x = torch.arange(width, device=device) grid = torch.meshgrid(time, y, x, indexing="ij") return torch.stack(grid, dim=-1).reshape(-1, 3).float() class Rotary3D(nn.Module): def __init__(self, head_size: int, base: float = 10_000.0) -> None: super().__init__() self.rotary_size = (head_size // 6) * 6 self.axis_size = self.rotary_size // 3 frequencies = torch.exp( -math.log(base) * torch.arange(self.axis_size // 2, dtype=torch.float32) / (self.axis_size // 2) ) self.register_buffer("frequencies", frequencies, persistent=False) @staticmethod def _rotate_half(inputs: torch.Tensor) -> torch.Tensor: pairs = inputs.reshape(*inputs.shape[:-1], -1, 2) first, second = pairs.unbind(dim=-1) return torch.stack((-second, first), dim=-1).flatten(-2) def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: angles = torch.cat( [positions[:, axis, None] * self.frequencies[None] for axis in range(3)], dim=-1, ) cosine = angles.cos().repeat_interleave(2, dim=-1)[None, None] sine = angles.sin().repeat_interleave(2, dim=-1)[None, None] rotary, remainder = inputs.split( (self.rotary_size, inputs.shape[-1] - self.rotary_size), dim=-1, ) rotated = rotary * cosine.to(rotary.dtype) rotated = rotated + self._rotate_half(rotary) * sine.to(rotary.dtype) return torch.cat((rotated, remainder), dim=-1) class ScalarEmbedding(nn.Module): def __init__(self, hidden_size: int) -> None: super().__init__() self.hidden_size = hidden_size self.mlp = nn.Sequential( nn.Linear(hidden_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size), ) def forward(self, value: torch.Tensor) -> torch.Tensor: value = value.reshape(-1).float() half = self.hidden_size // 2 frequencies = torch.exp( -math.log(10_000.0) * torch.arange(half, device=value.device, dtype=torch.float32) / half ) angles = value[:, None] * frequencies[None] * 1_000.0 embedding = torch.cat((angles.cos(), angles.sin()), dim=-1) return self.mlp(embedding.to(dtype=self.mlp[0].weight.dtype)) class SwiGLU(nn.Module): def __init__(self, hidden_size: int, ratio: float) -> None: super().__init__() inner_size = int(hidden_size * ratio * 2 / 3) inner_size = ((inner_size + 63) // 64) * 64 self.input = nn.Linear(hidden_size, inner_size * 2, bias=False) self.output = nn.Linear(inner_size, hidden_size, bias=False) def forward(self, inputs: torch.Tensor) -> torch.Tensor: values, gates = self.input(inputs).chunk(2, dim=-1) return self.output(values * F.silu(gates)) class SelfAttention(nn.Module): def __init__(self, hidden_size: int, num_heads: int, qk_norm: bool) -> None: super().__init__() self.num_heads = num_heads self.head_size = hidden_size // num_heads self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False) self.output = nn.Linear(hidden_size, hidden_size, bias=False) self.q_norm = ( nn.RMSNorm(self.head_size, eps=1e-6) if qk_norm else nn.Identity() ) self.k_norm = ( nn.RMSNorm(self.head_size, eps=1e-6) if qk_norm else nn.Identity() ) self.rotary = Rotary3D(self.head_size) def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: batch, tokens, channels = inputs.shape query, key, value = self.qkv(inputs).reshape( batch, tokens, 3, self.num_heads, self.head_size, ).unbind(dim=2) query = self.rotary(self.q_norm(query).transpose(1, 2), positions) key = self.rotary(self.k_norm(key).transpose(1, 2), positions) value = value.transpose(1, 2) attended = F.scaled_dot_product_attention(query, key, value) return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels)) class PrefixAttention(nn.Module): """Target attention over the joint clean-prefix and noisy-target sequence.""" def __init__(self, hidden_size: int, num_heads: int, qk_norm: bool) -> None: super().__init__() self.num_heads = num_heads self.head_size = hidden_size // num_heads self.target_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False) self.prefix_key_value = nn.Linear(hidden_size, hidden_size * 2, bias=False) self.output = nn.Linear(hidden_size, hidden_size, bias=False) self.q_norm = ( nn.RMSNorm(self.head_size, eps=1e-6) if qk_norm else nn.Identity() ) self.k_norm = ( nn.RMSNorm(self.head_size, eps=1e-6) if qk_norm else nn.Identity() ) self.rotary = Rotary3D(self.head_size) def project_prefix( self, prefix: torch.Tensor, positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: batch, tokens, _ = prefix.shape key, value = self.prefix_key_value(prefix).reshape( batch, tokens, 2, self.num_heads, self.head_size, ).unbind(dim=2) key = self.rotary(self.k_norm(key).transpose(1, 2), positions) return key, value.transpose(1, 2) def forward_cached( self, target: torch.Tensor, target_positions: torch.Tensor, prefix_key_value: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: batch, tokens, channels = target.shape query, key, value = self.target_qkv(target).reshape( batch, tokens, 3, self.num_heads, self.head_size, ).unbind(dim=2) query = self.rotary(self.q_norm(query).transpose(1, 2), target_positions) key = self.rotary(self.k_norm(key).transpose(1, 2), target_positions) value = value.transpose(1, 2) prefix_key, prefix_value = prefix_key_value key = torch.cat((prefix_key, key), dim=2) value = torch.cat((prefix_value, value), dim=2) attended = F.scaled_dot_product_attention(query, key, value) return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels)) class PrefixDiTBlock(nn.Module): def __init__(self, config: ModelConfig) -> None: super().__init__() hidden_size = config.hidden_size self.prefix_norm1 = nn.RMSNorm(hidden_size, eps=1e-6) self.prefix_attention = SelfAttention( hidden_size, config.num_heads, config.qk_norm, ) self.prefix_norm2 = nn.RMSNorm(hidden_size, eps=1e-6) self.prefix_mlp = SwiGLU(hidden_size, config.mlp_ratio) self.prefix_projection_norm = nn.RMSNorm(hidden_size, eps=1e-6) self.target_norm1 = nn.RMSNorm( hidden_size, eps=1e-6, elementwise_affine=False, ) self.target_attention = PrefixAttention( hidden_size, config.num_heads, config.qk_norm, ) self.target_norm2 = nn.RMSNorm( hidden_size, eps=1e-6, elementwise_affine=False, ) self.target_mlp = SwiGLU(hidden_size, config.mlp_ratio) self.modulation = nn.Linear(hidden_size, hidden_size * 6) @staticmethod def _modulate( inputs: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor, ) -> torch.Tensor: return inputs * (1.0 + scale[:, None]) + shift[:, None] def advance_prefix( self, prefix: torch.Tensor, positions: torch.Tensor, ) -> torch.Tensor: prefix = prefix + self.prefix_attention(self.prefix_norm1(prefix), positions) return prefix + self.prefix_mlp(self.prefix_norm2(prefix)) def prefix_key_value( self, prefix: torch.Tensor, positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: return self.target_attention.project_prefix( self.prefix_projection_norm(prefix), positions, ) def advance_target( self, target: torch.Tensor, condition: torch.Tensor, target_positions: torch.Tensor, prefix_key_value: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: shift1, scale1, gate1, shift2, scale2, gate2 = self.modulation( F.silu(condition) ).chunk(6, dim=-1) normalized = self._modulate(self.target_norm1(target), shift1, scale1) target = target + gate1[:, None] * self.target_attention.forward_cached( normalized, target_positions, prefix_key_value, ) normalized = self._modulate(self.target_norm2(target), shift2, scale2) return target + gate2[:, None] * self.target_mlp(normalized) def forward( self, prefix: torch.Tensor, target: torch.Tensor, condition: torch.Tensor, prefix_positions: torch.Tensor, target_positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: prefix = self.advance_prefix(prefix, prefix_positions) key_value = self.prefix_key_value(prefix, prefix_positions) target = self.advance_target( target, condition, target_positions, key_value, ) return prefix, target class CausalLatentVideoDiT(nn.Module): """Multiscale clean-prefix/noisy-suffix latent video transformer.""" def __init__( self, model_config: ModelConfig, latent_channels: int, history_frames: int = 12, ) -> None: super().__init__() self.model_config = model_config self.latent_channels = latent_channels self.history_frames = history_frames self.gradient_checkpointing = False hidden_size = model_config.hidden_size target_patch = model_config.target_patch_size coarse_patch = model_config.memory_patch_size fine_patch = model_config.fine_memory_patch_size self.target_patch_embed = nn.Conv3d( latent_channels, hidden_size, kernel_size=(1, target_patch, target_patch), stride=(1, target_patch, target_patch), ) self.coarse_prefix_embed = nn.Conv3d( latent_channels, hidden_size, kernel_size=(1, coarse_patch, coarse_patch), stride=(1, coarse_patch, coarse_patch), ) self.fine_prefix_embed = nn.Conv3d( latent_channels, hidden_size, kernel_size=(1, fine_patch, fine_patch), stride=(1, fine_patch, fine_patch), ) self.motion_prefix_embed = nn.Conv3d( latent_channels, hidden_size, kernel_size=(1, fine_patch, fine_patch), stride=(1, fine_patch, fine_patch), ) self.coarse_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size)) self.fine_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size)) self.t_embedding = ScalarEmbedding(hidden_size) self.r_embedding = ScalarEmbedding(hidden_size) self.delta_embedding = ScalarEmbedding(hidden_size) self.time_projection = nn.Sequential( nn.Linear(hidden_size * 3, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size), ) self.prefix_condition = nn.Linear(hidden_size, hidden_size) self.action_projection = ( nn.Sequential( nn.Linear(model_config.action_dim, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size), ) if model_config.action_dim > 0 else None ) self.context_projection = ( nn.Linear(model_config.context_dim, hidden_size, bias=False) if model_config.context_dim > 0 else None ) self.blocks = nn.ModuleList( PrefixDiTBlock(model_config) for _ in range(model_config.depth) ) self.output_norm = nn.RMSNorm( hidden_size, eps=1e-6, elementwise_affine=False, ) self.output_modulation = nn.Linear(hidden_size, hidden_size * 2) self.output_projection = nn.Linear( hidden_size, target_patch * target_patch * latent_channels, ) self.outcome_norm = nn.RMSNorm(hidden_size, eps=1e-6) self.reward_head = nn.Linear(hidden_size, 1) self.terminal_head = nn.Linear(hidden_size, 1) self.reset_parameters() def enable_activation_checkpointing(self, enabled: bool = True) -> None: self.gradient_checkpointing = enabled def reset_parameters(self) -> None: for module in self.modules(): if isinstance(module, nn.Linear): nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) for embedding in ( self.target_patch_embed, self.coarse_prefix_embed, self.fine_prefix_embed, self.motion_prefix_embed, ): nn.init.normal_(embedding.weight, std=0.02) nn.init.normal_(self.coarse_scale_embedding, std=0.02) nn.init.normal_(self.fine_scale_embedding, std=0.02) for block in self.blocks: nn.init.zeros_(block.modulation.weight) nn.init.zeros_(block.modulation.bias) nn.init.zeros_(self.output_modulation.weight) nn.init.zeros_(self.output_modulation.bias) nn.init.zeros_(self.output_projection.weight) nn.init.zeros_(self.output_projection.bias) nn.init.zeros_(self.reward_head.weight) nn.init.zeros_(self.reward_head.bias) nn.init.zeros_(self.terminal_head.weight) nn.init.zeros_(self.terminal_head.bias) @staticmethod def _patch( embedding: nn.Conv3d, video: torch.Tensor, time_offset: int, ) -> tuple[torch.Tensor, torch.Tensor]: tokens = embedding(video.transpose(1, 2)) _, _, frames, height, width = tokens.shape positions = _video_positions( frames, height, width, time_offset, video.device, ) return tokens.flatten(2).transpose(1, 2), positions def _encode_prefix( self, history: torch.Tensor, context: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if history.shape[1] != self.history_frames: raise ValueError( f"Expected {self.history_frames} history frames, got {history.shape[1]}." ) coarse, coarse_positions = self._patch( self.coarse_prefix_embed, history, 0, ) fine_frames = self.model_config.fine_history_frames fine_history = history[:, -fine_frames:] previous = history[:, -fine_frames - 1 : -1] if previous.shape[1] != fine_frames: previous = torch.cat((fine_history[:, :1], fine_history[:, :-1]), dim=1) motion = fine_history - previous fine, fine_positions = self._patch( self.fine_prefix_embed, fine_history, self.history_frames - fine_frames, ) motion_tokens, _ = self._patch( self.motion_prefix_embed, motion, self.history_frames - fine_frames, ) coarse = coarse + self.coarse_scale_embedding fine = fine + motion_tokens + self.fine_scale_embedding prefix = torch.cat((coarse, fine), dim=1) positions = torch.cat((coarse_positions, fine_positions), dim=0) if self.context_projection is not None and context is not None: projected = self.context_projection(context) prefix = torch.cat((prefix, projected), dim=1) positions = torch.cat( ( positions, torch.zeros(projected.shape[1], 3, device=history.device), ), dim=0, ) return prefix, positions def _condition( self, prefix: torch.Tensor, r: torch.Tensor, t: torch.Tensor, ) -> torch.Tensor: time = self.time_projection( torch.cat( ( self.t_embedding(t), self.r_embedding(r), self.delta_embedding(t - r), ), dim=-1, ) ) return time + self.prefix_condition(prefix.mean(dim=1)) def build_memory_cache( self, history: torch.Tensor, context: torch.Tensor | None = None, ) -> PrefixCache: prefix, positions = self._encode_prefix(history, context) pooled = self.prefix_condition(prefix.mean(dim=1)) key_values: list[tuple[torch.Tensor, torch.Tensor]] = [] for block in self.blocks: prefix = block.advance_prefix(prefix, positions) key_values.append(block.prefix_key_value(prefix, positions)) return PrefixCache(tuple(key_values), pooled) def predict_outcomes( self, history: torch.Tensor, actions: torch.Tensor | None = None, frame_count: int = 1, memory_cache: PrefixCache | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if memory_cache is None: prefix, _ = self._encode_prefix(history) pooled = self.prefix_condition(prefix.mean(dim=1)) else: pooled = memory_cache.pooled hidden = pooled[:, None].expand(-1, frame_count, -1) if self.action_projection is not None: if actions is None or actions.shape[1] != frame_count: raise ValueError("Outcome actions must match frame_count.") hidden = hidden + self.action_projection(actions) hidden = self.outcome_norm(hidden) return self.reward_head(hidden).squeeze(-1), self.terminal_head(hidden).squeeze(-1) def _unpatchify( self, patches: torch.Tensor, frames: int, height: int, width: int, ) -> torch.Tensor: batch = patches.shape[0] patch = self.model_config.target_patch_size patches = patches.reshape( batch, frames, height, width, patch, patch, self.latent_channels, ) return patches.permute(0, 1, 6, 2, 4, 3, 5).reshape( batch, frames, self.latent_channels, height * patch, width * patch, ) def forward( self, noisy_target: torch.Tensor, history: torch.Tensor, r: torch.Tensor, t: torch.Tensor, actions: torch.Tensor | None = None, context: torch.Tensor | None = None, memory_cache: PrefixCache | None = None, ) -> torch.Tensor: target_volume = self.target_patch_embed(noisy_target.transpose(1, 2)) batch, _, frames, height, width = target_volume.shape target = target_volume.flatten(2).transpose(1, 2) target_positions = _video_positions( frames, height, width, self.history_frames, noisy_target.device, ) if memory_cache is None: prefix, prefix_positions = self._encode_prefix(history, context) condition = self._condition(prefix, r, t) else: prefix = None prefix_positions = None condition = self.time_projection( torch.cat( ( self.t_embedding(t), self.r_embedding(r), self.delta_embedding(t - r), ), dim=-1, ) ) + memory_cache.pooled if self.action_projection is not None: if actions is None or actions.shape[:2] != (batch, frames): raise ValueError("Actions must match [batch, target_frames, action_dim].") action_embedding = self.action_projection(actions) target = target + action_embedding.repeat_interleave(height * width, dim=1) condition = condition + action_embedding.mean(dim=1) for block_index, block in enumerate(self.blocks): if memory_cache is None: if prefix is None or prefix_positions is None: raise RuntimeError("Prefix state is unavailable.") if self.training and self.gradient_checkpointing: prefix, target = checkpoint( block, prefix, target, condition, prefix_positions, target_positions, use_reentrant=False, ) else: prefix, target = block( prefix, target, condition, prefix_positions, target_positions, ) else: target = block.advance_target( target, condition, target_positions, memory_cache.key_values[block_index], ) shift, scale = self.output_modulation(F.silu(condition)).chunk(2, dim=-1) target = self.output_norm(target) * (1.0 + scale[:, None]) + shift[:, None] return self._unpatchify( self.output_projection(target), frames, height, width, ) CausalVideoDiT = CausalLatentVideoDiT @torch.inference_mode() def sample_next_latent( model: nn.Module, history: torch.Tensor, steps: int = 2, actions: torch.Tensor | None = None, generator: torch.Generator | None = None, window_frames: int = 1, memory_cache: object | None = None, instantaneous: bool = False, ) -> torch.Tensor: if steps < 1: raise ValueError("steps must be positive.") batch, _, channels, height, width = history.shape state = torch.randn( (batch, window_frames, channels, height, width), device=history.device, dtype=history.dtype, generator=generator, ) boundaries = torch.linspace(1.0, 0.0, steps + 1, device=history.device) for index in range(steps): t_value = boundaries[index] r_value = boundaries[index + 1] t = t_value.expand(batch) r = r_value.expand(batch) model_r = t if instantaneous else r average_velocity = model( state, history, model_r, t, actions=actions, memory_cache=memory_cache, ) state = state - (t_value - r_value).to(state.dtype) * average_velocity return state def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Live Diffusion Pong Base (unconditional).") parser.add_argument("--repo", default=DEFAULT_REPO) parser.add_argument("--revision", default="main") parser.add_argument("--local-dir", default=None) parser.add_argument("--steps", type=int, default=2) parser.add_argument("--frames", type=int, default=0, help="0 = until window closed") parser.add_argument("--seed", type=int, default=None) parser.add_argument("--window-scale", type=int, default=6) parser.add_argument("--fps-cap", type=float, default=0.0) return parser.parse_args() def _download(repo: str, revision: str, local_dir: str | None) -> Path: if local_dir: root = Path(local_dir) root.mkdir(parents=True, exist_ok=True) for name in ("config.json", "ema.safetensors"): hf_hub_download( repo_id=repo, filename=name, revision=revision, local_dir=str(root), ) return root config_path = Path( hf_hub_download(repo_id=repo, filename="config.json", revision=revision) ) hf_hub_download(repo_id=repo, filename="ema.safetensors", revision=revision) return config_path.parent def _load_settings(config_path: Path) -> tuple[VideoConfig, ModelConfig, CodecSettings]: raw = json.loads(config_path.read_text(encoding="utf-8")) video = VideoConfig(**{key: raw["video"][key] for key in VideoConfig.__dataclass_fields__}) model = ModelConfig(**{key: raw["model"][key] for key in ModelConfig.__dataclass_fields__}) codec_raw = raw["codec"] codec = CodecSettings( model_id=codec_raw.get("model_id", "madebyollin/sdxl-vae-fp16-fix"), subfolder=codec_raw.get("subfolder", "") or "", revision=codec_raw.get("revision", "main"), latent_channels=int(codec_raw.get("latent_channels", 4)), spatial_compression=int(codec_raw.get("spatial_compression", 8)), frame_batch_size=int(codec_raw.get("frame_batch_size", 256)), sample_posterior=bool(codec_raw.get("sample_posterior", False)), enable_slicing=bool(codec_raw.get("enable_slicing", True)), enable_tiling=bool(codec_raw.get("enable_tiling", False)), ) return video, model, codec def _frames_to_tensor(frames: np.ndarray, device: torch.device) -> torch.Tensor: tensor = torch.from_numpy(frames.copy()).permute(0, 3, 1, 2).float() return tensor.div_(127.5).sub_(1.0).unsqueeze(0).to(device, non_blocking=True) def _tensor_to_image(frame: torch.Tensor) -> Image.Image: array = frame.float().clamp(-1.0, 1.0).add(1.0).mul(127.5) return Image.fromarray(array.byte().permute(1, 2, 0).cpu().numpy(), mode="RGB") def _inject_ball(image: Image.Image, position: tuple[float, float]) -> Image.Image: image = image.copy() draw = ImageDraw.Draw(image) radius = max(3.0, min(image.width, image.height) * 0.025) x, y = position draw.ellipse( (x - radius, y - radius, x + radius, y + radius), fill=(255, 209, 72), outline=(255, 247, 196), width=max(1, round(radius * 0.25)), ) return image def _image_to_video_tensor(image: Image.Image, device: torch.device) -> torch.Tensor: array = np.asarray(image, dtype=np.uint8).copy() tensor = torch.from_numpy(array).permute(2, 0, 1).float() return ( tensor.div_(127.5) .sub_(1.0) .reshape(1, 1, 3, image.height, image.width) .to(device, non_blocking=True) ) class LiveWindow: def __init__(self, scale: int) -> None: self.scale = max(1, scale) self.is_open = True self.root = tk.Tk() self.root.title("Diffusion Pong Base") self.label = tk.Label(self.root, background="black") self.label.pack() self.photo: ImageTk.PhotoImage | None = None self.generated_width = 0 self.generated_height = 0 self.pending_click: tuple[float, float] | None = None self.pressed: set[str] = set() self.label.bind("", 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 Pong Base — frame {frame_index + 1} {latency_ms:.0f} ms{held_text}" ) try: self.root.update_idletasks() self.root.update() except tk.TclError: self.is_open = False return self.is_open @torch.inference_mode() def main() -> None: args = _parse_args() if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): raise SystemExit("Need a CUDA GPU with BF16 support.") print(f"Downloading {args.repo} …") checkpoint = _download(args.repo, args.revision, args.local_dir) video_cfg, model_cfg, codec_cfg = _load_settings(checkpoint / "config.json") weights = checkpoint / "ema.safetensors" if not weights.exists(): raise SystemExit(f"Missing {weights}") device = torch.device("cuda") torch.set_float32_matmul_precision("high") seed = args.seed if args.seed is not None else random.SystemRandom().randrange(2**31) simulation = PongSimulation.from_seed(video_cfg, seed) for _ in range(random.Random(seed).randint(0, 400)): simulation.step() history_rgb = _frames_to_tensor( simulation.sample_frames(video_cfg.history_frames), device, ) print("Loading VAE + DiT …") codec = FrameAutoencoderCodec(codec_cfg, device, torch.bfloat16) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): latents = codec.encode(history_rgb) history: deque[torch.Tensor] = deque( (latents[:, index] for index in range(video_cfg.history_frames)), maxlen=video_cfg.history_frames, ) model = ( CausalLatentVideoDiT( model_cfg, latent_channels=codec.channels, history_frames=codec.latent_frame_count(video_cfg.history_frames), ) .to(device) .eval() ) model.load_state_dict(load_file(str(weights)), strict=True) steps = max(1, args.steps) frame_limit = args.frames if args.frames > 0 else 10_000_000 window = LiveWindow(args.window_scale) generator = torch.Generator(device=device).manual_seed(seed + 7) inject_pos: tuple[float, float] | None = None inject_frames = 0 print("Ready. Close the window to stop.") for frame_index in range(frame_limit): click = window.consume_click() if click is not None: inject_pos = click inject_frames = 3 history_tensor = torch.stack(tuple(history), dim=1) action = torch.tensor( [window.action()], device=device, dtype=torch.float32, ).unsqueeze(0) torch.cuda.synchronize() started = time.perf_counter() with torch.autocast(device_type="cuda", dtype=torch.bfloat16): memory_cache = model.build_memory_cache(history_tensor) prediction = sample_next_latent( model, history_tensor, steps, actions=action if model_cfg.action_dim > 0 else None, generator=generator, window_frames=1, memory_cache=memory_cache, instantaneous=True, ) decoded = codec.decode(prediction) torch.cuda.synchronize() latency_ms = (time.perf_counter() - started) * 1000 image = _tensor_to_image(decoded[0, 0]) frame_latent = prediction[:, 0] if inject_pos is not None and inject_frames > 0: image = _inject_ball(image, inject_pos) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): frame_latent = codec.encode(_image_to_video_tensor(image, device))[:, 0] inject_frames -= 1 if inject_frames == 0: inject_pos = None history.append(frame_latent) if not window.show(image, frame_index, latency_ms): break if args.fps_cap > 0: target = 1.0 / args.fps_cap elapsed = time.perf_counter() - started if elapsed < target: time.sleep(target - elapsed) if window.is_open: window.close() print("Done.") if __name__ == "__main__": main()