| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def _init_linear_identity(linear): |
| if linear.weight.shape[0] != linear.weight.shape[1]: |
| raise ValueError( |
| "Identity initialization requires a square Linear layer. " |
| f"Got weight shape {tuple(linear.weight.shape)}." |
| ) |
| nn.init.eye_(linear.weight) |
| if linear.bias is not None: |
| nn.init.constant_(linear.bias, 0) |
|
|
|
|
| def _zero_linear(linear): |
| nn.init.constant_(linear.weight, 0) |
| if linear.bias is not None: |
| nn.init.constant_(linear.bias, 0) |
|
|
|
|
| def _zero_last_linear(module): |
| for layer in reversed(module): |
| if isinstance(layer, nn.Linear): |
| _zero_linear(layer) |
| return |
| raise ValueError("Expected at least one Linear layer.") |
|
|
|
|
| def _init_layer_norm_identity(norm): |
| nn.init.constant_(norm.weight, 1) |
| nn.init.constant_(norm.bias, 0) |
|
|
|
|
| class RotaryEmbedding(nn.Module): |
| """Rotary positional embedding for temporal self-attention.""" |
|
|
| def __init__(self, dim, max_seq_len=2048): |
| super().__init__() |
| if dim % 2 != 0: |
| raise ValueError(f"RoPE head dim must be even, got {dim}.") |
|
|
| self.dim = dim |
| self.max_seq_len = max_seq_len |
|
|
| inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq, persistent=True) |
| self._set_cos_sin_cache(max_seq_len, device=torch.device("cpu")) |
|
|
| def _set_cos_sin_cache(self, seq_len, device): |
| self.max_seq_len = seq_len |
| t = torch.arange(seq_len, device=device, dtype=torch.float32) |
| freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(device)) |
| emb = torch.cat((freqs, freqs), dim=-1) |
|
|
| self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False) |
| self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False) |
|
|
| @staticmethod |
| def _rotate_half(x): |
| x1 = x[..., : x.shape[-1] // 2] |
| x2 = x[..., x.shape[-1] // 2 :] |
| return torch.cat((-x2, x1), dim=-1) |
|
|
| def forward(self, x, seq_len): |
| if seq_len > self.max_seq_len or self.cos_cached.device != x.device: |
| self._set_cos_sin_cache(seq_len=max(seq_len, self.max_seq_len), device=x.device) |
|
|
| cos = self.cos_cached[..., :seq_len, :].to(dtype=x.dtype) |
| sin = self.sin_cached[..., :seq_len, :].to(dtype=x.dtype) |
| return (x * cos) + (self._rotate_half(x) * sin) |
|
|
|
|
| class TransformerLayer(nn.Module): |
| """Pre-LN temporal self-attention layer.""" |
|
|
| def __init__(self, token_dim, num_heads, dim_feedforward, dropout): |
| super().__init__() |
| if token_dim % num_heads != 0: |
| raise ValueError(f"token_dim={token_dim} must be divisible by num_heads={num_heads}.") |
|
|
| self.num_heads = num_heads |
| self.head_dim = token_dim // num_heads |
| self.dropout = dropout |
|
|
| self.qkv = nn.Linear(token_dim, token_dim * 3) |
| self.out_proj = nn.Linear(token_dim, token_dim) |
| self.norm1 = nn.LayerNorm(token_dim) |
| self.norm2 = nn.LayerNorm(token_dim) |
| self.ffn = nn.Sequential( |
| nn.Linear(token_dim, dim_feedforward), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(dim_feedforward, token_dim), |
| nn.Dropout(dropout), |
| ) |
| self._init_identity_weights() |
|
|
| def _init_identity_weights(self): |
| _zero_linear(self.out_proj) |
| _zero_last_linear(self.ffn) |
| _init_layer_norm_identity(self.norm1) |
| _init_layer_norm_identity(self.norm2) |
|
|
| def forward(self, x, rope=None, mask=None): |
| residual = x |
| x = self.norm1(x) |
| batch_size, seq_len, token_dim = x.shape |
|
|
| q, k, v = self.qkv(x).chunk(3, dim=-1) |
| q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() |
| k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() |
| v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() |
|
|
| if rope is not None: |
| q = rope(q, seq_len=seq_len) |
| k = rope(k, seq_len=seq_len) |
|
|
| x = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| attn_mask=mask, |
| dropout_p=self.dropout if self.training else 0.0, |
| ) |
|
|
| x = x.transpose(1, 2).reshape(batch_size, seq_len, token_dim).contiguous() |
| x = residual + self.out_proj(x) |
| x = x + self.ffn(self.norm2(x)) |
| return x |
|
|
|
|
| class TemporalTrajectoryPredictor(nn.Module): |
| """ |
| Temporal predictor without text cross-attention. |
| |
| Input: history_frames, shape (B, T, N, D), representing frames 0..T-1. |
| Output: predictions, shape (B, T, N, D), representing next frames 1..T. |
| Each output step only attends to current and previous input steps through a causal mask. |
| """ |
|
|
| def __init__( |
| self, |
| num_tokens=512, |
| token_dim=768, |
| max_seq_len=64, |
| num_blocks=1, |
| num_heads=8, |
| dim_feedforward=2048, |
| dropout=0.05, |
| ): |
| super().__init__() |
| if token_dim % num_heads != 0: |
| raise ValueError(f"token_dim={token_dim} must be divisible by num_heads={num_heads}.") |
|
|
| self.num_tokens = num_tokens |
| self.token_dim = token_dim |
| self.max_seq_len = max_seq_len |
| self.num_heads = num_heads |
| self.head_dim = token_dim // num_heads |
|
|
| self.rope = RotaryEmbedding(self.head_dim, max_seq_len) |
| self.layers = nn.ModuleList( |
| [ |
| TransformerLayer(token_dim, num_heads, dim_feedforward, dropout) |
| for _ in range(num_blocks) |
| ] |
| ) |
| self.output_projection = nn.Linear(token_dim, token_dim) |
| self._init_weights() |
|
|
| def _init_weights(self): |
| _init_linear_identity(self.output_projection) |
|
|
| @staticmethod |
| def _normalize_temporal_mask(mask, seq_len, device): |
| if mask is None: |
| return torch.ones((seq_len, seq_len), device=device, dtype=torch.bool).tril().unsqueeze(0).unsqueeze(1) |
| if mask.dtype != torch.bool: |
| return (mask < 0) if mask.dtype.is_floating_point else mask.bool() |
| return mask |
|
|
| def forward(self, history_frames, mask=None, predict_next_only=False): |
| target_dtype = self.output_projection.weight.dtype |
| history_frames = history_frames.to(dtype=target_dtype) |
|
|
| batch_size, seq_len, num_tokens, token_dim = history_frames.shape |
| if token_dim != self.token_dim: |
| raise ValueError(f"Expected token_dim={self.token_dim}, got {token_dim}.") |
|
|
| x = history_frames.transpose(1, 2).reshape(batch_size * num_tokens, seq_len, token_dim).contiguous() |
| mask = self._normalize_temporal_mask(mask, seq_len, x.device) |
|
|
| for layer in self.layers: |
| x = layer(x, rope=self.rope, mask=mask) |
|
|
| if predict_next_only: |
| x = self.output_projection(x[:, -1:, :]) |
| return x.reshape(batch_size, num_tokens, 1, token_dim).transpose(1, 2).contiguous() |
|
|
| x = self.output_projection(x) |
| return x.reshape(batch_size, num_tokens, seq_len, token_dim).transpose(1, 2).contiguous() |
|
|
|
|
| class SpatialGraphRefiner(nn.Module): |
| """Spatial graph refiner without text cross-attention.""" |
|
|
| def __init__( |
| self, |
| num_tokens=512, |
| token_dim=768, |
| max_graph_dist=10, |
| num_blocks=3, |
| num_heads=8, |
| dim_feedforward=2048, |
| dropout=0.05, |
| ): |
| super().__init__() |
| if token_dim % num_heads != 0: |
| raise ValueError(f"token_dim={token_dim} must be divisible by num_heads={num_heads}.") |
|
|
| self.num_tokens = num_tokens |
| self.token_dim = token_dim |
| self.max_graph_dist = max_graph_dist |
| self.num_heads = num_heads |
| self.head_dim = token_dim // num_heads |
| self.dropout = dropout |
|
|
| distances = torch.arange(max_graph_dist + 1, dtype=torch.float32) |
| scales = torch.linspace(0.1, 1.0, num_heads) |
| self.register_buffer("graph_bias_table", -1.0 * (distances.unsqueeze(1) * scales.unsqueeze(0))) |
|
|
| self.blocks = nn.ModuleList( |
| [ |
| nn.ModuleDict( |
| { |
| "qkv1": nn.Linear(token_dim, token_dim * 3), |
| "out1": nn.Linear(token_dim, token_dim), |
| "ln1": nn.LayerNorm(token_dim), |
| "qkv2": nn.Linear(token_dim, token_dim * 3), |
| "out2": nn.Linear(token_dim, token_dim), |
| "ln2": nn.LayerNorm(token_dim), |
| "ffn": nn.Sequential( |
| nn.Linear(token_dim, dim_feedforward), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(dim_feedforward, token_dim), |
| nn.Dropout(dropout), |
| ), |
| "ln3": nn.LayerNorm(token_dim), |
| } |
| ) |
| for _ in range(num_blocks) |
| ] |
| ) |
| self._init_identity_weights() |
|
|
| def _init_identity_weights(self): |
| for block in self.blocks: |
| _zero_linear(block["out1"]) |
| _zero_linear(block["out2"]) |
| _zero_last_linear(block["ffn"]) |
| _init_layer_norm_identity(block["ln1"]) |
| _init_layer_norm_identity(block["ln2"]) |
| _init_layer_norm_identity(block["ln3"]) |
|
|
| def _compute_graph_bias(self, shortest_path_matrix, batch_size, seq_len): |
| num_tokens = shortest_path_matrix.shape[-1] |
| dist_clipped = torch.clamp(shortest_path_matrix, 0, self.max_graph_dist).long() |
| bias = F.embedding(dist_clipped, self.graph_bias_table.to(shortest_path_matrix.device)) |
|
|
| if bias.shape[0] == 1 and batch_size > 1: |
| bias = bias.expand(batch_size, -1, -1, -1) |
| elif bias.shape[0] != batch_size: |
| raise ValueError( |
| "shortest_path_matrix must have batch size 1 or B. " |
| f"Got {bias.shape[0]} while input batch size is {batch_size}." |
| ) |
|
|
| return ( |
| bias.unsqueeze(1) |
| .expand(-1, seq_len, -1, -1, -1) |
| .permute(0, 1, 4, 2, 3) |
| .reshape(batch_size * seq_len, self.num_heads, num_tokens, num_tokens) |
| ) |
|
|
| def _spatial_attention(self, x, qkv, out, graph_attn_bias): |
| batch_time, num_tokens, token_dim = x.shape |
| q, k, v = qkv(x).chunk(3, dim=-1) |
| q = q.view(batch_time, num_tokens, self.num_heads, self.head_dim).transpose(1, 2).contiguous() |
| k = k.view(batch_time, num_tokens, self.num_heads, self.head_dim).transpose(1, 2).contiguous() |
| v = v.view(batch_time, num_tokens, self.num_heads, self.head_dim).transpose(1, 2).contiguous() |
|
|
| attn = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| attn_mask=graph_attn_bias, |
| dropout_p=self.dropout if self.training else 0.0, |
| ) |
| return out(attn.transpose(1, 2).reshape(batch_time, num_tokens, token_dim)) |
|
|
| def forward(self, next_frame_init, shortest_path_matrix): |
| target_dtype = self.blocks[0]["qkv1"].weight.dtype |
| next_frame_init = next_frame_init.to(dtype=target_dtype) |
| if shortest_path_matrix.dim() == 2: |
| shortest_path_matrix = shortest_path_matrix.unsqueeze(0) |
|
|
| batch_size, seq_len, num_tokens, token_dim = next_frame_init.shape |
| if token_dim != self.token_dim: |
| raise ValueError(f"Expected token_dim={self.token_dim}, got {token_dim}.") |
| if shortest_path_matrix.shape[-2:] != (num_tokens, num_tokens): |
| raise ValueError( |
| "shortest_path_matrix must have shape (B, N, N) or (1, N, N). " |
| f"Got {tuple(shortest_path_matrix.shape)} for N={num_tokens}." |
| ) |
|
|
| graph_attn_bias = self._compute_graph_bias(shortest_path_matrix, batch_size, seq_len).to(dtype=target_dtype) |
| x = next_frame_init.reshape(batch_size * seq_len, num_tokens, token_dim) |
|
|
| for block in self.blocks: |
| x = x + self._spatial_attention(block["ln1"](x), block["qkv1"], block["out1"], graph_attn_bias) |
| x = x + self._spatial_attention(block["ln2"](x), block["qkv2"], block["out2"], graph_attn_bias) |
| x = x + block["ffn"](block["ln3"](x)) |
|
|
| return x.reshape(batch_size, seq_len, num_tokens, token_dim) |
|
|
|
|
| class TextConditionedCrossAttentionGate(nn.Module): |
| """ |
| Independent text-conditioning module. |
| |
| It applies two text cross-attention layers, then uses text embeddings to produce |
| a per-output-dimension gate over the refined features. |
| """ |
|
|
| def __init__( |
| self, |
| token_dim=768, |
| text_dim=1024, |
| num_heads=8, |
| dim_feedforward=2048, |
| dropout=0.05, |
| gate_hidden_dim=None, |
| gate_residual=True, |
| ): |
| super().__init__() |
| if token_dim % num_heads != 0: |
| raise ValueError(f"token_dim={token_dim} must be divisible by num_heads={num_heads}.") |
|
|
| self.token_dim = token_dim |
| self.text_dim = text_dim |
| self.num_heads = num_heads |
| self.head_dim = token_dim // num_heads |
| self.dropout = dropout |
| self.gate_residual = gate_residual |
|
|
| self.cross_attn_1 = self._make_cross_attention_block(token_dim, text_dim, dim_feedforward, dropout) |
| self.cross_attn_2 = self._make_cross_attention_block(token_dim, text_dim, dim_feedforward, dropout) |
|
|
| hidden_dim = gate_hidden_dim or token_dim |
| self.gate_norm = nn.LayerNorm(text_dim) |
| self.gate_mlp = nn.Sequential( |
| nn.Linear(text_dim, hidden_dim), |
| nn.GELU(), |
| nn.Linear(hidden_dim, token_dim), |
| ) |
| self.gate_scale = nn.Parameter(torch.zeros(1)) |
| self._init_identity_weights() |
|
|
| def _init_identity_weights(self): |
| for block in (self.cross_attn_1, self.cross_attn_2): |
| _zero_linear(block["out"]) |
| _zero_last_linear(block["ffn"]) |
| _init_layer_norm_identity(block["ln1"]) |
| _init_layer_norm_identity(block["ln2"]) |
| _init_layer_norm_identity(self.gate_norm) |
| _zero_last_linear(self.gate_mlp) |
|
|
| @staticmethod |
| def _make_cross_attention_block(token_dim, text_dim, dim_feedforward, dropout): |
| return nn.ModuleDict( |
| { |
| "q": nn.Linear(token_dim, token_dim), |
| "kv": nn.Linear(text_dim, token_dim * 2), |
| "out": nn.Linear(token_dim, token_dim), |
| "ln1": nn.LayerNorm(token_dim), |
| "ln2": nn.LayerNorm(token_dim), |
| "ffn": nn.Sequential( |
| nn.Linear(token_dim, dim_feedforward), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(dim_feedforward, token_dim), |
| nn.Dropout(dropout), |
| ), |
| } |
| ) |
|
|
| @staticmethod |
| def _expand_text_for_tokens(text_embed, num_tokens): |
| batch_size, text_len, text_dim = text_embed.shape |
| return ( |
| text_embed.unsqueeze(1) |
| .expand(batch_size, num_tokens, text_len, text_dim) |
| .reshape(batch_size * num_tokens, text_len, text_dim) |
| .contiguous() |
| ) |
|
|
| def _cross_attention(self, x, text_embed, block): |
| residual = x |
| x_norm = block["ln1"](x) |
| batch_tokens, seq_len, token_dim = x_norm.shape |
|
|
| q = block["q"](x_norm).view(batch_tokens, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() |
| k, v = block["kv"](text_embed).chunk(2, dim=-1) |
| k = k.view(batch_tokens, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous() |
| v = v.view(batch_tokens, -1, self.num_heads, self.head_dim).transpose(1, 2).contiguous() |
|
|
| x = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| attn_mask=None, |
| dropout_p=self.dropout if self.training else 0.0, |
| ) |
| x = x.transpose(1, 2).reshape(batch_tokens, seq_len, token_dim).contiguous() |
| x = residual + block["out"](x) |
| x = x + block["ffn"](block["ln2"](x)) |
| return x |
|
|
| def forward(self, features, text_embed): |
| target_dtype = self.gate_mlp[0].weight.dtype |
| features = features.to(dtype=target_dtype) |
| if text_embed.dim() == 2: |
| text_embed = text_embed.unsqueeze(1) |
| text_embed = text_embed.to(dtype=target_dtype) |
|
|
| batch_size, seq_len, num_tokens, token_dim = features.shape |
| if token_dim != self.token_dim: |
| raise ValueError(f"Expected token_dim={self.token_dim}, got {token_dim}.") |
| if text_embed.shape[0] != batch_size: |
| raise ValueError(f"text_embed batch size must be B={batch_size}, got {text_embed.shape[0]}.") |
|
|
| x = features.transpose(1, 2).reshape(batch_size * num_tokens, seq_len, token_dim).contiguous() |
| text_for_tokens = self._expand_text_for_tokens(text_embed, num_tokens) |
|
|
| x = self._cross_attention(x, text_for_tokens, self.cross_attn_1) |
| x = self._cross_attention(x, text_for_tokens, self.cross_attn_2) |
| x = x.reshape(batch_size, num_tokens, seq_len, token_dim).transpose(1, 2).contiguous() |
|
|
| pooled_text = text_embed.mean(dim=1) |
| gate = torch.sigmoid(self.gate_mlp(self.gate_norm(pooled_text))).view(batch_size, 1, 1, token_dim) |
| if self.gate_residual: |
| return features + self.gate_scale * gate * x |
| return gate * x |
|
|
|
|
| class TrajectorySpatialTextRefinementBlock(nn.Module): |
| """One group: TemporalTrajectoryPredictor -> SpatialGraphRefiner -> TextConditionedCrossAttentionGate.""" |
|
|
| def __init__(self, temporal_kwargs=None, spatial_kwargs=None, text_kwargs=None): |
| super().__init__() |
| temporal_kwargs = temporal_kwargs or {} |
| spatial_kwargs = spatial_kwargs or {} |
| text_kwargs = text_kwargs or {} |
|
|
| self.temporal_predictor = TemporalTrajectoryPredictor(**temporal_kwargs) |
| self.spatial_refiner = SpatialGraphRefiner(**spatial_kwargs) |
| self.text_refiner = TextConditionedCrossAttentionGate(**text_kwargs) |
|
|
| def forward(self, x, text_embed, shortest_path_matrix, mask=None): |
| x = self.temporal_predictor(x, mask=mask, predict_next_only=False) |
| x = self.spatial_refiner(x, shortest_path_matrix=shortest_path_matrix) |
| x = self.text_refiner(x, text_embed=text_embed) |
| return x |
|
|
|
|
| class StackedTrajectorySpatialTextRefiner(nn.Module): |
| """ |
| Composite autoregressive model. |
| |
| It stacks N groups. Each group contains: |
| TemporalTrajectoryPredictor + SpatialGraphRefiner + TextConditionedCrossAttentionGate. |
| |
| Args: |
| num_groups: number of repeated groups. |
| temporal_kwargs: either one dict shared by all groups, or a list of dicts. |
| spatial_kwargs: either one dict shared by all groups, or a list of dicts. |
| text_kwargs: either one dict shared by all groups, or a list of dicts. |
| |
| Forward: |
| history_frames: (B, T, N, D), frames 0..T-1. |
| text_embed: (B, L_text, text_dim). |
| shortest_path_matrix: (B, N, N) or (1, N, N). |
| return: (B, T, N, D), refined predictions for frames 1..T. |
| """ |
|
|
| def __init__(self, num_groups=1, temporal_kwargs=None, spatial_kwargs=None, text_kwargs=None): |
| super().__init__() |
| if num_groups < 1: |
| raise ValueError(f"num_groups must be >= 1, got {num_groups}.") |
|
|
| temporal_kwargs = self._expand_kwargs(temporal_kwargs, num_groups, "temporal_kwargs") |
| spatial_kwargs = self._expand_kwargs(spatial_kwargs, num_groups, "spatial_kwargs") |
| text_kwargs = self._expand_kwargs(text_kwargs, num_groups, "text_kwargs") |
|
|
| self.blocks = nn.ModuleList( |
| [ |
| TrajectorySpatialTextRefinementBlock( |
| temporal_kwargs=temporal_kwargs[i], |
| spatial_kwargs=spatial_kwargs[i], |
| text_kwargs=text_kwargs[i], |
| ) |
| for i in range(num_groups) |
| ] |
| ) |
|
|
| @staticmethod |
| def _expand_kwargs(kwargs, num_groups, name): |
| if kwargs is None: |
| return [{} for _ in range(num_groups)] |
| if isinstance(kwargs, dict): |
| return [dict(kwargs) for _ in range(num_groups)] |
| if isinstance(kwargs, (list, tuple)): |
| if len(kwargs) != num_groups: |
| raise ValueError(f"{name} length must equal num_groups={num_groups}, got {len(kwargs)}.") |
| return [dict(item) for item in kwargs] |
| raise TypeError(f"{name} must be None, a dict, or a list/tuple of dicts.") |
|
|
| def forward(self, history_frames, text_embed, shortest_path_matrix, mask=None): |
| x = history_frames |
| for block in self.blocks: |
| x = block(x, text_embed=text_embed, shortest_path_matrix=shortest_path_matrix, mask=mask) |
| return x |
|
|
