| |
| |
| """Adapter that loads the upstream HaMeR Lightning checkpoint into the |
| self-contained torch reference shipped in :mod:`hamer.reference.hamer`. |
| |
| The upstream ``mano_head.transformer.transformer.layers.{i}.{0|1|2}.{norm|fn}`` |
| naming (and the dropout-padded ``Sequential`` indices) are remapped onto the |
| reference module's flatter ``head.transformer.layers[i].{norm_sa|sa|...}`` |
| layout. No new dependencies — pure ``torch`` + ``numpy``. |
| """ |
| from __future__ import annotations |
|
|
| import re |
| from pathlib import Path |
| from typing import Dict |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| def _remap_key(k: str) -> str | None: |
| """Translate one upstream key to the reference module's name; return |
| ``None`` for keys we don't model (e.g. discriminator).""" |
| if k.startswith("backbone."): |
| return k |
| m = re.match(r"mano_head\.transformer\.transformer\.layers\.(\d+)\.(\d+)\.(.+)", k) |
| if m: |
| block, sub, rest = int(m.group(1)), int(m.group(2)), m.group(3) |
| slot = ("sa", "ca", "ff")[sub] |
| if rest.startswith("norm."): |
| slot_norm = ("norm_sa", "norm_ca", "norm_ff")[sub] |
| return f"head.transformer.layers.{block}.{slot_norm}." + rest[len("norm.") :] |
| if rest.startswith("fn."): |
| inner = rest[len("fn.") :] |
| |
| if slot == "ff": |
| inner = inner.replace("net.3.", "net.2.") |
| |
| inner = inner.replace("to_out.0.", "to_out.") |
| return f"head.transformer.layers.{block}.{slot}." + inner |
| return None |
| if k.startswith("mano_head.transformer.to_token_embedding."): |
| return "head.transformer.to_token_embedding." + k.split(".", 3)[3] |
| if k == "mano_head.transformer.pos_embedding": |
| return "head.transformer.pos_embedding" |
| if k.startswith("mano_head.decpose.") or k.startswith("mano_head.decshape.") or k.startswith("mano_head.deccam."): |
| return "head." + k.split(".", 1)[1] |
| if k.startswith("mano_head.init_"): |
| return "head." + k.split(".", 1)[1] |
| return None |
|
|
|
|
| def load_hamer_checkpoint( |
| ref_module, |
| ckpt_path: Path | str, |
| mean_params_path: Path | str | None = None, |
| weights_only: bool = True, |
| ) -> dict: |
| """Load the official HaMeR checkpoint into the reference ``Hamer`` module. |
| |
| Returns a dict reporting how many tensors were mapped, missing, and extra. |
| Also overwrites the ``init_*`` buffers with the upstream MANO mean params |
| if ``mean_params_path`` is provided. |
| |
| ``weights_only=True`` (default) refuses arbitrary pickle globals. The |
| released ``hamer.ckpt`` only references ``collections.OrderedDict`` and |
| torch storage/rebuild helpers, so it loads under the safe loader; pass |
| ``weights_only=False`` only for a checkpoint you trust that does not. |
| """ |
| ckpt = torch.load(str(ckpt_path), map_location="cpu", weights_only=weights_only) |
| src = ckpt["state_dict"] |
|
|
| remapped: Dict[str, torch.Tensor] = {} |
| skipped = [] |
| for k, v in src.items(): |
| new_k = _remap_key(k) |
| if new_k is None: |
| skipped.append(k) |
| continue |
| remapped[new_k] = v |
|
|
| missing, unexpected = ref_module.load_state_dict(remapped, strict=False) |
|
|
| if mean_params_path is not None: |
| mp = np.load(str(mean_params_path)) |
| with torch.no_grad(): |
| ref_module.head.init_hand_pose.copy_(torch.from_numpy(mp["pose"].astype("float32")).unsqueeze(0)) |
| ref_module.head.init_betas.copy_(torch.from_numpy(mp["shape"].astype("float32")).unsqueeze(0)) |
| ref_module.head.init_cam.copy_(torch.from_numpy(mp["cam"].astype("float32")).unsqueeze(0)) |
|
|
| return { |
| "loaded": len(remapped), |
| "skipped_upstream": len(skipped), |
| "missing_in_ref": len(missing), |
| "unexpected_in_ref": len(unexpected), |
| "missing_examples": missing[:5], |
| "unexpected_examples": unexpected[:5], |
| } |
|
|
