# SPDX-License-Identifier: Apache-2.0
"""Python API for SAM3 on one Tenstorrent Blackhole chip.

    from tt_sam3 import Sam3

    with Sam3.from_pretrained(device_id=0) as model:
        result = model("media/source_1.png", "car")
        result.masks   # (N, H, W) bool, original image size
        result.boxes   # (N, 4) float32, [x0, y0, x1, y1] in original image pixels
        result.scores  # (N,) float32, sorted high to low

The model runs the same code path as the HTTP server (``tt_sam3.server.app``): the same device traces, the same
host pre- and post-processing and the same instance selection, so the outputs are the same bits.  Full reference:
``code/PYTHON.md``.
"""
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
from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

from tt_sam3.imaging import INPUT_SIZE, DecodedImage, _png_fast_decode, to_pixel_values

log = logging.getLogger("tt_sam3")

DEFAULT_REPO = "facebook/sam3"
#: The weights revision every published number was measured with (``TT_WEIGHTS_REVISION`` of the server).
PINNED_REVISION = "3c879f39826c281e95690f02c7821c4de09afae7"
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
}

ImageLike = Union[str, os.PathLike, bytes, bytearray, memoryview, Image.Image, np.ndarray, torch.Tensor]


# ---------------------------------------------------------------------------------------------------------------------
# Result
# ---------------------------------------------------------------------------------------------------------------------
@dataclass
class Sam3Result:
    """Instances of one text prompt on one image, sorted by score (high to low).

    Attributes:
        prompt: the text prompt (stripped).
        scores: ``(N,)`` float32 in [0, 1]. ``sigmoid(logit) * presence``, the upstream ``Sam3Processor`` score.
        boxes: ``(N, 4)`` float32 ``[x0, y0, x1, y1]`` in pixels of the original image, clamped to the image.
        masks: ``(N, H, W)`` bool at the original image size, or None when ``return_masks=False``.
            Pixel ``(y, x)`` is True when the bilinear upsample of the mask logits is > 0 (``sigmoid > 0.5``).
        presence: float in [0, 1]. Probability that the concept is in the image at all.
        width, height: original image size in pixels.
        mask_logits: ``(N, 288, 288)`` float32 raw mask logits (1008x1008 squashed frame), only with
            ``return_mask_logits=True``.

    Dict-style access also works, as with the upstream ``Sam3Processor`` state: ``result["masks"]``.
    """

    prompt: str
    scores: np.ndarray
    boxes: np.ndarray
    masks: Optional[np.ndarray]
    presence: float
    width: int
    height: int
    mask_logits: Optional[np.ndarray] = field(default=None, repr=False)

    def __len__(self) -> int:
        return int(self.scores.shape[0])

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

    def keys(self):
        return list(self.__dataclass_fields__)

    def to_dict(self) -> Dict[str, Any]:
        """JSON-friendly summary (no masks): prompt, presence, scores, boxes and mask areas."""
        return {
            "prompt": self.prompt,
            "presence": float(self.presence),
            "num_instances": len(self),
            "scores": [float(s) for s in self.scores],
            "boxes_xyxy": [[float(v) for v in b] for b in self.boxes],
            "mask_areas": None if self.masks is None else [int(m.sum()) for m in self.masks],
        }

    def __repr__(self) -> str:
        top = ", ".join(f"{s:.3f}" for s in self.scores[:5]) + (", ..." if len(self) > 5 else "")
        return (f"Sam3Result(prompt={self.prompt!r}, num_instances={len(self)}, scores=[{top}], "
                f"presence={self.presence:.3f}, size={self.width}x{self.height})")


# ---------------------------------------------------------------------------------------------------------------------
# Input handling (no device)
# ---------------------------------------------------------------------------------------------------------------------
def _decode_bytes(raw: bytes) -> DecodedImage:
    """PNG / JPEG / any PIL format -> RGB uint8 (same pixels as the server's decode)."""
    arr = _png_fast_decode(raw)
    if arr is not None:
        return DecodedImage(arr[..., :3] if arr.shape[2] == 4 else arr)
    im = Image.open(io.BytesIO(raw))
    im.load()
    if im.mode in ("RGB", "RGBA"):
        arr = np.asarray(im)
        return DecodedImage(arr[..., :3] if im.mode == "RGBA" else arr)
    return DecodedImage(np.asarray(im.convert("RGB")))


