Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

model/__init__.py

@@ -0,0 +1,30 @@
+"""HandFlow model package — attention function with flash-attn fallback."""
+
+import torch
+from einops import rearrange
+from torch import Tensor
+
+FA3_FOUND = True
+try:
+    from flash_attn_interface import flash_attn_func  # noqa: F401
+except ImportError:
+    FA3_FOUND = False
+
+if FA3_FOUND:
+
+    def attention(q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        q = rearrange(q, "B H L D -> B L H D")
+        k = rearrange(k, "B H L D -> B L H D")
+        v = rearrange(v, "B H L D -> B L H D")
+        x = flash_attn_func(q, k, v)
+        if not isinstance(x, Tensor):
+            x = x[0]
+        x = rearrange(x, "B L H D -> B L (H D)")
+        return x
+
+else:
+
+    def attention(q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+        x = rearrange(x, "B H L D -> B L (H D)")
+        return x

model/feature_extractors/__init__.py

@@ -0,0 +1 @@
+"""HandFlow feature extractors package."""

model/feature_extractors/condition_builder.py

@@ -0,0 +1,67 @@
+"""
+Condition Builder (with cmask mechanism).
+
+Concatenates image token + skeleton token into the condition stream.
+cmask: learnable mask token; both training and inference use HaMeR confidence as a continuous mask.
+
+Inputs:
+  - image_tokens: (B, T, d_model) — from Image Refiner + Compression
+  - skeleton_tokens: (B, T, d_model) — from Skeleton Encoder
+  - hamer_confidence: (B, T) — HaMeR detection confidence
+
+Output: condition (B, 2*T, d_model)
+"""
+
+import torch
+import torch.nn as nn
+
+
+class ConditionBuilder(nn.Module):
+    def __init__(self, d_model: int, mask_ratio: float = 0.2):
+        super().__init__()
+        self.d_model = d_model
+        self.mask_ratio = mask_ratio
+
+        # Learnable mask token — initialized with truncated normal
+        self.c_mask = nn.Parameter(torch.empty(1, 1, d_model))
+        nn.init.trunc_normal_(self.c_mask, std=0.02)
+
+    def forward(
+        self,
+        image_tokens: torch.Tensor,
+        skeleton_tokens: torch.Tensor,
+        hamer_confidence: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            image_tokens: (B, T, d_model)
+            skeleton_tokens: (B, T, d_model)
+            hamer_confidence: (B, T) — used as m_t during inference
+
+        Returns:
+            condition: (B, 2*T, d_model)
+            m: (B, T) — actual mask values used (for logging)
+        """
+        B, T, D = image_tokens.shape
+
+        # concat: (B, T, 2, d_model) -> (B, 2*T, d_model)
+        tokens = torch.stack([image_tokens, skeleton_tokens], dim=2)  # (B, T, 2, D)
+        tokens = tokens.reshape(B, 2 * T, D)
+
+        # Use HaMeR confidence as a continuous mask (consistent between training and inference)
+        m = hamer_confidence.float()  # (B, T)
+
+        # Optional: randomly dropout some frames during training as regularization
+        if self.training and self.mask_ratio > 0:
+            keep = (torch.rand(B, T, device=tokens.device) >= self.mask_ratio).float()
+            m = m * keep
+
+        # Expand m to match 2*T tokens (the two tokens per frame share the same m)
+        # m: (B, T) -> (B, 2*T)
+        m_expanded = m.unsqueeze(-1).expand(-1, -1, 2).reshape(B, 2 * T)  # (B, 2*T)
+        m_expanded = m_expanded.unsqueeze(-1)  # (B, 2*T, 1)
+
+        # cmask: c'_t = m_t * c_t + (1 - m_t) * c_mask
+        condition = m_expanded * tokens + (1 - m_expanded) * self.c_mask
+
+        return condition, m

model/feature_extractors/hamer_extractor.py

@@ -0,0 +1,110 @@
+"""
+HaMeR ViT-Huge backbone wrapper (frozen).
+
+Provides image token extraction for the 2D condition stream.
+Input: (B, 3, 256, 256) float32 [0,1]
+Output: (B, 192, 1280) — 16x12 patches x 1280D
+"""
+
+import sys
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+# Add HaMeR repository to the import path
+_HAMER_ROOT = Path(__file__).resolve().parents[2] / "third_party" / "hamer"
+if _HAMER_ROOT.exists() and str(_HAMER_ROOT) not in sys.path:
+    sys.path.insert(0, str(_HAMER_ROOT))
+
+
+class HaMeRBackbone(nn.Module):
+    """
+    Frozen HaMeR ViT-Huge backbone.
+
+    Internal pipeline:
+      1. ImageNet normalization (mean/std)
+      2. Crop 32px from each horizontal side: (B,3,256,256) -> (B,3,256,192)
+      3. Patch embedding: 192/16 x 256/16 = 12x16 = 192 tokens x 1280D
+      4. 32-layer ViT-H Transformer
+      5. Reshape output -> (B, 192, 1280)
+    """
+
+    # ImageNet normalization parameters
+    _IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406])
+    _IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225])
+
+    def __init__(self, checkpoint_path: Optional[str] = None):
+        super().__init__()
+        # Import the ViT backbone directly, bypassing hamer/models/__init__.py (avoids smplx dependency)
+        import importlib.util
+        vit_path = _HAMER_ROOT / "hamer" / "models" / "backbones" / "vit.py"
+        spec = importlib.util.spec_from_file_location("hamer_vit", str(vit_path))
+        vit_mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(vit_mod)
+        ViT = vit_mod.ViT
+
+        self.backbone = ViT(
+            img_size=(256, 192),
+            patch_size=16,
+            embed_dim=1280,
+            depth=32,
+            num_heads=16,
+            ratio=1,
+            use_checkpoint=False,
+            mlp_ratio=4,
+            qkv_bias=True,
+            drop_path_rate=0.55,
+        )
+        self.backbone.eval()
+
+        # Load pretrained weights
+        if checkpoint_path is not None:
+            ckpt = torch.load(checkpoint_path, map_location="cpu")
+            if "state_dict" in ckpt:
+                state_dict = {
+                    k.replace("backbone.", ""): v
+                    for k, v in ckpt["state_dict"].items()
+                    if k.startswith("backbone.")
+                }
+            else:
+                state_dict = ckpt
+            self.backbone.load_state_dict(state_dict, strict=False)
+
+        # Freeze all parameters
+        self.backbone.requires_grad_(False)
+        for p in self.backbone.parameters():
+            p.requires_grad = False
+
+        # Register normalization buffers
+        self.register_buffer("_mean", self._IMAGENET_MEAN.view(1, 3, 1, 1))
+        self.register_buffer("_std", self._IMAGENET_STD.view(1, 3, 1, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (B, 3, 256, 256) float32 in range [0, 1]
+        Returns:
+            tokens: (B, 192, 1280) patch tokens
+        """
+        # ImageNet normalization
+        x = (x - self._mean) / self._std
+
+        # HaMeR internal crop: 32px from each side (256 -> 192)
+        x = x[:, :, :, 32:-32]
+
+        # Backbone forward
+        feat_map = self.backbone.forward_features(x)  # (B, 1280, Hp, Wp)
+        B, C, Hp, Wp = feat_map.shape
+
+        # Flatten into tokens
+        tokens = feat_map.flatten(2).transpose(1, 2)  # (B, Hp*Wp, 1280)
+
+        return tokens
+
+    def train(self, mode: bool = True):
+        """Always stay in eval mode."""
+        super().train(False)
+        self.backbone.eval()
+        return self

model/feature_extractors/image_refiner.py

@@ -0,0 +1,282 @@
+"""
+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)

model/feature_extractors/skeleton_encoder.py

@@ -0,0 +1,96 @@
+"""
+Ray Direction Skeleton Encoder.
+
+2D skeleton -> inverse intrinsics -> ray direction -> high-frequency sinusoidal PE -> flatten -> MLP -> embed_dim.
+Outputs one skeleton token per frame.
+
+Input: landmarks (B, T, 21, 2) + crop_intrinsics (B, T, 4)
+Output: (B, T, embed_dim)
+"""
+
+import torch
+import torch.nn as nn
+
+
+class RayDirectionSkeletonEncoder(nn.Module):
+    """
+    Encodes a 2D skeleton into one token per frame via ray direction.
+
+    Pipeline:
+      1. un-normalize landmarks -> pixel coordinates (in crop coordinate system)
+      2. ray = [(u-cx)/fx, (v-cy)/fy, 1] -> normalize -> (21, 3)
+      3. High-frequency sinusoidal PE: [sin(2^L * ray), cos(2^L * ray)] L=0..num_freqs-1
+      4. flatten: 21 joints x (3 x 2 x num_freqs) per joint
+      5. MLP -> embed_dim
+    """
+
+    def __init__(
+        self,
+        n_joints: int = 21,
+        num_freqs: int = 6,
+        embed_dim: int = 512,
+        crop_h: int = 256,
+        crop_w: int = 256,
+        mlp_hidden: int = 512,
+    ):
+        super().__init__()
+        self.n_joints = n_joints
+        self.num_freqs = num_freqs
+        self.crop_h = crop_h
+        self.crop_w = crop_w
+
+        # Per-joint dim after PE: 3 (ray) x 2 (sin+cos) x num_freqs
+        pe_dim = 3 * 2 * num_freqs  # = 36
+        flat_dim = n_joints * pe_dim  # = 756
+
+        self.mlp = nn.Sequential(
+            nn.Linear(flat_dim, mlp_hidden),
+            nn.SiLU(),
+            nn.Linear(mlp_hidden, embed_dim),
+        )
+
+    def forward(
+        self,
+        landmarks: torch.Tensor,
+        crop_intrinsics: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            landmarks: (B, T, 21, 2) normalized [0, 1]
+            crop_intrinsics: (B, T, 4) = [fx, fy, cx, cy] in crop coords
+        Returns:
+            (B, T, embed_dim)
+        """
+        B, T, J, _ = landmarks.shape
+        fx = crop_intrinsics[..., 0:1]  # (B, T, 1)
+        fy = crop_intrinsics[..., 1:2]
+        cx = crop_intrinsics[..., 2:3]
+        cy = crop_intrinsics[..., 3:4]
+
+        # un-normalize -> crop pixel coordinates
+        u = landmarks[..., 0:1] * self.crop_w  # (B, T, 21, 1)
+        v = landmarks[..., 1:2] * self.crop_h
+
+        # ray direction: [(u-cx)/fx, (v-cy)/fy, 1]
+        ray_x = (u - cx.unsqueeze(-2)) / (fx.unsqueeze(-2) + 1e-8)
+        ray_y = (v - cy.unsqueeze(-2)) / (fy.unsqueeze(-2) + 1e-8)
+        ray_z = torch.ones_like(ray_x)
+        ray = torch.cat([ray_x, ray_y, ray_z], dim=-1)  # (B, T, 21, 3)
+
+        # normalize ray direction
+        ray = ray / (ray.norm(dim=-1, keepdim=True) + 1e-8)
+
+        # High-frequency sinusoidal PE
+        pe_parts = []
+        for L in range(self.num_freqs):
+            freq = 2.0 ** L
+            pe_parts.append(torch.sin(freq * ray))
+            pe_parts.append(torch.cos(freq * ray))
+        # pe: (B, T, 21, 3 * 2 * num_freqs) = (B, T, 21, 36)
+        pe = torch.cat(pe_parts, dim=-1)
+
+        # flatten all joints: (B, T, 21 * 36) = (B, T, 756)
+        flat = pe.reshape(B, T, -1)
+
+        # MLP -> embed_dim
+        return self.mlp(flat)

model/flow_matching/__init__.py

@@ -0,0 +1 @@
+"""HandFlow flow matching package."""

model/flow_matching/denoiser.py

@@ -0,0 +1,345 @@
+"""
+HandFlow Rectified Flow Denoiser — all-frame direct regression.
+
+Pose (48D axis-angle) and Trans (3D) are both regressed per-frame directly,
+without velocity/cumsum. Beta (10D) serves as a global token.
+
+Flow Target layout (flat packing per window, managed by mano_slice):
+  x_1 = [beta(10) | pose_0(48) | ... | pose_{T-1}(48) | trans_0(3) | ... | trans_{T-1}(3)]
+  flat_dim = 10 + T × 51
+
+z-stream: 1 beta token + T pose tokens + T trans tokens = 1+2T tokens
+c-stream: 2 × T tokens (1 image + 1 skeleton per frame, with cmask)
+
+Condition pipeline:
+  HaMeR backbone (frozen) → FrameCompressor (8-layer cross-attention pooling) → image tokens
+  2D skeleton → ray direction → PE → MLP → skeleton tokens
+  image + skeleton → cmask → condition stream
+"""
+
+import torch
+import torch.nn as nn
+from omegaconf import DictConfig
+
+from model.flow_matching.dualstream_transformer import (
+    FlowMatchingTransformer,
+    MLPEmbedder,
+)
+from model.flow_matching.helpers.model_wrapper import ModelWrapper
+from model.flow_matching.helpers.scheduler import FluxTimeSampler
+from model.feature_extractors.hamer_extractor import HaMeRBackbone
+from model.feature_extractors.image_refiner import FrameCompressor
+from model.feature_extractors.skeleton_encoder import RayDirectionSkeletonEncoder
+from model.feature_extractors.condition_builder import ConditionBuilder
+from model.pose_embedding import RotaryEmbedding
+
+
+class ManoSlice(nn.Module):
+    """Manage slice positions and normalization of each component in the flow target flat vector.
+
+    Flat layout: [beta(10) | poses(T,48) | trans(T,3)]
+    flat_dim = 10 + T * 51
+    """
+
+    def __init__(self, beta_dim: int, pose_dim: int,
+                 trans_dim: int, window_size: int, stats_path: str | None = None):
+        super().__init__()
+        self.beta_dim = beta_dim       # 10
+        self.pose_dim = pose_dim       # 48
+        self.trans_dim = trans_dim     # 3
+        self.window_size = window_size
+
+        self.flat_dim = beta_dim + window_size * (pose_dim + trans_dim)
+
+        # Build slices
+        s = 0
+        self.beta = slice(s, s + beta_dim)                                # [0:10]
+        s += beta_dim
+        self.poses = slice(s, s + window_size * pose_dim)                 # [10:10+48T]
+        s += window_size * pose_dim
+        self.trans = slice(s, s + window_size * trans_dim)                # [10+48T:10+51T]
+
+        # Normalization statistics
+        flat_mean = torch.zeros(self.flat_dim)
+        flat_std = torch.ones(self.flat_dim)
+        if stats_path is not None:
+            import numpy as np
+            stats = np.load(stats_path)
+            beta_mean = torch.from_numpy(stats["beta_mean"])         # (10,)
+            beta_std = torch.from_numpy(stats["beta_std"])           # (10,)
+            pose_mean = torch.from_numpy(stats["pose_mean"])         # (48,)
+            pose_std = torch.from_numpy(stats["pose_std"])           # (48,)
+            trans_mean = torch.from_numpy(stats["trans_mean"])       # (3,)
+            trans_std = torch.from_numpy(stats["trans_std"])         # (3,)
+            flat_mean = torch.cat([
+                beta_mean,
+                pose_mean.repeat(window_size),
+                trans_mean.repeat(window_size),
+            ])
+            flat_std = torch.cat([
+                beta_std,
+                pose_std.repeat(window_size),
+                trans_std.repeat(window_size),
+            ])
+        self.register_buffer("flat_mean", flat_mean)
+        self.register_buffer("flat_std", flat_std)
+
+        # per-component statistics for auxiliary loss denormalization
+        self.register_buffer("beta_mean_", flat_mean[self.beta])
+        self.register_buffer("beta_std_", flat_std[self.beta])
+
+        if stats_path is not None:
+            self.register_buffer("pose_mean_", pose_mean)
+            self.register_buffer("pose_std_", pose_std)
+            self.register_buffer("trans_mean_", trans_mean)
+            self.register_buffer("trans_std_", trans_std)
+        else:
+            self.register_buffer("pose_mean_", torch.zeros(pose_dim))
+            self.register_buffer("pose_std_", torch.ones(pose_dim))
+            self.register_buffer("trans_mean_", torch.zeros(trans_dim))
+            self.register_buffer("trans_std_", torch.ones(trans_dim))
+
+    def normalize(self, x_flat: torch.Tensor) -> torch.Tensor:
+        """Normalize the flat vector to ~N(0,1)."""
+        return (x_flat - self.flat_mean) / self.flat_std
+
+    def denormalize(self, x_flat: torch.Tensor) -> torch.Tensor:
+        """Denormalize the flat vector back to original scale."""
+        return x_flat * self.flat_std + self.flat_mean
+
+    def denormalize_components(
+        self, beta: torch.Tensor, poses: torch.Tensor, trans: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Denormalize each component back to original scale."""
+        beta_orig = beta * self.beta_std_ + self.beta_mean_
+        poses_orig = poses * self.pose_std_ + self.pose_mean_    # (48,) broadcast
+        trans_orig = trans * self.trans_std_ + self.trans_mean_  # (3,) broadcast
+        return beta_orig, poses_orig, trans_orig
+
+    def unpack(self, x_flat: torch.Tensor):
+        """Unpack flat → (beta, poses, trans)"""
+        B = x_flat.shape[0]
+        beta = x_flat[:, self.beta]       # (B, 10)
+        poses = x_flat[:, self.poses].reshape(B, self.window_size, self.pose_dim)    # (B, T, 48)
+        trans = x_flat[:, self.trans].reshape(B, self.window_size, self.trans_dim)   # (B, T, 3)
+        return beta, poses, trans
+
+    def pack(self, beta: torch.Tensor, poses: torch.Tensor,
+             trans: torch.Tensor) -> torch.Tensor:
+        """Pack → flat"""
+        B = beta.shape[0]
+        return torch.cat([
+            beta.reshape(B, -1),
+            poses.reshape(B, -1),
+            trans.reshape(B, -1),
+        ], dim=-1)
+
+
+class HandPoseDenoiser(nn.Module):
+    """
+    Rectified Flow Denoiser — all-frame direct regression.
+    """
+
+    def __init__(self, config: DictConfig):
+        super().__init__()
+        fm_cfg = config.model.fm
+        cond_cfg = fm_cfg.condition
+        comp_cfg = fm_cfg.frame_compressor
+        tf_cfg = fm_cfg.transformer
+
+        self.d_model = tf_cfg.hidden_size
+        self.beta_dim = fm_cfg.beta_dim       # 10
+        self.pose_dim = fm_cfg.pose_dim       # 48
+        self.trans_dim = fm_cfg.trans_dim     # 3
+        self.window_size = config.dataset.window_size
+
+        # mano_slice (includes normalization statistics)
+        stats_path = fm_cfg.get("normalization_stats", None)
+        self.mano_slice = ManoSlice(
+            beta_dim=self.beta_dim,
+            pose_dim=self.pose_dim,
+            trans_dim=self.trans_dim,
+            window_size=self.window_size,
+            stats_path=stats_path,
+        )
+
+        # z-stream input projections: each component → d_model
+        self.beta_proj_in = nn.Linear(self.beta_dim, self.d_model)
+        self.pose_proj_in = nn.Linear(self.pose_dim, self.d_model)
+        self.trans_proj_in = nn.Linear(self.trans_dim, self.d_model)
+
+        # z-stream output projections: d_model → native dims
+        self.beta_proj_out = nn.Linear(self.d_model, self.beta_dim)
+        self.pose_proj_out = nn.Linear(self.d_model, self.pose_dim)
+        self.trans_proj_out = nn.Linear(self.d_model, self.trans_dim)
+
+        # HaMeR backbone (frozen)
+        hamer_ckpt = config.model.get("hamer", {}).get("checkpoint", None)
+        self.hamer_backbone = HaMeRBackbone(checkpoint_path=hamer_ckpt)
+
+        # Frame Compressor (directly compresses HaMeR patch tokens)
+        self.frame_compressor = FrameCompressor(
+            d_model=self.d_model,
+            input_dim=cond_cfg.hamer_embed_dim,
+            num_heads=tf_cfg.num_heads,
+            num_layers=comp_cfg.num_layers,
+        )
+
+        # Skeleton Encoder
+        self.skeleton_encoder = RayDirectionSkeletonEncoder(
+            n_joints=cond_cfg.n_joints,
+            num_freqs=cond_cfg.skeleton_pe_freqs,
+            embed_dim=self.d_model,
+        )
+
+        # Condition Builder (cmask; mask_ratio only used for optional regularization)
+        self.condition_builder = ConditionBuilder(
+            d_model=self.d_model,
+            mask_ratio=cond_cfg.get("cmask_ratio", 0.0),
+        )
+
+        # RoPE for Denoiser
+        self.rope = RotaryEmbedding(dim=self.d_model // tf_cfg.num_heads)
+
+        # Flow Matching Transformer
+        self.transformer = FlowMatchingTransformer(
+            in_channels=self.d_model,
+            out_channels=self.d_model,
+            use_context_in=True,
+            context_in_dim=self.d_model,
+            use_txt_in=False,
+            vec_in_dim=0,
+            use_pre_text_attn=False,
+            config=tf_cfg,
+        )
+
+    def _extract_image_features(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        Image feature extraction: HaMeR backbone → Frame Compressor.
+        images: (B, T, 3, 256, 256) uint8
+        Returns: (B, T, d_model) — 1 image token per frame
+        """
+        B, T = images.shape[:2]
+        imgs = images.reshape(B * T, 3, images.shape[-2], images.shape[-1]).float() / 255.0
+        with torch.no_grad():
+            patches = self.hamer_backbone(imgs)
+        patches = patches.reshape(B, T, 192, -1)
+        compressed = self.frame_compressor(patches)
+        return compressed
+
+    def _build_condition(self, batch: dict) -> tuple[torch.Tensor, torch.Tensor]:
+        """Build the condition stream and return (condition, m)."""
+        if "image_tokens" in batch:
+            image_tokens = batch["image_tokens"]
+        else:
+            image_tokens = self._extract_image_features(batch["images"])
+        skeleton_tokens = self.skeleton_encoder(
+            batch["hamer_landmarks"],
+            batch["crop_intrinsics"],
+        )
+        condition, m = self.condition_builder(
+            image_tokens,
+            skeleton_tokens,
+            batch["hamer_confidence"],
+        )
+        return condition, m
+
+    def _build_z_positions(self, B: int, T: int, device: torch.device) -> torch.Tensor:
+        """z-stream RoPE positions: [0, 1, 1, 2, 2, ..., T, T]"""
+        pos = torch.zeros(1 + 2 * T, dtype=torch.long, device=device)
+        for t in range(T):
+            pos[1 + t] = t + 1           # pose token
+            pos[1 + T + t] = t + 1       # trans token
+        return pos
+
+    def _build_cond_positions(self, B: int, T: int, device: torch.device) -> torch.Tensor:
+        """c-stream RoPE positions: the 2 tokens per frame share position t"""
+        pos = torch.zeros(2 * T, dtype=torch.long, device=device)
+        for t in range(T):
+            pos[2 * t] = t
+            pos[2 * t + 1] = t
+        return pos
+
+    def forward(self, x_t_flat: torch.Tensor, t: torch.Tensor, batch: dict):
+        """
+        Args:
+            x_t_flat: (B, flat_dim) — noisy flow target
+            t: (B,) — flow timestep
+            batch: dict with all condition data
+        Returns:
+            v_pred_flat: (B, flat_dim) — predicted flow velocity
+        """
+        B, T = x_t_flat.shape[0], self.window_size
+        device = x_t_flat.device
+
+        # 1. Unpack flat → components
+        beta_t, poses_t, trans_t = self.mano_slice.unpack(x_t_flat)
+
+        # 2. Project to d_model → z-stream tokens
+        beta_tok = self.beta_proj_in(beta_t).unsqueeze(1)           # (B, 1, D)
+        pose_tok = self.pose_proj_in(poses_t)                       # (B, T, D)
+        trans_tok = self.trans_proj_in(trans_t)                     # (B, T, D)
+        z_stream = torch.cat([beta_tok, pose_tok, trans_tok], dim=1)  # (B, 1+2T, D)
+
+        # 3. Build condition → c-stream
+        condition, _ = self._build_condition(batch)  # (B, 2*T, D)
+
+        # 4. RoPE positions
+        z_positions = self._build_z_positions(B, T, device)
+        cond_positions = self._build_cond_positions(B, T, device)
+
+        # 5. Transformer forward
+        pc_out, _ = self.transformer(
+            pc=z_stream,
+            dino=condition,
+            timesteps=t,
+            y=None,
+            txt_tokens=None,
+            z_positions=z_positions,
+            cond_positions=cond_positions,
+            rope_module=self.rope,
+        )
+
+        # 6. Project back to native dims
+        v_beta = self.beta_proj_out(pc_out[:, 0:1])          # (B, 1, 10)
+        v_pose = self.pose_proj_out(pc_out[:, 1:1 + T])      # (B, T, 48)
+        v_trans = self.trans_proj_out(pc_out[:, 1 + T:])     # (B, T, 3)
+
+        # 7. Pack flat
+        v_pred_flat = self.mano_slice.pack(
+            v_beta.squeeze(1), v_pose, v_trans
+        )
+        return v_pred_flat
+
+    @torch.no_grad()
+    def build_flow_target(self, batch: dict) -> torch.Tensor:
+        """Build the GT flow target x_1 (normalized space) from the batch."""
+        beta = batch["mano_betas"]       # (B, 10)
+        poses = batch["mano_params"]     # (B, T, 48)
+        trans = batch["mano_trans"]      # (B, T, 3)
+        flat = self.mano_slice.pack(beta, poses, trans)
+        return self.mano_slice.normalize(flat)
+
+    @torch.no_grad()
+    def reconstruct_params(self, x_1_flat: torch.Tensor):
+        """
+        Reconstruct MANO parameters (original scale) from the predicted x_1 (normalized space).
+
+        All values are used directly, without cumsum.
+
+        Returns:
+            params: (B, T, 51) = cat[pose(48), trans(3)]
+            beta: (B, 10)
+        """
+        x_1_orig = self.mano_slice.denormalize(x_1_flat)
+        beta, poses, trans = self.mano_slice.unpack(x_1_orig)
+        params = torch.cat([poses, trans], dim=-1)   # (B, T, 51)
+        return params, beta
+
+    def get_model_wrapper(self):
+        """Return a ModelWrapper for the ODE solver."""
+        return ModelWrapper(self)
+
+    def get_time_sampler(self):
+        """Return the time sampler."""
+        return FluxTimeSampler()

model/flow_matching/dualstream_transformer.py

@@ -0,0 +1,464 @@
+"""
+Dual-stream transformer for flow matching on hand pose sequences.
+
+Derived from ShapeR/model/flow_matching/dualstream_transformer.py with changes:
+  - Import path adjusted from `from model import attention` (same, works as-is)
+  - Renamed internal variables: img→z (latent stream), txt→cond (condition stream)
+  - Added RoPE support: position IDs + rope_module passed through block calls
+
+Architecture (Flux-style):
+  - DoubleStreamBlock: joint attention over (latent z_t, condition context)
+  - SingleStreamBlock: self-attention on concatenated tokens
+  - LastLayer: final projection to output dimension
+"""
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the CC BY-NC 4.0 license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This file is inspired by the Flux dual stream transformer architecture.
+# Original source: https://github.com/black-forest-labs/flux
+# Original license: Apache License 2.0
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from einops import rearrange
+from model import attention
+from model.pose_embedding import RotaryEmbedding
+from torch import nn, Tensor
+
+
+class FlowMatchingTransformer(nn.Module):
+    """Transformer model for flow matching on sequences."""
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        use_context_in,
+        context_in_dim,
+        use_txt_in,
+        vec_in_dim,
+        use_pre_text_attn,
+        config,
+    ):
+        super().__init__()
+
+        self.config = config
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.use_txt_in = use_txt_in
+        self.use_context_in = use_context_in
+        self.use_pre_text_attn = use_pre_text_attn
+        if config.hidden_size % config.num_heads != 0:
+            raise ValueError(
+                f"Hidden size {config.hidden_size} must be divisible by num_heads {config.num_heads}"
+            )
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_heads
+        self.pc_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
+        if self.use_txt_in:
+            self.vector_in = MLPEmbedder(vec_in_dim, self.hidden_size)
+
+        self.dropout = getattr(config, 'dropout', 0.1)
+
+        if self.use_pre_text_attn:
+            self.text_attn_in = nn.Linear(context_in_dim, self.hidden_size)
+            self.text_blocks = nn.ModuleList(
+                [
+                    DoubleStreamBlock(
+                        self.hidden_size,
+                        self.num_heads,
+                        mlp_ratio=config.mlp_ratio,
+                        qkv_bias=config.qkv_bias,
+                        dropout=self.dropout,
+                    )
+                    for _ in range(4)
+                ]
+            )
+
+        if self.use_context_in:
+            self.dino_in = nn.Linear(context_in_dim, self.hidden_size)
+            self.double_blocks = nn.ModuleList(
+                [
+                    DoubleStreamBlock(
+                        self.hidden_size,
+                        self.num_heads,
+                        mlp_ratio=config.mlp_ratio,
+                        qkv_bias=config.qkv_bias,
+                        dropout=self.dropout,
+                    )
+                    for _ in range(config.depth)
+                ]
+            )
+
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(
+                    self.hidden_size, self.num_heads, mlp_ratio=config.mlp_ratio,
+                    dropout=self.dropout,
+                )
+                for _ in range(
+                    config.depth_single_blocks
+                    + (config.depth if not self.use_context_in else 0)
+                )
+            ]
+        )
+
+        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
+
+    def forward(
+        self,
+        pc: Tensor,
+        dino: Tensor,
+        timesteps: Tensor,
+        y: Tensor,
+        txt_tokens: Tensor,
+        z_positions: Tensor = None,
+        cond_positions: Tensor = None,
+        rope_module=None,
+    ) -> Tensor:
+        if pc.ndim != 3 or (dino is not None and dino.ndim != 3):
+            raise ValueError("Input pc and dino tensors must have 3 dimensions.")
+
+        pc = self.pc_in(pc)
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        intermediates = None
+        if self.use_txt_in and y is not None:
+            vec = vec + self.vector_in(y)
+
+        if self.use_pre_text_attn:
+            txt_tokens = self.text_attn_in(txt_tokens)
+            for block in self.text_blocks:
+                pc, txt_tokens = block(z=pc, cond=txt_tokens, vec=vec)
+
+        cond_len = 0
+        if self.use_context_in:
+            dino = self.dino_in(dino)
+            cond_len = dino.shape[1]
+            for block in self.double_blocks:
+                pc, dino = block(
+                    z=pc, cond=dino, vec=vec,
+                    z_positions=z_positions,
+                    cond_positions=cond_positions,
+                    rope_module=rope_module,
+                )
+            # Build combined positions (matching concatenation order: [cond, z])
+            if z_positions is not None and cond_positions is not None:
+                combined_positions = torch.cat([cond_positions, z_positions], dim=0)
+            else:
+                combined_positions = None
+            pc = torch.cat((dino, pc), 1)
+
+        for block_idx, block in enumerate(self.single_blocks):
+            pc = block(pc, vec=vec, positions=combined_positions, rope_module=rope_module)
+            if block_idx == 3:
+                intermediates = pc
+
+        if self.use_context_in:
+            pc = pc[:, cond_len:, ...]
+
+        pc = self.final_layer(pc, vec)
+        return pc, intermediates
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+    """Create sinusoidal timestep embeddings."""
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, dtype=torch.float32)
+        / half
+    ).to(t.device)
+
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+class MLPEmbedder(nn.Module):
+    """Two-layer MLP for embedding."""
+
+    def __init__(self, in_dim: int, hidden_dim: int):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(torch.nn.Module):
+    """Root Mean Square Layer Normalization."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: 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 QKNorm(torch.nn.Module):
+    """Query-Key normalization for stable attention."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tuple[Tensor, Tensor]:
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+
+class SelfAttention(nn.Module):
+    """Standard self-attention with QK normalization."""
+
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.norm = QKNorm(head_dim)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        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, k = self.norm(q, k, v)
+        x = attention(q, k, v)
+        x = self.proj(x)
+        return x
+
+
+@dataclass
+class ModulationOut:
+    shift: Tensor
+    scale: Tensor
+    gate: Tensor
+
+
+class Modulation(nn.Module):
+    """Adaptive layer norm modulation from timestep embedding."""
+
+    def __init__(self, dim: int, double: bool):
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> Tuple[ModulationOut, Optional[ModulationOut]]:
+        out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(
+            self.multiplier, dim=-1
+        )
+        return (
+            ModulationOut(*out[:3]),
+            ModulationOut(*out[3:]) if self.is_double else None,
+        )
+
+
+class DoubleStreamBlock(nn.Module):
+    """Double-stream attention block with separate z (latent) and cond (condition) streams."""
+
+    def __init__(
+        self, hidden_size: int, num_heads: int, mlp_ratio: float,
+        qkv_bias: bool = False, dropout: float = 0.1,
+    ):
+        super().__init__()
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        # z stream (latent) modules
+        self.z_mod = Modulation(hidden_size, double=True)
+        self.z_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.z_attn = SelfAttention(
+            dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
+        )
+        self.z_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.z_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+        self.z_attn_drop = nn.Dropout(dropout)
+        self.z_mlp_drop = nn.Dropout(dropout)
+
+        # cond stream (condition) modules
+        self.cond_mod = Modulation(hidden_size, double=True)
+        self.cond_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.cond_attn = SelfAttention(
+            dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias
+        )
+        self.cond_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.cond_mlp = nn.Sequential(
+            nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
+        )
+        self.cond_attn_drop = nn.Dropout(dropout)
+        self.cond_mlp_drop = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        z: Tensor,
+        cond: Tensor,
+        vec: Tensor,
+        z_positions: Tensor = None,
+        cond_positions: Tensor = None,
+        rope_module=None,
+    ) -> Tuple[Tensor, Tensor]:
+        z_mod1, z_mod2 = self.z_mod(vec)
+        cond_mod1, cond_mod2 = self.cond_mod(vec)
+
+        # prepare z stream for attention
+        z_modulated = self.z_norm1(z)
+        z_modulated = (1 + z_mod1.scale) * z_modulated + z_mod1.shift
+        z_qkv = self.z_attn.qkv(z_modulated)
+        z_q, z_k, z_v = rearrange(
+            z_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+        )
+        z_q, z_k = self.z_attn.norm(z_q, z_k, z_v)
+
+        # prepare cond stream for attention
+        cond_modulated = self.cond_norm1(cond)
+        cond_modulated = (1 + cond_mod1.scale) * cond_modulated + cond_mod1.shift
+        cond_qkv = self.cond_attn.qkv(cond_modulated)
+        cond_q, cond_k, cond_v = rearrange(
+            cond_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads
+        )
+        cond_q, cond_k = self.cond_attn.norm(cond_q, cond_k, cond_v)
+
+        # Apply RoPE (after QK-norm, before concatenation)
+        if rope_module is not None and z_positions is not None:
+            sin, cos = RotaryEmbedding.get_sin_cos_from_positions(
+                z_positions, rope_module.inv_freq
+            )
+            z_q = RotaryEmbedding.apply_rope(z_q, sin, cos)
+            z_k = RotaryEmbedding.apply_rope(z_k, sin, cos)
+        if rope_module is not None and cond_positions is not None:
+            sin, cos = RotaryEmbedding.get_sin_cos_from_positions(
+                cond_positions, rope_module.inv_freq
+            )
+            cond_q = RotaryEmbedding.apply_rope(cond_q, sin, cos)
+            cond_k = RotaryEmbedding.apply_rope(cond_k, sin, cos)
+
+        # run actual attention (concat [cond, z] streams)
+        q = torch.cat((cond_q, z_q), dim=2)
+        k = torch.cat((cond_k, z_k), dim=2)
+        v = torch.cat((cond_v, z_v), dim=2)
+
+        attn = attention(q, k, v)
+        cond_attn, z_attn = attn[:, : cond.shape[1]], attn[:, cond.shape[1] :]
+
+        # calculate the z stream blocks
+        z = z + z_mod1.gate * self.z_attn_drop(self.z_attn.proj(z_attn))
+        z = z + z_mod2.gate * self.z_mlp_drop(self.z_mlp(
+            (1 + z_mod2.scale) * self.z_norm2(z) + z_mod2.shift
+        ))
+
+        # calculate the cond stream blocks
+        cond = cond + cond_mod1.gate * self.cond_attn_drop(self.cond_attn.proj(cond_attn))
+        cond = cond + cond_mod2.gate * self.cond_mlp_drop(self.cond_mlp(
+            (1 + cond_mod2.scale) * self.cond_norm2(cond) + cond_mod2.shift
+        ))
+        return z, cond
+
+
+class SingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers as described in
+    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qk_scale: Optional[float] = None,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
+
+        self.norm = QKNorm(head_dim)
+
+        self.hidden_size = hidden_size
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.mlp_act = nn.GELU(approximate="tanh")
+        self.modulation = Modulation(hidden_size, double=False)
+        self.attn_drop = nn.Dropout(dropout)
+        self.mlp_drop = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: Tensor,
+        vec: Tensor,
+        positions: Tensor = None,
+        rope_module=None,
+    ) -> Tensor:
+        mod, _ = self.modulation(vec)
+        x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+        qkv, mlp = torch.split(
+            self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1
+        )
+
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+
+        # Apply RoPE (after QK-norm, before attention)
+        if rope_module is not None and positions is not None:
+            sin, cos = RotaryEmbedding.get_sin_cos_from_positions(
+                positions, rope_module.inv_freq
+            )
+            q = RotaryEmbedding.apply_rope(q, sin, cos)
+            k = RotaryEmbedding.apply_rope(k, sin, cos)
+
+        attn = self.attn_drop(attention(q, k, v))
+        output = self.mlp_drop(self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)))
+        return x + mod.gate * output
+
+
+class LastLayer(nn.Module):
+    """Final projection layer with adaptive layer norm."""
+
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(
+            hidden_size, patch_size * patch_size * out_channels, bias=True
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+        x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+        x = self.linear(x)
+        return x

model/flow_matching/helpers/__init__.py

@@ -0,0 +1,20 @@
+"""Flow matching helpers package — direct reuse from ShapeR."""
+
+from model.flow_matching.helpers.path import AffineProbPath, PathSample
+from model.flow_matching.helpers.scheduler import (
+    CondOTScheduler,
+    FluxTimeSampler,
+    TimeSampler,
+)
+from model.flow_matching.helpers.solver import ODESolver
+from model.flow_matching.helpers.model_wrapper import ModelWrapper
+
+__all__ = [
+    "AffineProbPath",
+    "PathSample",
+    "CondOTScheduler",
+    "FluxTimeSampler",
+    "TimeSampler",
+    "ODESolver",
+    "ModelWrapper",
+]

model/flow_matching/helpers/model_wrapper.py

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the CC BY-NC 4.0 license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This file is derived from https://github.com/facebookresearch/flow_matching
+# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
+
+# pyre-unsafe
+
+from abc import ABC
+
+from torch import nn, Tensor
+
+
+class ModelWrapper(ABC, nn.Module):
+    """
+    This class is used to wrap around another model, adding custom forward pass logic.
+    """
+
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.model = model
+
+    def forward(self, x: Tensor, t: Tensor, **extras) -> Tensor:
+        r"""
+        This method defines how inputs should be passed through the wrapped model.
+        Here, we're assuming that the wrapped model takes both :math:`x` and :math:`t` as input,
+        along with any additional keyword arguments.
+
+        Optional things to do here:
+            - check that t is in the dimensions that the model is expecting.
+            - add a custom forward pass logic.
+            - call the wrapped model.
+
+        | given x, t
+        | returns the model output for input x at time t, with extra information `extra`.
+
+        Args:
+            x (Tensor): input data to the model (batch_size, ...).
+            t (Tensor): time (batch_size).
+            **extras: additional information forwarded to the model, e.g., text condition.
+
+        Returns:
+            Tensor: model output.
+        """
+        return self.model(x=x, t=t, **extras)

model/flow_matching/helpers/path.py

@@ -0,0 +1,134 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the CC BY-NC 4.0 license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This file is derived from https://github.com/facebookresearch/flow_matching
+# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
+
+# pyre-unsafe
+
+from dataclasses import dataclass, field
+
+from torch import Tensor
+
+
+def expand_tensor_like(input_tensor: Tensor, expand_to: Tensor) -> Tensor:
+    """`input_tensor` is a 1d vector of length equal to the batch size of `expand_to`,
+    expand `input_tensor` to have the same shape as `expand_to` along all remaining dimensions.
+
+    Args:
+        input_tensor (Tensor): (batch_size,).
+        expand_to (Tensor): (batch_size, ...).
+
+    Returns:
+        Tensor: (batch_size, ...).
+    """
+    assert input_tensor.ndim == 1, "Input tensor must be a 1d vector."
+    assert (
+        input_tensor.shape[0] == expand_to.shape[0]
+    ), f"The first (batch_size) dimension must match. Got shape {input_tensor.shape} and {expand_to.shape}."
+
+    dim_diff = expand_to.ndim - input_tensor.ndim
+
+    t_expanded = input_tensor.clone()
+    t_expanded = t_expanded.reshape(-1, *([1] * dim_diff))
+
+    return t_expanded.expand_as(expand_to)
+
+
+@dataclass
+class PathSample:
+    r"""Represents a sample of a conditional-flow generated probability path.
+
+    Attributes:
+        x_1 (Tensor): the target sample :math:`X_1`.
+        x_0 (Tensor): the source sample :math:`X_0`.
+        t (Tensor): the time sample :math:`t`.
+        x_t (Tensor): samples :math:`X_t \sim p_t(X_t)`, shape (batch_size, ...).
+        dx_t (Tensor): conditional target :math:`\frac{\partial X}{\partial t}`, shape: (batch_size, ...).
+
+    """
+
+    x_1: Tensor = field(metadata={"help": "target samples X_1 (batch_size, ...)."})
+    x_0: Tensor = field(metadata={"help": "source samples X_0 (batch_size, ...)."})
+    t: Tensor = field(metadata={"help": "time samples t (batch_size, ...)."})
+    x_t: Tensor = field(
+        metadata={"help": "samples x_t ~ p_t(X_t), shape (batch_size, ...)."}
+    )
+    dx_t: Tensor = field(
+        metadata={"help": "conditional target dX_t, shape: (batch_size, ...)."}
+    )
+
+
+class AffineProbPath:
+    r"""The ``AffineProbPath`` class represents a specific type of probability path where the transformation between distributions is affine.
+    An affine transformation can be represented as:
+
+    .. math::
+
+        X_t = \alpha_t X_1 + \sigma_t X_0,
+
+    where :math:`X_t` is the transformed data point at time `t`. :math:`X_0` and :math:`X_1` are the source and target data points, respectively. :math:`\alpha_t` and :math:`\sigma_t` are the parameters of the affine transformation at time `t`.
+
+    The scheduler is responsible for providing the time-dependent parameters :math:`\alpha_t` and :math:`\sigma_t`, as well as their derivatives, which define the affine transformation at any given time `t`.
+
+    Using ``AffineProbPath`` in the flow matching framework:
+
+    .. code-block:: python
+
+        # Instantiates a probability path
+        my_path = AffineProbPath(...)
+        mse_loss = torch.nn.MSELoss()
+
+        for x_1 in dataset:
+            # Sets x_0 to random noise
+            x_0 = torch.randn()
+
+            # Sets t to a random value in [0,1]
+            t = torch.rand()
+
+            # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
+            path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
+
+            # Computes the MSE loss w.r.t. the velocity
+            loss = mse_loss(path_sample.dx_t, my_model(x_t, t))
+            loss.backward()
+
+    Args:
+        scheduler (Scheduler): An instance of a scheduler that provides the parameters :math:`\alpha_t`, :math:`\sigma_t`, and their derivatives over time.
+
+    """
+
+    def __init__(self, scheduler):
+        self.scheduler = scheduler
+
+    def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor):
+        r"""Sample from the affine probability path:
+
+        | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`.
+        | return :math:`X_0, X_1, X_t = \alpha_t X_1 + \sigma_t X_0`, and the conditional velocity at :math:`X_t, \dot{X}_t = \dot{\alpha}_t X_1 + \dot{\sigma}_t X_0`.
+
+        Args:
+            x_0 (Tensor): source data point, shape (batch_size, ...).
+            x_1 (Tensor): target data point, shape (batch_size, ...).
+            t (Tensor): times in [0,1], shape (batch_size).
+
+        Returns:
+            PathSample: a conditional sample at :math:`X_t \sim p_t`.
+        """
+        scheduler_output = self.scheduler(t)
+        alpha_t = expand_tensor_like(
+            input_tensor=scheduler_output.alpha_t, expand_to=x_1
+        )
+        sigma_t = expand_tensor_like(
+            input_tensor=scheduler_output.sigma_t, expand_to=x_1
+        )
+        d_alpha_t = expand_tensor_like(
+            input_tensor=scheduler_output.d_alpha_t, expand_to=x_1
+        )
+        d_sigma_t = expand_tensor_like(
+            input_tensor=scheduler_output.d_sigma_t, expand_to=x_1
+        )
+        x_t = sigma_t * x_0 + alpha_t * x_1
+        dx_t = d_sigma_t * x_0 + d_alpha_t * x_1
+        return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)

