# SPDX-FileCopyrightText: (c) 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

"""Python API: MoGe-2 monocular geometry on one Tenstorrent Blackhole chip.

    from tt_moge import MoGeModel

    with MoGeModel.from_pretrained() as model:        # opens device 0, loads the weights, captures the trace
        out = model("photo.jpg")                       # path, PIL image, numpy array or torch tensor
        out.depth        # (H, W) float32, metres; inf where out.mask is False
        out.points       # (H, W, 3) float32, metres, camera space (OpenCV axes: x right, y down, z forward)
        out.normal       # (H, W, 3) float32, unit vectors in camera space; 0 where invalid
        out.mask         # (H, W) bool, True = valid pixel
        out.intrinsics   # (3, 3) float32, normalized (row 0 x W and row 1 x H give pixels)

The call runs the same steps as ``POST /predict`` of the HTTP server (``tt_moge/server/app.py``): the
image is letterboxed onto the fixed canvas the device trace was captured for (1920x1080 by default),
the device runs the traced forward, and the host recovers focal length and z-shift, crops the
letterbox away and resizes every map back to the input resolution. For the same image and the same
options, the arrays are bit-identical to the server's ``npz`` response (the server stores ``normal``
as float16; here it is float32).

``MoGeModel.infer`` is a drop-in for the upstream ``moge.model.v2.MoGeModel.infer`` (same keys,
torch tensors), and ``MoGeModel.forward`` is the raw canvas-resolution forward (the call that the
published per-call latency measures).
"""
from __future__ import annotations

import atexit
import contextlib
import math
import os
import threading
import time
import warnings
import weakref
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import numpy as np
import torch

DEFAULT_REPO = "Ruicheng/moge-2-vitl-normal"
DEFAULT_REVISION = "cb0e8bbd6b1e243589717c78e750b1ba4c093acf"   # the revision the published numbers use
DEFAULT_RESOLUTION = (1920, 1080)                                  # canvas (width, height)
DEFAULT_NUM_TOKENS = 1800                                          # -> 32 x 57 ViT token grid at 16:9
NUM_TOKENS_RANGE = (1200, 3600)
MIN_IMAGE_SIDE = 8

ImageLike = Union[str, os.PathLike, "Image.Image", np.ndarray, torch.Tensor]  # noqa: F821

# --------------------------------------------------------------------------- warm-up variants
#
# The device graph has ONE variant per model object: the trace is captured for the canvas and the token
# budget at load time (both launch-fixed). What changes per request is the HOST path: how the image comes
# in (decoder), its size (canvas fit, resize back), and the options of the focal / shift solve and of the
# finish pass. A warm-up variant runs one full call of that path with a synthetic image of the right size
# and kind, so the lazy imports, decoder set-up, C++ helper first use, scratch buffers and allocator growth
# of that path are paid in from_pretrained, not by the first real request. No result is cached.

_INPUT_KINDS = ("array", "float", "png", "jpeg", "pil")
_CALL_KINDS = ("call", "predict_many", "infer", "forward")
_VARIANT_DEFAULTS = {"input": "array", "size": None, "fit": "pad", "fov_x": None, "apply_mask": True,
                     "force_projection": True, "call": "call"}

#: Named warm-up specs for ``from_pretrained(warmup_variants=...)``. Each entry is a variant (see
#: :meth:`MoGeModel.warmup`); ``size`` None = the canvas size.
WARMUP_PRESETS: Dict[str, Tuple[Dict[str, Any], ...]] = {
    "none": (),
    "minimal": ({},),
    "default": (
        {},                                                   # uint8 array at the canvas size, default options
        {"input": "float"},                                   # float tensor at the canvas size (upstream style)
        {"input": "png"},                                     # file path, PNG decoder (canvas size)
        {"input": "jpeg", "size": (800, 600)},                # file path, JPEG decoder; non-canvas fit + resize back
        {"size": (800, 600), "fit": "stretch"},
        {"fov_x": 60.0},
        {"apply_mask": False, "force_projection": False},
        {"input": "float", "call": "forward"},                # the raw forward on a float canvas
    ),
}
WARMUP_PRESETS["all"] = WARMUP_PRESETS["default"] + (
    {"input": "pil", "size": (800, 600)},
    {"input": "float", "size": (800, 600)},
    {"input": "float", "size": (800, 600), "call": "infer"},
    {"input": "png", "size": (800, 600)},
    {"input": "png", "call": "predict_many"},
    {"apply_mask": False},
    {"force_projection": False},
    {"call": "forward"},
)
DEFAULT_WARMUP_VARIANTS = "default"