def _float_to_uint8(a: np.ndarray) -> np.ndarray:
    a = a.astype(np.float32, copy=False)
    if a.size and float(np.nanmax(a)) > 1.5:  # looks like 0..255 values (a [0, 1] image never gets here)
        return np.clip(np.rint(a), 0, 255).astype(np.uint8)
    # [0, 1] floats: the upstream Sam3Processor's ToDtype(torch.uint8, scale=True) (x * 255.999, truncated)
    return (np.clip(a, 0.0, 1.0) * np.float32(255.999)).astype(np.uint8)


def _array_to_rgb(a: np.ndarray) -> np.ndarray:
    if a.ndim == 4 and a.shape[0] == 1:
        a = a[0]
    if a.ndim == 3 and a.shape[0] in (1, 3, 4) and a.shape[-1] not in (1, 3, 4):
        a = np.moveaxis(a, 0, -1)  # CHW (torch / upstream convention) -> HWC
    if a.ndim == 2:
        a = a[..., None]
    if a.ndim != 3 or a.shape[-1] not in (1, 3, 4) or a.shape[0] == 0 or a.shape[1] == 0:
        raise ValueError(f"image array must be (H, W), (H, W, 3|4) or (3|4, H, W); got shape {tuple(a.shape)}")
    if a.dtype == np.bool_:
        a = a.astype(np.uint8) * np.uint8(255)
    elif a.dtype != np.uint8:
        a = _float_to_uint8(a) if np.issubdtype(a.dtype, np.floating) else np.clip(a, 0, 255).astype(np.uint8)
    if a.shape[-1] == 1:
        a = np.repeat(a, 3, axis=-1)
    return a[..., :3]


def load_image(image: ImageLike) -> np.ndarray:
    """Any supported image input -> ``(H, W, 3)`` uint8 RGB numpy array (possibly a strided view).

    Accepts a file path (``str`` / ``os.PathLike``), encoded file bytes (PNG / JPEG / ...), a ``PIL.Image``
    (any mode; alpha is dropped, not composited, as the server does), or a numpy array / torch tensor in
    ``(H, W)``, ``(H, W, 3|4)`` or ``(3|4, H, W)`` layout.  uint8 arrays are used as they are; float arrays in
    [0, 1] are scaled like upstream ``Sam3Processor`` (x * 255.999, truncated); float arrays with values > 1.5 are
    taken as 0..255 and rounded.
    """
    return _load(image)[0]


def _load(image: ImageLike) -> Tuple[np.ndarray, Any, str]:
    """-> (rgb, buffer to hash for the image cache, cache tag)."""
    if isinstance(image, (str, os.PathLike)):
        with open(os.fspath(image), "rb") as f:
            raw = f.read()
        return _decode_bytes(raw).rgb, raw, "enc"
    if isinstance(image, (bytes, bytearray, memoryview)):
        raw = bytes(image)
        return _decode_bytes(raw).rgb, raw, "enc"
    if isinstance(image, Image.Image):
        if image.mode in ("RGB", "RGBA"):
            arr = np.asarray(image)
            return (arr[..., :3] if image.mode == "RGBA" else arr), arr, f"pil{image.mode}{arr.shape}"
        arr = np.asarray(image.convert("RGB"))
        return arr, arr, f"pil{arr.shape}"
    if isinstance(image, torch.Tensor):
        image = image.detach().cpu().numpy()
    if isinstance(image, np.ndarray):
        rgb = _array_to_rgb(image)
        return rgb, None, f"rgb{rgb.shape}"
    raise TypeError(f"unsupported image type {type(image).__name__}: pass a path, bytes, PIL.Image, numpy array or torch tensor")


def _check_prompts(prompts) -> Tuple[List[str], bool]:
    single = isinstance(prompts, str)
    plist = [prompts] if single else list(prompts)
    if not plist:
        raise ValueError("at least one prompt is needed")
    out = []
    for p in plist:
        if not isinstance(p, str) or not p.strip():
            raise ValueError(f"prompts must be non-empty strings, got {p!r}")
        out.append(p.strip())
    return out, single


