# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Python API: HaMeR hand mesh recovery on one Tenstorrent Blackhole chip.

    from hamer import HamerModel

    with HamerModel.from_pretrained(device_id=0) as model:
        hand = model("media/sample1.jpg", bbox=[20, 170, 295, 260], is_right=True)
        print(hand.rotmats.shape, hand.betas, hand.cam_t_full)

``from_pretrained`` loads the released checkpoint from the Hugging Face cache, opens the chip the way the published
numbers were measured (1x1 mesh, ETH dispatch, 12x10 compute grid, the ``tt-model.yaml`` serve configuration), uploads
the weights, captures one ttnn trace and checks the device output against the torch CPU reference.  Every call then
runs the same steps as ``POST /predict`` of the HTTP server: host crop + resize, one trace replay on the chip, host
post-processing (and the MANO mesh when ``MANO_RIGHT.pkl`` is available).  The outputs are the same numbers.

Nothing in this module imports ``ttnn`` at import time; the device stack loads inside ``from_pretrained``.
Full reference: ``PYTHON.md`` at the repository root.
"""
from __future__ import annotations

import functools
import gc
import os
import sys
import tempfile
import threading
import time
import warnings
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union

import numpy as np
import torch
from PIL import Image

from .preprocess import (
    IMAGE_HEIGHT, IMAGE_WIDTH, W_SLICE, center_square_bbox, hamer_bbox_size, normalize_raw, preprocess_pil,
)
from .rendering import FOCAL_LENGTH, MODEL_IMAGE_SIZE, cam_crop_to_full, project_mesh
from .served_config import (
    CKPT_FILE, MANO_FILE, MEAN_FILE, SERVED_DISPATCH, SERVED_ENV, SERVED_GRID, SERVED_WORKER_L1_SIZE,
    TRACE_REGION_BYTES, WEIGHTS_REPO, WEIGHTS_REVISION,
)

__all__ = ["HamerModel", "HandPrediction", "open_device", "load_reference", "apply_served_config",
           "DEFAULT_WARMUP_VARIANTS"]

ImageLike = Union[str, os.PathLike, Image.Image, np.ndarray, torch.Tensor]
BBox = Sequence[float]

MIN_CROP_PIXELS = 2.0  # same limit as the server: a crop side below this is not a hand
GRID_OVERRIDE_ENV = "TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE"


def _log(verbose: bool, msg: str) -> None:
    if verbose:
        print(f"[hamer] {msg}", flush=True)


# tt-metal / ttnn console logging, set before ``import ttnn`` when ``verbose=False`` (a value you set yourself stays):
# TT_LOGGER_LEVEL is the C++ logger (UMD / device open / JIT build lines), LOGURU_LEVEL the ttnn Python logger.
QUIET_LOG_ENV = {"TT_LOGGER_LEVEL": "error", "LOGURU_LEVEL": "WARNING"}


def _quiet_logs(verbose: bool) -> Dict[str, str]:
    """``verbose=False``: errors only from tt-metal / ttnn (when ``ttnn`` is not imported yet) and no ``[hamer]`` boot
    lines.  ``verbose=True``: the logging configuration is not changed.  Returns the variables that were set."""
    from ._verbosity import set_quiet  # noqa: WPS433

    set_quiet(not verbose)
    done: Dict[str, str] = {}
    if verbose or "ttnn" in sys.modules:
        return done
    for k, v in QUIET_LOG_ENV.items():
        if k not in os.environ:
            os.environ[k] = v
            done[k] = v
    return done


# --------------------------------------------------------------------------- #
# Warm-up variants
# --------------------------------------------------------------------------- #
# The device graph has ONE shape (a batch-1 uint8 256x192 crop); it is compiled, traced and replayed in from_pretrained.
# What differs between calls is host code: the input type and file format (decoder plugin), left-hand mirroring, the
# list form, the centre-square box, forward_crops / run_crop with uint8 or float crops, MANO.  Each variant below runs
# that code once on a dummy input of the real size, so that its one-time costs (lazy imports, decoder init, first
# native-helper call on a frame, first MANO kernels, first allocations) do not land on your first request.
WARMUP_FRAME_SIZE = (1920, 1080)  # (W, H) of the dummy frames: a 1080p camera frame (smaller frames are then warm too)
WARMUP_DEFAULTS: Dict[str, Dict[str, Any]] = {
    "frame": {"input": "path", "fmt": "jpeg", "hands": 1, "is_right": True, "bbox": "box",
              "frame_size": WARMUP_FRAME_SIZE},
    "crops": {"crop_dtype": "uint8", "crop_width": 256, "batch": 2},
    "run_crop": {"crop_dtype": "uint8"},
}
WARMUP_CHOICES: Dict[str, Tuple[Any, ...]] = {
    "input": ("path", "pil", "array", "tensor"), "fmt": ("jpeg", "png"), "bbox": ("box", "center"),
    "crop_dtype": ("uint8", "float"), "crop_width": (256, IMAGE_WIDTH), "is_right": (True, False, "mixed"),
}
#: ``from_pretrained(warmup_variants="default")``: every host code path of the API (13 device calls, about 0.25 s).
#: The quickstart call (a JPEG file, one right hand) runs last, so its code is the most recently used at the end.
DEFAULT_WARMUP_VARIANTS: Tuple[Dict[str, Any], ...] = (
    {"call": "run_crop", "crop_dtype": "float"},
    {"call": "run_crop", "crop_dtype": "uint8"},
    {"call": "crops", "crop_dtype": "float", "crop_width": 256},
    {"call": "crops", "crop_dtype": "uint8", "crop_width": 256},
    {"call": "frame", "input": "tensor"},
    {"call": "frame", "input": "array"},
    {"call": "frame", "input": "pil", "bbox": "center"},
    {"call": "frame", "input": "path", "fmt": "png", "hands": 2, "frame_size": (640, 480)},  # PNG encode is slow
    {"call": "frame", "input": "path", "fmt": "jpeg", "is_right": False},
    {"call": "frame", "input": "path", "fmt": "jpeg", "is_right": True},
)


def normalize_warmup_variant(call: str = "frame", **variant: Any) -> Tuple[Tuple[str, Any], ...]:
    """One warm-up variant as a sorted ``((key, value), ...)`` tuple with the defaults filled in (the idempotency key).
    ``call``: ``"frame"`` (``model(image, bbox)``), ``"crops"`` (:meth:`HamerModel.forward_crops`) or ``"run_crop"``.
    Raises ``ValueError`` for an unknown key or value."""
    if call not in WARMUP_DEFAULTS:
        raise ValueError(f"warm-up call must be one of {sorted(WARMUP_DEFAULTS)}, got {call!r}")
    allowed = WARMUP_DEFAULTS[call]
    bad = sorted(set(variant) - set(allowed))
    if bad:
        raise ValueError(f"unknown warm-up keys for call={call!r}: {bad} (allowed: {sorted(allowed)})")
    v = dict(allowed, **variant)
    if call == "frame":
        if "is_right" not in variant and int(v["hands"]) > 1:
            v["is_right"] = "mixed"
        v["hands"] = int(v["hands"])
        if v["hands"] < 1:
            raise ValueError(f"hands must be >= 1, got {v['hands']}")
        if v["hands"] == 1 and v["is_right"] == "mixed":
            raise ValueError("is_right='mixed' needs hands >= 2")
        fs = tuple(int(x) for x in v["frame_size"])
        if len(fs) != 2 or min(fs) < 16 or fs[0] * fs[1] > 50_000_000:
            raise ValueError(f"frame_size must be (W, H) with sides >= 16 and at most 50 MP, got {v['frame_size']}")
        v["frame_size"] = fs
        if v["input"] != "path":
            v.pop("fmt")
    if call == "crops":
        v["batch"] = int(v["batch"])
        if v["batch"] < 1:
            raise ValueError(f"batch must be >= 1, got {v['batch']}")
    for k, val in v.items():
        if k in WARMUP_CHOICES and val not in WARMUP_CHOICES[k]:
            raise ValueError(f"warm-up {k}={val!r} is not one of {WARMUP_CHOICES[k]}")
    return (("call", call),) + tuple(sorted(v.items()))


def _warmup_spec(spec: Any) -> List[Dict[str, Any]]:
    """``warmup_variants`` -> list of variant dicts.  ``"default"`` / ``"all"``: :data:`DEFAULT_WARMUP_VARIANTS`;
    ``None`` / ``"none"`` / ``False`` / ``[]``: none; a dict: one variant; a list of dicts: those."""
    if spec is None or spec is False or (isinstance(spec, str) and spec.lower() == "none"):
        return []
    if isinstance(spec, str):
        if spec.lower() in ("default", "all"):
            return [dict(v) for v in DEFAULT_WARMUP_VARIANTS]
        raise ValueError(f"warmup_variants must be 'default', 'none', a dict or a list of dicts, got {spec!r}")
    if isinstance(spec, dict):
        return [dict(spec)]
    out = []
    for v in spec:
        if not isinstance(v, dict):
            raise ValueError(f"warmup_variants list items must be dicts (model.warmup kwargs), got {v!r}")
        out.append(dict(v))
    return out


@functools.lru_cache(maxsize=4)
def _dummy_frame_cached(W: int, H: int) -> np.ndarray:
    a = _dummy_frame(W, H, seed=W * 7 + H)
    a.setflags(write=False)
    return a


def _dummy_frame(W: int, H: int, seed: int) -> np.ndarray:
    """A synthetic ``(H, W, 3)`` uint8 frame with photo-like statistics (smooth shading + texture + noise), so the
    decoders and the resize run at the cost of a real frame (a flat image decodes much faster)."""
    g = np.random.default_rng(seed)
    x = np.arange(W, dtype=np.float32)[None, :, None]
    y = np.arange(H, dtype=np.float32)[:, None, None]
    c = np.array([0.0, 1.1, 2.3], np.float32)
    a = 128 + 90 * np.sin(x / (W / 3.1) + c) * np.cos(y / (H / 2.3) - c) + 25 * np.sin(x / 3.7) * np.cos(y / 5.3)
    a += g.integers(-24, 25, (H, W, 1), dtype=np.int8)
    return np.clip(a, 0, 255).astype(np.uint8)


def _dummy_boxes(W: int, H: int, n: int) -> List[List[float]]:
    """n hand-sized boxes; the second one crosses the frame edge (exercises the grey padding)."""
    boxes = []
    for i in range(n):
        if i % 2 == 0:
            boxes.append([0.30 * W + i, 0.35 * H, 0.55 * W + i, 0.70 * H])
        else:
            boxes.append([-0.05 * W, 0.55 * H, 0.20 * W, 0.95 * H + i])
    return boxes


# --------------------------------------------------------------------------- #
# Output
# --------------------------------------------------------------------------- #
@dataclass
class HandPrediction:
    """One hand.  Field names follow the HTTP ``/predict`` response; :attr:`mano_params` gives the upstream
    HaMeR ``pred_mano_params`` layout.

    Attributes:
        rotmats:      ``(16, 3, 3)`` float32 rotation matrices; index 0 = global orientation, 1..15 = MANO hand pose.
                      Right-hand (model) frame: for a left hand they describe the mirrored crop, as upstream HaMeR.
        betas:        ``(10,)`` float32 MANO shape coefficients.
        cam:          ``(3,)`` float32 weak-perspective camera ``[s, tx, ty]`` in crop space (as the network emits it).
        cam_t_full:   ``(3,)`` float64 camera translation ``[tx, ty, tz]`` for the full frame (upstream
                      ``cam_crop_to_full``, focal length :attr:`focal_length`); x mirrored back for a left hand.
        focal_length: pixels, ``5000 / 256 * max(W, H)``.
        bbox:         ``(x1, y1, x2, y2)`` in original pixels (the request box, or the centre square).
        bbox_source:  ``"request"`` or ``"center_square"``.
        bbox_center:  ``(2,)`` crop centre in original pixels.
        bbox_size:    side of the square crop in original pixels (``rescale_factor * max(w, h)``).
        is_right, rescale_factor: the request values.
        image_size:   ``(W, H)`` of the frame.
        vertices:     ``(778, 3)`` float64 MANO mesh in metres (MANO frame, x mirrored back for a left hand), or
                      ``None`` without ``MANO_RIGHT.pkl`` or with ``return_vertices=False``.
        joints:       ``(21, 3)`` float64 joints in metres (same frame), or ``None`` without MANO.
        keypoints_2d: ``(21, 2)`` float64 joints projected into the original image, pixels, or ``None`` without MANO.
        timing_ms:    ``{"preprocess", "model", "mano", "total"}`` wall times of this call.
    """

    rotmats: np.ndarray
    betas: np.ndarray
    cam: np.ndarray
    cam_t_full: np.ndarray
    focal_length: float
    bbox: Tuple[float, float, float, float]
    bbox_source: str
    bbox_center: np.ndarray
    bbox_size: float
    is_right: bool
    rescale_factor: float
    image_size: Tuple[int, int]
    vertices: Optional[np.ndarray] = None
    joints: Optional[np.ndarray] = None
    keypoints_2d: Optional[np.ndarray] = None
    timing_ms: Dict[str, float] = field(default_factory=dict)

    @property
    def global_orient(self) -> np.ndarray:
        """``(1, 3, 3)`` global orientation (= ``rotmats[:1]``)."""
        return self.rotmats[:1]

    @property
    def hand_pose(self) -> np.ndarray:
        """``(15, 3, 3)`` MANO hand pose (= ``rotmats[1:]``)."""
        return self.rotmats[1:]

    @property
    def mano_params(self) -> Dict[str, np.ndarray]:
        """Upstream layout: ``{"global_orient": (1,3,3), "hand_pose": (15,3,3), "betas": (10,)}``."""
        return {"global_orient": self.global_orient, "hand_pose": self.hand_pose, "betas": self.betas}

    @property
    def mano_available(self) -> bool:
        return self.joints is not None

    def to_dict(self, rounded: bool = True) -> Dict[str, Any]:
        """JSON-ready dict with the keys and rounding of the HTTP ``/predict`` response (``timing_ms`` and
        ``trace_active`` aside).  ``rounded=False`` keeps full precision."""
        def r(a: Any, nd: int = 6) -> Any:
            a = np.asarray(a, dtype=np.float64)
            return (np.round(a, nd) if rounded else a).tolist()

        out: Dict[str, Any] = {
            "image_size": list(self.image_size),
            "bbox": r(self.bbox, 3),
            "bbox_source": self.bbox_source,
            "is_right": self.is_right,
            "rescale_factor": self.rescale_factor,
            "bbox_center": r(self.bbox_center, 3),
            "bbox_size": round(float(self.bbox_size), 3) if rounded else float(self.bbox_size),
            "focal_length": round(float(self.focal_length), 3) if rounded else float(self.focal_length),
            "rotmats": r(self.rotmats),
            "betas": r(self.betas),
            "cam": r(self.cam),
            "cam_t_full": r(self.cam_t_full),
            "mano_available": self.mano_available,
        }
        if self.joints is not None:
            out["joints"] = r(self.joints)
            out["keypoints_2d"] = r(self.keypoints_2d, 3)
        if self.vertices is not None:
            out["vertices"] = r(self.vertices)
        return out


# --------------------------------------------------------------------------- #
# Configuration / weights / device
# --------------------------------------------------------------------------- #
def apply_served_config(override: bool = False) -> Dict[str, Tuple[str, str]]:
    """Put the published ``tt-model.yaml`` serve configuration (:data:`hamer.served_config.SERVED_ENV`) into
    ``os.environ``.  The knobs are read when the device graph is built, so call this before building.  With
    ``override=False`` (default) a value already in the environment is kept; returns ``{name: (kept, served)}``
    for every kept value that differs from the published one."""
    differ = {}
    for k, v in SERVED_ENV.items():
        cur = os.environ.get(k)
        if cur is None or override:
            os.environ[k] = v
        elif cur != v:
            differ[k] = (cur, v)
    return differ


def resolve_weights(weights_dir: Optional[Union[str, os.PathLike]] = None, repo_id: str = WEIGHTS_REPO,
                    revision: Optional[str] = WEIGHTS_REVISION) -> Tuple[Path, Path]:
    """``(hamer.ckpt, mano_mean_params.npz)``: from ``weights_dir`` if given, else from the Hugging Face cache
    (downloaded on first use, 2.7 GB).  ``HAMER_WEIGHTS_DIR`` is honoured when ``weights_dir`` is ``None``."""
    weights_dir = weights_dir or os.environ.get("HAMER_WEIGHTS_DIR") or None
    if weights_dir:
        d = Path(weights_dir).expanduser()
        ck, mean = d / CKPT_FILE, d / MEAN_FILE
        for f in (ck, mean):
            if not f.is_file():
                raise FileNotFoundError(f"{d} does not contain {f.name}")
        return ck, mean
    from huggingface_hub import hf_hub_download  # noqa: WPS433

    ck = Path(hf_hub_download(repo_id=repo_id, filename=CKPT_FILE, revision=revision))
    mean = Path(hf_hub_download(repo_id=repo_id, filename=MEAN_FILE, revision=revision))
    return ck, mean


def resolve_mano_path(mano_path: Optional[Union[str, os.PathLike]] = None) -> Optional[Path]:
    """``MANO_RIGHT.pkl`` (licence-bound, not redistributed): ``mano_path`` (file or directory), else
    ``HAMER_MANO_PATH``, else ``$TT_DIT_CACHE_DIR/MANO_RIGHT.pkl`` or ``~/.cache/tt-model/hamer-p150/weights/``.
    An explicit path that does not exist raises; the default locations return ``None`` when absent."""
    explicit = mano_path or os.environ.get("HAMER_MANO_PATH")
    if explicit:
        p = Path(explicit).expanduser()
        p = p / MANO_FILE if p.is_dir() else p
        if not p.is_file():
            raise FileNotFoundError(f"MANO model not found at {p}")
        return p
    cache = os.environ.get("TT_DIT_CACHE_DIR")
    base = Path(cache) if cache else Path.home() / ".cache" / "tt-model" / "hamer-p150" / "weights"
    p = base / MANO_FILE
    return p if p.is_file() else None


def load_reference(weights_dir: Optional[Union[str, os.PathLike]] = None, repo_id: str = WEIGHTS_REPO,
                   revision: Optional[str] = WEIGHTS_REVISION) -> torch.nn.Module:
    """The torch CPU reference (``hamer.reference.Hamer``) with the released checkpoint loaded."""
    from .checkpoint import load_hamer_checkpoint  # noqa: WPS433
    from .reference import Hamer  # noqa: WPS433

    ck, mean = resolve_weights(weights_dir, repo_id, revision)
    ref = Hamer().eval()
    report = load_hamer_checkpoint(ref, ck, mean)
    if report["loaded"] == 0 or report["missing_in_ref"] != 0 or report["unexpected_in_ref"] != 0:
        raise RuntimeError(f"checkpoint {ck} does not map cleanly onto the reference: {report}")
    if not bool(ref.head.init_cam.abs().sum() > 0):
        raise RuntimeError(f"{mean} did not populate the MANO mean parameters")
    return ref


def open_device(device_id: int = 0, dispatch: str = "auto", verbose: bool = False) -> Tuple[Any, str]:
    """Open one Blackhole chip the way the published numbers were measured: a 1x1 ``MeshDevice`` with a 200 MiB
    trace region, ETH dispatch (the full 12x10 Tensix grid computes) and the larger kernel-config ring.

    ``dispatch``: ``"auto"`` (ETH, falling back to the default worker dispatch with a warning when the tt-metal build
    cannot place dispatch on Ethernet cores), ``"eth"`` or ``"worker"``.  When ``TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE``
    is not set and ``ttnn`` is not imported yet, it is set to ``11,9`` (grid capped at 12x10, as on the measured
    hosts; tt-metal reads it at ``import ttnn``).

    Returns ``(device, dispatch_used)``.  Close it with ``ttnn.close_mesh_device`` (or pass it to
    :meth:`HamerModel.from_pretrained` as ``device=`` and close it yourself afterwards)."""
    if dispatch not in ("auto", "eth", "worker"):
        raise ValueError(f"dispatch must be 'auto', 'eth' or 'worker', got {dispatch!r}")
    order = ["eth", "worker"] if dispatch == "auto" else [dispatch]
    # tt-metal reads the grid override when ttnn is imported: set it first (ETH dispatch on a Galaxy Blackhole chip
    # would otherwise compute on 13x10; a p150 with ETH dispatch has 12x10 anyway)
    if "eth" in order and GRID_OVERRIDE_ENV not in os.environ and "ttnn" not in sys.modules:
        os.environ[GRID_OVERRIDE_ENV] = f"{SERVED_GRID[0] - 1},{SERVED_GRID[1] - 1}"
    import ttnn  # noqa: WPS433

    extra_kb = int(os.environ.get("HAMER_KCFG_EXTRA_KB", "128"))
    last: Optional[BaseException] = None
    for i, d in enumerate(order):
        kw: Dict[str, Any] = {}
        if d == "eth":
            kw["dispatch_core_config"] = ttnn.DispatchCoreConfig(ttnn.DispatchCoreType.ETH)
        if extra_kb > 0:
            kw["worker_l1_size"] = 1461376 - extra_kb * 1024
        try:
            dev = ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(1, 1), physical_device_ids=[device_id],
                                        trace_region_size=TRACE_REGION_BYTES, num_command_queues=1, **kw)
        except Exception as e:  # noqa: BLE001
            last = e
            if i + 1 < len(order):
                warnings.warn(
                    f"hamer: ETH dispatch is not available on this tt-metal build / chip ({e}); falling back to the "
                    "default worker dispatch.  The published numbers were measured with ETH dispatch on the 12x10 "
                    "grid; worker dispatch leaves a smaller grid on a p150 and runs slower.", RuntimeWarning, stacklevel=2)
            continue
        g = dev.compute_with_storage_grid_size()
        if (g.x, g.y) != SERVED_GRID:
            warnings.warn(
                f"hamer: compute grid is {g.x}x{g.y}; the device graph is tuned and validated on "
                f"{SERVED_GRID[0]}x{SERVED_GRID[1]} (to cap a larger grid at 12x10, set {GRID_OVERRIDE_ENV}=11,9 before ttnn is "
                "imported).",
                RuntimeWarning, stacklevel=2)
        _log(verbose, f"device {device_id} open: dispatch={d} grid={g.x}x{g.y}")
        return dev, d
    raise RuntimeError(f"hamer: could not open device {device_id}: {last}") from last


def _close_device(dev: Any) -> None:
    import ttnn  # noqa: WPS433

    if isinstance(dev, getattr(ttnn, "MeshDevice", ())):
        ttnn.close_mesh_device(dev)
    else:
        ttnn.close_device(dev)


def _pcc(a: torch.Tensor, b: torch.Tensor) -> float:
    a = a.reshape(-1).double() - a.double().mean()
    b = b.reshape(-1).double() - b.double().mean()
    return float((a @ b) / (a.norm() * b.norm()).clamp_min(1e-12))


# --------------------------------------------------------------------------- #
# Inputs
# --------------------------------------------------------------------------- #
def to_pil(image: ImageLike) -> Image.Image:
    """A frame as an RGB ``PIL.Image``.  Accepts a file path, a ``PIL.Image`` (any mode), an ``(H, W, 3)`` uint8
    RGB array, or a torch tensor ``(3, H, W)`` / ``(H, W, 3)`` (uint8, or float in ``[0, 1]``)."""
    if isinstance(image, Image.Image):
        return image if image.mode == "RGB" else image.convert("RGB")
    if isinstance(image, (str, os.PathLike)):
        im = Image.open(image)
        im.load()
        return im if im.mode == "RGB" else im.convert("RGB")
    if isinstance(image, torch.Tensor):
        t = image.detach().cpu()
        if t.dim() == 4 and t.shape[0] == 1:
            t = t[0]
        if t.dim() == 3 and t.shape[0] in (1, 3) and t.shape[-1] not in (1, 3):
            t = t.permute(1, 2, 0)
        image = t.numpy()
    if isinstance(image, np.ndarray):
        a = image
        if a.ndim == 2:
            a = np.stack([a] * 3, -1)
        if a.ndim != 3 or a.shape[-1] not in (1, 3, 4):
            raise ValueError(f"expected an (H, W, 3) RGB image array, got shape {a.shape}")
        if a.shape[-1] == 1:
            a = np.repeat(a, 3, -1)
        a = a[..., :3]
        if a.dtype != np.uint8:
            if np.issubdtype(a.dtype, np.floating):
                if a.min() < 0.0 or a.max() > 1.0:
                    raise ValueError("float image arrays must be RGB in [0, 1]; pass uint8 for 0..255 pixels")
                a = np.round(a * 255.0)
            a = a.astype(np.uint8)
        return Image.fromarray(np.ascontiguousarray(a), "RGB")
    raise TypeError(f"unsupported image type {type(image).__name__}")


def _check_bbox(bbox: Optional[BBox], img: Image.Image, rescale_factor: float) -> Tuple[Tuple[float, ...], str]:
    """Same rules as the server's ``_validate_bbox`` (ValueError instead of HTTP 400)."""
    if bbox is None:
        return center_square_bbox(img), "center_square"
    try:
        x1, y1, x2, y2 = (float(v) for v in bbox)
    except (TypeError, ValueError):
        raise ValueError(f"bbox must be four numbers [x1, y1, x2, y2], got {bbox!r}") from None
    if not all(np.isfinite([x1, y1, x2, y2])) or x2 <= x1 or y2 <= y1:
        raise ValueError(f"bbox {list(bbox)} is not a valid [x1, y1, x2, y2] with x2 > x1, y2 > y1")
    if x2 < 0 or y2 < 0 or x1 > img.width or y1 > img.height:
        raise ValueError(f"bbox {list(bbox)} lies entirely outside the {img.width}x{img.height} image")
    side = hamer_bbox_size((x1, y1, x2, y2), rescale_factor)
    if side < MIN_CROP_PIXELS:
        raise ValueError(f"bbox {list(bbox)} is degenerate: crop side {side:.3f} px is below {MIN_CROP_PIXELS:g} px")
    return (x1, y1, x2, y2), "request"