def _normalize_variant(v: Dict[str, Any], canvas: Tuple[int, int]) -> Tuple:
    """A warm-up variant dict -> its hashable key (input, (w, h), fit, fov_x given, apply_mask,
    force_projection, call). Raises ValueError / TypeError for a bad variant."""
    if not isinstance(v, dict):
        raise TypeError(f"a warm-up variant is a dict of options, got {type(v).__name__}")
    unknown = set(v) - set(_VARIANT_DEFAULTS)
    if unknown:
        raise ValueError(f"unknown warm-up variant keys {sorted(unknown)}; valid: {sorted(_VARIANT_DEFAULTS)}")
    o = {**_VARIANT_DEFAULTS, **v}
    if o["input"] not in _INPUT_KINDS:
        raise ValueError(f"warm-up input must be one of {_INPUT_KINDS}, got {o['input']!r}")
    if o["call"] not in _CALL_KINDS:
        raise ValueError(f"warm-up call must be one of {_CALL_KINDS}, got {o['call']!r}")
    if o["fit"] not in ("pad", "stretch"):
        raise ValueError(f"fit must be 'pad' or 'stretch', got {o['fit']!r}")
    if o["fov_x"] is not None and not 1.0 < float(o["fov_x"]) < 179.0:
        raise ValueError(f"fov_x must be between 1 and 179 degrees, got {o['fov_x']}")
    size = canvas if o["size"] is None else tuple(int(s) for s in o["size"])
    if len(size) != 2 or min(size) < MIN_IMAGE_SIDE:
        raise ValueError(f"warm-up size must be (width, height), each at least {MIN_IMAGE_SIDE}; got {o['size']}")
    if o["call"] == "forward" and size != tuple(canvas):
        raise ValueError(f"call='forward' takes canvas-sized images ({canvas[0]}x{canvas[1]}); got size {size}")
    return (o["input"], size, o["fit"], o["fov_x"] is not None, bool(o["apply_mask"]),
            bool(o["force_projection"]), o["call"])


def _expand_warmup_spec(spec) -> List[Dict[str, Any]]:
    """``warmup_variants`` -> a list of variant dicts. Accepts None / a preset name / a variant dict / a
    list of preset names and variant dicts."""
    if spec is None or spec is False:
        return []
    if isinstance(spec, (str, dict)):
        spec = [spec]
    out: List[Dict[str, Any]] = []
    for s in spec:
        if isinstance(s, str):
            if s not in WARMUP_PRESETS:
                raise ValueError(f"unknown warm-up preset {s!r}; presets: {sorted(WARMUP_PRESETS)}")
            out.extend(dict(v) for v in WARMUP_PRESETS[s])
        elif isinstance(s, dict):
            out.append(dict(s))
        else:
            raise TypeError(f"warmup_variants entries are preset names or dicts, got {type(s).__name__}")
    return out