model/flow_matching/helpers/scheduler.py

@@ -0,0 +1,146 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the CC BY-NC 4.0 license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This file is derived from https://github.com/facebookresearch/flow_matching
+# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
+
+# pyre-unsafe
+
+import math
+from dataclasses import dataclass, field
+
+import torch
+from torch import Tensor
+from torch.distributions import LogisticNormal, Uniform
+
+
+def get_lin_function(
+    x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
+):
+    m = (y2 - y1) / (x2 - x1)
+    b = y1 - m * x1
+    return lambda x: m * x + b
+
+
+def time_shift(mu: float, sigma: float, t: Tensor):
+    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+
+@dataclass
+class SchedulerOutput:
+    r"""Represents a sample of a conditional-flow generated probability path.
+
+    Attributes:
+        alpha_t (Tensor): :math:`\alpha_t`, shape (...).
+        sigma_t (Tensor): :math:`\sigma_t`, shape (...).
+        d_alpha_t (Tensor): :math:`\frac{\partial}{\partial t}\alpha_t`, shape (...).
+        d_sigma_t (Tensor): :math:`\frac{\partial}{\partial t}\sigma_t`, shape (...).
+
+    """
+
+    alpha_t: Tensor = field(metadata={"help": "alpha_t"})
+    sigma_t: Tensor = field(metadata={"help": "sigma_t"})
+    d_alpha_t: Tensor = field(metadata={"help": "Derivative of alpha_t."})
+    d_sigma_t: Tensor = field(metadata={"help": "Derivative of sigma_t."})
+
+
+class CondOTScheduler:
+    """CondOT Scheduler."""
+
+    def __call__(self, t: Tensor) -> SchedulerOutput:
+        r"""
+        Args:
+            t (Tensor): times in [0,1], shape (...).
+
+        Returns:
+            SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t`
+        """
+        return SchedulerOutput(
+            alpha_t=t,
+            sigma_t=1 - t,
+            d_alpha_t=torch.ones_like(t),
+            d_sigma_t=-torch.ones_like(t),
+        )
+
+    def snr_inverse(self, snr: Tensor) -> Tensor:
+        r"""
+        Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`.
+
+        Args:
+            snr (Tensor): The signal-to-noise, shape (...)
+
+        Returns:
+            Tensor: t, shape (...)
+        """
+        kappa_t = snr / (1.0 + snr)
+
+        return self.kappa_inverse(kappa=kappa_t)
+
+    def kappa_inverse(self, kappa: Tensor) -> Tensor:
+        """
+        Computes :math:`t` from :math:`kappa_t`.
+
+        Args:
+            kappa (Tensor): :math:`kappa`, shape (...)
+
+        Returns:
+            Tensor: t, shape (...)
+        """
+        return kappa
+
+
+class TimeSampler:
+    def __init__(self, dist):
+        self.dist_name = dist
+        if dist == "uniform":
+            self.distribution = Uniform(0.0, 1.0)
+        elif dist == "lognormal":
+            self.distribution = LogisticNormal(0.0, 1.0)
+        elif dist == "flux":
+            self.distribution = FluxTimeSampler(mode="train")
+        else:
+            raise ValueError(f"Unknown distribution {dist}")
+
+    def __call__(self, shape, num_tokens, device):
+        if self.dist_name == "uniform":
+            return self.distribution.sample(sample_shape=(shape,)).to(device).squeeze(0)
+        elif self.dist_name == "lognormal":
+            return self.distribution.sample(sample_shape=(shape,))[:, 0].to(device)
+        elif self.dist_name == "flux":
+            return self.distribution(shape, num_tokens, device)
+
+
+class FluxTimeSampler:
+    def __init__(self, **kwargs):
+        self.mode = kwargs.get("mode", "train")
+
+    def __call__(self, bs, num_tokens, device):
+        if self.mode == "train":
+            logits_norm = torch.randn(bs, device=device)
+            timesteps = logits_norm.sigmoid()
+            mu = get_lin_function(y1=0.5, y2=1.15)(num_tokens)
+            timesteps = time_shift(mu, 1.0, timesteps)
+            timesteps = 1 - timesteps
+        elif self.mode == "inference":
+            timesteps = torch.linspace(1, 0, bs + 1).to(device)
+            mu = get_lin_function(y1=0.5, y2=1.15)(num_tokens)
+            timesteps = time_shift(mu, 1.0, timesteps)
+            timesteps = 1 - timesteps
+        return timesteps
+
+
+def visualize_sampling():
+    import matplotlib.pyplot as plt
+
+    flux_sampler = FluxTimeSampler(mode="inference")
+    timesteps = flux_sampler(25, 2048, torch.device("cuda"))
+    x = torch.linspace(0, 1, 25 + 1)
+    plt.scatter(x, y=timesteps.cpu().numpy())
+    plt.savefig("scatter_plot.png")
+
+    flux_sampler = FluxTimeSampler(mode="train")
+    timesteps = flux_sampler(5000, 4096, torch.device("cuda"))
+    plt.hist(timesteps.cpu().numpy(), bins=50, edgecolor="black")
+    plt.savefig("scatter_plot_4096.png")
+    print(timesteps)

