111 lines
3.5 KiB
Python
111 lines
3.5 KiB
Python
"""
|
|
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
|