def _is_box_list(bbox: Any) -> bool:
    """True for several boxes (``[[x1, y1, x2, y2], ...]`` or an ``(N, 4)`` array)."""
    if bbox is None:
        return False
    a = np.asarray(bbox, dtype=object)
    return a.ndim == 2 or (a.ndim == 1 and len(a) > 0 and isinstance(a[0], (list, tuple, np.ndarray)))


# --------------------------------------------------------------------------- #
# Model
# --------------------------------------------------------------------------- #
class HamerModel:
    """HaMeR (ViT-H/16 + MANO head) on one Blackhole chip.  Build with :meth:`from_pretrained`.

    Call it with a frame and a hand box::

        hand  = model(image, bbox=[x1, y1, x2, y2], is_right=True)            # -> HandPrediction
        hands = model(image, bbox=[box_a, box_b], is_right=[True, False])     # -> list[HandPrediction]

    ``image``: file path, ``PIL.Image``, ``(H, W, 3)`` uint8 RGB numpy array, or torch tensor (see :func:`to_pil`).
    One hand per device call (the device graph is batch 1); several boxes on one frame decode the frame once.
    Thread-safe (one device call at a time).  Release the chip with :meth:`close` or a ``with`` block.
    """

    def __init__(self, runner: Callable[[torch.Tensor], torch.Tensor], *, ref: torch.nn.Module,
                 raw_input: bool, mano: Optional[torch.nn.Module] = None, tt: Any = None,
                 owned_device: Any = None, info: Optional[Dict[str, Any]] = None) -> None:
        """Low-level constructor; use :meth:`from_pretrained` or :meth:`from_reference`."""
        self._runner = runner
        self.ref = ref
        self.raw_input = raw_input   # the device graph takes the uint8 crop (TT_FUSED_INPUT=u8)
        self.mano = mano
        self.tt = tt                 # hamer.model.TtHamer, or None for the CPU backend
        self._owned_device = owned_device
        self.info: Dict[str, Any] = dict(info or {})
        self._lock = threading.Lock()
        self._closed = False
        self._warm: Dict[Tuple[Tuple[str, Any], ...], Dict[str, Any]] = {}

    # ---------------- construction ----------------
    @classmethod
    def from_pretrained(
        cls,
        device_id: int = 0,
        *,
        device: Any = None,
        weights_dir: Optional[Union[str, os.PathLike]] = None,
        repo_id: str = WEIGHTS_REPO,
        revision: Optional[str] = WEIGHTS_REVISION,
        mano_path: Optional[Union[str, os.PathLike]] = None,
        dispatch: str = "auto",
        use_trace: bool = True,
        check_accuracy: bool = True,
        min_pcc: float = 0.99,
        served_config: bool = True,
        warmup_variants: Any = "default",
        verbose: bool = False,
    ) -> "HamerModel":
        """Load the released checkpoint and build the model on a chip, then warm up every call variant in
        ``warmup_variants`` so that your first real call runs at the steady speed.

        Args:
            device_id: chip to open (``TT_VISIBLE_DEVICES`` / ``TT_METAL_VISIBLE_DEVICES`` order).
            device: an already-opened ttnn device to use instead (not closed by :meth:`close`).  Open it with
                :func:`open_device` to get the published configuration; any 1x1 ``MeshDevice`` with a trace region of
                at least 200 MiB works.
            weights_dir: directory with ``hamer.ckpt`` + ``mano_mean_params.npz``; default: the Hugging Face cache
                (``repo_id`` at ``revision``, downloaded on first use).
            mano_path: ``MANO_RIGHT.pkl`` (file or directory) for vertices / joints / 2-D keypoints; default: see
                :func:`resolve_mano_path`.  Without it the outputs are the MANO parameters and cameras.
            dispatch: ``"auto"`` (ETH, else worker with a warning), ``"eth"`` or ``"worker"``.
            use_trace: capture and replay one ttnn trace (the published path).  ``False`` runs eagerly (slow).
            check_accuracy: run the server's warm-up gate: eager, trace replay and a fresh input vs the torch CPU
                reference, PCC >= ``min_pcc`` (about 1 s of CPU work).
            served_config: apply the ``tt-model.yaml`` serve configuration (:func:`apply_served_config`) first.
                Values already in the environment win and are reported in a warning.
            warmup_variants: host call variants to prepare after the device warm-up: ``"default"`` (all of
                :data:`DEFAULT_WARMUP_VARIANTS`, about 0.25 s), ``None`` / ``"none"``, or a list of
                :meth:`warmup` kwargs dicts.  The device graph (one shape) is always compiled and traced.
            verbose: print the boot steps and keep the tt-metal / ttnn log output.  ``False`` (default): errors only
                (``TT_LOGGER_LEVEL=error``, ``LOGURU_LEVEL=WARNING`` when not set and ``ttnn`` is not imported yet).
        """
        t_boot = time.perf_counter()
        _quiet_logs(verbose)
        if served_config:
            differ = apply_served_config()
            if differ:
                warnings.warn("hamer: environment overrides the published configuration: "
                              + ", ".join(f"{k}={a!r} (published {b!r})" for k, (a, b) in differ.items()),
                              RuntimeWarning, stacklevel=2)
        torch.set_grad_enabled(False)
        t = time.perf_counter()
        ref = load_reference(weights_dir, repo_id, revision)
        _log(verbose, f"checkpoint loaded in {time.perf_counter() - t:.1f}s")
        mano = _load_mano(mano_path, verbose)
        info = {"weights": {"repo": None if weights_dir else repo_id, "revision": None if weights_dir else revision,
                            "dir": str(weights_dir) if weights_dir else None}}
        return cls.from_reference(ref, device_id=device_id, device=device, mano=mano, dispatch=dispatch,
                                  use_trace=use_trace, check_accuracy=check_accuracy, min_pcc=min_pcc,
                                  warmup_variants=warmup_variants, verbose=verbose, _info=info, _t_boot=t_boot)

    @classmethod
    def from_reference(
        cls,
        ref: torch.nn.Module,
        *,
        backend: str = "tt",
        device_id: int = 0,
        device: Any = None,
        mano: Optional[torch.nn.Module] = None,
        dispatch: str = "auto",
        use_trace: bool = True,
        check_accuracy: bool = True,
        min_pcc: float = 0.99,
        warmup_variants: Any = "auto",
        verbose: bool = False,
        _info: Optional[Dict[str, Any]] = None,
        _t_boot: Optional[float] = None,
    ) -> "HamerModel":
        """Build from a torch reference module you already hold (``hamer.reference.Hamer`` with weights loaded).
        ``backend="cpu"`` runs that module on the host instead of the chip (no ttnn needed; slow; same API).
        ``warmup_variants``: as in :meth:`from_pretrained`; ``"auto"`` = ``"default"`` on the chip, none on the CPU."""
        t_boot = time.perf_counter() if _t_boot is None else _t_boot
        info = dict(_info or {})
        ref.eval()
        if isinstance(warmup_variants, str) and warmup_variants == "auto":
            warmup_variants = "default" if backend == "tt" else None
        variants = _warmup_spec(warmup_variants)
        for v in variants:
            normalize_warmup_variant(**v)  # fail before the device is opened
        if backend == "cpu":
            model = cls(lambda x: ref(x), ref=ref, raw_input=False, mano=mano, info=info)
            model._warmup_host(variants, verbose)
            model.info["boot_seconds"] = round(time.perf_counter() - t_boot, 1)
            return model
        if backend != "tt":
            raise ValueError(f"backend must be 'tt' or 'cpu', got {backend!r}")

        owned = None
        if device is None:
            device, used = open_device(device_id, dispatch, verbose)
            owned = device
        else:
            used = "caller"
        from .model import build_from_reference  # noqa: WPS433  (after open_device: it may import ttnn)
        tt = None
        try:
            t = time.perf_counter()
            tt = build_from_reference(ref, device_id=device_id, strict=True, use_trace=use_trace, device=device)
            if tt._vit_params is None or tt._head_params is None:  # noqa: SLF001
                raise RuntimeError("tt-nn parameters were not uploaded")
            _log(verbose, f"weights uploaded in {time.perf_counter() - t:.1f}s")
            raw = getattr(tt.fused_opts, "input", "norm") == "u8"
            g = device.compute_with_storage_grid_size()
            info.update(backend="tt", device_id=device_id if used != "caller" else None, dispatch=used,
                        num_command_queues=1 if used != "caller" else None,  # the model itself uses only CQ0
                        grid=f"{g.x}x{g.y}", fused=tt.fused_opts.describe())
            model = cls(tt, ref=ref, raw_input=raw, mano=mano, tt=tt, owned_device=owned, info=info)
            model.info["warmup"] = model._warmup(check_accuracy, min_pcc, verbose)
        except BaseException:
            if tt is not None:
                tt.close()
            if owned is not None:
                _close_device(owned)
            raise
        if use_trace and not tt.trace_active:
            model.close()
            raise RuntimeError("ttnn trace was not captured during warm-up")
        try:
            from . import native  # noqa: WPS433

            native.load()
            model.info["native_preprocess"] = native.STATUS
        except Exception:  # noqa: BLE001 - optional (PIL path is bit-identical)
            pass
        try:
            model._warmup_host(variants, verbose)
        except BaseException:
            model.close()
            raise
        model.info["boot_seconds"] = round(time.perf_counter() - t_boot, 1)
        _log(verbose, f"ready in {model.info['boot_seconds']}s ({model.info['grid']}, dispatch={used})")
        return model

    def _warmup(self, check: bool, min_pcc: float, verbose: bool) -> Dict[str, Any]:
        """The server's warm-up (``hamer.server.app.lifespan``): forward 1 eager (JIT), forward 2 trace capture +
        replay on the same input, forward 3 a different input through the copy-into-trace-buffer path."""
        from .reference import sample_input  # noqa: WPS433

        if self.raw_input:
            g = torch.Generator().manual_seed(1)
            u_a = torch.randint(0, 256, (1, 3, IMAGE_HEIGHT, IMAGE_WIDTH), generator=g, dtype=torch.uint8)
            u_b = torch.randint(0, 256, (1, 3, IMAGE_HEIGHT, IMAGE_WIDTH), generator=g, dtype=torch.uint8)
            x_a, x_b, d_a, d_b = normalize_raw(u_a), normalize_raw(u_b), u_a, u_b
        else:
            x_a, x_b = sample_input(batch=1, seed=0), sample_input(batch=1, seed=1)
            d_a, d_b = x_a, x_b
        outs = []
        with torch.inference_mode():
            for i, d in enumerate((d_a, d_a, d_b)):
                t = time.perf_counter()
                out = self._runner(d)
                if tuple(out.shape) != (1, 157) or not torch.isfinite(out).all():
                    raise RuntimeError(f"warm-up forward {i + 1} produced an invalid output {tuple(out.shape)}")
                outs.append(out.detach().clone())
                _log(verbose, f"warm-up forward {i + 1}/3: {(time.perf_counter() - t) * 1e3:.0f} ms")
            if not check:
                return {"checked": False}
            ref_a, ref_b = self.ref(x_a), self.ref(x_b)
        res = {"eager_vs_ref_pcc": round(_pcc(outs[0], ref_a), 5),
               "replay_vs_eager_pcc": round(_pcc(outs[1], outs[0]), 5),
               "fresh_input_vs_ref_pcc": round(_pcc(outs[2], ref_b), 5), "min_pcc": min_pcc, "checked": True}
        _log(verbose, f"warm-up PCC {res}")
        for k in ("eager_vs_ref_pcc", "replay_vs_eager_pcc", "fresh_input_vs_ref_pcc"):
            if res[k] < min_pcc:
                raise RuntimeError(f"warm-up gate failed: {k} = {res[k]} < {min_pcc}")
        return res

    def _warmup_host(self, variants: List[Dict[str, Any]], verbose: bool) -> None:
        """Run the requested host variants (after the device warm-up and the native-helper load)."""
        self.info.setdefault("warmup_variants", [])
        if not variants:
            return
        t = time.perf_counter()
        # the boot garbage (checkpoint dicts, reference outputs) is collected now, not in your first call; before the
        # variants, so that the full collection does not leave the CPU caches cold for the first call
        gc.collect()
        for v in variants:
            rec = self.warmup(**v)
            _log(verbose, f"warm-up {rec['variant']}: {rec['ms']:.1f} ms")
        self.info["warmup_host_seconds"] = round(time.perf_counter() - t, 3)

    def warmup(self, call: str = "frame", **variant: Any) -> Dict[str, Any]:
        """Prepare one call variant: run it once on a dummy input of the real size, so that its one-time host costs
        (lazy imports, image-decoder init, first native preprocess call on a frame, first MANO kernels, first
        allocations) are paid now and not in your first real call.  Idempotent: a variant that is already warm
        returns its record without running again.  The outputs of the dummy call are checked and dropped (nothing is
        cached; real calls are not changed).

        Args (the defaults are the quickstart call):
            call: ``"frame"`` (``model(image, bbox)``), ``"crops"`` (:meth:`forward_crops`), ``"run_crop"``.
            input (frame): ``"path"`` (an image file), ``"pil"``, ``"array"`` (uint8 ``(H, W, 3)``), ``"tensor"``
                (uint8 ``(3, H, W)``).
            fmt (frame, path): ``"jpeg"`` or ``"png"`` (the decoder that the first file of that format loads).
            hands (frame): boxes per call; 1 = single-box form, >= 2 = list form.
            is_right (frame): ``True``, ``False`` (mirrored crop) or ``"mixed"`` (with ``hands >= 2``).
            bbox (frame): ``"box"`` or ``"center"`` (``bbox=None``).
            frame_size (frame): ``(W, H)`` of the dummy frame, default ``(640, 480)``.
            crop_dtype (crops, run_crop): ``"uint8"`` or ``"float"`` (ImageNet-normalized).
            crop_width (crops): 256 (upstream crop) or 192.
            batch (crops): crops per call (one device call each), default 2.

        Returns ``{"variant": {...}, "ms": wall time, "warmed": True/False (False: it was already warm)}``.
        """
        self._check_open()
        key = normalize_warmup_variant(call, **variant)
        done = self._warm.get(key)
        if done is not None:
            return dict(done, warmed=False)
        v = dict(key)
        t = time.perf_counter()
        self._run_warmup_variant(v)
        rec = {"variant": {k: (list(x) if isinstance(x, tuple) else x) for k, x in v.items()},
               "ms": round((time.perf_counter() - t) * 1e3, 2)}
        self._warm[key] = rec
        self.info.setdefault("warmup_variants", []).append(rec)
        return dict(rec, warmed=True)

    def _run_warmup_variant(self, v: Dict[str, Any]) -> None:
        call = v["call"]
        if call == "frame":
            W, H = v["frame_size"]
            a = _dummy_frame_cached(W, H).copy()
            n = v["hands"]
            boxes = _dummy_boxes(W, H, n) if v["bbox"] == "box" else None
            if n == 1:
                bbox: Any = None if boxes is None else boxes[0]
                right: Any = bool(v["is_right"])
            else:
                bbox = boxes if boxes is not None else [None] * n
                right = [bool(i % 2 == 0) for i in range(n)] if v["is_right"] == "mixed" else [bool(v["is_right"])] * n
            inp = v["input"]
            with tempfile.TemporaryDirectory(prefix="hamer-warmup-") as d:
                if inp == "path":
                    image: Any = Path(d) / ("frame.jpg" if v["fmt"] == "jpeg" else "frame.png")
                    Image.fromarray(a).save(image, **({"quality": 90} if v["fmt"] == "jpeg" else {"compress_level": 1}))
                elif inp == "pil":
                    image = Image.fromarray(a)
                elif inp == "array":
                    image = a
                else:
                    image = torch.from_numpy(a).permute(2, 0, 1)
                if n > 1 and boxes is None:  # the centre square, n times (the list form needs boxes)
                    img = to_pil(image)
                    bbox = [list(center_square_bbox(img))] * n
                out = self(image, bbox, is_right=right)
            hands = out if isinstance(out, list) else [out]
            if len(hands) != n:
                raise RuntimeError(f"warm-up {v}: {len(hands)} outputs for {n} boxes")
            for h in hands:
                if not (np.isfinite(h.rotmats).all() and np.isfinite(h.cam_t_full).all()):
                    raise RuntimeError(f"warm-up {v}: non-finite output")
                h.to_dict()
            # host allocator growth: the crop buffers scale with the crop side (rescale_factor * box side).  One
            # host-only preprocess with a crop side of max(W, H) grows the heap so that later crops up to that side
            # do not pay first-time page faults.
            side = max(W, H) / 2.0
            preprocess_pil(Image.fromarray(a), (W / 2 - side / 2, H / 2 - side / 2, W / 2 + side / 2, H / 2 + side / 2),
                           rescale_factor=2.0, is_right=True, raw=self.raw_input, contiguous=not self.raw_input)
            return
        width = v.get("crop_width", IMAGE_WIDTH)
        B = v.get("batch", 1)
        u = torch.from_numpy(_dummy_frame(width, IMAGE_HEIGHT * B, seed=width + B).copy())
        u = u.reshape(B, IMAGE_HEIGHT, width, 3).permute(0, 3, 1, 2).contiguous()
        x = u if v["crop_dtype"] == "uint8" else normalize_raw(u)
        if call == "crops":
            res = self.forward_crops(x)
            if not torch.isfinite(res["pred_cam"]).all():
                raise RuntimeError(f"warm-up {v}: non-finite output")
            return
        out = self.run_crop(x)
        if tuple(out.shape) != (1, 157) or not torch.isfinite(out).all():
            raise RuntimeError(f"warm-up {v}: invalid output {tuple(out.shape)}")

    @property
    def warmed_variants(self) -> List[Dict[str, Any]]:
        """The variants prepared so far (by :meth:`from_pretrained` and :meth:`warmup`)."""
        return [dict(k) for k in self._warm]

    # ---------------- properties ----------------
    @property
    def device(self) -> Any:
        """The ttnn device (``None`` for the CPU backend or after :meth:`close`)."""
        return None if self.tt is None else self.tt.device

    @property
    def trace_active(self) -> bool:
        return bool(getattr(self.tt, "trace_active", False))

    @property
    def mano_available(self) -> bool:
        return self.mano is not None

    @property
    def faces(self) -> Optional[np.ndarray]:
        """``(1538, 3)`` int64 MANO triangle indices (``None`` without ``MANO_RIGHT.pkl``)."""
        return None if self.mano is None else np.asarray(self.mano.faces, dtype=np.int64)

    # ---------------- inference ----------------
    def __call__(
        self,
        image: ImageLike,
        bbox: Optional[Union[BBox, Sequence[BBox], np.ndarray]] = None,
        *,
        is_right: Union[bool, Sequence[bool]] = True,
        rescale_factor: float = 2.0,
        return_vertices: bool = True,
    ) -> Union[HandPrediction, List[HandPrediction]]:
        """Hand mesh recovery for one hand (``bbox`` = ``[x1, y1, x2, y2]``) or several hands of the same frame
        (``bbox`` = a list / ``(N, 4)`` array of boxes, ``is_right`` = one bool or one per box).

        Args:
            image: the FULL frame (path, PIL, uint8 ``(H, W, 3)`` RGB array, or torch tensor; see :func:`to_pil`).
            bbox: hand box(es) in original pixels.  ``None``: the centre square of the frame (the camera is then
                geometrically meaningless; for smoke tests).
            is_right: ``False`` for a left hand: the crop is mirrored, the outputs are mirrored back.
            rescale_factor: square crop side = ``rescale_factor * max(box w, h)``, in ``(0.5, 5.0]``; default 2.0.
            return_vertices: include the 778 mesh vertices when MANO is available.

        Returns:
            :class:`HandPrediction` for one box, ``list[HandPrediction]`` for a list of boxes.
        """
        t0 = time.perf_counter()
        img = to_pil(image)
        if not (0.5 < float(rescale_factor) <= 5.0):
            raise ValueError(f"rescale_factor must be in (0.5, 5.0], got {rescale_factor}")
        if _is_box_list(bbox):
            boxes = [list(b) for b in np.asarray(bbox, dtype=np.float64).reshape(-1, 4)]
            rights = list(is_right) if isinstance(is_right, (list, tuple, np.ndarray)) else [is_right] * len(boxes)
            if len(rights) != len(boxes):
                raise ValueError(f"{len(boxes)} boxes but {len(rights)} is_right values")
            return [self._predict_one(img, b, bool(r), float(rescale_factor), return_vertices, t0 if i == 0 else None)
                    for i, (b, r) in enumerate(zip(boxes, rights))]
        if isinstance(is_right, (list, tuple, np.ndarray)):
            raise ValueError("is_right must be a single bool for a single box")
        return self._predict_one(img, bbox, bool(is_right), float(rescale_factor), return_vertices, t0)

    predict = __call__

    def _predict_one(self, img: Image.Image, bbox: Optional[BBox], is_right: bool, rescale_factor: float,
                     return_vertices: bool, t0: Optional[float]) -> HandPrediction:
        """The server's ``/predict`` body after the image decode (same calls, same order, no rounding)."""
        t0 = time.perf_counter() if t0 is None else t0
        box, source = _check_bbox(bbox, img, rescale_factor)
        W, H = img.size
        raw_in = self.raw_input
        x, (cx, cy, bbox_size) = preprocess_pil(img, box, rescale_factor=rescale_factor, is_right=is_right,
                                                raw=raw_in, contiguous=not raw_in)
        if tuple(x.shape) != (1, 3, IMAGE_HEIGHT, IMAGE_WIDTH):
            raise RuntimeError(f"preprocess produced {tuple(x.shape)}")
        t1 = time.perf_counter()
        out = self.run_crop(x if raw_in else x.contiguous())
        t2 = time.perf_counter()

        out = out[0].float().cpu()
        rotmats = out[:144].reshape(16, 3, 3)
        betas = out[144:154]
        cam = out[154:157]
        focal = FOCAL_LENGTH / MODEL_IMAGE_SIZE * max(W, H)
        cam_full = cam.numpy().astype(np.float64).copy()
        if not is_right:
            cam_full[1] *= -1.0  # mirror the camera x back for a left hand (upstream demo.py)
        cam_t = cam_crop_to_full(cam_full, np.array([cx, cy]), bbox_size, W, H, focal)
        pred = HandPrediction(
            rotmats=rotmats.numpy().copy(), betas=betas.numpy().copy(), cam=cam.numpy().copy(), cam_t_full=cam_t,
            focal_length=float(focal), bbox=tuple(float(v) for v in box), bbox_source=source,
            bbox_center=np.array([cx, cy], dtype=np.float64), bbox_size=float(bbox_size), is_right=is_right,
            rescale_factor=rescale_factor, image_size=(W, H),
        )
        t3 = t2
        if self.mano is not None:
            with torch.inference_mode():
                verts, joints = (v[0] for v in self.mano.forward_with_joints(rotmats[None], betas[None]))
            kp2d, _ = project_mesh(joints, cam, (cx, cy), bbox_size, W, H, is_right=is_right)
            verts_np = verts.numpy().astype(np.float64).copy()
            joints_np = joints.numpy().astype(np.float64).copy()
            if not is_right:
                verts_np[:, 0] *= -1.0
                joints_np[:, 0] *= -1.0
            pred.joints, pred.keypoints_2d = joints_np, np.asarray(kp2d, dtype=np.float64)
            if return_vertices:
                pred.vertices = verts_np
            t3 = time.perf_counter()
        pred.timing_ms = {"preprocess": (t1 - t0) * 1e3, "model": (t2 - t1) * 1e3, "mano": (t3 - t2) * 1e3,
                          "total": (time.perf_counter() - t0) * 1e3}
        return pred

    def run_crop(self, crop: torch.Tensor) -> torch.Tensor:
        """One device call on one network-input crop -> ``(1, 157)`` float32
        ``[rotmats (16*9) | betas (10) | cam (3)]``.

        ``crop``: ``(1, 3, 256, 192)``, either uint8 pixels (``hamer.preprocess.preprocess_pil(..., raw=True)``) or
        the ImageNet-normalized float crop (``raw=False``).  With the published uint8 device input a float crop must
        be an exactly normalized 8-bit image (else ``ValueError``)."""
        self._check_open()
        if self.tt is None and crop.dtype == torch.uint8:
            crop = normalize_raw(crop)
        with self._lock, torch.inference_mode():
            return self._runner(crop)

    def forward_crops(self, img: torch.Tensor) -> Dict[str, Any]:
        """Upstream-style batch forward on pre-cropped hands (``hamer.models.HAMER.forward(batch)`` conventions).

        Args:
            img: ``(B, 3, 256, 256)`` (the upstream ``ViTDetDataset`` crop; the model uses columns 32:224) or
                ``(B, 3, 256, 192)``; uint8 pixels or ImageNet-normalized float (``(x / 255 - mean) / std``).
                Left hands must already be mirrored, as upstream.

        Returns (torch tensors, batch first; B device calls, one per crop):
            ``pred_cam`` (B, 3), ``pred_cam_t`` (B, 3) (crop camera, focal 5000 px for the 256 crop),
            ``focal_length`` (B, 2), ``pred_mano_params`` = ``{"global_orient": (B, 1, 3, 3),
            "hand_pose": (B, 15, 3, 3), "betas": (B, 10)}``, and with MANO ``pred_vertices`` (B, 778, 3) and
            ``pred_keypoints_3d`` (B, 21, 3) in metres.
        """
        if img.dim() != 4 or img.shape[1] != 3 or img.shape[2] != IMAGE_HEIGHT or img.shape[3] not in (256, IMAGE_WIDTH):
            raise ValueError(f"expected (B, 3, 256, 256) or (B, 3, 256, 192), got {tuple(img.shape)}")
        if img.shape[3] == 256:
            img = img[:, :, :, W_SLICE:-W_SLICE]
        outs = [self.run_crop(img[b:b + 1].contiguous())[0].float().clone() for b in range(img.shape[0])]
        out = torch.stack(outs)
        B = out.shape[0]
        rot = out[:, :144].reshape(B, 16, 3, 3)
        betas, cam = out[:, 144:154], out[:, 154:157]
        cam_t = torch.stack([cam[:, 1], cam[:, 2], 2 * FOCAL_LENGTH / (MODEL_IMAGE_SIZE * cam[:, 0] + 1e-9)], dim=-1)
        res: Dict[str, Any] = {
            "pred_cam": cam, "pred_cam_t": cam_t,
            "focal_length": torch.full((B, 2), FOCAL_LENGTH, dtype=torch.float32),
            "pred_mano_params": {"global_orient": rot[:, :1], "hand_pose": rot[:, 1:], "betas": betas},
        }
        if self.mano is not None:
            with torch.inference_mode():
                verts, joints = self.mano.forward_with_joints(rot, betas)
            res["pred_vertices"], res["pred_keypoints_3d"] = verts, joints
        return res

    # ---------------- lifetime ----------------
    def _check_open(self) -> None:
        if self._closed:
            raise RuntimeError("HamerModel is closed")

    def close(self) -> None:
        """Release the trace and the device tensors, and close the chip if this object opened it.  Idempotent."""
        if self._closed:
            return
        self._closed = True
        with self._lock:
            if self.tt is not None:
                self.tt.close()
            if self._owned_device is not None:
                try:
                    _close_device(self._owned_device)
                finally:
                    self._owned_device = None

    def __enter__(self) -> "HamerModel":
        return self

    def __exit__(self, *exc: Any) -> None:
        self.close()

    def __del__(self) -> None:  # noqa: D401
        try:
            self.close()
        except Exception:  # noqa: BLE001
            pass

    def __repr__(self) -> str:
        state = "closed" if self._closed else f"backend={self.info.get('backend')}"
        extra = "" if self.tt is None else f", grid={self.info.get('grid')}, dispatch={self.info.get('dispatch')}"
        return f"HamerModel({state}{extra}, mano={'yes' if self.mano is not None else 'no'})"


def _load_mano(mano_path: Optional[Union[str, os.PathLike]], verbose: bool) -> Optional[torch.nn.Module]:
    p = resolve_mano_path(mano_path)
    if p is None:
        _log(verbose, "MANO_RIGHT.pkl not found: MANO parameters and cameras only (pass mano_path= for the mesh)")
        return None
    from .mano import MANO  # noqa: WPS433

    m = MANO(p)
    m.eval()
    _log(verbose, f"MANO loaded from {p}")
    return m