model/flow_matching/helpers/solver.py

@@ -0,0 +1,78 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the CC BY-NC 4.0 license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This file is derived from https://github.com/facebookresearch/flow_matching
+# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
+
+# pyre-unsafe
+
+import torch
+from torch import Tensor
+from torchdiffeq import odeint
+
+
+class ODESolver:
+    """A class to solve ordinary differential equations (ODEs) using a specified velocity model.
+
+    This class utilizes a velocity field model to solve ODEs over a given time grid using numerical ode solvers.
+
+    Args:
+        velocity_model (Union[ModelWrapper, Callable]): a velocity field model receiving :math:`(x,t)` and returning :math:`u_t(x)`
+    """
+
+    def __init__(self, velocity_model):
+        super().__init__()
+        self.velocity_model = velocity_model
+
+    def sample(
+        self,
+        x_init: Tensor,
+        # pyre-fixme[9]: step_size has type `float`; used as `None`.
+        step_size: float = None,
+        method: str = "euler",
+        atol: float = 1e-5,
+        rtol: float = 1e-5,
+        time_grid: Tensor = torch.tensor([0.0, 1.0]),
+        return_intermediates: bool = False,
+        enable_grad: bool = False,
+        **model_extras,
+    ):
+        r"""Solve the ODE with the velocity field.
+
+        Args:
+            x_init (Tensor): initial conditions (e.g., source samples :math:`X_0 \sim p`). Shape: [batch_size, ...].
+            step_size (Optional[float]): The step size. Must be None for adaptive step solvers.
+            method (str): A method supported by torchdiffeq. Defaults to "euler". Other commonly used solvers are "dopri5", "midpoint" and "heun3". For a complete list, see torchdiffeq.
+            atol (float): Absolute tolerance, used for adaptive step solvers.
+            rtol (float): Relative tolerance, used for adaptive step solvers.
+            time_grid (Tensor): The process is solved in the interval [min(time_grid, max(time_grid)] and if step_size is None then time discretization is set by the time grid. May specify a descending time_grid to solve in the reverse direction. Defaults to torch.tensor([0.0, 1.0]).
+            return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Defaults to False.
+            enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False.
+            **model_extras: Additional input for the model.
+
+        Returns:
+          The last timestep when return_intermediates=False, otherwise all values specified in time_grid.
+        """
+
+        def ode_func(t, x):
+            return self.velocity_model(x=x, t=t, **model_extras)
+
+        ode_opts = {"step_size": step_size} if step_size is not None else {}
+
+        with torch.set_grad_enabled(enable_grad):
+            # Approximate ODE solution with numerical ODE solver
+            sol = odeint(
+                ode_func,
+                x_init,
+                time_grid,
+                method=method,
+                options=ode_opts,
+                atol=atol,
+                rtol=rtol,
+            )
+
+        if return_intermediates:
+            return sol
+        else:
+            return sol[-1]

