283 lines
12 KiB
Python
283 lines
12 KiB
Python
"""
|
|
Image Spatio-Temporal Refiner + Frame Compression.
|
|
|
|
Step 1: 8+8 alternating spatio-temporal attention over patch tokens of all frames
|
|
Step 2: 4 layers of cross-attention pooling compress 192 tokens per frame -> 1 token
|
|
|
|
Input: HaMeR backbone output (B, T, 192, 1280)
|
|
Output: (B, T, d_model) — one condensed image condition token per frame
|
|
"""
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from einops import rearrange
|
|
from model import attention
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# 2D RoPE (consistent with the legacy TemporalRefiner)
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class RotaryEmbedding1D(nn.Module):
|
|
"""1D RoPE: used for temporal positional encoding."""
|
|
|
|
def __init__(self, dim: int, base: int = 10000):
|
|
super().__init__()
|
|
assert dim % 2 == 0
|
|
self.dim = dim
|
|
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
|
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
|
|
|
def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
x: (B, H, L, D), positions: (L,)
|
|
"""
|
|
freqs = torch.einsum("i,j->ij", positions.float(), self.inv_freq.to(positions))
|
|
sin = freqs.sin().unsqueeze(0).unsqueeze(0)
|
|
cos = freqs.cos().unsqueeze(0).unsqueeze(0)
|
|
x1, x2 = x[..., 0::2], x[..., 1::2]
|
|
return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
|
|
|
|
|
|
class RotaryEmbedding2D(nn.Module):
|
|
"""2D RoPE: split head_dim in half and apply 1D RoPE with y/x coordinates respectively."""
|
|
|
|
def __init__(self, dim: int, base: int = 10000):
|
|
super().__init__()
|
|
assert dim % 2 == 0
|
|
self.dim = dim
|
|
self.dim_per_axis = dim // 2
|
|
inv_freq = 1.0 / (base ** (torch.arange(0, self.dim_per_axis, 2).float() / self.dim_per_axis))
|
|
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
|
|
|
def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
x: (B, H, L, D), positions: (L, 2) (y, x)
|
|
"""
|
|
x_y, x_x = x.chunk(2, dim=-1)
|
|
x_y = self._apply_1d(x_y, positions[:, 0])
|
|
x_x = self._apply_1d(x_x, positions[:, 1])
|
|
return torch.cat([x_y, x_x], dim=-1)
|
|
|
|
def _apply_1d(self, x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
|
|
freqs = torch.einsum("i,j->ij", pos.float(), self.inv_freq.to(pos))
|
|
sin = freqs.sin().unsqueeze(0).unsqueeze(0)
|
|
cos = freqs.cos().unsqueeze(0).unsqueeze(0)
|
|
x1, x2 = x[..., 0::2], x[..., 1::2]
|
|
return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Building blocks
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, dim: int):
|
|
super().__init__()
|
|
self.scale = nn.Parameter(torch.ones(dim))
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x_dtype = x.dtype
|
|
x = x.float()
|
|
rrms = torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + 1e-6)
|
|
return (x * rrms).to(dtype=x_dtype) * self.scale
|
|
|
|
|
|
class RefinerBlock(nn.Module):
|
|
"""Pre-LN Transformer block with optional 2D RoPE."""
|
|
|
|
def __init__(self, d_model: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = d_model // num_heads
|
|
self.norm1 = nn.LayerNorm(d_model)
|
|
self.qkv = nn.Linear(d_model, d_model * 3)
|
|
self.q_norm = RMSNorm(self.head_dim)
|
|
self.k_norm = RMSNorm(self.head_dim)
|
|
self.proj = nn.Linear(d_model, d_model)
|
|
self.norm2 = nn.LayerNorm(d_model)
|
|
mlp_hidden = int(d_model * mlp_ratio)
|
|
self.mlp = nn.Sequential(nn.Linear(d_model, mlp_hidden), nn.GELU(), nn.Linear(mlp_hidden, d_model))
|
|
self.drop = nn.Dropout(dropout)
|
|
|
|
def forward(self, x: torch.Tensor, rope: Optional[RotaryEmbedding2D] = None, positions: Optional[torch.Tensor] = None) -> torch.Tensor:
|
|
residual = x
|
|
x = self.norm1(x)
|
|
qkv = self.qkv(x)
|
|
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
|
q = self.q_norm(q)
|
|
k = self.k_norm(k)
|
|
if rope is not None and positions is not None:
|
|
q = rope(q, positions)
|
|
k = rope(k, positions)
|
|
x = attention(q, k, v)
|
|
x = self.proj(x)
|
|
x = residual + self.drop(x)
|
|
x = x + self.drop(self.mlp(self.norm2(x)))
|
|
return x
|
|
|
|
|
|
class CrossAttentionPoolBlock(nn.Module):
|
|
"""Cross-attention pooling: 1 learnable query token cross-attends K patches."""
|
|
|
|
def __init__(self, d_model: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = d_model // num_heads
|
|
self.norm_q = nn.LayerNorm(d_model)
|
|
self.norm_kv = nn.LayerNorm(d_model)
|
|
self.q_proj = nn.Linear(d_model, d_model)
|
|
self.kv_proj = nn.Linear(d_model, d_model * 2)
|
|
self.q_norm = RMSNorm(self.head_dim)
|
|
self.k_norm = RMSNorm(self.head_dim)
|
|
self.out_proj = nn.Linear(d_model, d_model)
|
|
self.norm_ff = nn.LayerNorm(d_model)
|
|
mlp_hidden = int(d_model * mlp_ratio)
|
|
self.ff = nn.Sequential(nn.Linear(d_model, mlp_hidden), nn.GELU(), nn.Linear(mlp_hidden, d_model))
|
|
self.drop = nn.Dropout(dropout)
|
|
|
|
def forward(self, query: torch.Tensor, patches: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
query: (B, 1, D), patches: (B, N, D)
|
|
Returns: (B, 1, D)
|
|
"""
|
|
residual = query
|
|
q = self.q_proj(self.norm_q(query))
|
|
kv = self.kv_proj(self.norm_kv(patches))
|
|
k, v = rearrange(kv, "B N (K H D) -> K B H N D", K=2, H=self.num_heads)
|
|
q = rearrange(q, "B L (H D) -> B H L D", H=self.num_heads)
|
|
q = self.q_norm(q)
|
|
k = self.k_norm(k)
|
|
out = attention(q, k, v) # (B, 1, H*D) = (B, 1, D)
|
|
out = self.out_proj(out)
|
|
query = residual + self.drop(out)
|
|
query = query + self.drop(self.ff(self.norm_ff(query)))
|
|
return query
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Spatio-Temporal Refiner
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class SpatioTemporalRefiner(nn.Module):
|
|
"""
|
|
8+8 alternating spatio-temporal attention.
|
|
|
|
Even layers: intra-frame spatial attention (B*T, 192, dim)
|
|
Odd layers: cross-frame same-position attention (B, T, dim) per spatial position
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
d_model: int,
|
|
hamer_dim: int = 1280,
|
|
spatial_layers: int = 8,
|
|
temporal_layers: int = 8,
|
|
num_heads: int = 8,
|
|
mlp_ratio: float = 4.0,
|
|
dropout: float = 0.1,
|
|
n_patches_h: int = 16,
|
|
n_patches_w: int = 12,
|
|
):
|
|
super().__init__()
|
|
self.d_model = d_model
|
|
self.n_patches = n_patches_h * n_patches_w
|
|
self.depth = spatial_layers + temporal_layers
|
|
self.spatial_layers = spatial_layers
|
|
self.temporal_layers = temporal_layers
|
|
|
|
self.input_proj = nn.Linear(hamer_dim, d_model)
|
|
|
|
# 2D spatial positions
|
|
y = torch.arange(n_patches_h)
|
|
x = torch.arange(n_patches_w)
|
|
self.register_buffer("spatial_positions", torch.cartesian_prod(y, x))
|
|
|
|
self.rope = RotaryEmbedding2D(dim=d_model // num_heads)
|
|
self.temporal_rope = RotaryEmbedding1D(dim=d_model // num_heads)
|
|
self.blocks = nn.ModuleList([
|
|
RefinerBlock(d_model, num_heads, mlp_ratio, dropout)
|
|
for _ in range(self.depth)
|
|
])
|
|
self.use_checkpoint = False
|
|
|
|
def gradient_checkpointing_enable(self):
|
|
self.use_checkpoint = True
|
|
|
|
def forward(self, hamer_features: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
hamer_features: (B, T, 192, 1280)
|
|
Returns:
|
|
(B, T, 192, d_model)
|
|
"""
|
|
B, T, N, _ = hamer_features.shape
|
|
x = self.input_proj(hamer_features)
|
|
|
|
spatial_pos = self.spatial_positions # (N, 2)
|
|
temporal_pos = torch.arange(T, device=x.device) # (T,)
|
|
temporal_pos_expanded = temporal_pos.unsqueeze(1).expand(T, N).reshape(T * N) # (T*N,)
|
|
|
|
for i, blk in enumerate(self.blocks):
|
|
if i % 2 == 0:
|
|
# Intra-frame spatial attention (2D RoPE: y, x)
|
|
x = x.reshape(B * T, N, self.d_model)
|
|
x = blk(x, rope=self.rope, positions=spatial_pos)
|
|
x = x.reshape(B, T, N, self.d_model)
|
|
else:
|
|
# Cross-frame global attention (1D temporal RoPE: frame index t)
|
|
x = x.reshape(B, T * N, self.d_model)
|
|
x = blk(x, rope=self.temporal_rope, positions=temporal_pos_expanded)
|
|
x = x.reshape(B, T, N, self.d_model)
|
|
|
|
return x
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Frame Compressor (Cross-Attention Pooling)
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class FrameCompressor(nn.Module):
|
|
"""
|
|
Compress N patch tokens of each frame into 1 condensed token.
|
|
Uses a learnable query token with multi-layer cross-attention.
|
|
Includes input_proj: supports arbitrary input dims (e.g. HaMeR 1280D -> d_model).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
d_model: int,
|
|
input_dim: int | None = None,
|
|
num_heads: int = 8,
|
|
num_layers: int = 8,
|
|
mlp_ratio: float = 4.0,
|
|
dropout: float = 0.1,
|
|
):
|
|
super().__init__()
|
|
self.input_proj = nn.Linear(input_dim, d_model) if input_dim else nn.Identity()
|
|
self.query = nn.Parameter(torch.randn(1, 1, d_model))
|
|
nn.init.trunc_normal_(self.query, std=0.02)
|
|
self.layers = nn.ModuleList([
|
|
CrossAttentionPoolBlock(d_model, num_heads, mlp_ratio, dropout)
|
|
for _ in range(num_layers)
|
|
])
|
|
|
|
def forward(self, patches: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
patches: (B, T, N, input_dim or d_model)
|
|
Returns: (B, T, d_model)
|
|
"""
|
|
B, T, N, _ = patches.shape
|
|
patches = self.input_proj(patches) # (B, T, N, d_model)
|
|
patches = patches.reshape(B * T, N, -1) # (B*T, N, D)
|
|
q = self.query.expand(B * T, -1, -1) # (B*T, 1, D)
|
|
for layer in self.layers:
|
|
q = layer(q, patches)
|
|
return q.reshape(B, T, -1)
|