Python API: warm-up so the first real call is fast (warmup_variants, model.warmup()), quiet logs, install extras

from_pretrained() now warms up the common per-request variants with dummy inputs (device and host), so the first real call is within about 10% of later calls (independently verified in fresh processes; outputs bit-identical; steady-state latency unchanged). Adds model.warmup(...) for other variants, verbose=False by default, and documents the [server,test] extras and repo-root paths. tt-model.yaml and SERVING.md are unchanged.

PYTHON.md

@@ -10,13 +10,36 @@ You need an environment that already has `ttnn` (tt-metal `v0.78.0-dev20260820`,
 Then install the `code/` directory of this repo:
 
 ```bash
-pip install -e code/            # or: pip install code/
-pip install -e "code/[pose]"    # also installs OpenCV for return_pose=True
+pip install -e code/                         # or: pip install code/
+pip install -e "code/[pose]"                 # also installs OpenCV for return_pose=True
+pip install -e "code/[server]"               # also installs FastAPI and uvicorn for the HTTP server (SERVING.md)
+pip install -e "code/[pose,server,test]"     # everything the host tests and the device tests need
 ```
 
 The package declares only its own dependencies (numpy<2, torch, pillow, safetensors, huggingface_hub).
 It does not install ttnn. You do not need `sys.path` changes or a specific working directory.
 
+| extra | installs | you need it for |
+|---|---|---|
+| `pose` | `opencv-python-headless` | `return_pose=True` (PnP-RANSAC) |
+| `server` | `fastapi`, `uvicorn`, `pydantic>=2` | the HTTP server `models/server/app.py`, and the host tests that compare with the server |
+| `test` | `pytest`, `httpx` | all tests (`httpx` is for FastAPI's `TestClient`) |
+| `viz` | `matplotlib` | optional figures |
+
+The host tests do not open a chip and do not need ttnn. Run only the host tests:
+
+```bash
+cd code
+TT_VISIBLE_DEVICES=none python -m pytest -q models/tests/test_api_host.py models/tests/test_server_host.py models/tests/test_fused_host.py
+```
+
+The device tests need a chip, ttnn and the weights:
+
+```bash
+cd code
+timeout -s INT 1200 python -m pytest -q -s models/tests/test_api_device.py models/tests/test_warmup_device.py
+```
+
 ## Quickstart
 
 ```python
@@ -29,8 +52,12 @@ with Mast3rP150.from_pretrained(device_id=0) as model:
     out.save_ply("pointcloud.ply")
 ```
 
-`code/examples/quickstart.py` runs this snippet on the demo pair. It writes `pointmaps.npz`, `pointcloud.ply`,
-`depth.png` and `summary.json`.
+The paths `media/source_1.png` and `media/source_2.png` are relative to the repo root.
+Run the snippet from the repo root, or give absolute paths.
+
+`code/examples/quickstart.py` runs this snippet on the demo pair. It finds the demo images relative to its own
+file, so you can run it from any directory. It writes `pointmaps.npz`, `pointcloud.ply`, `depth.png` and
+`summary.json` into `--out-dir` (default `quickstart_out`, relative to the current directory).
 
 ## `Mast3rP150.from_pretrained(...)`
 
@@ -41,10 +68,13 @@ with Mast3rP150.from_pretrained(device_id=0) as model:
 | `dispatch` | `"auto"` | `auto`: if the tt-metal tree has the ETH-dispatch patch, use ETH dispatch and a 12x10 grid (the published configuration). If not, use the default Tensix dispatch and show a warning. `eth` or `worker` forces one mode. Env: `MAST3R_DISPATCH`. |
 | `repo_id`, `revision`, `weights_dir` | `naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt` at `61c57447` | The source of `model.safetensors`: the HF cache (downloaded on first use), or a local directory. Env: `HF_MODEL`, `TT_WEIGHTS_REVISION`, `MAST3R_WEIGHTS_DIR`. |
 | `preprocess` | `"pad"` | Makes each view square. `pad` adds gray padding. `crop` cuts a centre square. Env: `MAST3R_PREPROC`. |
-| `warmup` | `1` | Number of warm-up forwards. The first one compiles the kernels and captures the metal trace. |
-| `capture_pose` | `True` | Also captures the symmetric graph that `return_pose=True` uses. |
+| `warmup` | `1` | Number of warm-up forwards on a zero pair. The first one compiles the kernels and captures the metal trace. |
+| `warmup_variants` | `("pair", "pose", "batch")` | The call variants to prepare before `from_pretrained` returns, so that the first real call of each variant is fast. See [Warm-up](#warm-up). |
+| `capture_pose` | `None` | Legacy switch. `False` removes `"pose"` from `warmup_variants` (no symmetric graph). `True` adds it. |
+| `verbose` | `False` | `False`: tt-metal prints only errors and ttnn prints only warnings and errors. `True`: the full tt-metal log and this package's INFO log. Env: `MAST3R_VERBOSE=1`. See [Console output](#console-output). |
 
-Startup takes about 65 s with an empty kernel cache and about 7 s with a warm cache (weights in the HF cache).
+Startup takes about 67 s with an empty kernel cache and about 9 s with a warm cache (weights in the HF cache).
+Of the 9 s, the default warm-up takes about 7 s (device graphs about 6 s, host about 1 s).
 The device-graph knobs (`TT_FUSED`, `MAST3R_OPT`, `MAST3R_CQS`, ...) come from the environment, as in the server.
 Their defaults are the published configuration.
 You can open only one model at a time in a process. `model.info` shows the dispatch mode, the grid and the weights path.
@@ -106,6 +136,109 @@ output = model.inference(pairs)                           # DUSt3R inference() o
 `inference` returns `{'view1', 'view2', 'pred1', 'pred2', 'loss'}` with CPU torch tensors, in DUSt3R's layout.
 All views are 512x512 squares. Upstream DUSt3R keeps the aspect ratio, but this port does not.
 
+## Warm-up
+
+`from_pretrained` prepares the call variants in `warmup_variants` before it returns.
+Thus the first real call of each prepared variant is as fast as the later calls.
+The device graph is fixed (512x512, one pair), so the variants are few:
+
+| variant | call it prepares | device | host first-use costs it pays |
+|---|---|---|---|
+| `"pair"` | `model(v1, v2)`, `forward_raw` | pair graph: kernel compile, trace capture, persistent input | PNG and JPEG decode, pad resize, numpy-array input, host thread pool, readback and activation buffers, `point_cloud` |
+| `"pose"` | `model(v1, v2, return_pose=True)` | symmetric graph (trace capture) | OpenCV import (about 65 ms), focal vote, one PnP-RANSAC |
+| `"batch"` | `predict_pairs`, `inference`, `load_images` / `make_pairs` | none (uses the pair graph) | the pipelined submit / collect path with 2 and 4 pairs |
+| `"crop"` | `model(..., preprocess="crop")` | none (uses the pair graph) | the crop resize |
+
+The default is `("pair", "pose", "batch")`. `"all"` adds `"crop"`.
+A variant can also be a dict of call options: `{"return_pose": True, "preprocess": "crop"}` (keys `return_pose`, `batch`, `preprocess`).
+
+```python
+model = Mast3rP150.from_pretrained(warmup_variants="all")         # also prepare preprocess="crop"
+model = Mast3rP150.from_pretrained(warmup_variants=["pair"])      # faster start-up; no symmetric graph
+model.warmup(return_pose=True)        # prepare a variant later; returns {"pose": ms}
+model.warmup("pose")                  # again: {} (idempotent, no device work)
+model.warmed_variants                 # ['pair', 'pose']
+model.info["warmup_ms"]               # time of each variant warm-up
+```
+
+How the warm-up works:
+
+- The device graph of each variant is compiled and captured on a zero pair, as the server does.
+- Then each variant makes 2 complete calls through the public API with realistic dummy images.
+  The images are a 640x480 photo-like texture (view 1 as PNG bytes, then as a numpy array) and the same scene
+  moved by 24 px (view 2 as JPEG bytes). They have real image sizes, so the decoders, the resize, the readback
+  and the activations do the same work as for a photo.
+- The warm-up does not keep results. Each later call computes its outputs from its own inputs.
+  The outputs are bit-identical with and without the warm-up (27 arrays of 3 reference pairs, plain, pose and
+  `predict_pairs`, compared between the old and the new code; `models/tests/test_warmup_device.py` also compares them).
+- Without OpenCV, the `"pose"` warm-up captures the symmetric graph and skips the host PnP (an INFO log line tells you).
+
+Cost (warm kernel cache, weights in the HF cache): the default variants add about 1.0 s of host warm-up
+(`pair` 0.15 s, `pose` 0.7 s, `batch` 0.16 s) to the about 6 s device warm-up. `from_pretrained` takes about 9 s.
+The trace region and the DRAM use do not change: the pair graph and the symmetric graph were already captured
+at start-up before this API. `warmup_variants=["pair"]` does not capture the symmetric graph.
+Then the first `return_pose=True` call captures it and imports OpenCV (about 0.6 s more on that call).
+
+What the warm-up cannot prepare:
+
+- Input image size has no limit, but it changes only the host resize (no device shape, no compile).
+  A new image size costs nothing extra on its first use.
+- A variant that is not prepared pays its cost on its first call.
+  Measured without the host warm-up (`warmup_variants=()`, `test_warmup_device.py`): the first `pair` call took 74 ms
+  (later calls 44 ms). The first `return_pose=True` call took 870 ms (later calls 259 ms), because it captures the
+  symmetric graph and imports OpenCV.
+- Code outside this package: DUSt3R's `global_aligner`, your own OpenCV or matplotlib calls.
+- The first `from_pretrained` on a chip with an empty kernel cache compiles every kernel (about 67 s).
+
+### First-call measurements
+
+Fresh process per row; each call gets a NEW real pair (crops of the demo photos at different windows and sizes,
+read from PNG files), never the warm-up input. Median of 3 processes (8 for `pose` and `pairs`), Galaxy Blackhole chip,
+ETH dispatch, 12x10 grid, warm kernel cache. The host is shared with other jobs, so a single call varies by about 10 ms.
+"Same input" is the median of 3 later calls on the input of call 1. The PnP-RANSAC time depends on the data,
+so that is the fair reference for call 1. Script: `code/tools_prof/first_call_bench.py`.
+
+Before (HEAD `09c7d11`: device graphs warmed, host not warmed):
+
+| variant | `from_pretrained` | 1st call | 2nd | 10th | steady median (calls 11-30) | same input | 1st-call penalty |
+|---|---|---|---|---|---|---|---|
+| `model(path, path)` | 8.0 s | 46.6 ms | 41.0 | 43.6 | 46.8 | 46.7 | -0.1 ms |
+| `return_pose=True` | 7.9 s | 240.4 ms | 234.2 | 246.3 | 253.8 | 227.4 | +18.2 ms |
+| `predict_pairs` (4 pairs) | 8.1 s | 126.7 ms | 123.0 | 124.1 | 124.8 | 129.2 | -2.5 ms |
+| `load_images` + `inference` (3 views, 6 pairs) | 7.7 s | 219.1 ms | 227.4 | 220.5 | 225.9 | 225.9 | -6.6 ms |
+| `preprocess="crop"` | 7.5 s | 43.5 ms | 40.9 | 48.7 | 49.9 | 40.4 | +3.1 ms |
+| `model(ndarray, ndarray)` | 8.0 s | 53.3 ms | 42.8 | 52.6 | 51.6 | 55.5 | -1.4 ms |
+
+After (default `warmup_variants`):
+
+| variant | `from_pretrained` | 1st call | 2nd | 10th | steady median (calls 11-30) | same input | 1st call / steady median |
+|---|---|---|---|---|---|---|---|
+| `model(path, path)` | 9.2 s | 51.8 ms | 47.8 | 44.0 | 49.0 | 51.9 | 1.06 |
+| `return_pose=True` | 9.0 s | 240.9 ms | 241.0 | 254.3 | 255.6 | 230.8 | 0.94 |
+| `predict_pairs` (4 pairs) | 9.1 s | 132.4 ms | 126.8 | 122.2 | 124.2 | 128.2 | 1.07 |
+| `load_images` + `inference` (3 views, 6 pairs) | 9.3 s | 229.0 ms | 236.9 | 219.8 | 222.3 | 216.9 | 1.03 |
+| `preprocess="crop"` (not in the defaults) | 8.9 s | 55.7 ms | 46.2 | 55.0 | 50.7 | 53.8 | 1.10 |
+| `model(ndarray, ndarray)` | 8.7 s | 65.3 ms | 53.0 | 64.6 | 59.4 | 70.3 | 1.10 |
+
+The first call of every variant is within 10 % of the steady median. The steady state did not change:
+the device call (`timing_ms['forward']`, 30 calls) is 27.0 ms before and 26.9 / 27.1 ms after, in back-to-back runs.
+The pose first call has a timing breakdown that is the same as the later calls (focal vote about 15 ms, RANSAC 120-140 ms).
+The `test_warmup_device.py` run shows the effect in one process: the first call is 0.92x (`pair`), 1.08x (`pose`)
+and 0.99x (`batch`) of the same-input warm calls, against 1.70x, 3.36x and 1.10x without the host warm-up.
+
+## Console output
+
+`from_pretrained(verbose=False)` (the default) keeps the console readable.
+Before it imports ttnn, it sets `TT_LOGGER_LEVEL=error` (tt-metal's C++ log: no info lines such as
+`act_block_h_override ...`, `Opening user mode device driver`, `JIT cache stats`) and `LOGURU_LEVEL=WARNING`
+(ttnn's Python log: no `Initial ttnn.CONFIG` dump). Errors still print, and so do this package's warnings
+(for example, the ETH-dispatch fallback). It changes a variable only if you did not set it,
+and only if ttnn (or loguru) is not imported yet. The variables stay set in the process and its child processes.
+
+- `verbose=True` (or `MAST3R_VERBOSE=1`) keeps tt-metal's normal log and prints this package's INFO log
+  (weights, warm-up times, ready line) on stderr.
+- Set `TT_LOGGER_LEVEL` or `LOGURU_LEVEL` yourself for a different level.
+
 ## Lifetime
 
 Use `with Mast3rP150.from_pretrained() as model:` or call `model.close()`.

README.md

@@ -38,6 +38,7 @@ In that environment, download this repo and install its `code/` directory:
 ```bash
 hf download changh95/mast3r-p150 --exclude "image/*" --local-dir mast3r-p150 && cd mast3r-p150
 pip install -e "code/[pose]"      # the [pose] extra adds OpenCV for return_pose=True
+pip install -e "code/[pose,server,test]"   # also the HTTP server and the tests (see PYTHON.md)
 ```
 
 ```python
@@ -50,7 +51,12 @@ with Mast3rP150.from_pretrained(device_id=0) as model:
     out.save_ply("pointcloud.ply")
 ```
 
-The first call of `from_pretrained` downloads the NAVER weights to your HF cache. It also compiles the kernels and captures the metal trace (about 65 s). With a warm cache, startup takes about 7 s. The package does not install ttnn. You do not need the HTTP server.
+The paths `media/source_1.png` and `media/source_2.png` are relative to the repo root.
+Run the snippet from the repo root, or give absolute paths.
+
+The first call of `from_pretrained` downloads the NAVER weights to your HF cache. It also compiles the kernels and captures the metal trace (about 65 s). The package does not install ttnn. You do not need the HTTP server.
+
+`Mast3rP150.from_pretrained()` prepares the plain call, the pose call (`return_pose=True`) and `predict_pairs` before it returns. Thus the first call is as fast as the later calls (within 10 %). Start-up takes about 9 s with a warm kernel cache. Use `warmup_variants=["pair"]` for a faster start-up, or `model.warmup("crop")` to prepare more variants later.
 
 | | name | type / default | meaning |
 |---|---|---|---|

code/examples/quickstart.py

@@ -4,7 +4,8 @@
 
     python examples/quickstart.py [--out-dir quickstart_out] [--pose]
 
-Reads media/source_1.png + media/source_2.png, runs the pair on the chip and writes
+Reads media/source_1.png + media/source_2.png (found relative to this file, so any working directory
+works), runs the pair on the chip and writes (relative to the current directory)
 <out-dir>/pointmaps.npz (the server's npz format), <out-dir>/pointcloud.ply (fused coloured
 cloud) and <out-dir>/depth.png (inputs + view-1 depth + confidence).
 """

code/mast3r_p150/__init__.py

@@ -11,8 +11,9 @@
 See PYTHON.md for the full API. ``import mast3r_p150`` needs no device and does not import ttnn.
 """
 from .inputs import load_images, make_pairs, to_pil
-from .model import DEFAULT_REPO_ID, DEFAULT_REVISION, Mast3rP150, PairResult, Pose
+from .model import (DEFAULT_REPO_ID, DEFAULT_REVISION, DEFAULT_WARMUP_VARIANTS, WARMUP_PRESETS, Mast3rP150,
+                    PairResult, Pose)
 
 __version__ = "0.1.0"
 __all__ = ["Mast3rP150", "PairResult", "Pose", "load_images", "make_pairs", "to_pil",
-           "DEFAULT_REPO_ID", "DEFAULT_REVISION", "__version__"]
+           "DEFAULT_REPO_ID", "DEFAULT_REVISION", "DEFAULT_WARMUP_VARIANTS", "WARMUP_PRESETS", "__version__"]

code/mast3r_p150/model.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 import io
 import logging
 import os
+import sys
 import threading
 import time
 import warnings
@@ -36,7 +37,7 @@ from models.demos.mast3r.postprocess import (
     preprocess_mode,
 )
 
-from .inputs import ImageLike, View, prepare_view
+from .inputs import ImageLike, View, load_images, make_pairs, prepare_view
 
 log = logging.getLogger("mast3r_p150")
 
@@ -73,6 +74,98 @@ def _activate(raw: torch.Tensor):
     return activate_pts3d(raw)[0].numpy(), activate_conf(raw)[0].numpy()
 
 
+# ---------------------------------------------------------------- warm-up variants
+
+#: Named warm-up variants. A variant is the set of call options that changes what runs on the first
+#: call: ``return_pose`` (the symmetric device graph + OpenCV PnP), ``batch`` (the pipelined
+#: ``predict_pairs`` / ``inference`` path) and ``preprocess`` (``pad`` / ``crop``; ``None`` = the model's).
+WARMUP_PRESETS: Dict[str, Dict[str, Any]] = {
+    "pair": {"return_pose": False, "batch": False, "preprocess": None},
+    "pose": {"return_pose": True, "batch": False, "preprocess": None},
+    "batch": {"return_pose": False, "batch": True, "preprocess": None},
+    "crop": {"return_pose": False, "batch": False, "preprocess": "crop"},
+}
+#: ``from_pretrained(warmup_variants=...)`` default: the plain call, the pose call and the batch helpers.
+DEFAULT_WARMUP_VARIANTS: Tuple[str, ...] = ("pair", "pose", "batch")
+_VARIANT_KEYS = ("return_pose", "batch", "preprocess")
+
+
+def normalize_warmup_variants(spec) -> List[Dict[str, Any]]:
+    """``warmup_variants`` -> list of variant dicts ``{return_pose, batch, preprocess}`` (duplicates removed).
+
+    ``spec``: ``None`` / ``()`` (no variant), ``"all"``, one name of :data:`WARMUP_PRESETS`, one dict of
+    call options, or a list of names / dicts (e.g. ``["pair", {"return_pose": True, "preprocess": "crop"}]``).
+    """
+    if spec is None:
+        return []
+    if isinstance(spec, (str, dict)):
+        spec = [spec]
+    out: List[Dict[str, Any]] = []
+    for item in spec:
+        if isinstance(item, str):
+            if item == "all":
+                out.extend(normalize_warmup_variants(list(WARMUP_PRESETS)))
+                continue
+            if item not in WARMUP_PRESETS:
+                raise ValueError(f"unknown warm-up variant {item!r}; use one of {sorted(WARMUP_PRESETS)}, "
+                                 "'all', or a dict of {return_pose, batch, preprocess}")
+            v = dict(WARMUP_PRESETS[item])
+        elif isinstance(item, dict):
+            bad = set(item) - set(_VARIANT_KEYS)
+            if bad:
+                raise ValueError(f"unknown warm-up variant option(s) {sorted(bad)}; allowed: {_VARIANT_KEYS}")
+            v = {"return_pose": bool(item.get("return_pose", False)), "batch": bool(item.get("batch", False)),
+                 "preprocess": item.get("preprocess")}
+        else:
+            raise TypeError(f"a warm-up variant is a name or a dict, got {type(item).__name__}")
+        if v["preprocess"] not in (None, "pad", "crop"):
+            raise ValueError(f"warm-up preprocess must be 'pad', 'crop' or None, got {v['preprocess']!r}")
+        if v not in out:
+            out.append(v)
+    return out
+
+
+def variant_name(v: Dict[str, Any]) -> str:
+    """Short label of a normalized variant (``pair``, ``pose``, ``batch``, ``pair/crop``, ...)."""
+    name = "pose" if v["return_pose"] else "batch" if v["batch"] else "pair"
+    if v["return_pose"] and v["batch"]:
+        name = "pose+batch"
+    return name + (f"/{v['preprocess']}" if v.get("preprocess") else "")
+
+
+def warmup_images(seed: int = 0) -> Tuple[bytes, bytes, np.ndarray]:
+    """Realistic dummy inputs for the warm-up: a 640x480 photo-like texture (smooth random blobs plus
+    fine noise) as view 1 (PNG bytes) and the same scene shifted by 24 px as view 2 (JPEG bytes), plus
+    the view-1 pixels as a ``(480, 640, 3)`` uint8 array. Non-square, so ``pad`` and ``crop`` do real work;
+    the two decoders, the bicubic resize and the network see inputs of the same sizes as real photos."""
+    from PIL import Image
+    rng = np.random.default_rng(seed)
+    coarse = rng.integers(0, 256, size=(12, 17, 3), dtype=np.uint8)
+    base = np.asarray(Image.fromarray(coarse).resize((664, 480), Image.BICUBIC), dtype=np.int16)
+    base = np.clip(base + rng.integers(-12, 13, size=base.shape), 0, 255).astype(np.uint8)
+    v1, v2 = base[:, :640], base[:, 24:664]
+    b1, b2 = io.BytesIO(), io.BytesIO()
+    Image.fromarray(v1).save(b1, format="PNG")
+    Image.fromarray(v2).save(b2, format="JPEG", quality=92)
+    return b1.getvalue(), b2.getvalue(), np.ascontiguousarray(v1)
+
+
+def quiet_native_logs(verbose: bool) -> Dict[str, str]:
+    """Before ttnn is imported: unless ``verbose``, make tt-metal's C++ logger (``TT_LOGGER_LEVEL``) print
+    errors only and ttnn's loguru logger (``LOGURU_LEVEL``) warnings and up. Values the caller already set win
+    (``os.environ.setdefault``); nothing changes once ttnn (or loguru) is imported. Returns what was set."""
+    if verbose or "ttnn" in sys.modules:
+        return {}
+    done = {}
+    for k, v in (("TT_LOGGER_LEVEL", "error"), ("LOGURU_LEVEL", "WARNING")):
+        if k == "LOGURU_LEVEL" and "loguru" in sys.modules:
+            continue
+        if k not in os.environ:
+            os.environ[k] = v
+            done[k] = v
+    return done
+
+
 # ---------------------------------------------------------------- outputs
 
 @dataclass
@@ -245,6 +338,7 @@ class _Engine:
         self.traced = bool(fused.enabled and fused.trace)
         self.lock = threading.RLock()
         self._pending = None
+        self._eager_warm: set = set()
 
     def _inp(self, v: View) -> torch.Tensor:
         return v.pixels if self.uint8_input else v.float_input()
@@ -297,24 +391,52 @@ class _Engine:
             self._pending = None
             self.lock.release()
 
-    def warmup(self, runs: int, capture_pose: bool) -> None:
-        """The server's warm-up: ``runs`` forwards on a zero pair (the first compiles and captures),
-        a capturing forward when needed, and the symmetric (pose) graph."""
+    def _dummy(self) -> View:
         z = torch.zeros(1, 3, IMG_SIZE, IMG_SIZE, dtype=torch.uint8 if self.uint8_input else torch.float32)
         dummy = View(pixels=z, rgb=np.zeros((IMG_SIZE, IMG_SIZE, 3), np.uint8), record={})
         if not self.uint8_input:
             dummy._float = z
-        for i in range(max(runs, 1 if self.traced else 0)):
+        return dummy
+
+    def graph_ready(self, pose: bool) -> bool:
+        """True when the device graph of this variant is compiled (and, when traced, captured)."""
+        if pose and not self.sym_graph:
+            pose = False                     # pose = two replays of the pair graph
+        if not self.traced:
+            return ("pose" if pose else "pair") in self._eager_warm
+        m = self.port.get_model(self.state, self.device)
+        return (m.input_kind(self._dummy().pixels) + ("_sym" if pose else "")) in m._kinds
+
+    def ensure_graph(self, pose: bool, runs: int = 1) -> float:
+        """Compile and capture the device graph of one variant on a zero pair (the server's warm-up input)
+        unless it is ready; idempotent. ``runs`` warm forwards for the pair graph (the first compiles and
+        captures). Returns the time spent in ms."""
+        if self.graph_ready(pose) and runs <= 1:
+            return 0.0
+        t_start = time.perf_counter()
+        dummy = self._dummy()
+        if pose and self.sym_graph:
             t0 = time.perf_counter()
-            self.forward(dummy, dummy)
-            log.info("warm-up forward %d: %.0f ms", i + 1, (time.perf_counter() - t0) * 1e3)
-        if self.traced:
-            if not self.port.get_model(self.state, self.device).trace_captured:
-                raise RuntimeError("the warm-up did not capture a metal trace")
-            if capture_pose and self.sym_graph:
+            self.forward_pose(dummy, dummy)
+            log.info("symmetric (pose) graph captured: %.0f ms", (time.perf_counter() - t0) * 1e3)
+        else:
+            for i in range(max(runs, 1)):
                 t0 = time.perf_counter()
-                self.forward_pose(dummy, dummy)
-                log.info("symmetric (pose) graph captured: %.0f ms", (time.perf_counter() - t0) * 1e3)
+                self.forward(dummy, dummy)
+                log.info("warm-up forward %d: %.0f ms", i + 1, (time.perf_counter() - t0) * 1e3)
+            if self.traced and not self.port.get_model(self.state, self.device).trace_captured:
+                raise RuntimeError("the warm-up did not capture a metal trace")
+        if not self.traced:
+            self._eager_warm.add("pose" if (pose and self.sym_graph) else "pair")
+        return (time.perf_counter() - t_start) * 1e3
+
+    def warmup(self, runs: int, capture_pose: bool) -> None:
+        """The server's warm-up: ``runs`` forwards on a zero pair (the first compiles and captures),
+        a capturing forward when needed, and the symmetric (pose) graph."""
+        if runs > 0 or self.traced:
+            self.ensure_graph(False, runs=runs)
+        if capture_pose and self.traced and self.sym_graph:
+            self.ensure_graph(True)
 
     def close(self) -> None:
         with self.lock:
@@ -355,6 +477,8 @@ class Mast3rP150:
     def __init__(self, engine, preprocess: str = "pad"):
         self._engine = engine
         self.preprocess = preprocess
+        self._warmed: List[Dict[str, Any]] = []     # normalized variants already warmed (host + device)
+        self._warm_lock = threading.Lock()
         self._finalizer = weakref.finalize(self, _finalize, engine.close) if engine is not None else None
         Mast3rP150._open.add(self)
 
@@ -363,7 +487,8 @@ class Mast3rP150:
     def from_pretrained(cls, device_id: int = 0, *, device=None, dispatch: Optional[str] = None,
                         repo_id: Optional[str] = None, revision: Optional[str] = None,
                         weights_dir: Optional[str] = None, preprocess: Optional[str] = None,
-                        warmup: int = 1, capture_pose: bool = True) -> "Mast3rP150":
+                        warmup: int = 1, capture_pose: Optional[bool] = None,
+                        warmup_variants=DEFAULT_WARMUP_VARIANTS, verbose: Optional[bool] = None) -> "Mast3rP150":
         """Load the weights, open the chip, build the device graph and capture its traces.
 
         ``device_id``     chip to open (``ttnn.open_device(device_id=...)``).
@@ -378,8 +503,20 @@ class Mast3rP150:
                           (download on first use) at the pinned commit, or a local directory.
                           Env: ``HF_MODEL``, ``TT_WEIGHTS_REVISION``, ``MAST3R_WEIGHTS_DIR``.
         ``preprocess``    default view preprocessing, ``"pad"`` or ``"crop"`` (``$MAST3R_PREPROC``).
-        ``warmup``        warm-up forwards (the first compiles the kernels and captures the trace).
-        ``capture_pose``  also capture the symmetric graph that ``return_pose=True`` uses.
+        ``warmup``        warm-up forwards on a zero pair (the first compiles the kernels and captures the trace).
+        ``warmup_variants``  the call variants to prepare before this returns, so the FIRST real call of each
+                          is as fast as the later ones: device graph compiled and captured, plus the host
+                          first-use costs (image decoders, resize, thread pool, OpenCV import and PnP,
+                          readback and activation buffers), run on realistic dummy images. Default
+                          ``("pair", "pose", "batch")``; also ``"crop"``, ``"all"``, dicts of
+                          ``{return_pose, batch, preprocess}``, or ``()`` for the device graph only.
+                          :meth:`warmup` adds variants later. Outputs do not change.
+        ``capture_pose``  legacy switch: ``False`` drops ``"pose"`` from ``warmup_variants`` (no symmetric
+                          graph is captured), ``True`` adds it.
+        ``verbose``       ``False`` (default; ``$MAST3R_VERBOSE=1`` sets ``True``): tt-metal's C++ log prints
+                          errors only and ttnn's loguru log warnings and up (``TT_LOGGER_LEVEL`` /
+                          ``LOGURU_LEVEL``, set only if unset and only before ttnn is first imported).
+                          ``True``: tt-metal's normal log plus this package's INFO log on stderr.
 
         Device-graph knobs (``TT_FUSED``, ``MAST3R_OPT``, ``MAST3R_CQS``, ...) are read from the
         environment once, as in the server; the defaults are the published configuration.
@@ -392,6 +529,19 @@ class Mast3rP150:
         mode = preprocess or preprocess_mode()
         if mode not in ("pad", "crop"):
             raise ValueError(f"preprocess must be 'pad' or 'crop', got {mode!r}")
+        variants = normalize_warmup_variants(warmup_variants)
+        if capture_pose is False:
+            variants = [v for v in variants if not v["return_pose"]]
+        elif capture_pose is True and not any(v["return_pose"] for v in variants):
+            variants += normalize_warmup_variants("pose")
+        if verbose is None:
+            verbose = os.environ.get("MAST3R_VERBOSE", "0") not in ("", "0", "false", "False")
+        quiet_native_logs(verbose)
+        if verbose and not log.handlers:
+            h = logging.StreamHandler()
+            h.setFormatter(logging.Formatter("%(asctime)s mast3r_p150 %(levelname)s: %(message)s"))
+            log.addHandler(h)
+            log.setLevel(logging.INFO)
         torch.set_grad_enabled(False)
         t0 = time.perf_counter()
         ckpt = resolve_checkpoint_path(repo_id=repo_id or os.environ.get("HF_MODEL") or DEFAULT_REPO_ID,
@@ -414,14 +564,24 @@ class Mast3rP150:
         info.update(weights=str(ckpt), fused_opts=list(fused.opts), num_cqs=fused.num_cqs,
                     traced=bool(fused.enabled and fused.trace))
         engine = _Engine(state, device, port, ttnn, owns, info)
+        model = None
         try:
-            engine.warmup(warmup, capture_pose)
+            t1 = time.perf_counter()
+            engine.warmup(warmup, any(v["return_pose"] for v in variants))
+            info["warmup_device_s"] = round(time.perf_counter() - t1, 2)
+            model = cls(engine, mode)
+            t1 = time.perf_counter()
+            model.warmup(variants)
+            info["warmup_host_s"] = round(time.perf_counter() - t1, 2)
         except BaseException:
-            engine.close()
+            if model is not None:
+                model.close()
+            else:
+                engine.close()
             raise
         info["startup_s"] = round(time.perf_counter() - t0, 1)
         log.info("ready: dispatch %s, grid %s, %.1f s", info["dispatch"], info["grid"], info["startup_s"])
-        return cls(engine, mode)
+        return model
 
     @property
     def info(self) -> dict:
@@ -453,6 +613,80 @@ class Mast3rP150:
             raise RuntimeError("this Mast3rP150 is closed")
         return self._engine
 
+    # ------------------------------------------------------------ warm-up
+    @property
+    def warmed_variants(self) -> List[str]:
+        """Labels of the variants prepared so far (``pair``, ``pose``, ``batch``, ``pair/crop``, ...)."""
+        return [variant_name(v) for v in self._warmed]
+
+    def warmup(self, variants=None, *, runs: int = 2, **options) -> Dict[str, float]:
+        """Prepare call variants so that their first real call is fast; idempotent.
+
+            model.warmup()                              # the defaults: "pair", "pose", "batch"
+            model.warmup("crop")                        # a preset name (see WARMUP_PRESETS), or "all"
+            model.warmup(return_pose=True, preprocess="crop")   # the options of one variant
+            model.warmup(["pair", {"batch": True}])     # several
+
+        For each variant not prepared yet: compile and capture its device graph (zero pair, as the
+        server does), then make ``runs`` complete calls through the public path on realistic dummy
+        images (:func:`warmup_images`: PNG + JPEG bytes, 640x480) so the host first-use costs are paid
+        here: PIL decoders, the pad/crop resize, the host thread pool, the readback and activation
+        buffers, and for ``return_pose`` the OpenCV import and one PnP-RANSAC. Nothing is cached:
+        a later call computes its outputs from its own inputs exactly as without the warm-up.
+        Returns ``{variant label: ms}`` for the variants prepared by this call (empty when all were ready).
+        """
+        if options:
+            if variants is not None:
+                raise TypeError("pass either variants or keyword options, not both")
+            variants = dict(options)
+        todo = normalize_warmup_variants(variants if variants is not None else DEFAULT_WARMUP_VARIANTS)
+        for v in todo:                       # explicit preprocess == the model's default -> same variant
+            if v["preprocess"] == self.preprocess:
+                v["preprocess"] = None
+        done: Dict[str, float] = {}
+        imgs = None
+        with self._warm_lock:
+            eng = self._eng()
+            for v in todo:
+                if v in self._warmed:
+                    continue
+                t0 = time.perf_counter()
+                eng.ensure_graph(v["return_pose"])
+                if imgs is None:
+                    imgs = warmup_images()
+                self._warm_host(v, max(int(runs), 0), imgs)
+                self._warmed.append(v)
+                done[variant_name(v)] = round((time.perf_counter() - t0) * 1e3, 1)
+                log.info("warm-up %s: %.0f ms", variant_name(v), done[variant_name(v)])
+        if done and self._engine is not None:
+            self._engine.info.setdefault("warmup_ms", {}).update(done)
+        return done
+
+    def _warm_host(self, v: Dict[str, Any], runs: int, imgs) -> None:
+        png, jpg, arr = imgs
+        mode = v["preprocess"] or self.preprocess
+        pose = v["return_pose"]
+        if pose:
+            try:
+                import cv2  # noqa: F401  (the PnP-RANSAC of the pose path imports it lazily)
+            except ImportError:
+                log.info("return_pose warm-up: OpenCV is not installed (pip install opencv-python-headless); "
+                            "the symmetric device graph is ready, the host PnP is not warmed")
+                pose = False
+        for i in range(runs):            # run 0: encoded PNG + JPEG; run 1: a numpy array + JPEG
+            a = png if i % 2 == 0 else arr
+            if v["batch"]:
+                if i == 0:
+                    views = load_images([png, jpg], preprocess=mode)
+                    self.inference(make_pairs(views), preprocess=mode)          # 2 pairs, pipelined
+                else:                    # 4 pairs: the pipeline's steady state (prepare i+1 + finish i-1)
+                    self.predict_pairs([(a, jpg), (jpg, a), (png, jpg), (jpg, png)], return_pose=pose,
+                                       preprocess=mode)
+            else:
+                r = self(a, jpg, return_pose=pose, preprocess=mode)
+                if i == 0:
+                    r.point_cloud()
+
     # ------------------------------------------------------------ inference
     def __call__(self, view1: ImageLike, view2: ImageLike, *, return_pose: bool = False,
                  intrinsics: Optional[Sequence[Sequence[float]]] = None, conf_pct: float = 50.0,

code/models/tests/test_api_host.py

@@ -168,6 +168,10 @@ class FakeEngine:
         self._pending = None
         return self.forward(v1, v2)
 
+    def ensure_graph(self, pose, runs=1):
+        self.calls.append(("graph", pose))
+        return 0.0
+
     def close(self):
         self.closed = True
         self.device = None
@@ -308,3 +312,98 @@ def test_lifetime():
     with pytest.raises(RuntimeError, match="closed"):
         m(IMG1, IMG2)
     m.close()                                            # idempotent
+
+
+# ---------------------------------------------------------------- warm-up API
+
+def test_warmup_variant_spec():
+    from mast3r_p150.model import DEFAULT_WARMUP_VARIANTS, WARMUP_PRESETS, normalize_warmup_variants, variant_name
+    assert DEFAULT_WARMUP_VARIANTS == ("pair", "pose", "batch") == mast3r_p150.DEFAULT_WARMUP_VARIANTS
+    assert [variant_name(v) for v in normalize_warmup_variants(DEFAULT_WARMUP_VARIANTS)] == ["pair", "pose", "batch"]
+    assert [variant_name(v) for v in normalize_warmup_variants("all")] == ["pair", "pose", "batch", "pair/crop"]
+    assert normalize_warmup_variants(None) == [] and normalize_warmup_variants(()) == []
+    assert normalize_warmup_variants("pose") == [WARMUP_PRESETS["pose"]]
+    assert normalize_warmup_variants({"return_pose": True, "preprocess": "crop"}) == \
+        [{"return_pose": True, "batch": False, "preprocess": "crop"}]
+    assert len(normalize_warmup_variants(["pair", {"return_pose": False}, "pair"])) == 1     # duplicates removed
+    with pytest.raises(ValueError, match="unknown warm-up variant"):
+        normalize_warmup_variants("huge")
+    with pytest.raises(ValueError, match="option"):
+        normalize_warmup_variants({"num_views": 3})
+    with pytest.raises(ValueError, match="preprocess"):
+        normalize_warmup_variants({"preprocess": "native"})
+    with pytest.raises(TypeError):
+        normalize_warmup_variants([3])
+
+
+def test_warmup_images_are_realistic():
+    from mast3r_p150.model import warmup_images
+    png, jpg, arr = warmup_images()
+    a, b = Image.open(io.BytesIO(png)), Image.open(io.BytesIO(jpg))
+    assert a.format == "PNG" and b.format == "JPEG" and a.size == b.size == (640, 480)
+    assert arr.shape == (480, 640, 3) and arr.dtype == np.uint8 and arr.std() > 20     # textured, not constant
+    assert np.array_equal(np.asarray(a.convert("RGB")), arr)
+    png2, _, _ = warmup_images()
+    assert png2 == png                                                               # deterministic
+
+
+def test_from_pretrained_signature():
+    import inspect
+    sig = inspect.signature(Mast3rP150.from_pretrained)
+    assert sig.parameters["warmup_variants"].default == ("pair", "pose", "batch")
+    assert sig.parameters["verbose"].default is None
+    assert sig.parameters["capture_pose"].default is None
+    assert "runs" in inspect.signature(Mast3rP150.warmup).parameters
+
+
+def test_warmup_is_idempotent_and_does_not_change_outputs():
+    eng = FakeEngine()
+    m = Mast3rP150(eng)
+    before = m(IMG1, IMG2)
+    n0 = len(eng.calls)
+    done = m.warmup(["pair", "batch"])
+    assert list(done) == ["pair", "batch"] and all(t >= 0 for t in done.values())
+    assert m.warmed_variants == ["pair", "batch"]
+    assert ("graph", False) in eng.calls[n0:]
+    fwd = [c for c in eng.calls[n0:] if c[0] == "fwd"]
+    assert len(fwd) >= 4                                         # real calls through the public path
+    # the warm-up input is the textured dummy, not a cached real input
+    real_px = eng.calls[0][1]
+    assert not any(torch.equal(c[1], real_px) for c in fwd)
+    n1 = len(eng.calls)
+    assert m.warmup(["pair", "batch"]) == {} and m.warmup("pair") == {} and m.warmup(return_pose=False) == {}
+    assert len(eng.calls) == n1                                  # idempotent: no device work
+    assert m.warmup(preprocess="pad") == {}                      # pad == the model default == "pair"
+    after = m(IMG1, IMG2)
+    for k in ("pts3d1", "pts3d2", "conf1", "conf2", "rgb1", "rgb2"):
+        assert np.array_equal(getattr(before, k), getattr(after, k)), k
+    assert m.info["warmup_ms"].keys() == {"pair", "batch"}
+    assert list(m.warmup("crop")) == ["pair/crop"]
+    with pytest.raises(TypeError):
+        m.warmup("pair", return_pose=True)
+    m.close()
+    with pytest.raises(RuntimeError, match="closed"):
+        m.warmup("pose")
+
+
+def test_warmup_pose_variant():
+    pytest.importorskip("cv2")
+    eng = FakeEngine()
+    m = Mast3rP150(eng)
+    assert list(m.warmup(return_pose=True)) == ["pose"]
+    assert ("graph", True) in eng.calls
+    import sys as _sys
+    assert "cv2" in _sys.modules
+    m.close()
+
+
+def test_quiet_native_logs():
+    code = ("import os, sys; from mast3r_p150.model import quiet_native_logs as q; "
+            "os.environ.pop('TT_LOGGER_LEVEL', None); os.environ.pop('LOGURU_LEVEL', None); "
+            "assert q(True) == {} and 'TT_LOGGER_LEVEL' not in os.environ; "
+            "assert q(False) == {'TT_LOGGER_LEVEL': 'error', 'LOGURU_LEVEL': 'WARNING'}, os.environ; "
+            "os.environ['TT_LOGGER_LEVEL'] = 'debug'; os.environ.pop('LOGURU_LEVEL'); "
+            "assert q(False) == {'LOGURU_LEVEL': 'WARNING'} and os.environ['TT_LOGGER_LEVEL'] == 'debug'; "
+            "sys.modules['ttnn'] = object(); os.environ.pop('LOGURU_LEVEL'); assert q(False) == {}")
+    env = dict(os.environ, TT_VISIBLE_DEVICES="none", PYTHONPATH=_CODE_ROOT)
+    subprocess.run([sys.executable, "-c", code], check=True, env=env, cwd="/")

code/models/tests/test_warmup_device.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Device test of the warm-up API: the FIRST real call after ``from_pretrained`` is already fast
+(needs a Blackhole chip, ttnn and the weights).
+
+    cd code && timeout -s INT 1200 python -m pytest -q -s models/tests/test_warmup_device.py
+
+1. A model opened with ``warmup_variants=()`` (device graph of the plain pair only) runs reference inputs:
+   the outputs, and the cold first-call cost of each variant (reported, not gated).
+2. A fresh model with the default ``warmup_variants`` (``pair``, ``pose``, ``batch``): the first call of each
+   variant on a NEW real input (crops of the demo photos, never the warm-up input) must be within
+   ``MAST3R_FIRST_TOL`` (10 %) + ``MAST3R_FIRST_SLACK_MS`` (5 ms, shared-host noise) of the median of 3 later
+   calls on the SAME input (the PnP-RANSAC time depends on the data, so the same input is the fair reference).
+3. The outputs of 2 are bit-identical to those of 1 (warm-up changes nothing), and ``warmup()`` is idempotent.
+``MAST3R_WARMUP_REPORT=<path>`` writes the numbers as JSON.
+"""
+from __future__ import annotations
+
+import json
+import os
+import statistics
+import sys
+import time
+
+import numpy as np
+import pytest
+
+_CODE_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+if _CODE_ROOT not in sys.path:
+    sys.path.insert(0, _CODE_ROOT)
+
+ttnn = pytest.importorskip("ttnn")
+pytest.importorskip("cv2")
+
+from PIL import Image  # noqa: E402
+
+from mast3r_p150 import Mast3rP150  # noqa: E402
+
+_REPO = os.path.dirname(_CODE_ROOT)
+TOL = float(os.environ.get("MAST3R_FIRST_TOL", "0.10"))
+SLACK_MS = float(os.environ.get("MAST3R_FIRST_SLACK_MS", "5"))
+REPEATS = 3
+
+
+def _crops(tmp_path, n=6):
+    """n real pairs: the same window cut from both demo photos, written as PNG files."""
+    a0 = Image.open(os.path.join(_REPO, "media", "source_1.png")).convert("RGB")
+    b0 = Image.open(os.path.join(_REPO, "media", "source_2.png")).convert("RGB")
+    W, H = a0.size
+    out = []
+    for i in range(n):
+        f = 0.6 + 0.05 * i
+        w, h = int(W * f), int(H * (0.95 - 0.05 * i))
+        x, y = (W - w) * (i % 3) // 2, (H - h) * ((i + 1) % 3) // 2
+        pa, pb = tmp_path / f"r{i}_1.png", tmp_path / f"r{i}_2.png"
+        a0.crop((x, y, x + w, y + h)).save(pa)
+        b0.crop((x, y, x + w, y + h)).save(pb)
+        out.append((str(pa), str(pb)))
+    return out
+
+
+def _variants(pairs):
+    """name -> callable(model) running one real call of that variant (each on its own new input)."""
+    return {
+        "pair": lambda m: m(*pairs[0]),
+        "pose": lambda m: m(*pairs[1], return_pose=True),
+        "batch": lambda m: m.predict_pairs([pairs[2], pairs[3]]),
+    }
+
+
+def _arrays(res):
+    rs = res if isinstance(res, list) else [res]
+    out = []
+    for r in rs:
+        out += [r.raw1.numpy(), r.raw2.numpy(), r.pts3d1, r.conf2]
+        if r.pose is not None:
+            out += [r.pose.R, r.pose.t]
+    return out
+
+
+def _timed(fn, m):
+    t0 = time.perf_counter()
+    r = fn(m)
+    return (time.perf_counter() - t0) * 1e3, r
+
+
+def test_first_call_is_fast(tmp_path):
+    pairs = _crops(tmp_path)
+    calls = _variants(pairs)
+    report = {"tol": TOL, "slack_ms": SLACK_MS}
+
+    # 1. reference outputs + the cold first-call cost without the host/pose warm-up
+    t0 = time.perf_counter()
+    m = Mast3rP150.from_pretrained(device_id=0, warmup_variants=())
+    report["startup_no_variants_s"] = round(time.perf_counter() - t0, 2)
+    ref, cold = {}, {}
+    try:
+        for name, fn in calls.items():
+            first, r = _timed(fn, m)
+            ref[name] = _arrays(r)
+            same = statistics.median(_timed(fn, m)[0] for _ in range(REPEATS))
+            cold[name] = {"first_ms": round(first, 1), "same_input_warm_ms": round(same, 1)}
+    finally:
+        m.close()
+    report["cold"] = cold
+
+    # 2. default warm-up: the first real call of each variant is fast
+    t0 = time.perf_counter()
+    m = Mast3rP150.from_pretrained(device_id=0)
+    report["startup_default_s"] = round(time.perf_counter() - t0, 2)
+    report["warmup_ms"] = m.info.get("warmup_ms")
+    warm, failures = {}, []
+    try:
+        assert sorted(m.warmed_variants) == ["batch", "pair", "pose"]
+        for name, fn in calls.items():
+            first, r = _timed(fn, m)
+            got = _arrays(r)
+            same = statistics.median(_timed(fn, m)[0] for _ in range(REPEATS))
+            warm[name] = {"first_ms": round(first, 1), "same_input_warm_ms": round(same, 1),
+                          "ratio": round(first / same, 3)}
+            # 3. bit-identical to the model without the warm-up
+            assert len(got) == len(ref[name]), name
+            for a, b in zip(got, ref[name]):
+                assert np.array_equal(a, b), f"{name}: outputs differ from the un-warmed model"
+            if first > same * (1.0 + TOL) + SLACK_MS:
+                failures.append(f"{name}: first {first:.1f} ms vs warm {same:.1f} ms")
+        t0 = time.perf_counter()
+        assert m.warmup() == {}                     # idempotent: everything is ready
+        assert (time.perf_counter() - t0) < 0.05
+    finally:
+        m.close()
+    report["warm"] = warm
+    print("\n# warm-up report: " + json.dumps(report, indent=1))
+    path = os.environ.get("MAST3R_WARMUP_REPORT")
+    if path:
+        with open(path, "w") as f:
+            json.dump(report, f, indent=1)
+    assert not failures, "first real call slower than the warm calls: " + "; ".join(failures)

code/pyproject.toml

@@ -22,7 +22,7 @@ dependencies = [
 pose = ["opencv-python-headless>=4.9,<4.12"]     # return_pose=True (PnP-RANSAC)
 viz = ["matplotlib"]                             # examples/quickstart.py figure
 server = ["fastapi", "uvicorn", "pydantic>=2"]   # models/server/app.py
-test = ["pytest"]
+test = ["pytest", "httpx"]                    # host tests (fastapi TestClient needs httpx); add [server,pose] too
 
 [project.urls]
 "Model card" = "https://huggingface.co/changh95/mast3r-p150"

code/tools_prof/first_call_bench.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""First-call latency of the Python API, one variant per FRESH process (needs a chip).
+
+    python tools_prof/first_call_bench.py --variant pair  [--json out.json] [--dump out.npz] [--from-kw '{...}']
+
+Variants (the request shapes a user hits): ``pair`` = model(path, path); ``pose`` = model(..., return_pose=True);
+``pairs`` = model.predict_pairs(4 pairs); ``inference`` = load_images(3 views) + make_pairs + model.inference;
+``crop`` = model(..., preprocess="crop"); ``array`` = model(np.ndarray, np.ndarray).
+
+Times ``import mast3r_p150``, ``from_pretrained(**from_kw)``, then calls 1..10 and ``--steady`` more calls. EVERY call
+gets a NEW real input pair (crops of the two demo photos at different windows/sizes, written once as PNG files under
+``--img-dir``), never the warm-up input. At the end the inputs of call 1 run 3 more times (warm): the
+first-call penalty is call 1 minus the median of those (same input, so the data-dependent PnP-RANSAC time cancels). ``--dump`` stores the outputs of fixed reference pairs (demo pair + 2 crops,
+plain and with pose) for a bit-identity check between code versions.
+"""
+import argparse
+import json
+import os
+import statistics
+import sys
+import time
+from pathlib import Path
+
+T_PROC = time.perf_counter()
+CODE = Path(__file__).resolve().parents[1]
+REPO = CODE.parent
+sys.path.insert(0, str(CODE))
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--variant", required=True, choices=["pair", "pose", "pairs", "inference", "crop", "array"])
+ap.add_argument("--calls", type=int, default=10)
+ap.add_argument("--steady", type=int, default=20)
+ap.add_argument("--img-dir", default=os.environ.get("MAST3R_BENCH_IMGS", "/tmp/mast3r_first_call_imgs"))
+ap.add_argument("--from-kw", default="{}", help="extra from_pretrained kwargs (JSON)")
+ap.add_argument("--json", default=None)
+ap.add_argument("--dump", default=None)
+args = ap.parse_args()
+
+
+def make_images(d: Path, n: int):
+    """n crop pairs of the demo photos (same window in both views), saved as PNG; deterministic."""
+    import numpy as np
+    from PIL import Image
+    d.mkdir(parents=True, exist_ok=True)
+    out = []
+    a0, b0 = Image.open(REPO / "media" / "source_1.png").convert("RGB"), Image.open(REPO / "media" / "source_2.png").convert("RGB")
+    rng = np.random.default_rng(1234)
+    for i in range(n):
+        pa, pb = d / f"c{i:03d}_1.png", d / f"c{i:03d}_2.png"
+        if not (pa.exists() and pb.exists()):
+            W, H = a0.size
+            f = rng.uniform(0.55, 0.95)
+            w, h = int(W * f), int(H * rng.uniform(0.55, 0.95))
+            x, y = int(rng.uniform(0, W - w)), int(rng.uniform(0, H - h))
+            flip = bool(rng.integers(2))
+            for im, p in ((a0, pa), (b0, pb)):
+                c = im.crop((x, y, x + w, y + h))
+                if flip:
+                    c = c.transpose(Image.FLIP_LEFT_RIGHT)
+                c.save(p)
+        out.append((str(pa), str(pb)))
+    return out
+
+
+need = (args.calls + args.steady) * (4 if args.variant == "pairs" else 2 if args.variant == "inference" else 1) + 4
+IMGS = make_images(Path(args.img_dir), need)
+REF = [(str(REPO / "media" / "source_1.png"), str(REPO / "media" / "source_2.png")), IMGS[-1], IMGS[-2]]
+POOL = IMGS[:-4]
+
+t0 = time.perf_counter()
+import mast3r_p150  # noqa: E402
+from mast3r_p150 import Mast3rP150  # noqa: E402
+t_import = time.perf_counter() - t0
+from_kw = json.loads(args.from_kw)
+t0 = time.perf_counter()
+model = Mast3rP150.from_pretrained(device_id=0, **from_kw)
+t_from = time.perf_counter() - t0
+
+it = iter(POOL)
+_first = []          # the inputs of call 1, re-run warm at the end (same-input reference for the first call)
+
+
+def one_call():
+    v = args.variant
+    if not _first:
+        need1 = {"pairs": 4, "inference": 2}.get(v, 1)
+        _first.extend(POOL[:need1])
+    if v == "pair":
+        a, b = next(it)
+        return model(a, b)
+    if v == "crop":
+        a, b = next(it)
+        return model(a, b, preprocess="crop")
+    if v == "array":
+        import numpy as np
+        from PIL import Image
+        a, b = next(it)
+        return model(np.asarray(Image.open(a).convert("RGB")), np.asarray(Image.open(b).convert("RGB")))
+    if v == "pose":
+        a, b = next(it)
+        return model(a, b, return_pose=True)
+    if v == "pairs":
+        return model.predict_pairs([next(it) for _ in range(4)])
+    if v == "inference":
+        p1, p2 = next(it), next(it)
+        views = mast3r_p150.load_images([p1[0], p1[1], p2[0]])
+        return model.inference(mast3r_p150.make_pairs(views))
+
+
+lat, fwd = [], []
+for i in range(args.calls + args.steady):
+    t0 = time.perf_counter()
+    r = one_call()
+    lat.append((time.perf_counter() - t0) * 1e3)
+    if hasattr(r, "timing_ms"):
+        fwd.append(r.timing_ms.get("forward"))
+it = iter(_first * 3)
+same = []
+for _ in range(3):
+    t0 = time.perf_counter()
+    one_call()
+    same.append((time.perf_counter() - t0) * 1e3)
+steady = lat[args.calls:]
+med = statistics.median(steady) if steady else float("nan")
+res = {"variant": args.variant, "from_kw": from_kw, "import_s": round(t_import, 2), "from_pretrained_s": round(t_from, 2),
+       "first_ms": round(lat[0], 1), "second_ms": round(lat[1], 1), "tenth_ms": round(lat[min(9, len(lat) - 1)], 1),
+       "steady_median_ms": round(med, 1), "steady_min_ms": round(min(steady), 1) if steady else None,
+       "first_over_median": round(lat[0] / med, 3),
+       "same_input_warm_ms": round(statistics.median(same), 1),
+       "first_penalty_ms": round(lat[0] - statistics.median(same), 1), "first_forward_ms": fwd[0] if fwd else None,
+       "lat_ms": [round(x, 1) for x in lat], "info": {k: v for k, v in model.info.items() if k != "open_kwargs"}}
+
+if args.dump:
+    import numpy as np
+    arrs = {}
+    for k, (a, b) in enumerate(REF):
+        r = model(a, b)
+        arrs[f"r{k}_raw1"], arrs[f"r{k}_raw2"] = r.raw1.numpy(), r.raw2.numpy()
+        arrs[f"r{k}_pts1"], arrs[f"r{k}_conf2"] = r.pts3d1, r.conf2
+        rp = model(a, b, return_pose=True)
+        arrs[f"r{k}_pose_raw1"], arrs[f"r{k}_pose_pts2"] = rp.raw1.numpy(), rp.pts3d2
+        if rp.pose is not None:
+            arrs[f"r{k}_pose_R"], arrs[f"r{k}_pose_t"] = rp.pose.R, rp.pose.t
+    out = model.predict_pairs(REF)
+    for k, r in enumerate(out):
+        arrs[f"pp{k}_raw2"] = r.raw2.numpy()
+    np.savez(args.dump, **arrs)
+model.close()
+print("RESULT " + json.dumps({k: v for k, v in res.items() if k != "lat_ms"}))
+if args.json:
+    Path(args.json).write_text(json.dumps(res, indent=1, default=str))