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.

README.md

@@ -37,8 +37,11 @@ Prerequisite: a Python environment with tt-metal's `ttnn`, built at tt-metal [`8
 ```bash
 hf download changh95/sam3-p150 --exclude "image/*" --local-dir sam3-p150 && cd sam3-p150
 pip install -e code/      # into a Python env that already has tt-metal's ttnn
+pip install -e "code/[server,test]"     # also the HTTP server (fastapi, uvicorn, ...) and the tests (pytest, httpx)
 ```
 
+Run the Python example from the model repo root. The demo image path `media/source_1.png` is relative to the repo root. You can also give an absolute path.
+
 ```python
 from tt_sam3 import Sam3
 
@@ -63,6 +66,7 @@ result.scores          # (N,) float32, high to low
 | Output | `presence`, `width`, `height` | Probability that the concept is in the image, and the original image size. |
 
 - `from_pretrained()` opens the chip, loads the weights and captures the device traces. The first run compiles the kernels and can take some minutes.
+- Write `model = Sam3.from_pretrained(warmup_variants={"prompts": ["car", "person"]})` to prepare the model and the text of your prompts. The first call is then as fast as later calls. A prompt that you did not list adds about 80 to 175 ms to the first call that uses it.
 - Use `with ... as model:` or call `model.close()` to release the chip.
 - Give a list of prompts to segment one image with many prompts. The image part runs one time, and each prompt adds about 20 ms.
 - The model keeps the last image on the device. When the same image comes again, the call takes about 21 ms.

code/PYTHON.md

@@ -6,24 +6,44 @@ scores for each instance.
 
 ## Install
 
-Install the package into a Python environment that already has tt-metal's `ttnn` (and its `torch`):
+Install the package into a Python environment that already has tt-metal's `ttnn` (and its `torch`). Run the commands
+from the model repo root:
 
 ```bash
-pip install -e code/          # from the model repo root; or: pip install code/
+pip install -e code/                    # the Python API only
+pip install -e "code/[server,test]"     # also the HTTP server (fastapi, uvicorn, ...) and the tests (pytest, httpx)
 ```
 
 The package does not declare `ttnn` or `torch`, so pip does not replace the tt-metal build of torch.
 
+| you want to run | extras you need |
+|---|---|
+| `from tt_sam3 import Sam3` | none |
+| the HTTP server (`python -m uvicorn tt_sam3.server.app:app`) | `server` |
+| the host tests (no device, no weights) | `server` and `test`: `test_server_host.py` and `test_api_host.py` start the server app with fastapi's `TestClient` |
+| the device tests | `test`, and a Blackhole chip |
+
+Run only the host tests (about 1 minute; they do not open a device and do not need the weights):
+
+```bash
+TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_sam3/tests/test_*_host.py
+```
+
 The weights `facebook/sam3` are gated (Meta SAM License). Request access on the Hub and run `hf auth login`
 once. The first `from_pretrained` downloads `sam3.pt` (3.45 GB) to the HF cache.
 
+**Paths in the examples.** The examples use the demo image `media/source_1.png` with a path relative to the model repo
+root. Run them from the repo root, or give an absolute path. `examples/quickstart.py` finds the demo image relative to
+its own file, so you can run it from any directory.
+
 ## Quickstart
 
 ```python
 from tt_sam3 import Sam3
 
-with Sam3.from_pretrained(device_id=0) as model:
-    result = model("media/source_1.png", "car")
+# run from the model repo root (the demo image path is relative to it)
+with Sam3.from_pretrained(device_id=0, warmup_variants={"prompts": ["car"]}) as model:
+    result = model("media/source_1.png", "car")     # the first call is as fast as later calls
 
 print(result)                    # Sam3Result(prompt='car', num_instances=3, scores=[0.973, 0.817, 0.797], ...)
 result.masks                     # (N, H, W) bool, original image size
@@ -40,7 +60,8 @@ import numpy as np
 from PIL import Image
 from tt_sam3 import Sam3
 
-model = Sam3.from_pretrained()                       # device 0, weights from the HF cache, traces warmed up
+model = Sam3.from_pretrained()                       # device 0, weights from the HF cache, warmed up
+# "street.jpg", "a.jpg", ...: your own image files (paths relative to the current directory)
 
 # Several prompts on one image: the image part runs once, each prompt adds about 20 ms.
 car, person, tree = model(Image.open("street.jpg"), ["car", "person", "tree"])
@@ -63,11 +84,12 @@ model.close()                                        # or use "with ... as model
 
 ## API reference
 
-### `Sam3.from_pretrained(repo_id="facebook/sam3", *, revision=None, checkpoint_path=None, device_id=0, device=None, dispatch="auto", warmup=2, num_threads=None, cache_image=True) -> Sam3`
+### `Sam3.from_pretrained(repo_id="facebook/sam3", *, revision=None, checkpoint_path=None, device_id=0, device=None, dispatch="auto", warmup=2, warmup_variants=None, num_threads=None, cache_image=True, verbose=False) -> Sam3`
 
 Loads the weights, opens the chip, converts and uploads the weights, then runs the warm-up. The first warm-up
 compiles the kernels and captures the device traces (minutes with a cold kernel cache, cached in
-`TT_METAL_CACHE`). When it returns, the model is ready and every call runs the captured traces.
+`TT_METAL_CACHE`). Then it prepares the request variants of `warmup_variants` (see "Warm-up" below). When it
+returns, the model is ready and the first real call is as fast as a later call.
 
 | argument | meaning |
 |---|---|
