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

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

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)