| from __future__ import annotations
|
|
|
| import math
|
| from dataclasses import dataclass
|
|
|
| import torch
|
| import torch.nn.functional as F
|
| from torch import nn
|
| from torch.utils.checkpoint import checkpoint
|
|
|
| from livediffusion.config import ModelConfig
|
|
|
|
|
| @dataclass(slots=True)
|
| class PrefixCache:
|
| key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...]
|
| pooled: torch.Tensor
|
|
|
|
|
| MemoryCache = PrefixCache
|
|
|
|
|
| def _video_positions(
|
| frames: int,
|
| height: int,
|
| width: int,
|
| time_offset: int,
|
| device: torch.device,
|
| ) -> torch.Tensor:
|
| time = torch.arange(time_offset, time_offset + frames, device=device)
|
| y = torch.arange(height, device=device)
|
| x = torch.arange(width, device=device)
|
| grid = torch.meshgrid(time, y, x, indexing="ij")
|
| return torch.stack(grid, dim=-1).reshape(-1, 3).float()
|
|
|
|
|
| class Rotary3D(nn.Module):
|
| """3D RoPE with a shorter temporal base so time and space do not share a spectrum."""
|
|
|
| def __init__(
|
| self,
|
| head_size: int,
|
| time_base: float = 1_000.0,
|
| space_base: float = 10_000.0,
|
| ) -> None:
|
| super().__init__()
|
| self.rotary_size = (head_size // 6) * 6
|
| self.axis_size = self.rotary_size // 3
|
| axis_half = self.axis_size // 2
|
| steps = torch.arange(axis_half, dtype=torch.float32) / max(axis_half, 1)
|
| time_frequencies = torch.exp(-math.log(time_base) * steps)
|
| space_frequencies = torch.exp(-math.log(space_base) * steps)
|
| self.register_buffer(
|
| "frequencies",
|
| torch.stack((time_frequencies, space_frequencies, space_frequencies)),
|
| persistent=False,
|
| )
|
|
|
| @staticmethod
|
| def _rotate_half(inputs: torch.Tensor) -> torch.Tensor:
|
| pairs = inputs.reshape(*inputs.shape[:-1], -1, 2)
|
| first, second = pairs.unbind(dim=-1)
|
| return torch.stack((-second, first), dim=-1).flatten(-2)
|
|
|
| def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
| angles = torch.cat(
|
| [
|
| positions[:, axis, None] * self.frequencies[axis][None]
|
| for axis in range(3)
|
| ],
|
| dim=-1,
|
| )
|
| cosine = angles.cos().repeat_interleave(2, dim=-1)[None, None]
|
| sine = angles.sin().repeat_interleave(2, dim=-1)[None, None]
|
| rotary, remainder = inputs.split(
|
| (self.rotary_size, inputs.shape[-1] - self.rotary_size),
|
| dim=-1,
|
| )
|
| rotated = rotary * cosine.to(rotary.dtype)
|
| rotated = rotated + self._rotate_half(rotary) * sine.to(rotary.dtype)
|
| return torch.cat((rotated, remainder), dim=-1)
|
|
|
|
|
| class ScalarEmbedding(nn.Module):
|
| def __init__(self, hidden_size: int) -> None:
|
| super().__init__()
|
| self.hidden_size = hidden_size
|
| self.mlp = nn.Sequential(
|
| nn.Linear(hidden_size, hidden_size),
|
| nn.SiLU(),
|
| nn.Linear(hidden_size, hidden_size),
|
| )
|
|
|
| def forward(self, value: torch.Tensor) -> torch.Tensor:
|
| value = value.reshape(-1).float()
|
| half = self.hidden_size // 2
|
| frequencies = torch.exp(
|
| -math.log(10_000.0)
|
| * torch.arange(half, device=value.device, dtype=torch.float32)
|
| / half
|
| )
|
| angles = value[:, None] * frequencies[None] * 1_000.0
|
| embedding = torch.cat((angles.cos(), angles.sin()), dim=-1)
|
| return self.mlp(embedding.to(dtype=self.mlp[0].weight.dtype))
|
|
|
|
|
| class SwiGLU(nn.Module):
|
| def __init__(self, hidden_size: int, ratio: float) -> None:
|
| super().__init__()
|
| inner_size = int(hidden_size * ratio * 2 / 3)
|
| inner_size = ((inner_size + 63) // 64) * 64
|
| self.input = nn.Linear(hidden_size, inner_size * 2, bias=False)
|
| self.output = nn.Linear(inner_size, hidden_size, bias=False)
|
|
|
| def forward(self, inputs: torch.Tensor) -> torch.Tensor:
|
| values, gates = self.input(inputs).chunk(2, dim=-1)
|
| return self.output(values * F.silu(gates))
|
|
|
|
|
| class SelfAttention(nn.Module):
|
| def __init__(self, hidden_size: int, num_heads: int, qk_norm: bool) -> None:
|
| super().__init__()
|
| self.num_heads = num_heads
|
| self.head_size = hidden_size // num_heads
|
| self.qkv = nn.Linear(hidden_size, hidden_size * 3, bias=False)
|
| self.output = nn.Linear(hidden_size, hidden_size, bias=False)
|
| self.q_norm = (
|
| nn.RMSNorm(self.head_size, eps=1e-6)
|
| if qk_norm
|
| else nn.Identity()
|
| )
|
| self.k_norm = (
|
| nn.RMSNorm(self.head_size, eps=1e-6)
|
| if qk_norm
|
| else nn.Identity()
|
| )
|
| self.rotary = Rotary3D(self.head_size)
|
|
|
| def forward(self, inputs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
| batch, tokens, channels = inputs.shape
|
| query, key, value = self.qkv(inputs).reshape(
|
| batch,
|
| tokens,
|
| 3,
|
| self.num_heads,
|
| self.head_size,
|
| ).unbind(dim=2)
|
| query = self.rotary(self.q_norm(query).transpose(1, 2), positions)
|
| key = self.rotary(self.k_norm(key).transpose(1, 2), positions)
|
| value = value.transpose(1, 2)
|
| attended = F.scaled_dot_product_attention(query, key, value)
|
| return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels))
|
|
|
|
|
| class PrefixAttention(nn.Module):
|
| """Target attention over the joint clean-prefix and noisy-target sequence.
|
|
|
| Keys and values may use fewer heads than queries. Prefix and target keys are
|
| concatenated, so both sides share the same reduced head count.
|
| """
|
|
|
| def __init__(
|
| self,
|
| hidden_size: int,
|
| num_heads: int,
|
| qk_norm: bool,
|
| kv_heads: int = 0,
|
| ) -> None:
|
| super().__init__()
|
| self.num_heads = num_heads
|
| self.head_size = hidden_size // num_heads
|
| self.kv_heads = kv_heads or num_heads
|
| kv_size = self.kv_heads * self.head_size
|
| self.target_qkv = nn.Linear(hidden_size, hidden_size + kv_size * 2, bias=False)
|
| self.prefix_key_value = nn.Linear(hidden_size, kv_size * 2, bias=False)
|
| self.output = nn.Linear(hidden_size, hidden_size, bias=False)
|
| self.q_norm = (
|
| nn.RMSNorm(self.head_size, eps=1e-6)
|
| if qk_norm
|
| else nn.Identity()
|
| )
|
| self.k_norm = (
|
| nn.RMSNorm(self.head_size, eps=1e-6)
|
| if qk_norm
|
| else nn.Identity()
|
| )
|
| self.rotary = Rotary3D(self.head_size)
|
|
|
| def project_prefix(
|
| self,
|
| prefix: torch.Tensor,
|
| positions: torch.Tensor,
|
| ) -> tuple[torch.Tensor, torch.Tensor]:
|
| batch, tokens, _ = prefix.shape
|
| key, value = self.prefix_key_value(prefix).reshape(
|
| batch,
|
| tokens,
|
| 2,
|
| self.kv_heads,
|
| self.head_size,
|
| ).unbind(dim=2)
|
| key = self.rotary(self.k_norm(key).transpose(1, 2), positions)
|
| return key, value.transpose(1, 2)
|
|
|
| def forward_cached(
|
| self,
|
| target: torch.Tensor,
|
| target_positions: torch.Tensor,
|
| prefix_key_value: tuple[torch.Tensor, torch.Tensor],
|
| ) -> torch.Tensor:
|
| batch, tokens, channels = target.shape
|
| kv_size = self.kv_heads * self.head_size
|
| projected = self.target_qkv(target)
|
| query, key, value = projected.split((channels, kv_size, kv_size), dim=-1)
|
| query = query.reshape(batch, tokens, self.num_heads, self.head_size)
|
| key = key.reshape(batch, tokens, self.kv_heads, self.head_size)
|
| value = value.reshape(batch, tokens, self.kv_heads, self.head_size)
|
| query = self.rotary(self.q_norm(query).transpose(1, 2), target_positions)
|
| key = self.rotary(self.k_norm(key).transpose(1, 2), target_positions)
|
| value = value.transpose(1, 2)
|
| prefix_key, prefix_value = prefix_key_value
|
| key = torch.cat((prefix_key, key), dim=2)
|
| value = torch.cat((prefix_value, value), dim=2)
|
| attended = F.scaled_dot_product_attention(
|
| query,
|
| key,
|
| value,
|
| enable_gqa=self.kv_heads != self.num_heads,
|
| )
|
| return self.output(attended.transpose(1, 2).reshape(batch, tokens, channels))
|
|
|
|
|
| @dataclass(slots=True)
|
| class WorldMemoryState:
|
| plane: torch.Tensor
|
| registers: torch.Tensor
|
|
|
|
|
| class WorldMemory(nn.Module):
|
| """Gated recurrent memory at latent resolution, plus global register state.
|
|
|
| The plane is a convolutional GRU over one cell per latent position, so an
|
| irreversible world edit stays represented after it scrolls out of the
|
| attention prefix. Registers carry frame-global state the plane cannot hold.
|
| """
|
|
|
| def __init__(
|
| self,
|
| latent_channels: int,
|
| hidden_size: int,
|
| memory_channels: int,
|
| read_patch: int,
|
| register_tokens: int,
|
| ) -> None:
|
| super().__init__()
|
| self.memory_channels = memory_channels
|
| self.read_patch = read_patch
|
| self.register_tokens = register_tokens
|
| self.cell = nn.Conv2d(
|
| latent_channels + memory_channels,
|
| memory_channels * 3,
|
| kernel_size=3,
|
| padding=1,
|
| )
|
| self.read = nn.Conv2d(
|
| memory_channels,
|
| hidden_size,
|
| kernel_size=read_patch,
|
| stride=read_patch,
|
| )
|
| self.read_norm = nn.RMSNorm(hidden_size, eps=1e-6)
|
| self.scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
|
| self.register_update = (
|
| nn.Linear(memory_channels + hidden_size, hidden_size * 2)
|
| if register_tokens > 0
|
| else None
|
| )
|
|
|
| def initial_state(
|
| self,
|
| batch: int,
|
| height: int,
|
| width: int,
|
| registers: torch.Tensor,
|
| device: torch.device,
|
| dtype: torch.dtype,
|
| ) -> WorldMemoryState:
|
| plane = torch.zeros(
|
| batch,
|
| self.memory_channels,
|
| height,
|
| width,
|
| device=device,
|
| dtype=dtype,
|
| )
|
| return WorldMemoryState(plane, registers)
|
|
|
| def advance(
|
| self,
|
| state: WorldMemoryState,
|
| frames: torch.Tensor,
|
| ) -> WorldMemoryState:
|
| plane = state.plane
|
| for index in range(frames.shape[1]):
|
| latent = frames[:, index].to(plane.dtype)
|
| reset, update, candidate = self.cell(
|
| torch.cat((latent, plane), dim=1)
|
| ).chunk(3, dim=1)
|
| gate = torch.sigmoid(update)
|
| proposal = torch.tanh(candidate + torch.sigmoid(reset) * plane)
|
| plane = plane * (1.0 - gate) + proposal * gate
|
| registers = state.registers
|
| if self.register_update is not None:
|
| pooled = plane.mean(dim=(2, 3))
|
| summary = torch.cat(
|
| (pooled[:, None].expand(-1, registers.shape[1], -1), registers),
|
| dim=-1,
|
| )
|
| gate, proposal = self.register_update(summary).chunk(2, dim=-1)
|
| gate = torch.sigmoid(gate)
|
| registers = registers * (1.0 - gate) + torch.tanh(proposal) * gate
|
| return WorldMemoryState(plane, registers)
|
|
|
| def tokens(self, state: WorldMemoryState) -> torch.Tensor:
|
| read = self.read(state.plane).flatten(2).transpose(1, 2)
|
| return self.read_norm(read) + self.scale_embedding
|
|
|
|
|
| class SharedModulation(nn.Module):
|
| """One adaLN projection shared by every block, with per-block offsets.
|
|
|
| Per-block modulation costs 6 * hidden^2 each; sharing recovers that budget
|
| for depth while keeping each block independently steerable.
|
| """
|
|
|
| def __init__(self, hidden_size: int, blocks: int) -> None:
|
| super().__init__()
|
| self.projection = nn.Linear(hidden_size, hidden_size * 6)
|
| self.offsets = nn.Parameter(torch.zeros(blocks, hidden_size * 6))
|
|
|
| def forward(self, condition: torch.Tensor, block_index: int) -> torch.Tensor:
|
| return self.projection(F.silu(condition)) + self.offsets[block_index]
|
|
|
|
|
| class PrefixEncoderBlock(nn.Module):
|
| """Clean-context stack that runs before the joint prefix/target DiT."""
|
|
|
| def __init__(self, config: ModelConfig) -> None:
|
| super().__init__()
|
| hidden_size = config.hidden_size
|
| self.norm1 = nn.RMSNorm(hidden_size, eps=1e-6)
|
| self.attention = SelfAttention(
|
| hidden_size,
|
| config.num_heads,
|
| config.qk_norm,
|
| )
|
| self.norm2 = nn.RMSNorm(hidden_size, eps=1e-6)
|
| self.mlp = SwiGLU(hidden_size, config.mlp_ratio)
|
| self.attn_scale = nn.Parameter(torch.ones(hidden_size))
|
| self.mlp_scale = nn.Parameter(torch.ones(hidden_size))
|
|
|
| def forward(self, prefix: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
| prefix = prefix + self.attn_scale * self.attention(self.norm1(prefix), positions)
|
| return prefix + self.mlp_scale * self.mlp(self.norm2(prefix))
|
|
|
|
|
| class PrefixDiTBlock(nn.Module):
|
| def __init__(self, config: ModelConfig) -> None:
|
| super().__init__()
|
| hidden_size = config.hidden_size
|
| self.prefix_norm1 = nn.RMSNorm(hidden_size, eps=1e-6)
|
| self.prefix_attention = SelfAttention(
|
| hidden_size,
|
| config.num_heads,
|
| config.qk_norm,
|
| )
|
| self.prefix_norm2 = nn.RMSNorm(hidden_size, eps=1e-6)
|
| self.prefix_mlp = SwiGLU(hidden_size, config.mlp_ratio)
|
| self.prefix_projection_norm = nn.RMSNorm(hidden_size, eps=1e-6)
|
| self.target_norm1 = nn.RMSNorm(
|
| hidden_size,
|
| eps=1e-6,
|
| elementwise_affine=False,
|
| )
|
| self.target_attention = PrefixAttention(
|
| hidden_size,
|
| config.num_heads,
|
| config.qk_norm,
|
| config.kv_heads,
|
| )
|
| self.target_norm2 = nn.RMSNorm(
|
| hidden_size,
|
| eps=1e-6,
|
| elementwise_affine=False,
|
| )
|
| self.target_mlp = SwiGLU(hidden_size, config.mlp_ratio)
|
| self.modulation = (
|
| None
|
| if config.shared_modulation
|
| else nn.Linear(hidden_size, hidden_size * 6)
|
| )
|
|
|
| @staticmethod
|
| def _modulate(
|
| inputs: torch.Tensor,
|
| shift: torch.Tensor,
|
| scale: torch.Tensor,
|
| ) -> torch.Tensor:
|
| return inputs * (1.0 + scale[:, None]) + shift[:, None]
|
|
|
| def advance_prefix(
|
| self,
|
| prefix: torch.Tensor,
|
| positions: torch.Tensor,
|
| ) -> torch.Tensor:
|
| prefix = prefix + self.prefix_attention(self.prefix_norm1(prefix), positions)
|
| return prefix + self.prefix_mlp(self.prefix_norm2(prefix))
|
|
|
| def prefix_key_value(
|
| self,
|
| prefix: torch.Tensor,
|
| positions: torch.Tensor,
|
| ) -> tuple[torch.Tensor, torch.Tensor]:
|
| return self.target_attention.project_prefix(
|
| self.prefix_projection_norm(prefix),
|
| positions,
|
| )
|
|
|
| def advance_target(
|
| self,
|
| target: torch.Tensor,
|
| condition: torch.Tensor,
|
| target_positions: torch.Tensor,
|
| prefix_key_value: tuple[torch.Tensor, torch.Tensor],
|
| modulation: torch.Tensor | None = None,
|
| ) -> torch.Tensor:
|
| if modulation is None:
|
| if self.modulation is None:
|
| raise RuntimeError("Shared modulation must be supplied by the trunk.")
|
| modulation = self.modulation(F.silu(condition))
|
| shift1, scale1, gate1, shift2, scale2, gate2 = modulation.chunk(6, dim=-1)
|
| normalized = self._modulate(self.target_norm1(target), shift1, scale1)
|
| target = target + gate1[:, None] * self.target_attention.forward_cached(
|
| normalized,
|
| target_positions,
|
| prefix_key_value,
|
| )
|
| normalized = self._modulate(self.target_norm2(target), shift2, scale2)
|
| return target + gate2[:, None] * self.target_mlp(normalized)
|
|
|
| def forward(
|
| self,
|
| prefix: torch.Tensor,
|
| target: torch.Tensor,
|
| condition: torch.Tensor,
|
| prefix_positions: torch.Tensor,
|
| target_positions: torch.Tensor,
|
| modulation: torch.Tensor | None = None,
|
| ) -> tuple[torch.Tensor, torch.Tensor]:
|
| prefix = self.advance_prefix(prefix, prefix_positions)
|
| key_value = self.prefix_key_value(prefix, prefix_positions)
|
| target = self.advance_target(
|
| target,
|
| condition,
|
| target_positions,
|
| key_value,
|
| modulation,
|
| )
|
| return prefix, target
|
|
|
|
|
| class CausalLatentVideoDiT(nn.Module):
|
| """Multiscale clean-prefix/noisy-suffix latent video transformer."""
|
|
|
| def __init__(
|
| self,
|
| model_config: ModelConfig,
|
| latent_channels: int,
|
| history_frames: int = 12,
|
| ) -> None:
|
| super().__init__()
|
| self.model_config = model_config
|
| self.latent_channels = latent_channels
|
| self.history_frames = history_frames
|
| self.gradient_checkpointing = False
|
| hidden_size = model_config.hidden_size
|
| target_patch = model_config.target_patch_size
|
| coarse_patch = model_config.memory_patch_size
|
| fine_patch = model_config.fine_memory_patch_size
|
| self.target_patch_embed = nn.Conv3d(
|
| latent_channels,
|
| hidden_size,
|
| kernel_size=(1, target_patch, target_patch),
|
| stride=(1, target_patch, target_patch),
|
| )
|
| self.coarse_prefix_embed = nn.Conv3d(
|
| latent_channels,
|
| hidden_size,
|
| kernel_size=(1, coarse_patch, coarse_patch),
|
| stride=(1, coarse_patch, coarse_patch),
|
| )
|
| self.fine_prefix_embed = nn.Conv3d(
|
| latent_channels,
|
| hidden_size,
|
| kernel_size=(1, fine_patch, fine_patch),
|
| stride=(1, fine_patch, fine_patch),
|
| )
|
| self.motion_prefix_embed = nn.Conv3d(
|
| latent_channels,
|
| hidden_size,
|
| kernel_size=(1, fine_patch, fine_patch),
|
| stride=(1, fine_patch, fine_patch),
|
| )
|
| self.coarse_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
|
| self.fine_scale_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
|
| self.register_tokens = (
|
| nn.Parameter(torch.randn(1, model_config.register_tokens, hidden_size) * 0.02)
|
| if model_config.register_tokens > 0
|
| else None
|
| )
|
| self.memory_blocks = nn.ModuleList(
|
| PrefixEncoderBlock(model_config)
|
| for _ in range(model_config.prefix_encoder_depth)
|
| )
|
| self.t_embedding = ScalarEmbedding(hidden_size)
|
| self.r_embedding = ScalarEmbedding(hidden_size)
|
| self.delta_embedding = ScalarEmbedding(hidden_size)
|
| self.time_projection = nn.Sequential(
|
| nn.Linear(hidden_size * 3, hidden_size),
|
| nn.SiLU(),
|
| nn.Linear(hidden_size, hidden_size),
|
| )
|
| self.prefix_condition = nn.Linear(hidden_size, hidden_size)
|
| self.action_projection = (
|
| nn.Sequential(
|
| nn.Linear(model_config.action_dim, hidden_size),
|
| nn.SiLU(),
|
| nn.Linear(hidden_size, hidden_size),
|
| )
|
| if model_config.action_dim > 0
|
| else None
|
| )
|
| self.context_projection = (
|
| nn.Linear(model_config.context_dim, hidden_size, bias=False)
|
| if model_config.context_dim > 0
|
| else None
|
| )
|
| self.blocks = nn.ModuleList(
|
| PrefixDiTBlock(model_config) for _ in range(model_config.depth)
|
| )
|
| self.shared_modulation = (
|
| SharedModulation(hidden_size, model_config.depth)
|
| if model_config.shared_modulation
|
| else None
|
| )
|
| self.world_memory = (
|
| WorldMemory(
|
| latent_channels,
|
| hidden_size,
|
| model_config.memory_channels,
|
| coarse_patch,
|
| model_config.register_tokens,
|
| )
|
| if model_config.memory_channels > 0
|
| else None
|
| )
|
| self.output_norm = nn.RMSNorm(
|
| hidden_size,
|
| eps=1e-6,
|
| elementwise_affine=False,
|
| )
|
| self.output_modulation = nn.Linear(hidden_size, hidden_size * 2)
|
| self.output_projection = nn.Linear(
|
| hidden_size,
|
| target_patch * target_patch * latent_channels,
|
| )
|
| self.outcome_norm = nn.RMSNorm(hidden_size, eps=1e-6)
|
| self.reward_head = nn.Linear(hidden_size, 1)
|
| self.terminal_head = nn.Linear(hidden_size, 1)
|
| self.reset_parameters()
|
|
|
| def enable_activation_checkpointing(self, enabled: bool = True) -> None:
|
| self.gradient_checkpointing = enabled
|
|
|
| def reset_parameters(self) -> None:
|
| for module in self.modules():
|
| if isinstance(module, nn.Linear):
|
| nn.init.xavier_uniform_(module.weight)
|
| if module.bias is not None:
|
| nn.init.zeros_(module.bias)
|
| for embedding in (
|
| self.target_patch_embed,
|
| self.coarse_prefix_embed,
|
| self.fine_prefix_embed,
|
| self.motion_prefix_embed,
|
| ):
|
| nn.init.normal_(embedding.weight, std=0.02)
|
| nn.init.normal_(self.coarse_scale_embedding, std=0.02)
|
| nn.init.normal_(self.fine_scale_embedding, std=0.02)
|
| for block in self.blocks:
|
| if block.modulation is not None:
|
| nn.init.zeros_(block.modulation.weight)
|
| nn.init.zeros_(block.modulation.bias)
|
| if self.shared_modulation is not None:
|
| nn.init.zeros_(self.shared_modulation.projection.weight)
|
| nn.init.zeros_(self.shared_modulation.projection.bias)
|
| nn.init.zeros_(self.shared_modulation.offsets)
|
| if self.world_memory is not None:
|
| nn.init.zeros_(self.world_memory.scale_embedding)
|
| nn.init.zeros_(self.output_modulation.weight)
|
| nn.init.zeros_(self.output_modulation.bias)
|
| nn.init.zeros_(self.output_projection.weight)
|
| nn.init.zeros_(self.output_projection.bias)
|
| nn.init.zeros_(self.reward_head.weight)
|
| nn.init.zeros_(self.reward_head.bias)
|
| nn.init.zeros_(self.terminal_head.weight)
|
| nn.init.zeros_(self.terminal_head.bias)
|
|
|
| @staticmethod
|
| def _patch(
|
| embedding: nn.Conv3d,
|
| video: torch.Tensor,
|
| time_offset: int,
|
| ) -> tuple[torch.Tensor, torch.Tensor]:
|
| tokens = embedding(video.transpose(1, 2))
|
| _, _, frames, height, width = tokens.shape
|
| positions = _video_positions(
|
| frames,
|
| height,
|
| width,
|
| time_offset,
|
| video.device,
|
| )
|
| return tokens.flatten(2).transpose(1, 2), positions
|
|
|
| def _encode_prefix(
|
| self,
|
| history: torch.Tensor,
|
| context: torch.Tensor | None = None,
|
| history_actions: torch.Tensor | None = None,
|
| ) -> tuple[torch.Tensor, torch.Tensor]:
|
| if history.shape[1] < self.history_frames:
|
| raise ValueError(
|
| f"Expected at least {self.history_frames} history frames, "
|
| f"got {history.shape[1]}."
|
| )
|
| batch = history.shape[0]
|
| registers = (
|
| self.register_tokens.expand(batch, -1, -1)
|
| if self.register_tokens is not None
|
| else None
|
| )
|
| memory_tokens: torch.Tensor | None = None
|
| if self.world_memory is not None:
|
|
|
|
|
| state = self.world_memory.initial_state(
|
| batch,
|
| history.shape[-2],
|
| history.shape[-1],
|
| registers
|
| if registers is not None
|
| else history.new_zeros(batch, 0, self.model_config.hidden_size),
|
| history.device,
|
| history.dtype,
|
| )
|
| state = self.world_memory.advance(state, history)
|
| memory_tokens = self.world_memory.tokens(state)
|
| if registers is not None:
|
| registers = state.registers
|
| window = history[:, -self.history_frames :]
|
| coarse, coarse_positions = self._patch(
|
| self.coarse_prefix_embed,
|
| window,
|
| 0,
|
| )
|
| if self.action_projection is not None and history_actions is not None:
|
| window_actions = history_actions[:, -self.history_frames :]
|
| tokens_per_frame = coarse.shape[1] // self.history_frames
|
| coarse = coarse + self.action_projection(
|
| window_actions
|
| ).repeat_interleave(tokens_per_frame, dim=1)
|
| history = window
|
| fine_frames = self.model_config.fine_history_frames
|
| fine_history = history[:, -fine_frames:]
|
| previous = history[:, -fine_frames - 1 : -1]
|
| if previous.shape[1] != fine_frames:
|
| previous = torch.cat((fine_history[:, :1], fine_history[:, :-1]), dim=1)
|
| motion = fine_history - previous
|
| fine, fine_positions = self._patch(
|
| self.fine_prefix_embed,
|
| fine_history,
|
| self.history_frames - fine_frames,
|
| )
|
| motion_tokens, _ = self._patch(
|
| self.motion_prefix_embed,
|
| motion,
|
| self.history_frames - fine_frames,
|
| )
|
| coarse = coarse + self.coarse_scale_embedding
|
| fine = fine + motion_tokens + self.fine_scale_embedding
|
| prefix = torch.cat((coarse, fine), dim=1)
|
| positions = torch.cat((coarse_positions, fine_positions), dim=0)
|
| if memory_tokens is not None:
|
| coarse_patch = self.model_config.memory_patch_size
|
| memory_positions = _video_positions(
|
| 1,
|
| history.shape[-2] // coarse_patch,
|
| history.shape[-1] // coarse_patch,
|
| -2,
|
| history.device,
|
| )
|
| prefix = torch.cat((memory_tokens, prefix), dim=1)
|
| positions = torch.cat((memory_positions, positions), dim=0)
|
| if registers is not None:
|
| register_count = registers.shape[1]
|
| register_positions = torch.stack(
|
| (
|
| torch.full((register_count,), -1.0, device=history.device),
|
| torch.zeros(register_count, device=history.device),
|
| torch.arange(register_count, device=history.device, dtype=torch.float32),
|
| ),
|
| dim=-1,
|
| )
|
| prefix = torch.cat((registers, prefix), dim=1)
|
| positions = torch.cat((register_positions, positions), dim=0)
|
| if self.context_projection is not None and context is not None:
|
| projected = self.context_projection(context)
|
| prefix = torch.cat((prefix, projected), dim=1)
|
| positions = torch.cat(
|
| (
|
| positions,
|
| torch.zeros(projected.shape[1], 3, device=history.device),
|
| ),
|
| dim=0,
|
| )
|
| for block in self.memory_blocks:
|
| prefix = block(prefix, positions)
|
| return prefix, positions
|
|
|
| def _pool_prefix(self, prefix: torch.Tensor) -> torch.Tensor:
|
| register_count = self.model_config.register_tokens
|
| if register_count > 0:
|
| return prefix[:, :register_count].mean(dim=1)
|
| return prefix.mean(dim=1)
|
|
|
| def _condition(
|
| self,
|
| prefix: torch.Tensor,
|
| r: torch.Tensor,
|
| t: torch.Tensor,
|
| ) -> torch.Tensor:
|
| time = self.time_projection(
|
| torch.cat(
|
| (
|
| self.t_embedding(t),
|
| self.r_embedding(r),
|
| self.delta_embedding(t - r),
|
| ),
|
| dim=-1,
|
| )
|
| )
|
| return time + self.prefix_condition(self._pool_prefix(prefix))
|
|
|
| def build_memory_cache(
|
| self,
|
| history: torch.Tensor,
|
| context: torch.Tensor | None = None,
|
| history_actions: torch.Tensor | None = None,
|
| ) -> PrefixCache:
|
| prefix, positions = self._encode_prefix(history, context, history_actions)
|
| pooled = self.prefix_condition(self._pool_prefix(prefix))
|
| key_values: list[tuple[torch.Tensor, torch.Tensor]] = []
|
| for block in self.blocks:
|
| prefix = block.advance_prefix(prefix, positions)
|
| key_values.append(block.prefix_key_value(prefix, positions))
|
| return PrefixCache(tuple(key_values), pooled)
|
|
|
| def predict_outcomes(
|
| self,
|
| history: torch.Tensor,
|
| actions: torch.Tensor | None = None,
|
| frame_count: int = 1,
|
| memory_cache: PrefixCache | None = None,
|
| history_actions: torch.Tensor | None = None,
|
| ) -> tuple[torch.Tensor, torch.Tensor]:
|
| if memory_cache is None:
|
| prefix, _ = self._encode_prefix(history, None, history_actions)
|
| pooled = self.prefix_condition(self._pool_prefix(prefix))
|
| else:
|
| pooled = memory_cache.pooled
|
| hidden = pooled[:, None].expand(-1, frame_count, -1)
|
| if self.action_projection is not None:
|
| if actions is None or actions.shape[1] != frame_count:
|
| raise ValueError("Outcome actions must match frame_count.")
|
| hidden = hidden + self.action_projection(actions)
|
| hidden = self.outcome_norm(hidden)
|
| return self.reward_head(hidden).squeeze(-1), self.terminal_head(hidden).squeeze(-1)
|
|
|
| def _unpatchify(
|
| self,
|
| patches: torch.Tensor,
|
| frames: int,
|
| height: int,
|
| width: int,
|
| ) -> torch.Tensor:
|
| batch = patches.shape[0]
|
| patch = self.model_config.target_patch_size
|
| patches = patches.reshape(
|
| batch,
|
| frames,
|
| height,
|
| width,
|
| patch,
|
| patch,
|
| self.latent_channels,
|
| )
|
| return patches.permute(0, 1, 6, 2, 4, 3, 5).reshape(
|
| batch,
|
| frames,
|
| self.latent_channels,
|
| height * patch,
|
| width * patch,
|
| )
|
|
|
| def forward(
|
| self,
|
| noisy_target: torch.Tensor,
|
| history: torch.Tensor,
|
| r: torch.Tensor,
|
| t: torch.Tensor,
|
| actions: torch.Tensor | None = None,
|
| context: torch.Tensor | None = None,
|
| memory_cache: PrefixCache | None = None,
|
| history_actions: torch.Tensor | None = None,
|
| ) -> torch.Tensor:
|
| target_volume = self.target_patch_embed(noisy_target.transpose(1, 2))
|
| batch, _, frames, height, width = target_volume.shape
|
| target = target_volume.flatten(2).transpose(1, 2)
|
| target_positions = _video_positions(
|
| frames,
|
| height,
|
| width,
|
| self.history_frames,
|
| noisy_target.device,
|
| )
|
| if memory_cache is None:
|
| prefix, prefix_positions = self._encode_prefix(
|
| history,
|
| context,
|
| history_actions,
|
| )
|
| condition = self._condition(prefix, r, t)
|
| else:
|
| prefix = None
|
| prefix_positions = None
|
| condition = self.time_projection(
|
| torch.cat(
|
| (
|
| self.t_embedding(t),
|
| self.r_embedding(r),
|
| self.delta_embedding(t - r),
|
| ),
|
| dim=-1,
|
| )
|
| ) + memory_cache.pooled
|
| if self.action_projection is not None:
|
| if actions is None or actions.shape[:2] != (batch, frames):
|
| raise ValueError("Actions must match [batch, target_frames, action_dim].")
|
| action_embedding = self.action_projection(actions)
|
| target = target + action_embedding.repeat_interleave(height * width, dim=1)
|
| condition = condition + action_embedding.mean(dim=1)
|
| for block_index, block in enumerate(self.blocks):
|
| modulation = (
|
| self.shared_modulation(condition, block_index)
|
| if self.shared_modulation is not None
|
| else None
|
| )
|
| if memory_cache is None:
|
| if prefix is None or prefix_positions is None:
|
| raise RuntimeError("Prefix state is unavailable.")
|
| if self.training and self.gradient_checkpointing:
|
| prefix, target = checkpoint(
|
| block,
|
| prefix,
|
| target,
|
| condition,
|
| prefix_positions,
|
| target_positions,
|
| modulation,
|
| use_reentrant=False,
|
| )
|
| else:
|
| prefix, target = block(
|
| prefix,
|
| target,
|
| condition,
|
| prefix_positions,
|
| target_positions,
|
| modulation,
|
| )
|
| else:
|
| target = block.advance_target(
|
| target,
|
| condition,
|
| target_positions,
|
| memory_cache.key_values[block_index],
|
| modulation,
|
| )
|
| shift, scale = self.output_modulation(F.silu(condition)).chunk(2, dim=-1)
|
| target = self.output_norm(target) * (1.0 + scale[:, None]) + shift[:, None]
|
| return self._unpatchify(
|
| self.output_projection(target),
|
| frames,
|
| height,
|
| width,
|
| )
|
|
|
|
|
| CausalVideoDiT = CausalLatentVideoDiT
|
|
|