@@ -77,10 +99,23 @@ compiles the kernels and captures the device traces (minutes with a cold kernel
 | `device_id` | Chip to open. Use 0 when `TT_VISIBLE_DEVICES` pins one chip. |
 | `device` | An already opened `ttnn` device. Open it with `ttnn.open_device(**tt_sam3.device_open_kwargs())`. `close()` does not close a device that you gave. |
 | `dispatch` | `"auto"` (default): ETH dispatch and a 12x10 compute grid when the tt-metal tree has the Blackhole ETH-dispatch patch. This is how the published Galaxy numbers were measured. Without the patch, the model uses the default dispatch and gives a `RuntimeWarning`. The outputs are the same. `"eth"` or `"worker"` select the dispatch directly. |
-| `warmup` | Number of warm-up forwards (default 2). With 0, the first call compiles and captures. |
+| `warmup` | Number of trace warm-up forwards (default 2). With 0, there is no warm-up at all (also no `warmup_variants`): the first call compiles and captures. |
+| `warmup_variants` | The request variants to prepare (see "Warm-up"). `None` (default): `DEFAULT_WARMUP_VARIANTS`. A dict replaces the defaults of the keys that it names. `False`: only the trace warm-up. |
 | `num_threads` | Host torch threads. Default: `SAM3_NUM_THREADS`, else 17, capped at the CPU count. |
+| `verbose` | `False` (default): a quiet console (see "Console output"). `True`: tt-metal's INFO log and the `tt_sam3` steps. |
 | `cache_image` | Keep the device results of the last image. When the same image comes again (same pixels or same file bytes), the image part is skipped. The output bits are the same. This needs `xxhash`. |
 
+### `model.warmup(*, prompt_counts=(), image_sizes=(), max_instances=(), prompts=()) -> dict`
+
+Prepares more request variants after `from_pretrained` (same keys as `warmup_variants`). It is idempotent: a variant
+that is already prepared is skipped. It returns `{variant: seconds}` for the variants that ran.
+`model.warmed_variants` and `model.info["warmup"]` show what is prepared and the total warm-up time.
+
+```python
+model.warmup(prompts=["bicycle", "traffic light"])    # about 75 ms for each new prompt
+model.warmup(image_sizes=[(720, 1280)])               # about 0.1 s
+```
+
 ### `model(image, prompts, *, threshold=0.5, max_instances=50, return_masks=True, return_mask_logits=False)`
 
 | argument | meaning |
@@ -131,6 +166,108 @@ closed at exit. One process can use a chip for only one model at a time.
   `from_pretrained` uses (`l1_small_size`, `trace_region_size`, and `dispatch_core_config` for ETH dispatch).
 - `model.info`: the checkpoint path, the revision, the dispatch, the compute grid and the host paths in use.
 
+## Warm-up
+
+After `from_pretrained` returns, the first real call of a prepared variant is as fast as a later call (within about
+2 %, see the table). The warm-up uses real calls on generated photo-like images of the real sizes and dummy noun
+phrases. These calls go through the full path: the device traces, the host text-encoder threads, PNG and JPEG decode,
+the input kernels, the kept-instance readback, the mask upsample and its buffers. The warm-up keeps no result. The only
+exception is the text features of the prompts that you give in `prompts`, which go into the normal text cache. A call
+computes and caches the same values. The image cache is empty after the warm-up. The outputs do not change: the
+digests of all 20 x 9 measured calls were bit-identical before and after the change.
+
+### Variants
+
+| key | default | what it prepares |
+|---|---|---|
+| `prompt_counts` | `(1, 2, 3, 4)` | Calls with 1..k prompts (the server's `SAM3_MAX_PROMPTS` default is 4). |
+| `image_sizes` | `((1080, 1920), (1008, 1008))` | `(height, width)` inputs. 1008x1008 uses its own input kernel (no resample). |
+| `max_instances` | `(1, 2, 4, 8, 16, 32, 64, 200)` | The device readback bucket for the number of kept instances (powers of two up to 64; more than 64 = full mask readback). |
+| `prompts` | `()` | Your noun phrases. Their text features are computed now. |
+
+Cost of the defaults: about 3.2 s more in `from_pretrained` (15.4..16.3 s -> 18.7..19.5 s with a warm kernel cache). No more device
+memory, because all traces and readback buffers exist after the trace warm-up. Each prompt in `prompts` adds about
+75 ms.
+
+### What the defaults give, and what needs your prompts
+
+- The defaults prepare the host first-use costs (decoders, the first use of the text encoder) and the request
+  shapes (1..4 prompts, the two image sizes, all kept-instance buckets). Alone, they make call 1 faster for calls
+  with several prompts (2 prompts: 375 -> 238 ms, 4 prompts: 414 -> 363 ms), but not for one prompt (210 -> 207 ms).
+- For a fixed prompt set, call 1 still runs the text encoder for each of your prompts (+70..+175 ms with the
+  defaults). Give your prompts in `warmup_variants={"prompts": [...]}` to remove this: call 1 / steady goes from
+  1.5..2.0x to 0.96..1.02x.
+- The text encoder has `SAM3_TEXT_WORKERS` (default 4) threads. A new prompt goes to the first idle thread, so one
+  new prompt per call always uses the same thread, and 4 new prompts use 4 threads in parallel. The warm-up runs its
+  text encodes on its own threads and stops them at the end, so the text threads of your calls are the same as
+  without the warm-up. The time of a call with new prompts does not change.
+
+### What the warm-up cannot prepare
+
+- **A new prompt.** The prompts are open vocabulary. The host text encoder runs once for each prompt that the model
+  has not seen yet, and the result goes into the text cache (256 entries). This costs about 75 to 150 ms on the call
+  that uses the prompt first (measured: +65..+150 ms for one new prompt, about +225 ms for 4 new prompts, which encode
+  in parallel). A later call with a new prompt pays the same cost, so this is not a first-call cost. To remove it from
+  the first call, give your prompts in `warmup_variants={"prompts": [...]}` or call `model.warmup(prompts=[...])`.
+- **Other image sizes.** Any size works without a first-use cost above about 1 ms (only the resample tables of
+  the new size are computed). In the measurement, the JPEG set has 9 different sizes and none of them was slower on
+  first use.
+
+### Measured
+
+Chip 11 of a Galaxy Blackhole host (12x10 grid, ETH dispatch, shared host), one fresh process per row,
+`perf/warmup/first_call_bench.py`. Every call uses a NEW real image (18 photos, never the warm-up input, never the
+previous image). "steady" is the median of calls 3..20. The image set is 1920x1080 PNG unless stated (`sq`:
+1008x1008 PNG, `jpg`: JPEG files of 9 sizes, 333x500 to 1600x900). Columns: from_pretrained, then call
+1 / 2 / 10 / steady in ms.
+
+| variant | before | after, defaults | after, prompts listed |
+|---|---|---|---|
+| 1 prompt (`"car"`) | 15.9 s, 210 / 132 / 128 / 130 | 18.8 s, 207 / 128 / 128 / 129 | 19.6 s, 127 / 128 / 128 / 128 |
+| 2 prompts | 15.7 s, 375 / 157 / 165 / 152 | 19.1 s, 238 / 149 / 149 / 149 | 19.5 s, 147 / 154 / 149 / 149 |
+| 4 prompts | 16.2 s, 414 / 197 / 196 / 188 | 19.2 s, 363 / 197 / 198 / 188 | 19.7 s, 187 / 197 / 193 / 188 |
+| 1 prompt, numpy input | 15.6 s, 154 / 82 / 84 / 81 | 19.4 s, 165 / 83 / 83 / 82 | 20.2 s, 81 / 82 / 84 / 81 |
+| 1 prompt, `return_masks=False, return_mask_logits=True` | 16.0 s, 191 / 129 / 129 / 128 | 18.9 s, 208 / 129 / 128 / 128 | 19.8 s, 128 / 128 / 127 / 128 |
+| 1 prompt, `threshold=0.05, max_instances=200` | 21.4 s, 198 / 131 / 129 / 133 | 19.0 s, 207 / 129 / 128 / 133 | 19.9 s, 128 / 131 / 126 / 134 |
+| 1 prompt, `sq` | 16.1 s, 177 / 106 / 107 / 108 | 19.1 s, 193 / 106 / 106 / 106 | 20.2 s, 105 / 105 / 106 / 106 |
+| 1 prompt, `jpg` | 16.0 s, 157 / 85 / 82 / 83 | 19.2 s, 166 / 84 / 83 / 84 | 19.9 s, 86 / 84 / 83 / 84 |
+| a new prompt on every call | 15.9 s, 217 / 228 / 219 / 209 | 18.8 s, 208 / 233 / 223 / 222 | - |
+
+"Prompts listed" = `warmup_variants={"prompts": ["car", "person", "tree", "road"]}` (the other keys at their
+defaults). The first-call cost that remains with the defaults is the text encoder for prompts that are not listed
+(see above). The output digests of every call were bit-identical to "before". The JSON of every run is in
+`perf/warmup/results/` (`before`, `after`, `afterp`).
+
+Steady state when every call has NEW prompts (the host text encoder is most of the call), the same images and
+prompts, old (54d1f2a, no variant warm-up) and new interleaved in fresh processes
+(`perf/warmup/results/ab_54d1f2a/`):
+
+| workload | 54d1f2a steady (ms) | now steady (ms) |
+|---|---|---|
+| 1 new prompt per call, numpy input, 12 pairs | 175.4 (mean) | 176.7 (mean); difference per pair -16.5..+13.6, median +1.5 |
+| same, `warmup_variants=False`, 6 rounds | - | 172.9 (mean) |
+| 1 new prompt per call, JPEG files, 3 pairs | 181.8 | 185.4 |
+| 4 new prompts per call, 2 pairs | 354.7 | 360.9 |
+| 1 fixed prompt, numpy input (the published path), 4 pairs | 81.3 | 81.1 |
+
+The differences are inside the run-to-run noise of this shared host (about +-10 ms per run). An earlier version of the
+warm-up used every text thread, and a single new prompt then moved between 4 threads: 171.8 vs 160.7 ms, slower in
+8 of 8 pairs (`perf/warmup/results/ab_4df6bea/`). The text pool and the warm-up's own text threads (see above) remove
+this.
+
+## Console output
+
+`from_pretrained(verbose=False)` (the default) keeps the console quiet:
+
+- It sets `TT_LOGGER_LEVEL=error` (tt-metal's C++ log) and `LOGURU_LEVEL=WARNING` (ttnn's Python log). It does this
+  only when you have not set them and only before `ttnn` is imported (they are read once, at import).
+- It hides two known harmless start-up warnings: sam3's CUDA autocast decorators on a CPU-only torch, and numba's TBB
+  version check.
+- Errors and all other warnings (also the `RuntimeWarning`s of `tt_sam3`) are shown.
+
+With `verbose=True`, nothing is changed and the `tt_sam3` logger prints its steps (load, device open, warm-up) to
+stderr. To see tt-metal's log while `verbose=False`, set `TT_LOGGER_LEVEL=info` before you start Python.
+
 ## Notes
 
 - Same results as the HTTP server: `model(...)` runs the same device traces, host preprocessing, instance

code/pyproject.toml

@@ -40,7 +40,7 @@ dependencies = [
 
 [project.optional-dependencies]
 server = ["fastapi", "uvicorn", "pydantic>=2", "pybase64>=1.4", "httptools>=0.6"]
-test = ["pytest"]
+test = ["pytest", "httpx"]       # httpx: fastapi TestClient of the host tests
 
 [tool.setuptools.packages.find]
 include = ["tt_sam3*"]

code/tt_sam3/api.py

@@ -16,12 +16,15 @@ host pre- and post-processing and the same instance selection, so the outputs ar
 from __future__ import annotations
 
 import atexit
+import contextlib
 import functools
 import gc
 import io
 import logging
 import os
+import sys
 import threading
+import time
 import warnings
 import weakref
 from concurrent.futures import ThreadPoolExecutor
@@ -44,6 +47,19 @@ WEIGHTS_FILENAME = "sam3.pt"
 NUM_QUERIES = 200
 MASK_RES = 288
 WARMUP_PROMPT = "object"
+#: Dummy prompts of the variant warm-up (text-encoder cache misses on the host text workers; never a cached answer to
+#: a user prompt: a user prompt that equals one of them gets the same text features it would get anyway).
+_WARMUP_WORDS = ("region", "thing", "item", "shape", "surface", "part", "area", "piece", "material", "pattern")
+#: Kept-instance readback buckets of the device select path (powers of two up to 64; more than 64 = full readback).
+INSTANCE_BUCKETS = (1, 2, 4, 8, 16, 32, 64, 200)
+#: What ``from_pretrained(warmup_variants=None)`` prepares (see ``Sam3.warmup``).  About 3 s on top of the trace
+#: warm-up; no extra device memory (every trace already exists after the base warm-up).
+DEFAULT_WARMUP_VARIANTS: Dict[str, Tuple] = {
+    "prompt_counts": (1, 2, 3, 4),            # 1 .. 4 prompts per call (the server's SAM3_MAX_PROMPTS default)
+    "image_sizes": ((1080, 1920), (1008, 1008)),  # (H, W): the resampling and the identity input kernels
+    "max_instances": INSTANCE_BUCKETS,        # every kept-instance readback bucket
+    "prompts": (),                            # your own noun phrases: their text features are computed now
+}
 _GRID_OVERRIDE_ENV = "TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE"
 _GRID_12x10 = "11,9"  # logical end coordinate (inclusive) of a 12x10 compute grid
 
@@ -295,6 +311,101 @@ def _warmup_input() -> torch.Tensor:
     return torch.cat([x, y, torch.sin(x * 3.14) * torch.cos(y * 3.14)], dim=1).contiguous()
 
 
+def _dummy_rgb(height: int, width: int, seed: int) -> np.ndarray:
+    """A photo-like (H, W, 3) uint8 warm-up image: smooth gradients, blobs and noise (real sizes, real host work:
+    decode, resample, mask upsample).  Generated, never a cached output."""
+    rng = np.random.default_rng(seed)
+    yy, xx = np.meshgrid(np.linspace(0, 1, height, dtype=np.float32), np.linspace(0, 1, width, dtype=np.float32), indexing="ij")
+    img = np.stack([xx, yy, 0.5 + 0.5 * np.sin(6.28 * (xx + yy))], axis=-1) * 160.0
+    for _ in range(6):  # a few bright ellipses ("objects")
+        cy, cx, ry, rx = rng.uniform(0.1, 0.9), rng.uniform(0.1, 0.9), rng.uniform(0.05, 0.25), rng.uniform(0.05, 0.25)
+        m = ((yy - cy) / ry) ** 2 + ((xx - cx) / rx) ** 2 < 1.0
+        img[m] = rng.uniform(0, 255, size=3).astype(np.float32)
+    img += rng.normal(0, 8, size=(height, width, 1)).astype(np.float32)
+    return np.clip(img, 0, 255).astype(np.uint8)
+
+
+def _encode(rgb: np.ndarray, fmt: str) -> bytes:
+    buf = io.BytesIO()
+    Image.fromarray(rgb).save(buf, format=fmt, **({"quality": 90} if fmt == "JPEG" else {"compress_level": 1}))
+    return buf.getvalue()
+
+
+def _as_tuple(value, pair: bool = False) -> Tuple:
+    """One value or a sequence -> tuple; with ``pair``, one ``(h, w)`` pair is one value."""
+    if value is None:
+        return ()
+    if isinstance(value, (str, int)):
+        return (value,)
+    value = tuple(value)
+    if pair and len(value) == 2 and all(isinstance(v, int) for v in value):
+        return (value,)
+    return value
+
+
+def normalize_warmup_variants(spec) -> Dict[str, Tuple]:
+    """``warmup_variants`` -> a full spec.  ``None`` = :data:`DEFAULT_WARMUP_VARIANTS`; a dict replaces the defaults of the
+    keys it names (an empty value turns that key off); ``False`` / ``{}``-with-all-keys-empty = trace warm-up only."""
+    if spec is None or spec is True:
+        spec = {}
+        base = DEFAULT_WARMUP_VARIANTS
+    elif spec is False:
+        spec, base = {}, {k: () for k in DEFAULT_WARMUP_VARIANTS}
+    else:
+        base = DEFAULT_WARMUP_VARIANTS
+    if not isinstance(spec, dict):
+        raise TypeError("warmup_variants must be None, False or a dict with keys " + ", ".join(DEFAULT_WARMUP_VARIANTS))
+    unknown = set(spec) - set(DEFAULT_WARMUP_VARIANTS)
+    if unknown:
+        raise ValueError(f"unknown warmup_variants keys {sorted(unknown)}; known: {list(DEFAULT_WARMUP_VARIANTS)}")
+    out = {k: _as_tuple(spec[k], pair=k == "image_sizes") if k in spec else tuple(v) for k, v in base.items()}
+    for k in out["prompt_counts"]:
+        if not isinstance(k, int) or not 1 <= k <= 64:
+            raise ValueError(f"prompt_counts: ints in 1..64, got {k!r}")
+    sizes = []
+    for hw in out["image_sizes"]:
+        if (not isinstance(hw, (tuple, list)) or len(hw) != 2 or not all(isinstance(v, int) and 1 <= v <= 8192 for v in hw)):
+            raise ValueError(f"image_sizes: (height, width) int pairs in 1..8192, got {hw!r}")
+        sizes.append((int(hw[0]), int(hw[1])))
+    out["image_sizes"] = tuple(sizes)
+    for b in out["max_instances"]:
+        if not isinstance(b, int) or not 1 <= b <= NUM_QUERIES:
+            raise ValueError(f"max_instances: ints in 1..{NUM_QUERIES}, got {b!r}")
+    out["prompts"] = tuple(_check_prompts(list(out["prompts"]))[0]) if out["prompts"] else ()
+    return out
+
+
+def _quiet_start(verbose: bool) -> None:
+    """Console noise policy of ``from_pretrained``.  Quiet (default): tt-metal's C++ log and ttnn's loguru log show
+    errors / warnings only (``TT_LOGGER_LEVEL=error``, ``LOGURU_LEVEL=WARNING``), set only when the user has not set
+    them and only before tt-metal / loguru are loaded (they read them once).  Verbose: the ``tt_sam3`` logger prints
+    its INFO steps (load, device open, warm-up) to stderr when it has no handler yet."""
+    if verbose:
+        if not log.handlers and not logging.getLogger().handlers:
+            h = logging.StreamHandler()
+            h.setFormatter(logging.Formatter("%(asctime)s tt_sam3 %(levelname)s %(message)s"))
+            log.addHandler(h)
+            log.setLevel(logging.INFO)
+        return
+    if "ttnn" not in sys.modules:
+        os.environ.setdefault("TT_LOGGER_LEVEL", "error")
+    if "loguru" not in sys.modules:
+        os.environ.setdefault("LOGURU_LEVEL", "WARNING")
+
+
+@contextlib.contextmanager
+def _quiet_warnings(verbose: bool):
+    """Hide two known harmless start-up warnings (sam3's CUDA autocast decorators on a CPU torch; numba's TBB version
+    probe, which then uses its OpenMP layer).  Everything else, our RuntimeWarnings included, is shown."""
+    if verbose:
+        yield
+        return
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message="CUDA is not available or torch_xla is imported", category=UserWarning)
+        warnings.filterwarnings("ignore", message="The TBB threading layer requires TBB version")
+        yield
+
+
 _PIN_PATCHED = [False]
 
 
@@ -348,6 +459,9 @@ class Sam3:
         self._xxhash = xxhash
         self.cache_image = bool(cache_image and xxhash is not None and self._rgb_input)
         self._last_decoded: List[Any] = [None, None]
+        self._warmed: set = set()  # variant atoms prepared by warmup()
+        self._warm_seed = 0
+        self.info.setdefault("warmup", {"variants": {k: [] for k in DEFAULT_WARMUP_VARIANTS}, "seconds": 0.0})
         self._atexit = functools.partial(_close_at_exit, weakref.ref(self))
         atexit.register(self._atexit)
 
@@ -363,8 +477,10 @@ class Sam3:
         device=None,
         dispatch: str = "auto",
         warmup: int = 2,
+        warmup_variants: Optional[Union[Dict[str, Any], bool]] = None,
         num_threads: Optional[int] = None,
         cache_image: bool = True,
+        verbose: bool = False,
     ) -> "Sam3":
         """Load the weights, open the chip, build the device traces and warm them up (the server's startup).
 
@@ -379,11 +495,34 @@ class Sam3:
             dispatch: ``"auto"`` (ETH dispatch + 12x10 grid when the tt-metal tree has the ETH-dispatch patch, as
                 the published numbers; else the default dispatch with a warning), ``"eth"`` or ``"worker"``.
             warmup: warm-up forwards (the first one compiles the kernels and captures the traces; minutes when the
-                kernel cache is cold).  0 makes the first call do it.
+                kernel cache is cold).  0 skips all warm-up (also ``warmup_variants``): the first call does it.
+            warmup_variants: the request variants to prepare after the trace warm-up, so that the first real call of
+                each is as fast as a later one (see :meth:`warmup`).  ``None`` = :data:`DEFAULT_WARMUP_VARIANTS`
+                (1..4 prompts, 1920x1080 and 1008x1008 inputs, every kept-instance bucket; about 3 s).  A dict
+                replaces the defaults of the keys it names, e.g. ``{"prompts": ["car", "person"]}`` also computes
+                the text features of your prompts now.  ``False``: only the trace warm-up.
             num_threads: host torch threads (default ``SAM3_NUM_THREADS`` or 17, capped at the CPU count).
             cache_image: reuse the device results of the image part when the same image comes again (one entry;
                 same output bits).  Needs the optional ``xxhash`` package.
+            verbose: False (default) keeps the console quiet: tt-metal / ttnn print errors and warnings only (see
+                ``PYTHON.md``, "Console output").  True prints tt-metal's INFO log and the ``tt_sam3`` steps.
         """
+        _quiet_start(verbose)
+        variants = normalize_warmup_variants(warmup_variants)  # check the spec before the slow part
+        with _quiet_warnings(verbose):
+            model = cls._build(repo_id, revision, checkpoint_path, device_id, device, dispatch, warmup, num_threads,
+                               cache_image)
+            if warmup > 0:
+                try:
+                    model.warmup(**variants)
+                except BaseException:
+                    model.close()
+                    raise
+        log.info("SAM3 ready: %s", model.info)
+        return model
+
+    @classmethod
+    def _build(cls, repo_id, revision, checkpoint_path, device_id, device, dispatch, warmup, num_threads, cache_image):
         from tt_sam3.server import maskpost
         from tt_sam3.tt.fused_config import FusedConfig
         from tt_sam3.tt.ttnn_sam3_model import TtSam3ImagePipeline, _patch_pin_memory, build_sam3_model, resolve_bpe_path
@@ -443,7 +582,6 @@ class Sam3:
             "fast_mask": fast_mask,
             "torch_threads": torch.get_num_threads(),
         }
