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

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

utils/checkpoint_utils.py

@@ -0,0 +1,68 @@
+"""
+Checkpoint loading utilities.
+"""
+
+import gc
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import torch
+
+from model.flow_matching.denoiser import HandPoseDenoiser
+
+
+def _clean_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    return {
+        (k[len("module.") :] if k.startswith("module.") else k): v
+        for k, v in state_dict.items()
+    }
+
+
+def _load_state_dict(
+    ckpt_path: Path, device: torch.device
+) -> Tuple[Dict[str, torch.Tensor], Optional[int]]:
+    # Load to CPU first to avoid a large checkpoint filling GPU memory
+    raw = None
+    if ckpt_path.is_dir():
+        state_file = ckpt_path / "mp_rank_00_model_states.pt"
+        if not state_file.exists():
+            raise FileNotFoundError(f"DeepSpeed checkpoint not found: {state_file}")
+        raw = torch.load(state_file, map_location="cpu", weights_only=False)
+    else:
+        raw = torch.load(str(ckpt_path), map_location="cpu", weights_only=False)
+
+    state_dict = raw.get("module", raw) if isinstance(raw, dict) else raw
+    epoch = raw.get("epoch") if isinstance(raw, dict) else None
+    del raw
+    gc.collect()
+
+    # Move state_dict to the target device
+    state_dict = {k: v.to(device) for k, v in state_dict.items()}
+    return _clean_state_dict(state_dict), epoch
+
+
+def load_denoiser_from_ckpt(
+    config,
+    ckpt_path: str,
+    device: torch.device,
+) -> HandPoseDenoiser:
+    """Load HandPoseDenoiser, supporting either a DeepSpeed directory or a .pt file."""
+    denoiser = HandPoseDenoiser(config).to(device)
+    ckpt = Path(ckpt_path)
+    state_dict, epoch = _load_state_dict(ckpt, device)
+    result = denoiser.load_state_dict(state_dict, strict=False)
+    # The frozen HaMeR backbone is loaded separately from HAMER_CKPT by HaMeRBackbone
+    # and is intentionally absent from the released denoiser checkpoint; filter its keys
+    # out of the warning so the log stays focused on genuinely unexpected mismatches.
+    missing = [k for k in result.missing_keys if not k.startswith("hamer_backbone.")]
+    unexpected = [k for k in result.unexpected_keys if not k.startswith("hamer_backbone.")]
+    if missing:
+        print(f"[Denoiser] Missing keys: {missing}")
+    if unexpected:
+        print(f"[Denoiser] Unexpected keys: {unexpected}")
+    epoch_str = epoch if epoch is not None else "?"
+    print(f"[Denoiser] Loaded from {ckpt}  (epoch={epoch_str})")
+    del state_dict
+    gc.collect()
+    torch.cuda.empty_cache()
+    return denoiser