model/pose_embedding.py

@@ -0,0 +1,58 @@
+"""
+RoPE rotary positional encoding, reused by the Denoiser and Transformer.
+Extracted from model/vae/autoencoder.py.
+"""
+
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+
+class RotaryEmbedding(nn.Module):
+    """RoPE rotary positional encoding (applied to Q/K)."""
+
+    def __init__(self, dim: int, base: int = 10000):
+        super().__init__()
+        if dim % 2 != 0:
+            raise ValueError("RoPE dim must be even")
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def get_sin_cos(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
+        t = torch.arange(seq_len, device=device, dtype=dtype)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(device=device, dtype=dtype))
+        sin = freqs.sin()
+        cos = freqs.cos()
+        return sin, cos
+
+    @staticmethod
+    def get_sin_cos_from_positions(
+        positions: torch.Tensor,
+        inv_freq: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute sin/cos from arbitrary position indices (supports repeated indices).
+
+        Args:
+            positions: (L,) int — arbitrary position indices, duplicates allowed
+            inv_freq:  (D/2,)   — from RotaryEmbedding.inv_freq
+
+        Returns:
+            sin, cos: each (L, D/2)
+        """
+        freqs = torch.einsum("i,j->ij", positions.float(), inv_freq.to(positions))
+        return freqs.sin(), freqs.cos()
+
+    @staticmethod
+    def apply_rope(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
+        """
+        x: (B, H, L, D)
+        sin/cos: (L, D/2)
+        """
+        x1 = x[..., 0::2]
+        x2 = x[..., 1::2]
+        sin = sin.unsqueeze(0).unsqueeze(0)
+        cos = cos.unsqueeze(0).unsqueeze(0)
+        x_rot = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+        return x_rot.flatten(-2)