-        log.info("SAM3 ready: %s", info)
         return cls(pipeline, device, owns_device=owns, info=info, cache_image=cache_image)
 
     # ---------------------------------------------------------------- inference
@@ -500,6 +638,125 @@ class Sam3:
                 out.append(self(loaded, prompt_list[0] if single else prompt_list, **kwargs))
         return out
 
+    # ---------------------------------------------------------------- warm-up
+    def warmup(
+        self,
+        *,
+        prompt_counts: Union[int, Sequence[int]] = (),
+        image_sizes: Union[Tuple[int, int], Sequence[Tuple[int, int]]] = (),
+        max_instances: Union[int, Sequence[int]] = (),
+        prompts: Union[str, Sequence[str]] = (),
+    ) -> Dict[str, float]:
+        """Prepare request variants now, so that their first real call is as fast as a later one.  Idempotent: a
+        variant that is already prepared is skipped.  Called by :meth:`from_pretrained` with ``warmup_variants``.
+
+        Every variant runs real calls on generated photo-like images of the real sizes and real (dummy) prompts, through
+        the same path as ``model(...)``: the device traces, the host text-encoder worker threads, PNG / JPEG decode,
+        the threaded input kernels, the kept-instance readback, the mask upsample and its buffers.  Nothing of these
+        calls is kept except the text features of ``prompts`` (the model's normal text cache, the same values a call
+        computes); the image cache is cleared afterwards.  The outputs of later calls do not change.
+
+        Args:
+            prompt_counts: prompts per call (``model(image, [p1, .., pk])``), e.g. ``(1, 2, 3, 4)``.
+            image_sizes: ``(height, width)`` input sizes, e.g. ``[(1080, 1920)]``.  Any other size works without a
+                first-use cost above about 1 ms (only its resample tables are new); 1008x1008 uses its own kernel.
+            max_instances: kept-instance counts to prepare: each value prepares the readback bucket of that many kept
+                instances (1, 2, 4, .. 64 and > 64).
+            prompts: your noun phrases.  Their text features (the host text encoder, about 75 ms per new prompt) are
+                computed now, so the first call with them does not pay it.  A prompt that is not listed pays it once,
+                on the call that first uses it (every new prompt does, also later).
+
+        Returns:
+            ``{variant: seconds}`` of the variants that ran now.
+        """
+        req = normalize_warmup_variants({"prompt_counts": prompt_counts, "image_sizes": image_sizes,
+                                         "max_instances": max_instances, "prompts": prompts})
+        if not any(req.values()):
+            return {}
+        todo = [("host", None)]
+        todo += [("prompt_counts", k) for k in req["prompt_counts"]]
+        todo += [("image_sizes", hw) for hw in req["image_sizes"]]
+        todo += [("max_instances", NUM_QUERIES if b > 64 else 1 << (b - 1).bit_length()) for b in req["max_instances"]]  # bucket
+        todo += [("prompts", p) for p in req["prompts"]]
+        todo = [a for a in dict.fromkeys(todo) if a not in self._warmed]
+        done: Dict[str, float] = {}
+        if not todo:
+            return done
+        if self._closed:
+            raise RuntimeError("Sam3 model is closed")
+        t_all = time.perf_counter()
+        scratch = getattr(self.pipeline, "scratch_text_pool", None)
+        with self._lock, (scratch() if scratch is not None else contextlib.nullcontext()):
+            try:
+                for atom in todo:
+                    t0 = time.perf_counter()
+                    self._warm_one(*atom)
+                    self._warmed.add(atom)
+                    if atom[0] != "host":
+                        self.info["warmup"]["variants"][atom[0]].append(list(atom[1]) if isinstance(atom[1], tuple) else atom[1])
+                    done[f"{atom[0]}={atom[1]}"] = round(time.perf_counter() - t0, 3)
+                    log.info("warm-up %s=%s: %.2f s", atom[0], atom[1], done[f"{atom[0]}={atom[1]}"])
+            finally:
+                self._invalidate_cache()  # the last warm-up image is not kept as "the previous image"
+                gc.collect()
+        self.info["warmup"]["seconds"] = round(self.info["warmup"]["seconds"] + time.perf_counter() - t_all, 2)
+        return done
+
+    @property
+    def warmed_variants(self) -> Dict[str, list]:
+        """The variants :meth:`warmup` has prepared so far (``{"prompt_counts": [1, 2], ...}``)."""
+        return {k: list(v) for k, v in self.info["warmup"]["variants"].items()}
+
+    def _dummy_prompts(self, k: int) -> List[str]:
+        """k prompts that are text-encoder cache misses now (the miss path: host worker threads, as a new user prompt)."""
+        out = []
+        while len(out) < k:
+            self._warm_seed += 1
+            n = self._warm_seed
+            out.append(f"{_WARMUP_WORDS[n % len(_WARMUP_WORDS)]} {n}")
+        return out
+
+    def _warm_one(self, kind: str, value) -> None:
+        if kind == "host":
+            # host first-use costs that every request variant shares: the PNG (libdeflate fast path and PIL) and JPEG
+            # decoders, the process-wide first use of the host text encoder (on the warm-up's own text threads, see
+            # scratch_text_pool: the real calls start their text threads new)
+            rgb = _dummy_rgb(270, 480, 1)
+            for fmt in ("PNG", "JPEG"):
+                _decode_bytes(_encode(rgb, fmt))
+            encode = getattr(self.pipeline, "_cached_forward_text", None)
+            if encode is not None:
+                ps = self._dummy_prompts(max(1, int(os.environ.get("SAM3_TEXT_WORKERS", "4"))))
+                self.pipeline.prefetch_text(ps)  # all workers at once: one encode on each worker thread
+                for p in ps:
+                    encode([p], device="cpu")
+            return
+        if kind == "prompt_counts":  # a new image as encoded PNG bytes (decode + cache key), k new prompts
+            self._warm_seed += 1
+            png = _encode(_dummy_rgb(1080, 1920, 100 + self._warm_seed), "PNG")
+            self(png, self._dummy_prompts(value))
+            return
+        if kind == "image_sizes":
+            h, w = value
+            self._warm_seed += 1
+            self(_dummy_rgb(h, w, 200 + self._warm_seed), self._dummy_prompts(1)[0])
+            return
+        if kind == "max_instances":
+            # threshold 0 keeps every query (score > 0), so max_instances = n kept instances: the readback bucket of n
+            # (n > 64: the full mask readback); one image, the image trace runs once (image cache)
+            if not hasattr(self, "_warm_inst_img"):
+                self._warm_seed += 1
+                self._warm_inst_img = _dummy_rgb(1080, 1920, 300 + self._warm_seed)
+            self(self._warm_inst_img, WARMUP_PROMPT, threshold=0.0, max_instances=int(value), return_masks=value <= 8)
+            return
+        if kind == "prompts":  # the text features a call with this prompt computes (and caches) on its first use
+            encode = getattr(self.pipeline, "_cached_forward_text", None)
+            if encode is not None:
+                self.pipeline.prefetch_text([value])
+                encode([value], device="cpu")
+            return
+        raise ValueError(kind)
+
     # ---------------------------------------------------------------- internals
     def _load_cached(self, image) -> "_Loaded":
         if isinstance(image, _Loaded):