def boxes_cxcywh_to_xyxy_pixels(boxes_cxcywh: torch.Tensor, width: int, height: int) -> np.ndarray:
    """Normalised ``(N, 4)`` cx, cy, w, h -> ``(N, 4)`` float32 x0, y0, x1, y1 pixels clamped to the image (the
    server's arithmetic, without its 2-decimal rounding)."""
    if boxes_cxcywh.numel() == 0:
        return np.zeros((0, 4), dtype=np.float32)
    cx, cy, w, h = boxes_cxcywh.float().unbind(-1)
    xyxy = torch.stack([cx - 0.5 * w, cy - 0.5 * h, cx + 0.5 * w, cy + 0.5 * h], dim=-1)
    xyxy = xyxy * torch.tensor([width, height, width, height], dtype=torch.float32)
    xyxy[:, 0::2] = xyxy[:, 0::2].clamp(0, width)
    xyxy[:, 1::2] = xyxy[:, 1::2].clamp(0, height)
    return xyxy.numpy().astype(np.float32, copy=False)


# ---------------------------------------------------------------------------------------------------------------------
# Device open
# ---------------------------------------------------------------------------------------------------------------------
def _tt_metal_home() -> Optional[str]:
    from tt_sam3.tt.device_open import tt_metal_home

    return tt_metal_home()


def eth_dispatch_available() -> bool:
    """True when the tt-metal tree has the Blackhole ETH-dispatch patch (1 CQ without dispatch_s), the p150
    target configuration (ETH dispatch, 1 CQ, 12x10).  Detected from ``tt_metal/impl/dispatch/topology.cpp``."""
    from tt_sam3.tt.device_open import eth_dispatch_available as _avail

    return _avail()


def device_open_kwargs(device_id: int = 0, dispatch: str = "auto") -> Dict[str, Any]:
    """The ``ttnn.open_device`` keyword arguments ``Sam3.from_pretrained`` uses (``l1_small_size``,
    ``trace_region_size`` and, for ETH dispatch, ``dispatch_core_config`` + ``num_command_queues=1``).  Use them to
    open the device yourself and pass it as ``Sam3.from_pretrained(device=...)``; for ETH dispatch also set
    ``TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE=11,9`` (12x10 grid) before tt-metal starts."""
    from tt_sam3.tt.device_open import add_dispatch_kwargs
    from tt_sam3.tt.fused_config import FusedConfig
    from tt_sam3.tt.fused_config import device_open_kwargs as _kw

    kwargs = _kw(FusedConfig.from_env(), device_id=device_id, l1_small_size=int(os.environ.get("SAM3_L1_SMALL_SIZE", "0")))
    return add_dispatch_kwargs(kwargs, _resolve_dispatch(dispatch))


def _resolve_dispatch(dispatch: str) -> str:
    from tt_sam3.tt.device_open import resolve_dispatch

    return resolve_dispatch(dispatch, stacklevel=4)


def _open_device(device_id: int, dispatch: str):
    from tt_sam3.tt.device_open import open_device
    from tt_sam3.tt.fused_config import FusedConfig
    from tt_sam3.tt.fused_config import device_open_kwargs as _kw

    base = _kw(FusedConfig.from_env(), device_id=device_id, l1_small_size=int(os.environ.get("SAM3_L1_SMALL_SIZE", "0")))
    return open_device(base, dispatch)


def _warmup_input() -> torch.Tensor:
    """The deterministic structured pattern the server warms with (``app.warmup_input``)."""
    x = torch.linspace(-1, 1, INPUT_SIZE).view(1, 1, 1, INPUT_SIZE).expand(1, 1, INPUT_SIZE, INPUT_SIZE)
    y = torch.linspace(-1, 1, INPUT_SIZE).view(1, 1, INPUT_SIZE, 1).expand(1, 1, INPUT_SIZE, INPUT_SIZE)
    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]


def _resolve_checkpoint(repo_id: str, revision: Optional[str], checkpoint_path: Optional[str]) -> str:
    if checkpoint_path:
        if os.path.isdir(checkpoint_path):
            checkpoint_path = os.path.join(checkpoint_path, WEIGHTS_FILENAME)
        if not os.path.isfile(checkpoint_path):
            raise FileNotFoundError(f"checkpoint not found: {checkpoint_path!r}")
        return checkpoint_path
    weights_dir = os.environ.get("SAM3_WEIGHTS_DIR")
    if weights_dir:
        return _resolve_checkpoint(repo_id, revision, weights_dir)
    from huggingface_hub import hf_hub_download

    return hf_hub_download(repo_id=repo_id, filename=WEIGHTS_FILENAME, revision=revision)