def _synthetic_image(width: int, height: int, seed: int = 0) -> np.ndarray:
    """A deterministic photo-like uint8 (H, W, 3) scene (sky gradient, ground plane, textured boxes,
    sensor noise) for the warm-up calls: realistic sizes and statistics, so the decoders, the canvas fit,
    the focal / shift solve and the finish pass take their normal paths. Not a real input."""
    rng = np.random.default_rng(seed)
    yy = np.linspace(0.0, 1.0, height, dtype=np.float32)[:, None]
    xx = np.linspace(0.0, 1.0, width, dtype=np.float32)[None, :]
    horizon = 0.45
    sky = np.stack([0.45 + 0.3 * yy, 0.6 + 0.25 * yy, 0.9 - 0.1 * yy], -1) * np.ones((1, width, 1), np.float32)
    g = np.clip((yy - horizon) / (1.0 - horizon), 0.0, 1.0)
    stripes = 0.08 * np.sin(40.0 * xx * (0.3 + g) + 8.0 * g)
    ground = np.stack([0.35 + 0.25 * g + stripes, 0.3 + 0.2 * g + stripes, 0.2 + 0.1 * g + stripes], -1)
    img = np.where((yy < horizon)[..., None], sky, ground)
    for _ in range(6):                                     # boxes standing on the ground
        bw, bh = rng.uniform(0.05, 0.2), rng.uniform(0.1, 0.35)
        x0, yb = rng.uniform(0.0, 1.0 - bw), rng.uniform(horizon + 0.05, 0.95)
        ys, ye = int(max(0.0, yb - bh) * height), int(yb * height)
        xs, xe = int(x0 * width), int((x0 + bw) * width)
        if ye > ys and xe > xs:
            col = rng.uniform(0.1, 0.9, 3).astype(np.float32)
            shade = np.linspace(0.8, 1.1, xe - xs, dtype=np.float32)[None, :, None]
            img[ys:ye, xs:xe] = col * shade
    img += 0.02 * rng.standard_normal(img.shape, dtype=np.float32)
    return (np.clip(img, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)


# --------------------------------------------------------------------------- console noise

_LOG_ENV_QUIET = {"TT_LOGGER_LEVEL": "error",      # tt-metal / UMD C++ logger (read when ttnn loads)
                  "LOGURU_LEVEL": "WARNING"}       # ttnn's Python logger (loguru default handler, read at import)


@contextlib.contextmanager
def _quiet_tt_logs(verbose: bool):
    """With verbose=False: ttnn is imported and the device is opened with the tt-metal info / debug log
    lines off. Errors still print; an explicit TT_LOGGER_LEVEL / LOGURU_LEVEL in the environment wins.
    Both loggers read the variable once (loguru when ttnn imports it, the C++ logger at its first message,
    the device open), so they are set only inside this block and restored after it. No effect when ttnn
    is already imported (its loggers are set up then)."""
    import sys

    if verbose or "ttnn" in sys.modules:
        yield
        return
    saved = {k: os.environ.get(k) for k in _LOG_ENV_QUIET}
    for k, v in _LOG_ENV_QUIET.items():
        os.environ.setdefault(k, v)
    try:
        try:
            import ttnn  # noqa: F401
        except ImportError:
            pass        # reported by the code that needs ttnn
        yield
    finally:
        for k, v in saved.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v


# --------------------------------------------------------------------------- output

@dataclass
class MoGeOutput:
    """Result of ``model(image)`` at the resolution of the input image (H x W).

    Attributes:
        points: float32 (H, W, 3). Camera-space point map in metres (OpenCV axes: x right, y down,
            z forward). ``+inf`` where ``mask`` is False when ``apply_mask=True``.
        depth: float32 (H, W). Metric depth (z) in metres. ``+inf`` where ``mask`` is False when
            ``apply_mask=True``.
        normal: float32 (H, W, 3). Unit surface normals in camera space; ``0`` where ``mask`` is False
            when ``apply_mask=True``.
        mask: bool (H, W). True for valid pixels (sky, letterbox and invalid depth are False).
        intrinsics: float32 (3, 3). Normalized intrinsics: ``fx, cx`` are in units of the image width,
            ``fy, cy`` in units of the image height (see ``intrinsics_pixels``).
        metric_scale: float. The scale the model predicted (already applied to points and depth).
        fov_x, fov_y: float. Horizontal / vertical field of view in degrees.
        region: (x0, y0, x1, y1). Where the image was placed on the canvas (pixels).

    The object also acts like the upstream ``infer()`` dict: ``out["depth"]``, ``out.keys()``,
    ``dict(out.items())``.
    """

    points: np.ndarray
    depth: np.ndarray
    normal: np.ndarray
    mask: np.ndarray
    intrinsics: np.ndarray
    metric_scale: float
    fov_x: float
    fov_y: float
    region: Tuple[int, int, int, int] = field(default=(0, 0, 0, 0), repr=False)

    _KEYS = ("points", "depth", "normal", "mask", "intrinsics", "metric_scale")

    @property
    def height(self) -> int:
        return int(self.depth.shape[0])

    @property
    def width(self) -> int:
        return int(self.depth.shape[1])

    @property
    def intrinsics_pixels(self) -> np.ndarray:
        """(3, 3) float32 intrinsics in pixels of the input image."""
        k = self.intrinsics.astype(np.float32).copy()
        k[0] *= self.width
        k[1] *= self.height
        k[2] = (0.0, 0.0, 1.0)
        return k

    def keys(self):
        return self._KEYS

    def items(self):
        return ((k, getattr(self, k)) for k in self._KEYS)

    def __getitem__(self, key: str):
        if key not in self._KEYS:
            raise KeyError(key)
        return getattr(self, key)

    def __contains__(self, key) -> bool:
        return key in self._KEYS

    def to_dict(self) -> Dict[str, Any]:
        return {k: getattr(self, k) for k in self._KEYS}

    def to_torch(self) -> Dict[str, torch.Tensor]:
        """The upstream ``MoGeModel.infer()`` dict: torch tensors ``points``, ``depth``, ``normal``,
        ``mask``, ``intrinsics`` (no copies of the numpy arrays)."""
        return {k: torch.from_numpy(np.ascontiguousarray(getattr(self, k)))
                for k in ("points", "intrinsics", "depth", "mask", "normal")}

    def save_npz(self, path) -> None:
        """Save the arrays with the keys / dtypes of the server's ``npz`` response
        (``normal`` as float16, ``mask`` as uint8)."""
        np.savez_compressed(path, points=self.points, depth=self.depth, normal=self.normal.astype(np.float16),
                            mask=self.mask.astype(np.uint8), intrinsics=self.intrinsics,
                            metric_scale=np.float32(self.metric_scale))

    def __repr__(self) -> str:
        valid = int(self.mask.sum())
        d = self.depth[self.mask] if valid else np.zeros(0, np.float32)
        med = f"{float(np.median(d)):.3f} m" if valid else "n/a"
        return (f"MoGeOutput({self.width}x{self.height}, valid {valid / self.mask.size:.1%}, median depth {med}, "
                f"fov_x {self.fov_x:.1f} deg, metric_scale {self.metric_scale:.4f})")


# --------------------------------------------------------------------------- input conversion

def _load_rgb(image) -> Union[np.ndarray, torch.Tensor]:
    """Any supported image -> uint8 (H, W, 3) numpy array, or float32 (1, 3, H, W) torch tensor in [0, 1].

    uint8 inputs and files stay 8-bit (the served path). Float inputs stay float32.
    """
    from PIL import Image

    if isinstance(image, (str, os.PathLike)):
        with open(image, "rb") as f:
            raw = f.read()
        from tt_moge.server import fastpath
        if fastpath.enabled():
            arr = fastpath.decode_png_rgb(raw)     # the server's decoder: 8-bit RGB(A) PNG, == PIL's pixels
            if arr is not None:
                return arr
        import io
        with Image.open(io.BytesIO(raw)) as im:
            im.load()
            return np.array(im.convert("RGB"))
    if isinstance(image, Image.Image):
        return np.array(image if image.mode == "RGB" else image.convert("RGB"))
    if isinstance(image, torch.Tensor):
        t = image.detach().cpu()
        if t.dim() == 4:
            if t.shape[0] != 1:
                raise ValueError(f"one image per call; got a batch of {t.shape[0]} (pass a list or use predict_many)")
            t = t[0]
        if t.dim() == 3 and t.shape[0] in (1, 3) and t.shape[-1] not in (1, 3, 4):
            t = t.permute(1, 2, 0)          # (C, H, W) -> (H, W, C)
        image = t.numpy()
    if not isinstance(image, np.ndarray):
        raise TypeError(f"unsupported image type {type(image).__name__}: pass a path, PIL image, numpy array or "
                        "torch tensor")
    a = image
    if a.ndim == 3 and a.shape[0] in (1, 3) and a.shape[-1] not in (1, 3, 4):
        a = np.moveaxis(a, 0, -1)          # (C, H, W) -> (H, W, C)
    if a.ndim == 2:
        a = a[..., None]
    if a.ndim != 3 or a.shape[-1] not in (1, 3, 4):
        raise ValueError(f"expected an RGB image (H, W, 3) or (3, H, W); got shape {tuple(image.shape)}")
    if a.shape[-1] == 1:
        a = np.repeat(a, 3, axis=-1)
    elif a.shape[-1] == 4:
        a = a[..., :3]                     # PIL's RGBA -> RGB also drops alpha
    if a.dtype == np.uint8:
        return np.ascontiguousarray(a)
    if a.dtype == np.bool_ or not np.issubdtype(a.dtype, np.number):
        raise TypeError(f"unsupported image dtype {a.dtype}")
    if np.issubdtype(a.dtype, np.integer):
        raise TypeError(f"integer images must be uint8 (0..255); got {a.dtype}")
    t = torch.from_numpy(np.ascontiguousarray(a, dtype=np.float32)).permute(2, 0, 1)[None]
    return t


def _check_size(h: int, w: int) -> None:
    if h < MIN_IMAGE_SIDE or w < MIN_IMAGE_SIDE:
        raise ValueError(f"image too small: {w}x{h} (at least {MIN_IMAGE_SIDE}x{MIN_IMAGE_SIDE})")


# --------------------------------------------------------------------------- model

_OPEN_MODELS: "weakref.WeakSet[MoGeModel]" = weakref.WeakSet()


@atexit.register
def _close_open_models() -> None:      # best effort: release traces / close owned devices at exit
    for m in list(_OPEN_MODELS):
        try:
            m.close()
        except Exception:  # noqa: BLE001
            pass


class MoGeModel:
    """MoGe-2 (ViT-L, normal head) on one Tenstorrent Blackhole chip.

    Create it with :meth:`from_pretrained`; call it with an image; close it with :meth:`close` or a
    ``with`` block. One device trace is captured for a fixed canvas (``resolution``) and token budget
    (``num_tokens``); every image is fitted onto that canvas, so any input size works with the same trace.
    Calls are thread-safe (the device runs one image at a time).
    """

    def __init__(self, tt_model, device, *, resolution: Tuple[int, int] = DEFAULT_RESOLUTION,
                 num_tokens: int = DEFAULT_NUM_TOKENS, owns_device: bool = False,
                 device_info: Optional[Dict[str, Any]] = None, checkpoint: Optional[str] = None):
        """Internal: use :meth:`from_pretrained`. ``tt_model`` is a built ``tt_moge.tt.ttnn_moge.TtMoGe``."""
        from tt_moge.server import postprocess as pp

        self._tt = tt_model
        self._device = device
        self._owns_device = owns_device
        self.resolution = (int(resolution[0]), int(resolution[1]))
        self.num_tokens = int(num_tokens)
        self.token_grid = pp.token_grid(self.resolution[0], self.resolution[1], self.num_tokens)
        self.device_info = dict(device_info or {})
        self.checkpoint = checkpoint
        # host threads of the post-processing pass (results do not depend on it); serve.env uses 16
        self.host_threads = int(os.environ.get("MOGE_SERVE_THREADS", "0")) or min(16, os.cpu_count() or 1)
        self._lock = threading.Lock()
        self._warm_lock = threading.Lock()
        self._closed = False
        self._warmed: Dict[Tuple, float] = {}        # variant key -> warm-up time (s)
        _OPEN_MODELS.add(self)

    # ------------------------------------------------------------------ construction
    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, os.PathLike] = DEFAULT_REPO,
        *,
        device=None,
        device_id: int = 0,
        revision: Optional[str] = DEFAULT_REVISION,
        resolution: Tuple[int, int] = DEFAULT_RESOLUTION,
        num_tokens: int = DEFAULT_NUM_TOKENS,
        dispatch: str = "eth",
        num_command_queues: Optional[int] = None,
        warmup: int = 2,
        warmup_variants: Any = DEFAULT_WARMUP_VARIANTS,
        verbose: bool = False,
    ) -> "MoGeModel":
        """Load the weights, open the device, build the model and capture the device trace.

        Args:
            pretrained_model_name_or_path: Hugging Face repo id (default ``Ruicheng/moge-2-vitl-normal``,
                downloaded to the HF cache), a local directory that holds ``model.pt``, or the path of
                ``model.pt`` itself. ``$MOGE_WEIGHTS_DIR`` (a directory with ``model.pt``) overrides the repo.
            device: an already opened ttnn device. The model then uses it and does not close it.
                Pass ``num_command_queues=2`` if you opened it with 2 command queues.
            device_id: ttnn device id to open when ``device`` is None (default 0).
            revision: weights revision (commit sha) on the Hub; default: the pinned revision of the
                published numbers. None = the repo's default branch.
            resolution: canvas (width, height) that every image is fitted onto. Default (1920, 1080),
                the validated setting. Other canvases work but are not validated.
            num_tokens: ViT token budget (1200..3600, default 1800).
            dispatch: ``"eth"`` (default; ETH dispatch, 1 command queue, 12x10 compute grid: the p150
                configuration of the published numbers; needs the ETH-dispatch patch of tt-metal, else it falls
                back to worker dispatch with a RuntimeWarning) or ``"worker"`` (opt-in; worker dispatch, which
                leaves an 11x10 compute grid on a p150; 12x10 only on a Galaxy chip).
            num_command_queues: override the command-queue count (default 1). 2 is an opt-in for worker
                dispatch only (Galaxy-only speed-up: the output is read back on CQ1); outputs are
                bit-identical with 1 or 2 queues.
            warmup: device forwards to run at load (>= 1; the first one compiles and captures the trace).
            warmup_variants: the host paths to prepare at load, so that the first real call of each is
                already as fast as later calls. A preset name (``"default"``, ``"all"``, ``"minimal"``,
                ``"none"``; see ``WARMUP_PRESETS``), a variant dict (the keyword arguments of
                :meth:`warmup`), or a list of them. Each variant costs one full call with a synthetic
                image (about 30-60 ms). ``"none"``: only the device warm-up.
            verbose: print progress, and keep the tt-metal info log lines. With False (default), ttnn is
                imported with ``TT_LOGGER_LEVEL=error`` and ``LOGURU_LEVEL=WARNING`` unless you set them.
        """
        if not NUM_TOKENS_RANGE[0] <= int(num_tokens) <= NUM_TOKENS_RANGE[1]:
            raise ValueError(f"num_tokens={num_tokens} outside the model's range {NUM_TOKENS_RANGE}")
        cw, ch = (int(resolution[0]), int(resolution[1]))
        if cw < 224 or ch < 224:
            raise ValueError(f"resolution {cw}x{ch} is too small (each side at least 224)")
        variants = _expand_warmup_spec(warmup_variants)          # validated before the device is opened
        for v in variants:
            _normalize_variant(v, (cw, ch))
        if (cw, ch) != DEFAULT_RESOLUTION or int(num_tokens) != DEFAULT_NUM_TOKENS:
            warnings.warn(f"tt_moge: resolution {cw}x{ch} / num_tokens {num_tokens} is not the validated setting "
                          f"{DEFAULT_RESOLUTION[0]}x{DEFAULT_RESOLUTION[1]} / {DEFAULT_NUM_TOKENS}", stacklevel=2)

        def log(msg):
            if verbose:
                print(f"[tt_moge] {msg}", flush=True)

        # The vendored DINOv2 layers probe xformers unless this is set (serve.env sets it too).
        os.environ.setdefault("XFORMERS_DISABLED", "1")
        with _quiet_tt_logs(verbose):
            return cls._load(pretrained_model_name_or_path, device, device_id, revision, (cw, ch), num_tokens,
                             dispatch, num_command_queues, warmup, variants, log)

    @classmethod
    def _load(cls, pretrained_model_name_or_path, device, device_id, revision, resolution, num_tokens, dispatch,
              num_command_queues, warmup, variants, log) -> "MoGeModel":
        cw, ch = resolution
        import tt_moge  # noqa: F401  (vendored moge / utils3d on sys.path)
        from tt_moge.reference.load_pretrained import load_moge2, resolve_checkpoint

        t0 = time.perf_counter()
        ckpt = cls._resolve_weights(pretrained_model_name_or_path, revision, resolve_checkpoint)
        log(f"loading weights {ckpt}")
        ref = load_moge2(ckpt, verbose=False)
        log(f"weights loaded in {time.perf_counter() - t0:.1f} s")

        owns = device is None
        info: Dict[str, Any]
        if owns:
            from tt_moge.device import open_device
            device, info = open_device(device_id, dispatch=dispatch, num_command_queues=num_command_queues)
        else:
            g = device.compute_with_storage_grid_size()
            info = {"dispatch": "external", "num_command_queues": num_command_queues or 1, "grid": (g.x, g.y)}
        log(f"device open: {info}")

        tt = None
        try:
            from tt_moge.tt.ttnn_moge import TtMoGe

            two_cq = info["num_command_queues"] == 2
            prev = os.environ.get("MOGE_2CQ")
            os.environ["MOGE_2CQ"] = "1" if two_cq else "0"   # TtMoGe reads it once, at build
            try:
                t1 = time.perf_counter()
                tt = TtMoGe(ref_moge_model=ref, device=device, trace=True)
                log(f"model built in {time.perf_counter() - t1:.1f} s")
            finally:
                if prev is None:
                    os.environ.pop("MOGE_2CQ", None)
                else:
                    os.environ["MOGE_2CQ"] = prev
            model = cls(tt, device, resolution=(cw, ch), num_tokens=num_tokens, owns_device=owns,
                        device_info=info, checkpoint=ckpt)
            model._warmup(max(1, int(warmup)), log)
            t2 = time.perf_counter()
            model.warmup(variants, _log=log)
            log(f"{len(model._warmed)} host warm-up variants in {time.perf_counter() - t2:.2f} s")
            log(f"ready after {time.perf_counter() - t0:.1f} s")
            return model
        except BaseException:
            if tt is not None:
                try:
                    tt.release()
                except Exception:  # noqa: BLE001
                    pass
            if owns:
                from tt_moge.device import close_device
                close_device(device)
            raise

    @staticmethod
    def _resolve_weights(name_or_path, revision, resolve_checkpoint) -> str:
        p = Path(os.fspath(name_or_path)).expanduser()
        if p.is_file():
            return str(p)
        if p.is_dir():
            f = p / "model.pt"
            if not f.is_file():
                raise FileNotFoundError(f"{p} does not contain model.pt")
            return str(f)
        return resolve_checkpoint(repo_id=str(name_or_path), revision=revision)

    def _warmup(self, iters: int, log) -> None:
        cw, ch = self.resolution
        g = torch.Generator().manual_seed(0)
        dummy = torch.rand(1, 3, ch, cw, generator=g)     # the server's warm-up image
        out = None
        for i in range(iters):
            t = time.perf_counter()
            out = self.forward(dummy)
            log(f"warm-up forward {i + 1}/{iters}: {1e3 * (time.perf_counter() - t):.0f} ms")
        for k in ("points", "normal", "mask", "metric_scale"):
            if k not in out:
                raise RuntimeError(f"the checkpoint has no {k!r} head")
        if tuple(out["points"].shape[1:3]) != (ch, cw):
            raise RuntimeError(f"warm-up output {tuple(out['points'].shape)} does not match the canvas {(ch, cw)}")

    def warmup(self, variants=None, /, *, _log=None, **variant) -> float:
        """Prepare one or more host-path variants, so that their first real call is as fast as later calls.

        ``model.warmup(input="jpeg", size=(640, 480), fit="stretch")`` prepares one variant;
        ``model.warmup(["all"])`` or ``model.warmup([{...}, {...}])`` prepares several. Idempotent: a
        variant that is already prepared (same key) is skipped. Returns the seconds spent.

        Variant keys (all optional):
            input: ``"array"`` (uint8 numpy, default), ``"float"`` (float32 torch (1, 3, H, W) in [0, 1]),
                ``"png"`` / ``"jpeg"`` (a file path; the decoder), ``"pil"`` (a PIL image).
            size: ``(width, height)`` of the input image; default the canvas size.
            fit, fov_x, apply_mask, force_projection: as in :meth:`__call__` (for ``fov_x`` only
                "given or not" matters).
            call: ``"call"`` (``model(image)``, default), ``"predict_many"`` (2 images), ``"infer"``,
                ``"forward"`` (canvas size only).

        Each variant runs the real call once on a synthetic image of that size and kind (and a temporary
        file for png / jpeg). The device trace is the same for every variant, so this adds no device memory.
        Outputs of later calls do not change (no result is kept).
        """
        self._check_open()
        if variants is None:
            todo = [variant]
        else:
            if variant:
                raise TypeError("pass either a list / preset name or keyword options, not both")
            todo = _expand_warmup_spec(variants)
        keys = [(_normalize_variant(v, self.resolution), v) for v in todo]
        spent = 0.0
        images: Dict[Tuple[int, int, int], np.ndarray] = {}     # synthetic images of this warm-up, by size
        with self._warm_lock:
            tmp = None
            try:
                for key, v in keys:
                    if key in self._warmed:
                        continue
                    if key[0] in ("png", "jpeg") and tmp is None:
                        import tempfile
                        tmp = tempfile.TemporaryDirectory(prefix="tt_moge_warmup_")
                    t = time.perf_counter()
                    self._warm_one(key, v, tmp.name if tmp is not None else None, images)
                    dt = time.perf_counter() - t
                    self._warmed[key] = dt
                    spent += dt
                    if _log is not None:
                        _log(f"warm-up {self._variant_str(key)}: {1e3 * dt:.0f} ms")
            finally:
                if tmp is not None:
                    tmp.cleanup()
        return spent

    @property
    def warmed_variants(self) -> List[Dict[str, Any]]:
        """The host-path variants prepared so far, with the warm-up time of each (ms)."""
        names = ("input", "size", "fit", "fov_x_given", "apply_mask", "force_projection", "call")
        return [dict(zip(names, k), warmup_ms=round(1e3 * dt, 1)) for k, dt in self._warmed.items()]

    @staticmethod
    def _variant_str(key) -> str:
        inp, (w, h), fit, fov, am, fp, call = key
        extra = "".join([f" fit={fit}" if fit != "pad" else "", " fov_x" if fov else "",
                         " apply_mask=False" if not am else "", " force_projection=False" if not fp else ""])
        return f"{call}({inp} {w}x{h}{extra})"

    def _warm_one(self, key, v: Dict[str, Any], tmpdir: Optional[str], images: Dict) -> None:
        inp, (w, h), fit, _fov, apply_mask, force_projection, call = key
        n = 2 if call == "predict_many" else 1
        imgs = []
        for i in range(n):
            if (w, h, i) not in images:
                images[(w, h, i)] = _synthetic_image(w, h, seed=1 + i)
            a = images[(w, h, i)].copy()            # a fresh array per call, as a user passes
            if inp == "array":
                x = a
            elif inp == "float":
                x = torch.from_numpy(a).permute(2, 0, 1)[None].float().div_(255.0)
            elif inp == "pil":
                from PIL import Image
                x = Image.fromarray(a)
            else:
                from PIL import Image
                x = os.path.join(tmpdir, f"warmup_{w}x{h}_{i}.{'png' if inp == 'png' else 'jpg'}")
                if inp == "png":
                    Image.fromarray(a).save(x, format="PNG", compress_level=1)
                else:
                    Image.fromarray(a).save(x, format="JPEG", quality=90)
            imgs.append(x)
        opts = dict(fit=fit, fov_x=v.get("fov_x"), apply_mask=apply_mask, force_projection=force_projection)
        if call == "forward":
            self.forward(imgs[0])
        elif call == "infer":
            x = imgs[0]
            if isinstance(x, torch.Tensor) and x.dim() == 4:
                x = x[0]
            self.infer(x, **opts)
        elif call == "predict_many":
            self.predict_many(imgs, **opts)
        else:
            self(imgs[0], **opts)

    # ------------------------------------------------------------------ properties
    @property
    def device(self):
        """The ttnn device the model runs on."""
        return self._device

    @property
    def closed(self) -> bool:
        return self._closed

    def __repr__(self) -> str:
        state = "closed" if self._closed else "open"
        return (f"MoGeModel(resolution={self.resolution[0]}x{self.resolution[1]}, num_tokens={self.num_tokens}, "
                f"token_grid={self.token_grid}, device={self.device_info}, {state})")

    # ------------------------------------------------------------------ lifetime
    def close(self) -> None:
        """Release the device trace and buffers; close the device if this model opened it. Idempotent."""
        if self._closed:
            return
        with self._lock:
            if self._closed:
                return
            self._closed = True
            try:
                self._tt.release()
            finally:
                if self._owns_device:
                    from tt_moge.device import close_device
                    close_device(self._device)
                self._tt = None
                self._device = None
                _OPEN_MODELS.discard(self)

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

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

    def _check_open(self) -> None:
        if self._closed:
            raise RuntimeError("this MoGeModel is closed")

    # ------------------------------------------------------------------ steps
    def _prepare(self, image, fit: str):
        """image -> (model input on the canvas, region, (H, W) of the input)."""
        from tt_moge.server import fastpath
        from tt_moge.server import postprocess as pp

        if fit not in ("pad", "stretch"):
            raise ValueError(f"fit must be 'pad' or 'stretch', got {fit!r}")
        x = _load_rgb(image)
        cw, ch = self.resolution
        if isinstance(x, torch.Tensor):                  # float32 (1, 3, H, W)
            h, w = int(x.shape[2]), int(x.shape[3])
            _check_size(h, w)
            if (w, h) == (cw, ch):                       # canvas-sized float image: the model input as is
                return x, (0, 0, cw, ch), (h, w)
            # the canvas fit works on 8-bit images (as the server does)
            x = (x[0].permute(1, 2, 0).clamp(0, 1) * 255).round().to(torch.uint8).numpy()
        h, w = x.shape[:2]
        _check_size(h, w)
        if not x.flags.writeable and (w, h) == (cw, ch):
            x = x.copy()                                 # the canvas-sized array becomes the model input tensor
        if fastpath.enabled():   # C++ fit, uint8 canvas (== pp.fit_to_canvas * 255 bit for bit; the server's path)
            canvas, region = fastpath.fit_to_canvas_u8(x, self.resolution, fit)
        else:
            from PIL import Image
            canvas, region = pp.fit_to_canvas(Image.fromarray(x), self.resolution, fit)
        return canvas, region, (h, w)

    def _run(self, model_input: torch.Tensor) -> Dict[str, torch.Tensor]:
        self._check_open()
        with self._lock, torch.inference_mode():
            self._check_open()
            return self._tt(model_input, num_tokens=self.num_tokens)

    def _postprocess(self, out, region, size, fit, fov_x, force_projection, apply_mask) -> MoGeOutput:
        """Canvas forward output -> MoGeOutput at ``size`` (the server's /predict post-processing)."""
        from tt_moge.server import fastpath
        from tt_moge.server import postprocess as pp

        with torch.inference_mode():
            fov_canvas = pp.fov_x_to_canvas(fov_x, self.resolution, region) if fov_x is not None else None
            if fastpath.enabled():
                cam = fastpath.camera(out, region if fit == "pad" else None, fov_canvas)
                # == fastpath.finish (the server's pass), with this model's host thread count
                f = fastpath._Finish(out, cam, self.resolution, region, size, force_projection, apply_mask, False)
                f.rows(0, size[0], self.host_threads)
                res = f.result()
                depth, points, normal, mask = res["depth"], res["points"], res["normal"], res["mask"]
            else:
                valid = pp.region_mask(self.resolution, region) if fit == "pad" else None
                cam = pp.canvas_to_camera(out, valid, fov_canvas)
                res = pp.to_original(cam, self.resolution, region, size, force_projection=force_projection,
                                     apply_mask=apply_mask)
                depth, points, normal, mask = (res[k].numpy() for k in ("depth", "points", "normal", "mask"))
            fx, fy = pp.fov_deg(res["intrinsics"])
            return MoGeOutput(points=points, depth=depth, normal=normal, mask=mask,
                              intrinsics=res["intrinsics"].numpy().astype(np.float32),
                              metric_scale=float(cam["metric_scale"]), fov_x=fx, fov_y=fy, region=tuple(region))

    @staticmethod
    def _check_opts(fov_x, fit) -> None:
        if fov_x is not None and not 1.0 < float(fov_x) < 179.0:
            raise ValueError(f"fov_x must be between 1 and 179 degrees, got {fov_x}")
        if fit not in ("pad", "stretch"):
            raise ValueError(f"fit must be 'pad' or 'stretch', got {fit!r}")

    # ------------------------------------------------------------------ public calls
    def __call__(self, image, *, fit: str = "pad", fov_x: Optional[float] = None, apply_mask: bool = True,
                 force_projection: bool = True):
        """Estimate metric geometry for one image (or a list of images).

        Args:
            image: a file path, a PIL image, a numpy array (H, W, 3) uint8 or float in [0, 1] (also
                (3, H, W) and grayscale (H, W)), or a torch tensor (3, H, W) / (1, 3, H, W) float in [0, 1]
                (the upstream convention) or (H, W, 3) uint8. A list or tuple of images returns a list
                (see :meth:`predict_many`). Any size from 8x8 up; it is fitted onto the canvas.
            fit: ``"pad"`` (default) letterboxes the image onto the canvas, keeping its aspect ratio;
                ``"stretch"`` resizes it to the canvas.
            fov_x: known horizontal field of view of the image in degrees. Only the z-shift is then
                solved for. Default: estimated from the image.
            apply_mask: set invalid pixels to ``inf`` (depth, points) and 0 (normal). Default True.
            force_projection: compute the points from depth and intrinsics (upstream default True).

        Returns:
            :class:`MoGeOutput` at the input resolution (a list of them for a list input).

        Float images that are not canvas-sized are converted to 8 bit before the canvas fit (the
        canvas fit works on 8-bit images, as in the server).
        """
        if isinstance(image, (list, tuple)):
            return self.predict_many(image, fit=fit, fov_x=fov_x, apply_mask=apply_mask,
                                     force_projection=force_projection)
        if isinstance(image, torch.Tensor) and image.dim() == 4 and image.shape[0] > 1:
            return self.predict_many(list(image), fit=fit, fov_x=fov_x, apply_mask=apply_mask,
                                     force_projection=force_projection)
        self._check_opts(fov_x, fit)
        x, region, size = self._prepare(image, fit)
        out = self._run(x)
        return self._postprocess(out, region, size, fit, fov_x, force_projection, apply_mask)

    predict = __call__

    def predict_many(self, images: Iterable[ImageLike], *, fit: str = "pad", fov_x: Optional[float] = None,
                     apply_mask: bool = True, force_projection: bool = True) -> List[MoGeOutput]:
        """Run several images; returns one :class:`MoGeOutput` per image, in order.

        Same results as calling the model on each image. While the device runs image i, a host thread
        loads and fits image i + 1 onto the canvas (most useful for file paths: the PNG / JPEG decode
        is hidden). The device runs one image at a time (batch 1). Overlapping the post-processing too
        was measured slower (host contention with the device call's own host threads).
        """
        self._check_opts(fov_x, fit)
        images = list(images)
        results: List[MoGeOutput] = []
        if not images:
            return results
        with ThreadPoolExecutor(max_workers=1, thread_name_prefix="tt-moge-prefetch") as ex:
            prep = ex.submit(self._prepare, images[0], fit)
            for i in range(len(images)):
                x, region, size = prep.result()
                if i + 1 < len(images):
                    prep = ex.submit(self._prepare, images[i + 1], fit)
                out = self._run(x)
                results.append(self._postprocess(out, region, size, fit, fov_x, force_projection, apply_mask))
        return results

    def infer(self, image: torch.Tensor, num_tokens: Optional[int] = None, resolution_level: int = 9,
              force_projection: bool = True, apply_mask: bool = True, fov_x: Optional[float] = None,
              use_fp16: bool = True, *, fit: str = "pad") -> Dict[str, torch.Tensor]:
        """Drop-in for upstream ``moge.model.v2.MoGeModel.infer``.

        Args:
            image: float tensor (3, H, W) or (B, 3, H, W) in [0, 1] (anything ``__call__`` takes also works).
            num_tokens: must be None or the ``num_tokens`` the model was loaded with (the device trace
                is bound to it). ``resolution_level`` and ``use_fp16`` are accepted and ignored.
            force_projection, apply_mask, fov_x: as upstream.

        Returns:
            dict of CPU torch tensors: ``points`` (H, W, 3), ``intrinsics`` (3, 3), ``depth`` (H, W),
            ``mask`` (H, W) bool, ``normal`` (H, W, 3); with a leading batch dimension for a 4-D input.
        """
        if num_tokens is not None and int(num_tokens) != self.num_tokens:
            raise ValueError(f"num_tokens={num_tokens}: the device trace was captured for num_tokens="
                             f"{self.num_tokens} (pass it to from_pretrained)")
        if isinstance(image, torch.Tensor) and image.dim() == 4:
            outs = self.predict_many(list(image), fit=fit, fov_x=fov_x, apply_mask=apply_mask,
                                     force_projection=force_projection)
            ds = [o.to_torch() for o in outs]
            return {k: torch.stack([d[k] for d in ds]) for k in ds[0]}
        return self(image, fit=fit, fov_x=fov_x, apply_mask=apply_mask, force_projection=force_projection).to_torch()

    def forward(self, image, num_tokens: Optional[int] = None) -> Dict[str, torch.Tensor]:
        """Raw forward on a canvas-sized image, like upstream ``MoGeModel.forward``.

        Args:
            image: canvas-sized (``resolution``) image: float32 torch tensor (1, 3, H, W) in [0, 1], or
                uint8 (1, H, W, 3) / (H, W, 3), or anything ``__call__`` takes at the canvas size.

        Returns:
            dict of CPU torch tensors at the canvas resolution: ``points`` (1, H, W, 3) affine-invariant
            point map (before the focal / shift recovery), ``normal`` (1, H, W, 3), ``mask`` (1, H, W)
            probabilities, ``metric_scale`` (1,).

        This is the call the published per-call latency (``tt(image)``) measures.
        """
        if num_tokens is not None and int(num_tokens) != self.num_tokens:
            raise ValueError(f"num_tokens={num_tokens}: the device trace was captured for num_tokens={self.num_tokens}")
        cw, ch = self.resolution
        if isinstance(image, torch.Tensor) and image.dim() == 4 and image.dtype == torch.float32 \
                and image.shape[0] == 1 and image.shape[1] == 3:
            x = image
        elif isinstance(image, torch.Tensor) and image.dtype == torch.uint8 and image.dim() in (3, 4) \
                and image.shape[-1] == 3:
            x = image if image.dim() == 4 else image[None]
        else:
            x = _load_rgb(image)
            x = torch.from_numpy(x)[None] if isinstance(x, np.ndarray) else x
        hw = (x.shape[1], x.shape[2]) if x.dtype == torch.uint8 else (x.shape[2], x.shape[3])
        if tuple(hw) != (ch, cw):
            raise ValueError(f"forward() takes canvas-sized images ({cw}x{ch}); got {hw[1]}x{hw[0]}. "
                             "Use model(image) for other sizes.")
        return self._run(x)