code/tt_sam3/tests/test_warmup_device.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Device test of the warm-up: after ``from_pretrained`` the FIRST real call is already fast, and the warm-up does not
+change any output bit.
+
+    source tools/chipenv.sh <chip> <model>/.venv
+    timeout -s INT 1800 python -m pytest -q -s code/tt_sam3/tests/test_warmup_device.py
+
+1. ``from_pretrained(warmup_variants={"prompts": [...]})`` in this process (fresh model), then real calls on NEW images
+   (two real photos, alternated, so the one-entry image cache never hits) with 1, 2 and 4 prompts: the first call of each
+   variant must be within ``WARM_TOL`` (default 1.10) of the median of its later calls (+ 3 ms for timer noise).
+2. A second model with ``warmup_variants=False`` (trace warm-up only) gives the same output bits for the same calls.
+"""
+import hashlib
+import os
+import statistics
+import time
+
+import numpy as np
+import pytest
+
+pytestmark = pytest.mark.skipif(os.environ.get("TT_VISIBLE_DEVICES", "") == "none", reason="needs a device")
+
+MODEL_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+IMAGES = [os.path.join(MODEL_DIR, "media", f) for f in ("source_1.png", "target_1.png")]
+VARIANTS = {"k1": ["car"], "k2": ["car", "tree"], "k4": ["car", "person", "tree", "road"]}
+PROMPTS = sorted({p for v in VARIANTS.values() for p in v})
+TOL = float(os.environ.get("WARM_TOL", "1.10"))
+CALLS = 8
+
+
+def _digest(res):
+    h = hashlib.sha1()
+    for r in res:
+        for a in (r.scores, r.boxes, np.packbits(r.masks)):
+            h.update(np.ascontiguousarray(a).tobytes())
+        h.update(np.float64(r.presence).tobytes())
+    return h.hexdigest()
+
+
+def _series(model, prompts):
+    raws = [open(p, "rb").read() for p in IMAGES]
+    ms, digests = [], []
+    for i in range(CALLS):
+        t0 = time.perf_counter()
+        res = model(raws[i % 2], prompts)  # a different image than the previous call: the whole graph runs
+        ms.append((time.perf_counter() - t0) * 1e3)
+        digests.append(_digest(res))
+    return ms, digests
+
+
+def test_first_call_fast_and_outputs_unchanged():
+    from tt_sam3 import Sam3
+
+    t0 = time.perf_counter()
+    model = Sam3.from_pretrained(device_id=0, warmup_variants={"prompts": PROMPTS})
+    t_fp = time.perf_counter() - t0
+    info = model.info["warmup"]
+    print(f"\nfrom_pretrained {t_fp:.1f} s (variant warm-up {info['seconds']} s): {info['variants']}", flush=True)
+    assert info["variants"]["prompt_counts"] == [1, 2, 3, 4] and info["variants"]["prompts"] == PROMPTS
+    warm = {}
+    try:
+        for name, prompts in VARIANTS.items():
+            ms, dg = _series(model, prompts)
+            med = statistics.median(ms[2:])
+            print(f"{name}: first {ms[0]:.1f} ms, second {ms[1]:.1f} ms, steady median {med:.1f} ms "
+                  f"(first / steady = {ms[0] / med:.3f})", flush=True)
+            warm[name] = dg
+            assert ms[0] <= med * TOL + 3.0, (name, ms)
+        assert model.warmup(prompt_counts=(1, 2), prompts=PROMPTS) == {}  # idempotent: nothing runs again
+    finally:
+        model.close()
+
+    cold = Sam3.from_pretrained(device_id=0, warmup_variants=False)
+    try:
+        assert cold.info["warmup"]["variants"]["prompt_counts"] == []
+        for name, prompts in VARIANTS.items():
+            _, dg = _series(cold, prompts)
+            assert dg == warm[name], f"{name}: the variant warm-up changed outputs"
+    finally:
+        cold.close()

code/tt_sam3/tests/test_warmup_host.py

@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Host tests (no device, no weights) of the warm-up API: ``from_pretrained(warmup_variants=..., verbose=...)``,
+``model.warmup(**variant)`` (spec checks, idempotence, what it runs, no effect on outputs) and the quiet console
+policy.
+
+    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_sam3/tests/test_warmup_host.py
+"""
+import inspect
+import os
+import subprocess
+import sys
+
+import numpy as np
+import pytest
+import torch
+
+from tt_sam3 import api
+
+MODEL_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+
+class _Runner:
+    def __init__(self):
+        self.stats = {}
+        self.img_key = None
+
+
+class _Pipe:
+    """Deterministic fake of the device pipeline that records what the warm-up asks for."""
+
+    def __init__(self):
+        self._full = _Runner()
+        self.calls = []  # (shape, prompts, threshold, max_instances, cache hit)
+        self.text = []  # encoded prompts (cache misses)
+        self.text_cache = {}
+        self.prefetched = []
+        self.released = 0
+
+    def prefetch_text(self, prompts):
+        self.prefetched.append(list(prompts))
+
+    def _cached_forward_text(self, captions, device="cpu"):
+        key = tuple(captions)
+        if key not in self.text_cache:
+            self.text.append(captions[0])
+            self.text_cache[key] = {"language_features": torch.zeros(32, 1, 256)}
+        return self.text_cache[key]
+
+    def release(self):
+        self.released += 1
+
+    def forward_rgb_selected_multi(self, rgb, prompts, threshold, max_instances, image_key=None):
+        hit = image_key is not None and image_key == self._full.img_key
+        self._full.img_key = image_key
+        self.calls.append((rgb.shape, list(prompts), threshold, max_instances, hit))
+        for p in prompts:
+            self._cached_forward_text([p])
+        m = float(np.asarray(rgb, dtype=np.float64).mean()) / 255.0
+        out = []
+        for k, _p in enumerate(prompts):
+            n = 200
+            scores = torch.linspace(0.9, 0.01, n) * (0.5 + m)
+            keep = torch.nonzero(scores > threshold).flatten()[:max_instances]
+            g = torch.Generator().manual_seed(k)
+            logits = torch.randn(n, api.MASK_RES, api.MASK_RES, generator=g)
+            boxes = torch.rand(n, 4, generator=g) * 0.5 + 0.25
+            out.append((scores[keep], boxes[keep], logits[keep], 0.9))
+        return out
+
+
+def _model():
+    pytest.importorskip("xxhash")
+    pipe = _Pipe()
+    return api.Sam3(pipe, device=None, owns_device=False, info={"rgb_input": True, "fast_mask": True}), pipe
+
+
+# ------------------------------------------------------------------------------------------------
+# spec
+# ------------------------------------------------------------------------------------------------
+def test_from_pretrained_signature():
+    sig = inspect.signature(api.Sam3.from_pretrained)
+    assert sig.parameters["warmup_variants"].default is None
+    assert sig.parameters["verbose"].default is False
+    assert sig.parameters["warmup"].default == 2
+    assert "warmup_variants" in api.Sam3.from_pretrained.__doc__ and "verbose" in api.Sam3.from_pretrained.__doc__
+    assert callable(api.Sam3.warmup) and api.Sam3.warmup.__doc__
+    assert set(inspect.signature(api.Sam3.warmup).parameters) == {"self", *api.DEFAULT_WARMUP_VARIANTS}
+
+
+def test_spec_defaults_and_overrides():
+    d = api.normalize_warmup_variants(None)
+    assert d["prompt_counts"] == (1, 2, 3, 4)
+    assert d["image_sizes"] == ((1080, 1920), (1008, 1008))
+    assert d["max_instances"] == api.INSTANCE_BUCKETS and d["prompts"] == ()
+    o = api.normalize_warmup_variants({"prompts": ["car", " person "], "image_sizes": (720, 1280)})
+    assert o["prompts"] == ("car", "person") and o["image_sizes"] == ((720, 1280),)
+    assert o["prompt_counts"] == d["prompt_counts"]  # keys not named keep their defaults
+    assert api.normalize_warmup_variants({"prompt_counts": 2})["prompt_counts"] == (2,)
+    off = api.normalize_warmup_variants(False)
+    assert all(v == () for v in off.values())
+    assert api.normalize_warmup_variants({"max_instances": ()})["max_instances"] == ()
+
+
+@pytest.mark.parametrize("bad", [{"foo": 1}, {"prompt_counts": 0}, {"prompt_counts": [1, "2"]}, {"image_sizes": [(10,)]},
+                                 {"image_sizes": [(0, 10)]}, {"max_instances": 201}, {"prompts": [""]}, "all"])
+def test_spec_rejects(bad):
+    with pytest.raises((ValueError, TypeError)):
+        api.normalize_warmup_variants(bad)
+
+
+# ------------------------------------------------------------------------------------------------
+# warmup(): what it runs, idempotence, no effect on outputs
+# ------------------------------------------------------------------------------------------------
+def test_warmup_runs_each_variant_once():
+    model, pipe = _model()
+    done = model.warmup(prompt_counts=(1, 3), image_sizes=[(1080, 1920), (1008, 1008)], max_instances=(1, 3, 64, 100, 200),
+                        prompts=["car", "dog"])
+    assert "host=None" in done and len(done) == 1 + 2 + 2 + 4 + 2  # 3 -> bucket 4; 100 and 200 -> one bucket
+    shapes = [c[0] for c in pipe.calls]
+    assert (1080, 1920, 3) in shapes and (1008, 1008, 3) in shapes  # real sizes, real host work
+    assert sorted(len(c[1]) for c in pipe.calls if c[2] == 0.5)[:4] == [1, 1, 1, 3]
+    assert sorted(c[3] for c in pipe.calls if c[2] == 0.0) == [1, 4, 64, 200]  # threshold 0: one call per bucket
+    assert "car" in pipe.text and "dog" in pipe.text
+    assert len(pipe.prefetched[0]) == int(os.environ.get("SAM3_TEXT_WORKERS", "4"))  # every text worker thread
+    assert pipe._full.img_key is None and model._last_decoded == [None, None]  # nothing kept as "previous image"
+    w = model.warmed_variants
+    assert w["prompt_counts"] == [1, 3] and w["max_instances"] == [1, 4, 64, 200] and w["prompts"] == ["car", "dog"]
+    assert w["image_sizes"] == [[1080, 1920], [1008, 1008]]
+    n = len(pipe.calls)
+    assert model.warmup(prompt_counts=3, image_sizes=(1008, 1008), max_instances=3, prompts="car") == {}  # idempotent (3: bucket 4)
+    assert len(pipe.calls) == n
+    more = model.warmup(prompt_counts=2)
+    assert list(more) == ["prompt_counts=2"] and len(pipe.calls) == n + 1
+
+
+def test_warmup_prompts_do_not_collide_with_user_prompts():
+    model, pipe = _model()
+    model.warmup(prompt_counts=(1, 2, 3, 4))
+    dummies = {p for c in pipe.calls for p in c[1]}
+    assert len(dummies) == 10 and not dummies & {"car", "person", "tree", "dog", "object"}
+
+
+def test_warmup_does_not_change_outputs():
+    rng = np.random.default_rng(5)
+    img = rng.integers(0, 256, size=(300, 400, 3), dtype=np.uint8)
+    cold, _ = _model()
+    ref = cold(img, ["car", "tree"], threshold=0.3)
+    warm, _ = _model()
+    warm.warmup(**api.normalize_warmup_variants(None))
+    got = warm(img, ["car", "tree"], threshold=0.3)
+    for a, b in zip(ref, got):
+        assert np.array_equal(a.masks, b.masks) and np.array_equal(a.scores, b.scores) and np.array_equal(a.boxes, b.boxes)
+
+
+def test_warmup_after_close_raises():
+    model, _ = _model()
+    model.close()
+    with pytest.raises(RuntimeError):
+        model.warmup(prompt_counts=1)
+
+
+# ------------------------------------------------------------------------------------------------
+# quiet console
+# ------------------------------------------------------------------------------------------------
+def _run(code, env_extra=None):
+    env = {k: v for k, v in os.environ.items() if k not in ("TT_LOGGER_LEVEL", "LOGURU_LEVEL")}
+    env.update(TT_VISIBLE_DEVICES="none", **(env_extra or {}))
+    out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, env=env, timeout=300)
+    assert out.returncode == 0, out.stderr[-2000:]
+    return out.stdout.strip()
+
+
+def test_quiet_start_sets_log_levels_only_when_unset():
+    code = ("import os; from tt_sam3 import api; api._quiet_start(False); "
+            "print(os.environ.get('TT_LOGGER_LEVEL'), os.environ.get('LOGURU_LEVEL'))")
+    assert _run(code) == "error WARNING"
+    assert _run(code, {"TT_LOGGER_LEVEL": "debug", "LOGURU_LEVEL": "DEBUG"}) == "debug DEBUG"  # the user's choice wins
+    verbose = ("import os, logging; from tt_sam3 import api; api._quiet_start(True); "
+               "print(os.environ.get('TT_LOGGER_LEVEL'), api.log.level == logging.INFO, len(api.log.handlers))")
+    assert _run(verbose) == "None True 1"
+
+
+def test_quiet_warnings_hide_only_known_noise():
+    import warnings
+
+    with warnings.catch_warnings(record=True) as rec:
+        warnings.simplefilter("always")
+        with api._quiet_warnings(False):
+            warnings.warn("CUDA is not available or torch_xla is imported. Disabling autocast.", UserWarning)
+            warnings.warn("The TBB threading layer requires TBB version 2021 update 6 or later")
+            warnings.warn("something real", RuntimeWarning)
+        with api._quiet_warnings(True):
+            warnings.warn("CUDA is not available or torch_xla is imported. Disabling autocast.", UserWarning)
+    assert [str(w.message)[:14] for w in rec] == ["something real", "CUDA is not av"]
+
+
+# ------------------------------------------------------------------------------------------------
+# docs
+# ------------------------------------------------------------------------------------------------
+def test_python_md_documents_warmup_and_install_extras():
+    text = open(os.path.join(MODEL_DIR, "code", "PYTHON.md"), encoding="utf-8").read()
+    for needle in ("warmup_variants", "model.warmup(", "prompt_counts", "image_sizes", "max_instances", "prompts",
+                   'code/[server,test]', "test_*_host.py", "verbose", "repo root"):
+        assert needle in text, needle
+
+
+def test_text_pool_keeps_single_prompt_jobs_on_slot0():
+    """The text-encoder pool sends a job to the lowest idle slot: after the warm-up has used every slot, a stream of
+    one new prompt per call still runs on the same (slot 0) thread; k concurrent jobs use k different slots."""
+    import threading
+    import time
+
+    from tt_sam3.tt.ttnn_sam3_model import _TextPool
+
+    pool = _TextPool(4)
+    try:
+        def name(sleep=0.0):
+            time.sleep(sleep)
+            return threading.current_thread().name
+
+        burst = [pool.submit(name, 0.2) for _ in range(4)]  # the warm-up: one job on each slot at once
+        assert sorted(f.result() for f in burst) == [f"sam3-text{i}" for i in range(4)]
+        seq = [pool.submit(name).result() for _ in range(20)]  # one new prompt per call
+        assert set(seq) == {"sam3-text0"}
+        two = [pool.submit(name, 0.1) for _ in range(2)]
+        assert [f.result() for f in two] == ["sam3-text0", "sam3-text1"]
+        six = [pool.submit(name, 0.1) for _ in range(6)]  # more jobs than slots: the first free slots take the rest
+        names = [f.result() for f in six]
+        assert names[:4] == [f"sam3-text{i}" for i in range(4)] and set(names[4:]) <= set(names[:4])
+
+        def boom():
+            raise ValueError("x")
+
+        with pytest.raises(ValueError):
+            pool.submit(boom).result()
+        assert pool.submit(name).result() == "sam3-text0"
+    finally:
+        pool.shutdown()
+
+
+def test_scratch_text_pool_leaves_no_warm_text_threads():
+    """The warm-up's text encodes run on their own threads; afterwards the real pool is as before (none yet: the
+    first real prompt starts a new thread, as without the warm-up)."""
+    import threading
+    import types
+
+    from tt_sam3.tt.ttnn_sam3_model import TtSam3ImagePipeline
+
+    seen = []
+    fake = types.SimpleNamespace(_text_dev=None, _text_cache={}, _text_cache_max=256,
+                                 _orig_forward_text=lambda *a, **k: seen.append(threading.current_thread()))
+    prefetch = TtSam3ImagePipeline.prefetch_text.__get__(fake)
+    with TtSam3ImagePipeline.scratch_text_pool(fake):
+        assert prefetch(["a", "b"]) == 2
+        tmp = fake._text_pool
+        [f.result() for f in fake._text_futures.values()]
+    assert fake._text_pool is None and all(not t.is_alive() for t in seen) and len(seen) == 2
+    assert prefetch(["c"]) == 1 and fake._text_pool is not tmp
+    fake._text_futures[(("c",), False, False)].result()
+    fake._text_pool.shutdown()

