| |
| |
| """Self-contained torch CPU reference for the HaMeR per-frame regressor used |
| by HaMeR. |
| |
| Adapted from ``geopavlakos/hamer`` (Pavlakos et al., ICCV 2023) with upstream |
| dependencies on ``timm``, ``pytorch_lightning``, ``yacs``, ``smplx`` and |
| ``chumpy`` removed — this file needs only ``torch`` and ``einops`` so it runs |
| in the Tenstorrent venv without extra packages. Random weights, deterministic |
| under ``torch.manual_seed``; used as the accuracy anchor for the tt-nn port. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from functools import partial |
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from einops import rearrange |
|
|
|
|
| |
| |
| |
| @dataclass |
| class HamerConfig: |
| |
| img_size: tuple = (256, 192) |
| patch_size: int = 16 |
| in_chans: int = 3 |
| |
| embed_dim: int = 1280 |
| depth: int = 32 |
| num_heads: int = 16 |
| mlp_ratio: float = 4.0 |
| qkv_bias: bool = True |
| |
| num_hand_joints: int = 15 |
| joint_rep_dim: int = 6 |
| ief_iters: int = 1 |
| head_dim: int = 1024 |
| head_depth: int = 6 |
| head_heads: int = 8 |
| head_dim_head: int = 64 |
| head_mlp_dim: int = 1024 |
|
|
| @property |
| def npose(self) -> int: |
| return self.joint_rep_dim * (self.num_hand_joints + 1) |
|
|
|
|
| |
| |
| |
| class Mlp(nn.Module): |
| def __init__(self, dim: int, hidden_dim: int) -> None: |
| super().__init__() |
| self.fc1 = nn.Linear(dim, hidden_dim) |
| self.act = nn.GELU() |
| self.fc2 = nn.Linear(hidden_dim, dim) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.fc2(self.act(self.fc1(x))) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, dim: int, num_heads: int, qkv_bias: bool = True) -> None: |
| super().__init__() |
| self.num_heads = num_heads |
| self.head_dim = dim // num_heads |
| self.scale = self.head_dim ** -0.5 |
| self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| self.proj = nn.Linear(dim, dim) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| B, N, C = x.shape |
| qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) |
| q, k, v = qkv[0], qkv[1], qkv[2] |
| attn = (q * self.scale) @ k.transpose(-2, -1) |
| attn = attn.softmax(dim=-1) |
| out = (attn @ v).transpose(1, 2).reshape(B, N, C) |
| return self.proj(out) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, dim: int, num_heads: int, mlp_ratio: float, qkv_bias: bool) -> None: |
| super().__init__() |
| norm_layer = partial(nn.LayerNorm, eps=1e-6) |
| self.norm1 = norm_layer(dim) |
| self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias) |
| self.norm2 = norm_layer(dim) |
| self.mlp = Mlp(dim, int(dim * mlp_ratio)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = x + self.attn(self.norm1(x)) |
| x = x + self.mlp(self.norm2(x)) |
| return x |
|
|
|
|
| class PatchEmbed(nn.Module): |
| def __init__(self, img_size: tuple, patch_size: int, in_chans: int, embed_dim: int) -> None: |
| super().__init__() |
| self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, padding=2) |
| self.img_h = img_size[0] |
| self.img_w = img_size[1] |
|
|
| def forward(self, x: torch.Tensor): |
| x = self.proj(x) |
| Hp, Wp = x.shape[2], x.shape[3] |
| x = x.flatten(2).transpose(1, 2) |
| return x, (Hp, Wp) |
|
|
|
|
| class ViT(nn.Module): |
| def __init__(self, cfg: HamerConfig) -> None: |
| super().__init__() |
| self.patch_embed = PatchEmbed(cfg.img_size, cfg.patch_size, cfg.in_chans, cfg.embed_dim) |
| num_patches = ((cfg.img_size[0] + 2 * 2 - cfg.patch_size) // cfg.patch_size + 1) * \ |
| ((cfg.img_size[1] + 2 * 2 - cfg.patch_size) // cfg.patch_size + 1) |
| self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, cfg.embed_dim)) |
| self.blocks = nn.ModuleList([ |
| Block(cfg.embed_dim, cfg.num_heads, cfg.mlp_ratio, cfg.qkv_bias) |
| for _ in range(cfg.depth) |
| ]) |
| self.last_norm = nn.LayerNorm(cfg.embed_dim, eps=1e-6) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| B = x.shape[0] |
| x, (Hp, Wp) = self.patch_embed(x) |
| x = x + self.pos_embed[:, 1:] + self.pos_embed[:, :1] |
| for blk in self.blocks: |
| x = blk(x) |
| x = self.last_norm(x) |
| return x.permute(0, 2, 1).reshape(B, -1, Hp, Wp).contiguous() |
|
|
|
|
| |
| |
| |
| class _FeedForward(nn.Module): |
| def __init__(self, dim: int, hidden_dim: int) -> None: |
| super().__init__() |
| self.net = nn.Sequential(nn.Linear(dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.net(x) |
|
|
|
|
| class _SelfAttention(nn.Module): |
| def __init__(self, dim: int, heads: int, dim_head: int) -> None: |
| super().__init__() |
| inner = dim_head * heads |
| self.heads = heads |
| self.scale = dim_head ** -0.5 |
| self.to_qkv = nn.Linear(dim, inner * 3, bias=False) |
| self.to_out = nn.Linear(inner, dim) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| qkv = self.to_qkv(x).chunk(3, dim=-1) |
| q, k, v = (rearrange(t, "b n (h d) -> b h n d", h=self.heads) for t in qkv) |
| attn = (q @ k.transpose(-1, -2) * self.scale).softmax(dim=-1) |
| out = rearrange(attn @ v, "b h n d -> b n (h d)") |
| return self.to_out(out) |
|
|
|
|
| class _CrossAttention(nn.Module): |
| def __init__(self, dim: int, heads: int, dim_head: int, context_dim: Optional[int] = None) -> None: |
| super().__init__() |
| inner = dim_head * heads |
| context_dim = context_dim or dim |
| self.heads = heads |
| self.scale = dim_head ** -0.5 |
| self.to_q = nn.Linear(dim, inner, bias=False) |
| self.to_kv = nn.Linear(context_dim, inner * 2, bias=False) |
| self.to_out = nn.Linear(inner, dim) |
|
|
| def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor: |
| k, v = self.to_kv(context).chunk(2, dim=-1) |
| q = self.to_q(x) |
| q, k, v = (rearrange(t, "b n (h d) -> b h n d", h=self.heads) for t in (q, k, v)) |
| attn = (q @ k.transpose(-1, -2) * self.scale).softmax(dim=-1) |
| out = rearrange(attn @ v, "b h n d -> b n (h d)") |
| return self.to_out(out) |
|
|
|
|
| class _CrossAttnBlock(nn.Module): |
| def __init__(self, dim: int, heads: int, dim_head: int, mlp_dim: int, context_dim: int) -> None: |
| super().__init__() |
| self.norm_sa = nn.LayerNorm(dim) |
| self.sa = _SelfAttention(dim, heads, dim_head) |
| self.norm_ca = nn.LayerNorm(dim) |
| self.ca = _CrossAttention(dim, heads, dim_head, context_dim=context_dim) |
| self.norm_ff = nn.LayerNorm(dim) |
| self.ff = _FeedForward(dim, mlp_dim) |
|
|
| def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor: |
| x = x + self.sa(self.norm_sa(x)) |
| x = x + self.ca(self.norm_ca(x), context=context) |
| x = x + self.ff(self.norm_ff(x)) |
| return x |
|
|
|
|
| class TransformerDecoder(nn.Module): |
| def __init__(self, num_tokens: int, token_dim: int, dim: int, depth: int, |
| heads: int, dim_head: int, mlp_dim: int, context_dim: int) -> None: |
| super().__init__() |
| self.to_token_embedding = nn.Linear(token_dim, dim) |
| self.pos_embedding = nn.Parameter(torch.randn(1, num_tokens, dim) * 0.02) |
| self.layers = nn.ModuleList([ |
| _CrossAttnBlock(dim, heads, dim_head, mlp_dim, context_dim) for _ in range(depth) |
| ]) |
|
|
| def forward(self, tokens: torch.Tensor, context: torch.Tensor) -> torch.Tensor: |
| x = self.to_token_embedding(tokens) + self.pos_embedding[:, : tokens.shape[1]] |
| for blk in self.layers: |
| x = blk(x, context) |
| return x |
|
|
|
|
| def rot6d_to_rotmat(x: torch.Tensor) -> torch.Tensor: |
| """Convert 6-D continuous rotation to 3×3 rotation matrix (Zhou et al., CVPR'19).""" |
| x = x.reshape(-1, 2, 3).permute(0, 2, 1).contiguous() |
| a1 = x[..., 0] |
| a2 = x[..., 1] |
| b1 = F.normalize(a1, dim=-1) |
| b2 = F.normalize(a2 - (b1 * a2).sum(-1, keepdim=True) * b1, dim=-1) |
| b3 = torch.cross(b1, b2, dim=-1) |
| return torch.stack((b1, b2, b3), dim=-1) |
|
|
|
|
| class MANOHead(nn.Module): |
| def __init__(self, cfg: HamerConfig) -> None: |
| super().__init__() |
| self.cfg = cfg |
| self.transformer = TransformerDecoder( |
| num_tokens=1, |
| token_dim=1, |
| dim=cfg.head_dim, |
| depth=cfg.head_depth, |
| heads=cfg.head_heads, |
| dim_head=cfg.head_dim_head, |
| mlp_dim=cfg.head_mlp_dim, |
| context_dim=cfg.embed_dim, |
| ) |
| self.decpose = nn.Linear(cfg.head_dim, cfg.npose) |
| self.decshape = nn.Linear(cfg.head_dim, 10) |
| self.deccam = nn.Linear(cfg.head_dim, 3) |
| |
| |
| self.register_buffer("init_hand_pose", torch.zeros(1, cfg.npose)) |
| self.register_buffer("init_betas", torch.zeros(1, 10)) |
| self.register_buffer("init_cam", torch.zeros(1, 3)) |
|
|
| def forward(self, features: torch.Tensor) -> torch.Tensor: |
| B = features.shape[0] |
| context = rearrange(features, "b c h w -> b (h w) c") |
| token = torch.zeros(B, 1, 1, device=features.device, dtype=features.dtype) |
| pose = self.init_hand_pose.expand(B, -1) |
| betas = self.init_betas.expand(B, -1) |
| cam = self.init_cam.expand(B, -1) |
| for _ in range(self.cfg.ief_iters): |
| tok_out = self.transformer(token, context=context).squeeze(1) |
| pose = self.decpose(tok_out) + pose |
| betas = self.decshape(tok_out) + betas |
| cam = self.deccam(tok_out) + cam |
| rotmats = rot6d_to_rotmat(pose).view(B, self.cfg.num_hand_joints + 1, 3, 3) |
| |
| return torch.cat([rotmats.reshape(B, -1), betas, cam], dim=-1) |
|
|
|
|
| class Hamer(nn.Module): |
| def __init__(self, cfg: Optional[HamerConfig] = None) -> None: |
| super().__init__() |
| self.cfg = cfg or HamerConfig() |
| self.backbone = ViT(self.cfg) |
| self.head = MANOHead(self.cfg) |
|
|
| def forward(self, image: torch.Tensor) -> torch.Tensor: |
| feats = self.backbone(image) |
| return self.head(feats) |
|
|
|
|
| def build_reference(seed: int = 0) -> Hamer: |
| """Random-weight torch CPU reference; deterministic under the given seed.""" |
| torch.manual_seed(seed) |
| model = Hamer() |
| model.eval() |
| return model |
|
|
|
|
| def sample_input(batch: int = 1, seed: int = 0) -> torch.Tensor: |
| g = torch.Generator().manual_seed(seed + 1) |
| return torch.randn(batch, 3, 256, 192, generator=g) |
|
|
|
|
| def build_paired(tt_module, device_id: int = 0, seed: int = 0): |
| """Harness hook — returns (ref_model, tt_model, sample_tensor). |
| |
| ``tt_module`` is the imported ``models.experimental.hamer.tt.hamer`` |
| namespace; it must expose ``build_from_reference(ref_model, device_id)`` |
| returning a callable with the same output contract as ``Hamer.forward``. |
| """ |
| ref = build_reference(seed=seed) |
| sample = sample_input(seed=seed) |
| tt_model = tt_module.build_from_reference(ref, device_id=device_id) |
| return ref, tt_model, sample |
|
|