def _close_at_exit(ref):
    model = ref()
    if model is not None:
        try:
            model.close()
        except Exception:  # noqa: BLE001
            pass


# ---------------------------------------------------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------------------------------------------------
class Sam3:
    """SAM3 text-prompted instance segmentation on one Blackhole chip.

    Create with :meth:`from_pretrained`; call with an image and one or more noun-phrase prompts.  Use as a
    context manager or call :meth:`close` to release the traces and the device.
    """

    def __init__(self, pipeline, device, *, owns_device: bool, info: Dict[str, Any], cache_image: bool = True):
        self.pipeline = pipeline
        self.device = device
        self.info = info
        self._owns_device = owns_device
        self._lock = threading.RLock()
        self._closed = False
        self._rgb_input = bool(info.get("rgb_input"))
        self._fast_mask = bool(info.get("fast_mask"))
        try:
            import xxhash
        except ImportError:
            xxhash = None
        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)

    # ---------------------------------------------------------------- construction
    @classmethod
    def from_pretrained(
        cls,
        repo_id: str = DEFAULT_REPO,
        *,
        revision: Optional[str] = None,
        checkpoint_path: Optional[str] = None,
        device_id: int = 0,
        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).

        Args:
            repo_id: Hugging Face repo with ``sam3.pt`` (gated: accept the license and ``hf auth login`` first).
            revision: weights revision; default ``TT_WEIGHTS_REVISION`` or the pinned :data:`PINNED_REVISION`.
            checkpoint_path: local ``sam3.pt`` (or a directory with it) instead of the Hub / HF cache.
                ``SAM3_WEIGHTS_DIR`` does the same from the environment.
            device_id: chip to open (``0`` when ``TT_VISIBLE_DEVICES`` pins one chip).
            device: an already opened ``ttnn`` device to use instead (it must be opened with
                :func:`device_open_kwargs`; :meth:`close` then does not close it).
            dispatch: ``"auto"`` (default: ETH dispatch + 1 command queue + 12x10 grid, the p150 target, when the
                tt-metal tree has the ETH-dispatch patch; else the stock dispatch with a ``RuntimeWarning``),
                ``"eth"``, or ``"worker"`` (explicit opt-in: stock Tensix-column dispatch, 11x10 grid on a p150).
            warmup: warm-up forwards (the first one compiles the kernels and captures the traces; minutes when the
                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

        if revision is None:
            revision = os.environ.get("TT_WEIGHTS_REVISION") or (PINNED_REVISION if repo_id == DEFAULT_REPO else None)
        if not _PIN_PATCHED[0]:
            _patch_pin_memory()  # sam3's geometry encoder calls Tensor.pin_memory(), which raises on CPU-only torch
            _PIN_PATCHED[0] = True

        with torch.no_grad():
            ckpt = _resolve_checkpoint(repo_id, revision, checkpoint_path)
            bpe = resolve_bpe_path()
            log.info("Loading checkpoint %s", ckpt)
            sam3_model = build_sam3_model(use_pretrained=True, checkpoint_path=ckpt, bpe_path=bpe)

            owns = device is None
            dispatch_used, grid = "external", None
            if owns:
                device, dispatch_used, grid = _open_device(device_id, dispatch)
            else:
                g = device.compute_with_storage_grid_size()
                grid = (g.x, g.y)
            pipeline = None
            try:
                fused = FusedConfig.from_env()
                pipeline = TtSam3ImagePipeline(sam3_model, device, torch_compile=False, num_threads=num_threads,
                                               fused_config=fused)
                px = _warmup_input()
                for i in range(max(0, int(warmup))):
                    log.info("Warm-up %d/%d", i + 1, warmup)
                    with torch.inference_mode():
                        out = pipeline.forward_image(px, [WARMUP_PROMPT])
                    if out is None or "pred_masks" not in out:
                        raise RuntimeError("warm-up forward returned no predictions")
                if fused.enabled and fused.enc_trace and warmup > 0:
                    if not pipeline.encoder_trace_captured or pipeline.encoder_fallbacks != 0:
                        raise RuntimeError("warm-up did not run entirely through the device traces "
                                           f"(captured={pipeline.encoder_trace_captured}, "
                                           f"fallbacks={pipeline.encoder_fallbacks})")
                fast_mask = os.environ.get("SAM3_FAST_MASK", "1") == "1" and maskpost.self_check()
                rgb_input = _rgb_input_ok(pipeline, warm=warmup > 0)
            except BaseException:
                _release_pipeline(pipeline)
                if owns:
                    import ttnn

                    ttnn.close_device(device)
                raise
        info = {
            "checkpoint": ckpt,
            "revision": revision,
            "dispatch": dispatch_used,
            "grid": grid,
            "fused": pipeline.fused.describe(),
            "rgb_input": rgb_input,
            "fast_mask": fast_mask,
            "torch_threads": torch.get_num_threads(),
        }
        return cls(pipeline, device, owns_device=owns, info=info, cache_image=cache_image)

    # ---------------------------------------------------------------- inference
    def __call__(
        self,
        image: ImageLike,
        prompts: Union[str, Sequence[str]],
        *,
        threshold: float = 0.5,
        max_instances: int = 50,
        return_masks: bool = True,
        return_mask_logits: bool = False,
    ) -> Union[Sam3Result, List[Sam3Result]]:
        """Segment every instance of each text prompt in ``image``.

        Args:
            image: file path, encoded bytes, ``PIL.Image``, numpy array or torch tensor (see :func:`load_image`).
                Any size up to 4096 px per side is fine; the model squashes it to 1008x1008 internally.
            prompts: one noun phrase (``"car"``) or a list of them.  The image part runs once; each prompt adds
                about 20 ms.
            threshold: keep instances with score > threshold (upstream default 0.5).
            max_instances: per-prompt cap (1..200), highest scores first.
            return_masks: False skips the upsample to the original size (``result.masks`` is None).
            return_mask_logits: also return the raw ``(N, 288, 288)`` mask logits.

        Returns:
            One :class:`Sam3Result` for a string prompt, a list of them (same order) for a list of prompts.
        """
        prompt_list, single = _check_prompts(prompts)
        if not 0.0 <= float(threshold) <= 1.0:
            raise ValueError("threshold must be in [0, 1]")
        if not 1 <= int(max_instances) <= NUM_QUERIES:
            raise ValueError(f"max_instances must be in [1, {NUM_QUERIES}]")
        loaded = self._load_cached(image)
        results = self._run(loaded, prompt_list, float(threshold), int(max_instances), return_masks, return_mask_logits)
        return results[0] if single else results

    def predict_many(
        self,
        images: Iterable[ImageLike],
        prompts: Union[str, Sequence[str]],
        **kwargs,
    ) -> List[Union[Sam3Result, List[Sam3Result]]]:
        """Run :meth:`__call__` on several images with the same prompts.  The next image is read and decoded on a
        host thread while the chip works on the current one; the outputs are the same as separate calls."""
        prompt_list, single = _check_prompts(prompts)
        images = list(images)
        out = []
        with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sam3-load") as pool:
            nxt = pool.submit(self._load_cached, images[0]) if images else None
            for i in range(len(images)):
                loaded = nxt.result()
                nxt = pool.submit(self._load_cached, images[i + 1]) if i + 1 < len(images) else None
                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):
            return image
        if self.cache_image and isinstance(image, (str, os.PathLike, bytes, bytearray, memoryview)):
            if isinstance(image, (str, os.PathLike)):
                with open(os.fspath(image), "rb") as f:
                    raw = f.read()
            else:
                raw = bytes(image)
            key = ("enc", len(raw), self._xxhash.xxh3_128_digest(raw))
            last = self._last_decoded[:]
            if last[0] == key:
                return _Loaded(last[1], key)
            rgb = _decode_bytes(raw).rgb
            self._last_decoded[:] = [key, rgb]
            return _Loaded(rgb, key)
        rgb, buf, tag = _load(image)
        key = None
        if self.cache_image:
            b = np.ascontiguousarray(rgb if buf is None else buf)
            key = (tag, b.shape, self._xxhash.xxh3_128_digest(b))
        return _Loaded(rgb, key)

    def _run(self, loaded: "_Loaded", prompts, threshold, max_instances, return_masks, return_mask_logits):
        if self._closed:
            raise RuntimeError("Sam3 model is closed")
        pipeline = self.pipeline
        rgb = loaded.rgb
        height, width = int(rgb.shape[0]), int(rgb.shape[1])
        pipeline.prefetch_text(prompts)  # text-cache misses start on host threads now (as the server does)
        with self._lock:
            try:
                with torch.inference_mode():
                    if self._rgb_input:
                        selected = pipeline.forward_rgb_selected_multi(rgb, prompts, threshold, max_instances,
                                                                       image_key=loaded.key)
                    else:
                        px = to_pixel_values(DecodedImage(rgb))
                        selected = pipeline.forward_image_selected_multi(px, prompts, threshold, max_instances)
            except BaseException:
                self._invalidate_cache()
                raise
        return [self._result(p, sel, width, height, return_masks, return_mask_logits) for p, sel in zip(prompts, selected)]

    def _result(self, prompt, sel, width, height, return_masks, return_mask_logits) -> Sam3Result:
        scores, boxes, mask_logits, presence = sel
        masks = None
        if return_masks:
            masks = np.zeros((int(scores.numel()), height, width), dtype=bool)
            for i in range(masks.shape[0]):
                masks[i] = self._upsample(mask_logits[i], width, height)
        return Sam3Result(
            prompt=prompt,
            scores=scores.float().numpy().astype(np.float32, copy=False),
            boxes=boxes_cxcywh_to_xyxy_pixels(boxes, width, height),
            masks=masks,
            presence=float(presence),
            width=width,
            height=height,
            mask_logits=mask_logits.float().numpy() if return_mask_logits else None,
        )

    def _upsample(self, logits: torch.Tensor, width: int, height: int) -> np.ndarray:
        if self._fast_mask:  # bounding-box crop, bit-for-bit the full upsample (self-checked at startup)
            from tt_sam3.server import maskpost

            return maskpost.upsample_mask(logits.numpy(), width, height)
        up = F.interpolate(logits[None, None].float(), size=(height, width), mode="bilinear", align_corners=False)
        return (up[0, 0] > 0).numpy()

    def _invalidate_cache(self):
        self._last_decoded[:] = [None, None]
        runner = getattr(self.pipeline, "_full", None)
        if runner is not None:
            runner.img_key = None

    # ---------------------------------------------------------------- lifetime
    def close(self) -> None:
        """Release the device traces and buffers and close the device (if this object opened it).  Idempotent."""
        with self._lock:
            if self._closed:
                return
            self._closed = True
            atexit.unregister(self._atexit)
            _release_pipeline(self.pipeline)
            self.pipeline = None
            gc.collect()
            if self._owns_device and self.device is not None:
                import ttnn

                log.info("Closing device")
                ttnn.close_device(self.device)
            self.device = None

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

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

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

    def __repr__(self) -> str:
        state = "closed" if self._closed else f"dispatch={self.info.get('dispatch')}, grid={self.info.get('grid')}"
        return f"Sam3({state})"