code/tt_sam3/tt/ttnn_sam3_model.py

@@ -836,12 +836,12 @@ class TtSam3ImagePipeline:
         if os.environ.get("SAM3_TEXT_PREFETCH", "1") != "1" or getattr(self, "_text_dev", None) is not None:
             return 0
         import threading
-        from concurrent.futures import ThreadPoolExecutor
 
-        if getattr(self, "_text_pool", None) is None:
-            self._text_pool = ThreadPoolExecutor(max_workers=int(os.environ.get("SAM3_TEXT_WORKERS", "4")), thread_name_prefix="sam3-text")
+        if getattr(self, "_text_futures", None) is None:
             self._text_futures = {}
             self._text_lock = threading.Lock()
+        if getattr(self, "_text_pool", None) is None:
+            self._text_pool = _TextPool(int(os.environ.get("SAM3_TEXT_WORKERS", "4")))
         nthreads = torch.get_num_threads()
 
         def job(p):
@@ -861,6 +861,24 @@ class TtSam3ImagePipeline:
                 n += 1
         return n
 
+    def scratch_text_pool(self):
+        """Context manager: the text encodes started inside it run on their own worker threads, which stop at its end
+        (the warm-up uses it).  The real calls after it start their text threads new, as without the warm-up: a
+        measured single-prompt encode on a thread that the warm-up used was about 20 ms slower than on a new thread."""
+        import contextlib
+
+        @contextlib.contextmanager
+        def scope():
+            saved, self._text_pool = getattr(self, "_text_pool", None), None
+            try:
+                yield
+            finally:
+                tmp, self._text_pool = self._text_pool, saved
+                if tmp is not None:
+                    tmp.shutdown(wait=True)
+
+        return scope()
+
     def _cached_forward_text(self, captions, input_boxes=None, additional_text=None, device="cuda"):
         cache_key = (tuple(captions), input_boxes is not None, additional_text is not None)
         if cache_key in self._text_cache:
