hamer-p150 / code /tests /test_api_host.py
changh95's picture
Python API: warm-up so the first real call is fast (warmup_variants, model.warmup()), quiet logs, install extras
77e087f verified
Raw History Blame Contribute Delete
16.3 kB
# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Host-only tests of the Python API (``hamer.HamerModel``): no device, no ttnn tensors, no weight download.
The API runs with ``backend="cpu"`` on a seeded random-weight reference (and a synthetic MANO pkl with the real array
shapes). Its outputs are compared with the HTTP server's own ``/predict`` handler (``hamer.server.app.predict``,
called in-process with the same reference as its model): the published API must return the server's numbers.
cd models/hamer-p150 && TT_VISIBLE_DEVICES=none HAMER_USE_TT=0 .venv/bin/python -m pytest code/tests/test_api_host.py -q
"""
from __future__ import annotations
import base64
import os
import pickle
import subprocess
import sys
from pathlib import Path
os.environ.setdefault("HAMER_USE_TT", "0") # never open a chip from a host test
import numpy as np # noqa: E402
import pytest # noqa: E402
import torch # noqa: E402
from PIL import Image # noqa: E402
import hamer # noqa: E402
from hamer import HamerModel, HandPrediction # noqa: E402
from hamer.api import _check_bbox, to_pil # noqa: E402
from hamer.reference import build_reference # noqa: E402
REPO = Path(__file__).resolve().parents[2]
MEDIA = REPO / "media"
SAMPLE = MEDIA / "sample1.jpg"
BOX_R, BOX_L = [20, 170, 295, 260], [130, 340, 320, 510]
torch.set_grad_enabled(False)
@pytest.fixture(scope="module")
def ref():
return build_reference(seed=0)
@pytest.fixture(scope="module")
def mano_pkl(tmp_path_factory):
g = np.random.default_rng(0)
w = g.random((778, 16))
w /= w.sum(1, keepdims=True)
jr = np.zeros((16, 778))
for j in range(16):
jr[j, g.choice(778, 20, replace=False)] = 1 / 20
d = {"v_template": g.random((778, 3)) * 0.1, "shapedirs": g.normal(0, 1e-3, (778, 3, 10)),
"posedirs": g.normal(0, 1e-3, (778, 3, 135)), "J_regressor": jr, "weights": w,
"kintree_table": np.array([[4294967295, 0, 1, 2, 0, 4, 5, 0, 7, 8, 0, 10, 11, 0, 13, 14], list(range(16))]),
"f": g.integers(0, 778, (1538, 3))}
p = tmp_path_factory.mktemp("mano") / "MANO_RIGHT.pkl"
with open(p, "wb") as f:
pickle.dump(d, f)
return p
@pytest.fixture(scope="module")
def mano(mano_pkl):
from hamer.mano import MANO
return MANO(mano_pkl).eval()
@pytest.fixture(scope="module")
def cpu_model(ref):
m = HamerModel.from_reference(ref, backend="cpu")
yield m
m.close()
@pytest.fixture(scope="module")
def cpu_model_mano(ref, mano):
m = HamerModel.from_reference(ref, backend="cpu", mano=mano)
yield m
m.close()
# --------------------------------------------------------------------------- #
# Package surface
# --------------------------------------------------------------------------- #
def test_import_does_not_load_ttnn():
code = "import sys, hamer; from hamer import HamerModel, HandPrediction, Hamer; assert 'ttnn' not in sys.modules"
env = {k: v for k, v in os.environ.items() if k not in ("TT_METAL_HOME", "TT_METAL_RUNTIME_ROOT")}
r = subprocess.run([sys.executable, "-c", code], env=env, capture_output=True, text=True, timeout=120)
assert r.returncode == 0, r.stderr
def test_public_names():
for n in ("HamerModel", "HandPrediction", "open_device", "load_reference", "apply_served_config",
"Hamer", "HamerConfig", "build_reference", "sample_input", "build_from_reference",
"load_hamer_checkpoint"):
assert n in hamer.__all__ and getattr(hamer, n) is not None
with pytest.raises(AttributeError):
hamer.does_not_exist # noqa: B018
def test_package_data_present():
pkg = Path(hamer.__file__).parent
assert (pkg / "kernels" / "hgemv" / "compute.cpp").is_file()
assert (pkg / "native" / "resize.c").is_file()
def test_no_machine_paths_in_runtime_code():
pkg = Path(hamer.__file__).parent
bad = [str(p) for p in pkg.rglob("*.py") if "/home/" in p.read_text()]
assert not bad, bad
def test_served_config_matches_manifest():
yaml = pytest.importorskip("yaml")
man = REPO / "tt-model.yaml"
if not man.is_file():
pytest.skip("tt-model.yaml not shipped here")
from hamer import served_config as sc
env = dict(yaml.safe_load(man.read_text())["serve"]["env"])
assert env.pop("TT_WEIGHTS_REVISION") == sc.WEIGHTS_REVISION
assert env.pop("HAMER_DISPATCH") == sc.SERVED_DISPATCH
env.pop("TT_METAL_VISIBLE_DEVICES")
env.pop("HAMER_WARMUP_ITERS")
assert {k: str(v) for k, v in env.items()} == sc.SERVED_ENV
assert sc.SERVED_WORKER_L1_SIZE == 1461376 - int(sc.SERVED_ENV["HAMER_KCFG_EXTRA_KB"]) * 1024
def test_apply_served_config_keeps_caller_values(monkeypatch):
from hamer.served_config import SERVED_ENV
for k in SERVED_ENV:
monkeypatch.delenv(k, raising=False)
monkeypatch.setenv("TT_FUSED_SDPA", "default")
differ = hamer.apply_served_config()
assert differ == {"TT_FUSED_SDPA": ("default", "32x192")}
assert os.environ["TT_FUSED_SDPA"] == "default" and os.environ["TT_FUSED_INPUT"] == "u8"
# --------------------------------------------------------------------------- #
# Inputs
# --------------------------------------------------------------------------- #
def test_to_pil_accepts_common_types():
want = np.asarray(Image.open(SAMPLE).convert("RGB"))
cands = [str(SAMPLE), SAMPLE, Image.open(SAMPLE), want, want.astype(np.float32) / 255.0,
torch.from_numpy(want.copy()), torch.from_numpy(want.copy()).permute(2, 0, 1),
torch.from_numpy(want.copy()).permute(2, 0, 1)[None], Image.open(SAMPLE).convert("RGBA")]
for c in cands:
got = np.asarray(to_pil(c))
assert got.shape == want.shape and np.array_equal(got, want), type(c)
with pytest.raises(ValueError):
to_pil(np.zeros((4, 4, 3), np.float32) + 2.0)
with pytest.raises(TypeError):
to_pil(123)
def test_bbox_rules():
img = Image.new("RGB", (100, 80))
assert _check_bbox(None, img, 2.0)[1] == "center_square"
assert _check_bbox([1, 2, 30, 40], img, 2.0) == ((1.0, 2.0, 30.0, 40.0), "request")
for bad in ([5, 5, 5, 9], [0, 0, float("nan"), 3], [200, 200, 300, 300], [0, 0, 0.5, 0.5], "abcd"):
with pytest.raises(ValueError):
_check_bbox(bad, img, 2.0)
# --------------------------------------------------------------------------- #
# Outputs vs the HTTP server's handler (same reference model)
# --------------------------------------------------------------------------- #
def _server_predict(model_obj, mano_obj, **req):
pytest.importorskip("fastapi")
import hamer.server.app as app
app.STATE.clear()
app.STATE.update(model=model_obj, mano=mano_obj, ready=True)
try:
body = app.PredictRequest(image=base64.b64encode(SAMPLE.read_bytes()).decode(), **req)
out = app.predict(body)
finally:
app.STATE.clear()
out.pop("timing_ms")
out.pop("trace_active")
return out
@pytest.mark.parametrize("box,right", [(BOX_R, True), (BOX_L, False), (None, True)])
def test_matches_server_handler(ref, cpu_model, box, right):
want = _server_predict(ref, None, bbox=box, is_right=right)
hand = cpu_model(SAMPLE, bbox=box, is_right=right)
assert isinstance(hand, HandPrediction)
assert hand.to_dict() == want
assert hand.rotmats.shape == (16, 3, 3) and hand.betas.shape == (10,) and hand.cam.shape == (3,)
assert hand.cam_t_full.shape == (3,) and hand.vertices is None and not hand.mano_available
assert set(hand.timing_ms) == {"preprocess", "model", "mano", "total"}
@pytest.mark.parametrize("return_vertices", [True, False])
def test_matches_server_handler_with_mano(ref, mano, cpu_model_mano, return_vertices):
want = _server_predict(ref, mano, bbox=BOX_L, is_right=False, return_vertices=return_vertices)
hand = cpu_model_mano(SAMPLE, bbox=BOX_L, is_right=False, return_vertices=return_vertices)
assert hand.to_dict() == want
assert hand.joints.shape == (21, 3) and hand.keypoints_2d.shape == (21, 2)
assert (hand.vertices is not None) == return_vertices
assert cpu_model_mano.faces.shape == (1538, 3)
def test_several_hands_and_input_types(cpu_model):
frame = np.asarray(Image.open(SAMPLE).convert("RGB"))
hands = cpu_model(frame, bbox=np.array([BOX_R, BOX_L]), is_right=[True, False])
assert isinstance(hands, list) and len(hands) == 2
single = [cpu_model(SAMPLE, bbox=b, is_right=r) for b, r in ((BOX_R, True), (BOX_L, False))]
for a, b in zip(hands, single):
assert a.to_dict(rounded=False) == b.to_dict(rounded=False)
with pytest.raises(ValueError):
cpu_model(frame, bbox=[BOX_R, BOX_L], is_right=[True])
with pytest.raises(ValueError):
cpu_model(frame, bbox=BOX_R, rescale_factor=0.2)
def test_mano_params_layout(cpu_model):
h = cpu_model(SAMPLE, bbox=BOX_R)
mp = h.mano_params
assert mp["global_orient"].shape == (1, 3, 3) and mp["hand_pose"].shape == (15, 3, 3)
assert np.array_equal(np.concatenate([mp["global_orient"], mp["hand_pose"]]), h.rotmats)
r = h.rotmats.astype(np.float64)
assert np.allclose(r @ r.transpose(0, 2, 1), np.eye(3), atol=1e-5) # proper rotations
def test_forward_crops_upstream_layout(ref, cpu_model_mano):
from hamer.preprocess import crop_bbox, normalize_raw, preprocess_pil
img = Image.open(SAMPLE).convert("RGB")
u8, _ = preprocess_pil(img, BOX_R, raw=True) # (1, 3, 256, 192)
full, _ = crop_bbox(img, BOX_R) # 256 x 256 upstream-style crop
u256 = torch.from_numpy(np.asarray(full).copy()).permute(2, 0, 1)[None]
res = cpu_model_mano.forward_crops(torch.cat([u8, u8]))
assert res["pred_cam"].shape == (2, 3) and res["pred_cam_t"].shape == (2, 3)
assert res["pred_mano_params"]["global_orient"].shape == (2, 1, 3, 3)
assert res["pred_mano_params"]["hand_pose"].shape == (2, 15, 3, 3)
assert res["pred_mano_params"]["betas"].shape == (2, 10)
assert res["pred_vertices"].shape == (2, 778, 3) and res["pred_keypoints_3d"].shape == (2, 21, 3)
res_n = cpu_model_mano.forward_crops(normalize_raw(u8)) # uint8 == its normalized form
assert torch.equal(res["pred_cam"][0], res["pred_cam"][1]) and torch.equal(res_n["pred_cam"][0], res["pred_cam"][0])
want = ref(normalize_raw(u8))[0]
got = torch.cat([res["pred_mano_params"]["global_orient"][0].reshape(-1),
res["pred_mano_params"]["hand_pose"][0].reshape(-1),
res["pred_mano_params"]["betas"][0], res["pred_cam"][0]])
assert torch.equal(got, want)
r256 = cpu_model_mano.forward_crops(u256) # the [32:224] slice of a 256 crop
assert r256["pred_cam"].shape == (1, 3)
with pytest.raises(ValueError):
cpu_model_mano.forward_crops(torch.zeros(1, 3, 128, 128))
def test_close_and_context_manager(ref):
with HamerModel.from_reference(ref, backend="cpu") as m:
assert m(SAMPLE, bbox=BOX_R).rotmats.shape == (16, 3, 3)
with pytest.raises(RuntimeError):
m(SAMPLE, bbox=BOX_R)
m.close() # idempotent
assert "closed" in repr(m)
# --------------------------------------------------------------------------- #
# Warm-up API (from_pretrained(warmup_variants=...), model.warmup(**variant))
# --------------------------------------------------------------------------- #
def test_warmup_signature_and_defaults():
import inspect
from hamer.api import DEFAULT_WARMUP_VARIANTS, WARMUP_FRAME_SIZE, normalize_warmup_variant
sig = inspect.signature(HamerModel.from_pretrained)
assert sig.parameters["warmup_variants"].default == "default"
assert sig.parameters["verbose"].default is False
assert inspect.signature(HamerModel.from_reference).parameters["warmup_variants"].default == "auto"
assert "DEFAULT_WARMUP_VARIANTS" in hamer.api.__all__
keys = [normalize_warmup_variant(**v) for v in DEFAULT_WARMUP_VARIANTS]
assert len(set(keys)) == len(keys) # no duplicates
calls = {dict(k)["call"] for k in keys}
assert calls == {"frame", "crops", "run_crop"}
frame = [dict(k) for k in keys if dict(k)["call"] == "frame"]
assert {f["input"] for f in frame} == {"path", "pil", "array", "tensor"}
assert {f.get("fmt") for f in frame if f["input"] == "path"} == {"jpeg", "png"}
assert {f["is_right"] for f in frame} >= {True, False, "mixed"}
assert {f["bbox"] for f in frame} == {"box", "center"}
assert all(f["frame_size"] == WARMUP_FRAME_SIZE for f in frame if f.get("fmt") != "png")
def test_normalize_warmup_variant_rules():
from hamer.api import normalize_warmup_variant as nv
assert nv() == nv(call="frame", input="path", fmt="jpeg", hands=1, is_right=True, bbox="box")
assert dict(nv(hands=2))["is_right"] == "mixed"
assert "fmt" not in dict(nv(input="array", fmt="png")) # fmt only matters for files
assert dict(nv(frame_size=[512, 334]))["frame_size"] == (512, 334)
assert dict(nv(call="crops"))["batch"] == 2
for bad in (dict(call="nope"), dict(bogus=1), dict(input="url"), dict(hands=0), dict(is_right="mixed"),
dict(frame_size=(8, 8)), dict(call="crops", crop_width=128), dict(call="run_crop", batch=2),
dict(call="crops", batch=0), dict(call="run_crop", crop_dtype="fp16")):
with pytest.raises(ValueError):
nv(**bad)
def test_warmup_spec_parsing():
from hamer.api import DEFAULT_WARMUP_VARIANTS, _warmup_spec
assert _warmup_spec("default") == [dict(v) for v in DEFAULT_WARMUP_VARIANTS] == _warmup_spec("all")
for none in (None, "none", False, []):
assert _warmup_spec(none) == []
assert _warmup_spec({"call": "run_crop"}) == [{"call": "run_crop"}]
with pytest.raises(ValueError):
_warmup_spec("fast")
with pytest.raises(ValueError):
_warmup_spec(["run_crop"])
def test_bad_warmup_spec_fails_before_build(ref):
with pytest.raises(ValueError):
HamerModel.from_reference(ref, backend="cpu", warmup_variants=[{"call": "frame", "input": "url"}])
def test_warmup_is_idempotent_and_does_not_change_outputs(ref, mano):
small = (320, 240) # the CPU backend is slow: small dummy frames
variants = [{"call": "run_crop"}, {"call": "frame", "input": "array", "is_right": False, "frame_size": small}]
with HamerModel.from_reference(ref, backend="cpu", mano=mano) as plain:
want_l = plain(SAMPLE, bbox=BOX_L, is_right=False).to_dict(rounded=False)
want_r = plain(SAMPLE, bbox=BOX_R).to_dict(rounded=False)
assert plain.warmed_variants == [] and plain.info.get("warmup_variants") == []
with HamerModel.from_reference(ref, backend="cpu", mano=mano, warmup_variants=variants) as m:
assert len(m.warmed_variants) == 2 and len(m.info["warmup_variants"]) == 2
assert all(r["ms"] > 0 for r in m.info["warmup_variants"])
again = m.warmup(call="frame", input="array", is_right=False, frame_size=list(small))
assert again["warmed"] is False and len(m.warmed_variants) == 2
new = m.warmup(call="frame", input="path", fmt="png", hands=2, bbox="center", frame_size=small)
assert new["warmed"] is True and new["variant"]["is_right"] == "mixed" and len(m.warmed_variants) == 3
assert m(SAMPLE, bbox=BOX_L, is_right=False).to_dict(rounded=False) == want_l
assert m(SAMPLE, bbox=BOX_R).to_dict(rounded=False) == want_r
with pytest.raises(RuntimeError):
m.warmup(call="run_crop")
def test_quiet_logs(monkeypatch):
from hamer import _verbosity
from hamer.api import QUIET_LOG_ENV, _quiet_logs
for k in QUIET_LOG_ENV:
monkeypatch.delenv(k, raising=False)
monkeypatch.delitem(sys.modules, "ttnn", raising=False)
monkeypatch.setenv("LOGURU_LEVEL", "DEBUG") # a caller value stays
try:
assert _quiet_logs(verbose=False) == {"TT_LOGGER_LEVEL": "error"}
assert os.environ["TT_LOGGER_LEVEL"] == "error" and os.environ["LOGURU_LEVEL"] == "DEBUG"
assert _verbosity.quiet()
monkeypatch.delenv("TT_LOGGER_LEVEL")
assert _quiet_logs(verbose=True) == {} and "TT_LOGGER_LEVEL" not in os.environ
assert not _verbosity.quiet()
finally:
_verbosity.set_quiet(False)