@dataclass
class _Loaded:
    rgb: np.ndarray
    key: Optional[Tuple] = None


def _rgb_input_ok(pipeline, warm: bool) -> bool:
    """The server's ``_fused_input_ok``: the banded uint8 input path after its bit-for-bit self check."""
    if os.environ.get("SAM3_FUSED_INPUT", "1") != "1" or not pipeline.rgb_input_supported():
        return False
    from tt_sam3.server import hostinput

    hostinput.configure(torch.get_num_threads())
    ok = hostinput.self_check()
    if ok and warm:
        with torch.inference_mode():  # warm the band path through the traces once
            pipeline.forward_rgb_selected_multi(np.zeros((720, 1280, 3), dtype=np.uint8), [WARMUP_PROMPT], 0.5, 1)
    return ok


def _release_pipeline(pipeline) -> None:
    if pipeline is None:
        return
    try:
        pipeline.release()
    except Exception:  # noqa: BLE001 - close must reach close_device
        log.exception("pipeline.release() failed")
    for attr in ("backbone_params", "_enc_params", "_fused_encoder"):
        setattr(pipeline, attr, None)
    for pool in (getattr(pipeline, "_text_pool", None), getattr(getattr(pipeline, "_full", None), "_band_pool", None)):
        if pool is not None:
            pool.shutdown(wait=True)