@@ -1037,3 +1055,87 @@ def preprocess_image(image: torch.Tensor, target_size: int = 1008) -> torch.Tens
     if image.shape[-2:] != (target_size, target_size):
         image = F.interpolate(image, size=(target_size, target_size), mode="bilinear", align_corners=False)
     return (image - 0.5) / 0.5
+
+
+
+class _TextPool:
+    """The host text-encoder worker threads (``SAM3_TEXT_WORKERS`` slots, each thread started on its first job).
+
+    A new job goes to the LOWEST idle slot.  When no slot is idle it waits in one shared backlog, which the first slot
+    to finish takes (as a ThreadPoolExecutor does).  So one new prompt per call always runs on slot 0 (one warm thread
+    with its warm OpenMP pool and malloc arena), and k concurrent new prompts use slots 0..k-1.  A ThreadPoolExecutor
+    hands a job to ANY idle thread: once several threads exist (k > 1 prompts, or the warm-up) single-prompt jobs move
+    between them and pay their colder state (measured: about +11 ms per new prompt, 172 vs 161 ms per call)."""
+
+    def __init__(self, n: int):
+        import collections
+        import threading
+
+        self._n = max(1, int(n))
+        self._cond = threading.Condition()
+        self._inbox = [None] * self._n
+        self._idle = [True] * self._n
+        self._threads = [None] * self._n
+        self._backlog = collections.deque()
+        self._shutdown = False
+
+    def __len__(self):
+        return self._n
+
+    def submit(self, fn, *args):
+        from concurrent.futures import Future
+        import threading
+
+        fut = Future()
+        with self._cond:
+            if self._shutdown:
+                raise RuntimeError("text pool is shut down")
+            job = (fut, fn, args)
+            slot = next((i for i in range(self._n) if self._idle[i]), None)
+            if slot is None:
+                self._backlog.append(job)
+            else:
+                self._idle[slot] = False
+                self._inbox[slot] = job
+                if self._threads[slot] is None:
+                    t = threading.Thread(target=self._worker, args=(slot,), name=f"sam3-text{slot}", daemon=True)
+                    self._threads[slot] = t
+                    t.start()
+                self._cond.notify_all()
+        return fut
+
+    def _worker(self, i: int):
+        while True:
+            with self._cond:
+                while self._inbox[i] is None and not self._shutdown:
+                    self._cond.wait()
+                job, self._inbox[i] = self._inbox[i], None
+            if job is None:
+                return
+            while job is not None:
+                fut, fn, args = job
+                if not fut.set_running_or_notify_cancel():
+                    res, exc = None, None
+                else:
+                    try:
+                        res, exc = fn(*args), None
+                    except BaseException as e:  # noqa: BLE001 - handed to the caller through the future
+                        res, exc = None, e
+                with self._cond:  # the slot is idle BEFORE the result is visible: the next submit sees it idle
+                    job = self._backlog.popleft() if self._backlog else None
+                    if job is None:
+                        self._idle[i] = True
+                if fut.running():
+                    if exc is None:
+                        fut.set_result(res)
+                    else:
+                        fut.set_exception(exc)
+
+    def shutdown(self, wait=True):
+        with self._cond:
+            self._shutdown = True
+            self._cond.notify_all()
+        if wait:
+            for t in self._threads:
+                if t is not None:
+                    t.join()

examples/quickstart.py

@@ -5,6 +5,8 @@
     python examples/quickstart.py [image] [prompt ...]
 
 Writes examples/out/quickstart_overlay.png (masks and boxes on the image) and examples/out/quickstart.json.
+The default image and the output directory are found relative to this file, so it runs from any directory; an image
+path given on the command line is relative to the current directory.
 """
 import json
 import os
@@ -23,7 +25,7 @@ out_dir = os.path.join(HERE, "out")
 os.makedirs(out_dir, exist_ok=True)
 
 # --- the model card snippet -----------------------------------------------------------------------------
-with Sam3.from_pretrained(device_id=0) as model:
+with Sam3.from_pretrained(device_id=0, warmup_variants={"prompts": prompts}) as model:  # warm: call 1 is fast too
     results = model(image_path, prompts)            # one Sam3Result per prompt
     t0 = time.perf_counter()
     model(image_path, prompts)                      # warm call (same image: the image part is reused)