code/tt_sam3/server/app.py

# SPDX-License-Identifier: Apache-2.0
"""ASGI serving app for SAM3 (text-prompted open-vocabulary instance segmentation) on a
single Tenstorrent Blackhole p150a.

Served by tt-model-manager as ``kind: tt-dit-server``::

    runtime:
      app: tt_sam3.server.app:app

uvicorn runs this module with ``--lifespan on``. Everything that touches the chip or the
weights happens inside the ASGI **lifespan**, so uvicorn's ``Application startup complete``
-- the line ``tt-model serve`` waits for -- means the device is claimed, the checkpoint is
loaded and the pipeline has already run its warm-up forwards.

Import-time this module only needs numpy / torch / fastapi / pydantic / PIL and the
port's own side-effect-free ``tt_sam3.tt.ttnn_sam3_model`` (``ttnn`` and ``sam3`` are
imported lazily in the lifespan), so the image's build-time ``verify.sh`` can import it
with no device, no weights and no ``HF_MODEL``.

Recipe (the one the port validated -- ``tests/benchmark_sam3_e2e.py``):

    _patch_pin_memory()
    model = build_sam3_model(checkpoint_path=<sam3.pt from HF_MODEL @ TT_WEIGHTS_REVISION>)
    device = open_device(cfg)    # ETH dispatch + 1 CQ + 12x10 grid (SAM3_DISPATCH=auto), single chip
    pipeline = TtSam3ImagePipeline(model, device, torch_compile=False)
    pipeline.forward_image(px_1x3x1008x1008, ["object"])       # x SAM3_WARMUP_RUNS

Configuration (environment, read in the lifespan -- never at import):

    HF_MODEL              weights repo id            (default facebook/sam3; tt-model sets it)
    TT_WEIGHTS_REVISION   commit sha of the weights  (default None = repo default branch)
    SAM3_WEIGHTS_DIR      local dir holding sam3.pt  (offline / host override; skips the Hub)
    TT_MESH_SHAPE         "1x1" | "(1, 1)" | "1,1"   (anything but a 1-chip mesh is refused)
    TT_DEVICE_ID          chip to open               (default 0)
    SAM3_DISPATCH         "auto" (default) | "eth" | "worker". auto = ETH dispatch + 1 command queue + 12x10
                          compute grid (the p150 target) when tt-metal has the Blackhole ETH-dispatch patch, else
                          the stock dispatch with a warning. "worker" = explicit opt-in to the stock Tensix-column
                          dispatch (11x10 on a p150; 12x10 only on a Galaxy chip). Outputs are the same.
    SAM3_BPE_PATH         CLIP BPE vocab .txt.gz     (default: vendored tt_sam3/assets copy)
    SAM3_TORCH_COMPILE    "1" to torch.compile the decoder (default "0": no C++ compiler in the image)
    SAM3_NUM_THREADS      torch CPU threads          (default 17, capped at os.cpu_count())
    SAM3_WARMUP_RUNS      warm forwards before READY (default 2)
    SAM3_L1_SMALL_SIZE    ttnn l1_small_size         (default 0 = ttnn default, as validated)
    SAM3_MAX_PROMPTS      prompts per request        (default 4)
    SAM3_MAX_SIDE         longest accepted image side (default 4096)
    SAM3_IMG_CACHE        "1" (default; needs the optional xxhash package) = image-result cache: a request whose
                          image bytes equal the previous request's (xxh3_128 key) reuses the decoded image and the
                          image-trace outputs still resident on the device (no decode, upload or image trace; same
                          output bits). "0" = every request runs the whole forward. Read at import.
    TT_FUSED              fused device paths (default ON since the device pass of 2026-09-13; "0" =
                          the legacy graph, bit-for-bit). Read once here (device open needs the trace
                          region) and passed to the pipeline. With it: SAM3_TRACE_REGION_SIZE ->
                          ttnn.open_device(trace_region_size=), and the encoder trace is captured
                          during warm-up run 1, so READY implies a captured trace AND zero legacy
                          fallbacks in warm-up (startup fails otherwise; /info.runtime reports both).
                          Sub-knobs and their validated defaults: tt_sam3/tt/fused_config.py.
"""
from __future__ import annotations

import base64
import binascii
import gc
import io
import logging
import os
import struct
import threading
import time
import zlib
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any, Dict, List, Literal, Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as F
from fastapi import FastAPI, HTTPException, Query, Request
from fastapi.concurrency import run_in_threadpool
from PIL import Image
from pydantic import BaseModel, Field

from tt_sam3.server import maskpost as _maskpost
from tt_sam3.tt.fused_config import FusedConfig, device_open_kwargs
from tt_sam3.tt.ttnn_sam3_model import (
    TtSam3ImagePipeline,
    _patch_pin_memory,
    build_sam3_model,
    preprocess_image,
    resolve_bpe_path,
)

log = logging.getLogger("tt_sam3.server")
if not log.handlers:
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")

# ---------------------------------------------------------------------------------------
# Constants (what the port validated)
# ---------------------------------------------------------------------------------------
MODEL_NAME = "sam3-p150"
DEFAULT_WEIGHTS_REPO = "facebook/sam3"
WEIGHTS_FILENAME = "sam3.pt"
INPUT_SIZE = 1008  # fixed: pos_embed / SDPA chunking are precomputed for the 72x72 patch grid
NUM_QUERIES = 200
MASK_RES = 288  # pred_masks are 288x288 logits (stride 3.5 of the 1008 input)
SOURCE_REPO = "https://github.com/changh95/tt-sam3"
SOURCE_COMMIT = "cb728e85275334973b39d5acf6ae376ff2a94a68"
DEFAULT_WARMUP_PROMPT = "object"

STATE: Dict[str, Any] = {}
LOCK = threading.Lock()  # one chip, one globally monkey-patched torch model: serialize


# ---------------------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------------------
@dataclass
class Config:
    weights_repo: str
    weights_revision: Optional[str]
    weights_dir: Optional[str]
    mesh_shape: str
    device_id: int
    bpe_path: str
    torch_compile: bool
    warmup_runs: int
    l1_small_size: int
    max_prompts: int
    max_side: int
    fused: FusedConfig = FusedConfig.disabled()
    dispatch: str = "auto"


def _env_flag(name: str, default: str) -> bool:
    return os.environ.get(name, default).strip().lower() not in ("0", "false", "no", "off", "")


def read_config() -> Config:
    return Config(
        weights_repo=os.environ.get("HF_MODEL") or DEFAULT_WEIGHTS_REPO,
        weights_revision=os.environ.get("TT_WEIGHTS_REVISION") or None,
        weights_dir=os.environ.get("SAM3_WEIGHTS_DIR") or None,
        mesh_shape=os.environ.get("TT_MESH_SHAPE", "1x1"),
        device_id=int(os.environ.get("TT_DEVICE_ID", "0")),
        bpe_path=resolve_bpe_path(),
        torch_compile=_env_flag("SAM3_TORCH_COMPILE", "0"),
        warmup_runs=int(os.environ.get("SAM3_WARMUP_RUNS", "2")),
        l1_small_size=int(os.environ.get("SAM3_L1_SMALL_SIZE", "0")),
        max_prompts=int(os.environ.get("SAM3_MAX_PROMPTS", "4")),
        max_side=int(os.environ.get("SAM3_MAX_SIDE", "4096")),
        fused=FusedConfig.from_env(),  # TT_FUSED read once for the whole process
        dispatch=os.environ.get("SAM3_DISPATCH", "auto").strip().lower() or "auto",
    )


def parse_mesh_shape(value: str) -> Tuple[int, int]:
    """Accept the launcher's ``1x1`` as well as ``(1, 1)`` / ``1,1``; refuse anything else."""
    text = value.strip().lower().replace("(", "").replace(")", "").replace(" ", "")
    for sep in ("x", ","):
        if sep in text:
            parts = text.split(sep)
            if len(parts) == 2 and all(p.isdigit() for p in parts):
                return int(parts[0]), int(parts[1])
    raise RuntimeError(
        f"TT_MESH_SHAPE={value!r} is not understood; expected a shape like '1x1', '(1, 1)' or '1,1'."
    )


# ---------------------------------------------------------------------------------------
# Lifespan pieces
# ---------------------------------------------------------------------------------------
def resolve_checkpoint(cfg: Config) -> str:
    """Path of ``sam3.pt``: a local dir override, else the HF cache at the pinned revision.

    ``hf_hub_download(revision=<sha>)`` resolves from the cache without any network call when
    tt-model pre-downloaded the snapshot; without a sha it must resolve ``main`` on the Hub.
    """
    if cfg.weights_dir:
        path = os.path.join(cfg.weights_dir, WEIGHTS_FILENAME)
        if not os.path.isfile(path):
            raise FileNotFoundError(f"SAM3_WEIGHTS_DIR={cfg.weights_dir!r} does not contain {WEIGHTS_FILENAME}")
        return path
    from huggingface_hub import hf_hub_download

    if not cfg.weights_revision:
        log.warning("TT_WEIGHTS_REVISION is not set: resolving %s@main on the Hub (network needed)", cfg.weights_repo)
    return hf_hub_download(repo_id=cfg.weights_repo, filename=WEIGHTS_FILENAME, revision=cfg.weights_revision)


def open_device(cfg: Config):
    """Open the single chip: ``fused_config.device_open_kwargs`` (l1_small / trace region) + the dispatch of
    ``SAM3_DISPATCH`` (default auto = ETH dispatch, 1 CQ, 12x10 grid; ``tt_sam3.tt.device_open``)."""
    rows, cols = parse_mesh_shape(cfg.mesh_shape)
    if rows * cols != 1:
        raise RuntimeError(
            f"TT_MESH_SHAPE={cfg.mesh_shape!r} is a {rows * cols}-chip mesh; this port runs on exactly one "
            "Blackhole chip (mesh_device P150). Use ttnn.open_mesh_device and shard the model first."
        )
    # shared with benchmark.py / conftest.py: bare open for the legacy path, l1_small_size only when set
    # (SAM3_L1_SMALL_SIZE; the E2E runs used the default), trace_region_size when the fused encoder trace is on
    from tt_sam3.tt.device_open import open_device as _open

    kwargs: Dict[str, Any] = device_open_kwargs(cfg.fused, device_id=cfg.device_id, l1_small_size=cfg.l1_small_size)
    device, mode, grid = _open(kwargs, cfg.dispatch)
    STATE["dispatch"] = {"dispatch": mode, "num_command_queues": 1, "grid": f"{grid[0]}x{grid[1]}"}
    return device


def warmup_input() -> torch.Tensor:
    """The deterministic structured pattern the benchmark warms with (not random noise)."""
    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()


@asynccontextmanager
async def lifespan(_app: FastAPI):
    cfg = read_config()
    torch.set_grad_enabled(False)
    _patch_pin_memory()  # sam3's geometry encoder calls Tensor.pin_memory(), which raises on CPU-only torch

    # 1) weights + reference model on the host (no device yet)
    t0 = time.perf_counter()
    log.info("Loading weights %s@%s (%s)", cfg.weights_repo, cfg.weights_revision or "main", WEIGHTS_FILENAME)
    ckpt_path = resolve_checkpoint(cfg)
    if not os.path.isfile(cfg.bpe_path):
        raise FileNotFoundError(f"BPE vocabulary not found at {cfg.bpe_path!r} (set SAM3_BPE_PATH)")
    log.info("Loading checkpoint %s (bpe=%s)", ckpt_path, cfg.bpe_path)
    model = build_sam3_model(use_pretrained=True, checkpoint_path=ckpt_path, bpe_path=cfg.bpe_path)
    n_params = sum(p.numel() for p in model.parameters()) / 1e6
    t_load = time.perf_counter() - t0
    log.info("Reference model ready: %.1fM params in %.1fs", n_params, t_load)

    # 2) device + ttnn pipeline (weights converted to bfloat8_b/bf16 and uploaded)
    device = open_device(cfg)
    pipeline = None
    try:
        t1 = time.perf_counter()
        log.info("Loading pipeline (torch_compile=%s, fused=%s)", cfg.torch_compile, cfg.fused.describe())
        pipeline = TtSam3ImagePipeline(model, device, torch_compile=cfg.torch_compile, fused_config=cfg.fused)
        log.info("Pipeline built in %.1fs (torch threads=%d)", time.perf_counter() - t1, torch.get_num_threads())

        # 3) warm-up: JIT-compiles the ttnn kernels (minutes cold, cached under TT_METAL_CACHE)
        px = warmup_input()
        t2 = time.perf_counter()
        for i in range(max(0, cfg.warmup_runs)):
            log.info("Warming up (%d/%d) ...", i + 1, cfg.warmup_runs)
            with torch.inference_mode():
                out = pipeline.forward_image(px, [DEFAULT_WARMUP_PROMPT])
            if out is None or "pred_masks" not in out:
                raise RuntimeError("warm-up forward returned no predictions")
            log.info("Warmup run %d took %.2fs", i + 1, time.perf_counter() - t2)
            t2 = time.perf_counter()
        if cfg.fused.enabled and cfg.fused.enc_trace and cfg.warmup_runs > 0:
            # READY contract under TT_FUSED=1: the encoder trace is captured before startup completes AND
            # every warm-up forward ran through it (a fallback to the legacy eager encoder means the served
            # shape is not the traced one; refuse READY rather than serve the un-traced path silently).
            if not pipeline.encoder_trace_captured or pipeline.encoder_fallbacks != 0:
                raise RuntimeError(
                    "TT_FUSED=1: warm-up did not run entirely through the fused encoder trace "
                    f"(captured={pipeline.encoder_trace_captured}, legacy fallbacks={pipeline.encoder_fallbacks}); "
                    "refusing to report READY"
                )
            log.info("Encoder trace captured (%s)", pipeline.encoder_stats)
        _fast_mask_ok()  # one-time bit-for-bit self check of the crop post-processing against F.interpolate
        STATE["rgb_input"] = _fused_input_ok(pipeline)
        log.info("Warmup complete")

        STATE.update(
            rgb_input=STATE.get("rgb_input", False),
            config=cfg,
            device=device,
            pipeline=pipeline,
            checkpoint_path=ckpt_path,
            num_params_m=n_params,
            device_arch=str(device.arch()),
            ready_at=time.time(),
        )
        yield
    finally:
        import ttnn

        STATE.clear()
        # drop the device tensors held by the pipeline before closing the chip
        if pipeline is not None:
            try:
                pipeline.release()  # TT_FUSED=1: release the encoder trace + persistent buffers first
            except Exception:  # noqa: BLE001 - shutdown must reach close_device
                log.exception("pipeline.release() failed")
            for attr in ("backbone_params", "_enc_params", "_fused_encoder"):
                setattr(pipeline, attr, None)
        del pipeline
        gc.collect()
        log.info("Closing device")
        ttnn.close_device(device)


app = FastAPI(title="SAM3 on Tenstorrent Blackhole", lifespan=lifespan)


# ---------------------------------------------------------------------------------------
# Request / response
# ---------------------------------------------------------------------------------------
class PredictRequest(BaseModel):
    """One image, one to four noun-phrase prompts. Each prompt is an independent query."""

    image: str = Field(..., description="base64 PNG/JPEG (any size, RGB or RGBA; squashed to 1008x1008)")
    prompts: List[str] = Field(..., min_length=1, max_length=8, description="noun phrases, e.g. ['car']")
    threshold: float = Field(0.5, ge=0.0, le=1.0, description="keep instances with score > threshold (upstream default 0.5)")
    max_instances: int = Field(50, ge=1, le=NUM_QUERIES, description="per-prompt cap, highest scores first")
    mask_format: Literal["png", "rle", "none"] = Field("png", description="per-instance mask encoding at original resolution")


_B64_ALPHABET = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
# host worker pool for the per-instance mask upsample + PNG encode (torch / zlib release the GIL)
_POOL = ThreadPoolExecutor(max_workers=int(os.environ.get("SAM3_POST_THREADS", "4")), thread_name_prefix="sam3-post")


try:  # SIMD base64 (validate=True has the same strictness as base64.b64decode(validate=True)): 4.8 MB demo 0.2 ms
    import pybase64 as _pybase64
except ImportError:  # optional dependency
    _pybase64 = None


def _b64decode_strict(b64: str) -> bytes:
    """``base64.b64decode(b64, validate=True)`` semantics (only the base64 alphabet, '=' only as trailing padding,
    binascii padding checks) without the regex pass over the whole payload (4.8 MB demo: 11.9 -> ~8 ms with
    binascii, 0.2 ms with pybase64)."""
    if _pybase64 is not None:
        return _pybase64.b64decode(b64, validate=True)
    raw = b64.encode("ascii")  # non-ASCII -> UnicodeEncodeError (a client error)
    if raw.translate(None, _B64_ALPHABET):
        raise binascii.Error("Non-base64 digit found")
    eq = raw.find(b"=")
    if eq != -1 and (len(raw) - eq > 2 or raw[eq:] not in (b"=", b"==")):
        raise binascii.Error("Non-base64 digit found")
    return binascii.a2b_base64(raw)


# DecodedImage, the PNG fast path and the 1008x1008 preprocessing live in tt_sam3/imaging.py (shared with the
# Python API, which must not need fastapi); these are the same objects under their old names.
from tt_sam3.imaging import _PNG_SIG, DecodedImage, _libdeflate, _png_fast_decode  # noqa: E402,F401
from tt_sam3.imaging import to_pixel_values as _to_pixel_values  # noqa: E402

_PNG_FAST = os.environ.get("SAM3_FAST_PNG", "1") == "1"


def _decode_image(b64: str) -> DecodedImage:
    """base64 PNG/JPEG -> RGB uint8.  RGB / RGBA decode without PIL's convert pass: RGBA -> RGB in PIL drops the
    alpha channel (no compositing), so the RGB channels of the RGBA buffer are the same bytes."""
    try:
        raw = _b64decode_strict(b64)
    except Exception as e:  # noqa: BLE001 - any decode failure is a client error
        raise HTTPException(status_code=400, detail=f"bad image: {e}") from None
    return _decode_bytes(raw)


def _decode_bytes(raw: bytes) -> DecodedImage:
    """PNG/JPEG bytes -> RGB uint8 (see ``_decode_image``)."""
    try:
        arr = _png_fast_decode(raw) if _PNG_FAST else None
        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")))
    except Exception as e:  # noqa: BLE001 - any decode failure is a client error
        raise HTTPException(status_code=400, detail=f"bad image: {e}") from None



def _fused_input_ok(pipeline) -> bool:
    """SAM3_FUSED_INPUT=1 (default): the decoded image goes to the device runner as uint8 RGB and is resized,
    normalised and im2col'ed per window-row band straight into the upload buffers (``server/hostinput.py``), each
    band's upload overlapping the next band's host work.  Enabled only after a bit-for-bit self check against
    ``_to_pixel_values`` + im2col on this torch build (else the reference path)."""
    if os.environ.get("SAM3_FUSED_INPUT", "1") != "1" or not pipeline.rgb_input_supported():
        log.info("image input: reference path (_to_pixel_values + im2col)")
        return False
    from tt_sam3.server import hostinput

    hostinput.configure(torch.get_num_threads())
    ok = hostinput.self_check()
    if ok:
        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), [DEFAULT_WARMUP_PROMPT], 0.5, 1)
    log.info("image input: %s", "fused band path (hostinput)" if ok else "reference path (hostinput self check FAILED)")
    return ok


def _png_chunk(tag: bytes, data: bytes) -> bytes:
    return struct.pack(">I", len(data)) + tag + data + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF)


def _encode_mask_png(mask: np.ndarray) -> str:
    """(H, W) bool -> base64 8-bit grayscale PNG (0 / 255), written directly: filter type 0 on every row and zlib
    with the RLE strategy.  Decodes to exactly the same pixels as the previous PIL encode (``Image.fromarray(m *
    255, "L")``, checked in perf/dev/host_path_check.py); 1080p demo-size mask: PIL level 1 ~15 ms -> ~3 ms, and
    about half the bytes (the adaptive PNG filters PIL tries are what made it slow)."""
    h, w = mask.shape
    raw = np.empty((h, w + 1), dtype=np.uint8)
    raw[:, 0] = 0  # filter type None
    np.multiply(mask, np.uint8(255), out=raw[:, 1:], dtype=np.uint8, casting="unsafe")
    z = zlib.compressobj(1, zlib.DEFLATED, 15, 9, zlib.Z_RLE)
    idat = z.compress(raw.tobytes()) + z.flush()
    png = (b"\x89PNG\r\n\x1a\n" + _png_chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 0, 0, 0, 0))
           + _png_chunk(b"IDAT", idat) + _png_chunk(b"IEND", b""))
    return base64.b64encode(png).decode("ascii")


def _encode_mask_rle(mask: np.ndarray) -> Dict[str, Any]:
    """COCO-style uncompressed RLE: column-major runs, starting with the number of zeros."""
    flat = np.asarray(mask, dtype=np.uint8).flatten(order="F")
    if flat.size == 0:
        return {"size": [int(mask.shape[0]), int(mask.shape[1])], "counts": []}
    change = np.flatnonzero(np.diff(flat)) + 1
    starts = np.concatenate([[0], change])
    lengths = np.diff(np.concatenate([starts, [flat.size]])).tolist()
    if flat[0] == 1:
        lengths = [0] + lengths
    return {"size": [int(mask.shape[0]), int(mask.shape[1])], "counts": [int(v) for v in lengths]}


def select_instances(out: Dict[str, torch.Tensor], threshold: float, max_instances: int):
    """Upstream Sam3Processor scoring: sigmoid(logit) * sigmoid(presence), keep > threshold.

    Returns (scores (N,), boxes_cxcywh (N,4) normalized, mask_logits (N,288,288), presence float)
    sorted by score descending and capped at ``max_instances``. Runs under the lock, cheap.
    """
    logits = out["pred_logits"][0, :, 0].float()  # (200,)
    presence_logit = out["presence_logit_dec"].reshape(-1)[0].float()
    presence = torch.sigmoid(presence_logit)
    scores = torch.sigmoid(logits) * presence
    keep = torch.nonzero(scores > threshold, as_tuple=False).flatten()
    if keep.numel() > 0:
        order = torch.argsort(scores[keep], descending=True)
        keep = keep[order][:max_instances]
    return (
        scores[keep].clone(),
        out["pred_boxes"][0, keep].float().clone(),
        out["pred_masks"][0, keep].float().clone(),
        float(presence),
    )


def _boxes_to_xyxy_pixels(boxes_cxcywh: torch.Tensor, width: int, height: int) -> List[List[float]]:
    if boxes_cxcywh.numel() == 0:
        return []
    cx, cy, w, h = boxes_cxcywh.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 [[round(float(v), 2) for v in row] for row in xyxy.tolist()]


def _upsample_mask(mask_logits: torch.Tensor, width: int, height: int) -> np.ndarray:
    """(288, 288) logits -> (H, W) bool mask; bilinear, align_corners=False, > 0 (== sigmoid > 0.5).  One mask per
    call (the interpolation of each output pixel only reads its own mask, so batching does not change the values)."""
    up = F.interpolate(mask_logits[None, None], size=(height, width), mode="bilinear", align_corners=False)
    return (up[0, 0] > 0).numpy()


_FAST_MASK = os.environ.get("SAM3_FAST_MASK", "1") == "1"
_FAST_MASK_OK: Optional[bool] = None


def _fast_mask_ok() -> bool:
    """maskpost reproduces this torch build's bilinear arithmetic (bit-for-bit self check, once); else the full path."""
    global _FAST_MASK_OK
    if _FAST_MASK_OK is None:
        _FAST_MASK_OK = bool(_FAST_MASK and _maskpost.self_check())
        log.info("mask post-processing: %s", "bounding-box crop (maskpost)" if _FAST_MASK_OK else "full upsample")
    return _FAST_MASK_OK


def _mask_payload(mask_logits: torch.Tensor, width: int, height: int, mask_format: str) -> Dict[str, Any]:
    if _fast_mask_ok():  # same pixels as the full path below, computed only where the mask can be non-zero
        if mask_format == "png":
            png, area = _maskpost.mask_png_and_area(mask_logits.numpy(), width, height)
            return {"mask_area": area, "mask_png": png}
        m = _maskpost.upsample_mask(mask_logits.numpy(), width, height)
        return {"mask_area": int(np.count_nonzero(m)), "mask_rle": _encode_mask_rle(m)}
    m = _upsample_mask(mask_logits, width, height)
    out: Dict[str, Any] = {"mask_area": int(m.sum())}
    if mask_format == "png":
        out["mask_png"] = _encode_mask_png(m)
    else:
        out["mask_rle"] = _encode_mask_rle(m)
    return out


def _postprocess_all(prompts: List[str], selected: List[Any], width: int, height: int, mask_format: str) -> List[Dict[str, Any]]:
    """Per-instance work of every prompt (upsample to the original size, binarize, PNG / RLE encode).  The crop path
    (maskpost, ~0.3-1.5 ms per mask, mostly GIL-holding numpy) runs serially on the request thread, which measured
    faster than the pool; the full-upsample fallback runs on the pool with the request thread taking the first job
    (it keeps that core clocked up).  Outputs are byte-identical to a serial loop either way."""
    jobs = [(k, i) for k, sel in enumerate(selected) for i in range(sel[0].numel())] if mask_format != "none" else []
    job = lambda ki: _mask_payload(selected[ki[0]][2][ki[1]], width, height, mask_format)  # noqa: E731
    if _fast_mask_ok():
        payloads = [job(ki) for ki in jobs]
    else:
        futs = [_POOL.submit(job, ki) for ki in jobs[1:]]
        payloads = ([job(jobs[0])] if jobs else []) + [f.result() for f in futs]
    return _postprocess_results(prompts, selected, width, height, dict(zip(jobs, payloads)))


def _postprocess_results(prompts, selected, width: int, height: int, by_job) -> List[Dict[str, Any]]:
    results = []
    for k, (prompt, (scores, boxes, _mask_logits, presence)) in enumerate(zip(prompts, selected)):
        boxes_px = _boxes_to_xyxy_pixels(boxes, width, height)
        instances: List[Dict[str, Any]] = []
        for i in range(scores.numel()):
            inst: Dict[str, Any] = {"score": round(float(scores[i]), 4), "box_xyxy": boxes_px[i]}
            inst.update(by_job.get((k, i), {}))
            instances.append(inst)
        results.append({"prompt": prompt, "presence": round(presence, 4), "num_instances": len(instances), "instances": instances})
    return results


def _postprocess_prompt(prompt: str, selected, width: int, height: int, mask_format: str) -> Dict[str, Any]:
    return _postprocess_all([prompt], [selected], width, height, mask_format)[0]


# ---------------------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------------------
@app.get("/health")
def health() -> dict:
    ready = "pipeline" in STATE
    return {
        "status": "ok" if ready else "starting",
        "model": MODEL_NAME,
        "device": {"arch": STATE.get("device_arch"), "device_id": STATE["config"].device_id} if ready else None,
    }


@app.get("/info")
def info() -> dict:
    cfg: Optional[Config] = STATE.get("config")
    return {
        "model": MODEL_NAME,
        "description": "SAM3 (Segment Anything with Concepts, Meta) - text-prompted open-vocabulary instance segmentation",
        "task": "image-segmentation",
        "io": "image + text prompts -> per-prompt instance masks (original resolution), boxes and scores",
        "device_open": STATE.get("dispatch"),
        "hardware": "Tenstorrent Blackhole (single chip; default ETH dispatch, 1 CQ, 12x10 grid: see device_open) via tt-nn; whole image graph on device as two metal "
        "traces (image: ViT-H/14 + FPN neck + geometry encoder; per prompt: fusion encoder + decoder + seg head); "
        "text encoder on CPU (cached per prompt)",
        "weights": {
            "repo": cfg.weights_repo if cfg else DEFAULT_WEIGHTS_REPO,
            "revision": cfg.weights_revision if cfg else None,
            "file": WEIGHTS_FILENAME,
            "checkpoint_path": STATE.get("checkpoint_path"),
            "num_params_m": STATE.get("num_params_m"),
            "license": "Meta SAM License (gated repo; request access on the Hub)",
        },
        "source": {"repo": SOURCE_REPO, "commit": SOURCE_COMMIT, "package": "tt_sam3"},
        "limits": {
            "input_size": INPUT_SIZE,
            "resize": "bilinear squash to 1008x1008 (non-square images are squashed, as upstream)",
            "batch": 1,
            "max_prompts": cfg.max_prompts if cfg else None,
            "max_side": cfg.max_side if cfg else None,
            "num_queries": NUM_QUERIES,
            "mask_logit_resolution": MASK_RES,
        },
        "runtime": {
            "torch_compile": cfg.torch_compile if cfg else None,
            "torch_threads": torch.get_num_threads(),
            "warmup_runs": cfg.warmup_runs if cfg else None,
            "mesh_shape": cfg.mesh_shape if cfg else None,
            # the pipeline's copy: it may downgrade sub-knobs the loaded model cannot use (bf16 neck hand-off)
            "fused": (STATE["pipeline"].fused.describe() if "pipeline" in STATE else cfg.fused.describe()) if cfg else None,
            "encoder_trace_captured": bool(STATE["pipeline"].encoder_trace_captured) if "pipeline" in STATE else None,
            "encoder_fallbacks": int(STATE["pipeline"].encoder_fallbacks) if "pipeline" in STATE else None,
            "image_cache": bool(_IMG_CACHE and STATE.get("rgb_input")),
            "image_cache_hits": _image_cache_hits(),
        },
        "endpoints": {"predict": "POST /predict", "predict_raw": "POST /predict_raw", "health": "GET /health",
                      "models": "GET /v1/models"},
    }


@app.get("/v1/models")
def v1_models() -> dict:
    cfg: Optional[Config] = STATE.get("config")
    return {
        "object": "list",
        "data": [{"id": cfg.weights_repo if cfg else DEFAULT_WEIGHTS_REPO, "object": "model", "owned_by": "changh95"}],
    }


try:  # image-result cache key (SAM3_IMG_CACHE): xxh3_128 of the request's image bytes, 4.8 MB in ~0.2 ms
    import xxhash as _xxhash
except ImportError:  # optional dependency: no cache
    _xxhash = None
_IMG_CACHE = os.environ.get("SAM3_IMG_CACHE", "1") == "1" and _xxhash is not None
_LAST_DECODED: List[Any] = [None, None]  # [key, DecodedImage] of the last decoded image (one entry)


def _image_key(raw, tag: str) -> Optional[Tuple]:
    if not _IMG_CACHE:
        return None
    return (tag, len(raw), _xxhash.xxh3_128_digest(raw))


def _decode_cached(raw: bytes, key: Optional[Tuple], decode) -> DecodedImage:
    """Same bytes as the last request -> the same decoded pixels (skip the inflate)."""
    last = _LAST_DECODED[:]  # one consistent snapshot
    if key is not None and last[0] == key:
        return last[1]
    im = decode(raw)
    if key is not None:
        _LAST_DECODED[:] = [key, im]
    return im


def _invalidate_image_cache(pipeline) -> None:
    """Drop both cache levels (call with LOCK held): the decoded image and the device runner's resident-image key."""
    _LAST_DECODED[:] = [None, None]
    runner = getattr(pipeline, "_full", None)
    if runner is not None:
        runner.img_key = None


def _image_cache_hits() -> Optional[int]:
    runner = getattr(STATE.get("pipeline"), "_full", None)
    return int(runner.stats.get("image_cache_hits", 0)) if runner is not None else None


def _check_prompts(cfg, prompts_in: List[str]) -> List[str]:
    prompts = [p.strip() for p in prompts_in]
    if not prompts or any(not p for p in prompts):
        raise HTTPException(status_code=400, detail="prompts must be non-empty strings")
    if len(prompts) > cfg.max_prompts:
        raise HTTPException(status_code=400, detail=f"at most {cfg.max_prompts} prompts per request")
    return prompts


def _predict_core(raw: bytes, prompts_in: List[str], threshold: float, max_instances: int, mask_format: str,
                  decode, tag: str) -> dict:
    """raw: the image bytes (PNG/JPEG, or raw RGB uint8 for /predict_raw); decode(raw) -> DecodedImage."""
    pipeline: Optional[TtSam3ImagePipeline] = STATE.get("pipeline")
    cfg: Optional[Config] = STATE.get("config")
    if pipeline is None or cfg is None:
        raise HTTPException(status_code=503, detail="model is still starting")
    prompts = _check_prompts(cfg, prompts_in)

    t_start = time.perf_counter()
    # text-cache misses: the host text encoder starts now, overlapping the decode / preprocess / image trace
    pipeline.prefetch_text(prompts)
    rgb_input = STATE.get("rgb_input")
    key = _image_key(raw, tag) if rgb_input else None
    im = _decode_cached(raw, key, decode)
    width, height = im.size
    if max(width, height) > cfg.max_side:
        raise HTTPException(status_code=400, detail=f"image side {max(width, height)} exceeds SAM3_MAX_SIDE={cfg.max_side}")
    t_decoded = time.perf_counter()

    if rgb_input:  # resize / normalise / im2col run inside the forward, banded with the upload
        px, fwd = im.rgb, (lambda *a: pipeline.forward_rgb_selected_multi(*a, image_key=key))
    else:
        px, fwd = _to_pixel_values(im), pipeline.forward_image_selected_multi
    t_pre = time.perf_counter()

    selected = []
    try:
        with LOCK:
            try:
                with torch.inference_mode():
                    # full-device path: the image part runs once per request, the prompt part once per prompt, and
                    # only the kept instances' masks are read back from the device
                    selected = fwd(px, prompts, threshold, max_instances)
            except BaseException:
                _invalidate_image_cache(pipeline)  # a failed forward never leaves a cache entry behind
                raise
    except HTTPException:
        raise
    except Exception as e:  # noqa: BLE001 - surface the failure text to the client
        log.exception("inference failed")
        raise HTTPException(status_code=500, detail=f"inference failed: {type(e).__name__}: {e}") from None
    t_fwd = time.perf_counter()

    results = _postprocess_all(prompts, selected, width, height, mask_format)
    t_end = time.perf_counter()

    return {
        "width": width,
        "height": height,
        "input_size": INPUT_SIZE,
        "threshold": threshold,
        "mask_format": mask_format,
        "results": results,
        "timing_ms": {
            "decode": round((t_decoded - t_start) * 1000, 1),
            "preprocess": round((t_pre - t_decoded) * 1000, 1),
            "forward": round((t_fwd - t_pre) * 1000, 1),
            "forward_per_prompt": round((t_fwd - t_pre) * 1000 / max(1, len(prompts)), 1),
            "postprocess": round((t_end - t_fwd) * 1000, 1),
            "total": round((t_end - t_start) * 1000, 1),
        },
    }


@app.post("/predict")
def predict(req: PredictRequest) -> dict:
    if STATE.get("pipeline") is None or STATE.get("config") is None:
        raise HTTPException(status_code=503, detail="model is still starting")
    try:
        raw = _b64decode_strict(req.image)
    except Exception as e:  # noqa: BLE001 - any decode failure is a client error
        raise HTTPException(status_code=400, detail=f"bad image: {e}") from None
    return _predict_core(raw, req.prompts, req.threshold, req.max_instances, req.mask_format, _decode_bytes, "enc")


def _raw_rgb_decoder(width: int, height: int):
    def decode(raw: bytes) -> DecodedImage:
        if width <= 0 or height <= 0 or width * height * 3 != len(raw):
            raise HTTPException(status_code=400, detail=f"raw RGB body is {len(raw)} bytes, expected {width}x{height}x3")
        return DecodedImage(np.frombuffer(raw, dtype=np.uint8).reshape(height, width, 3))
    return decode


@app.post("/predict_raw")
async def predict_raw(
    request: Request,
    prompts: List[str] = Query(..., min_length=1, max_length=8),
    threshold: float = Query(0.5, ge=0.0, le=1.0),
    max_instances: int = Query(50, ge=1, le=NUM_QUERIES),
    mask_format: Literal["png", "rle", "none"] = Query("png"),
) -> dict:
    """Binary input, same pipeline and same response JSON as /predict without the JSON / base64 request layer.
    Body: the PNG/JPEG file (any Content-Type other than application/x-rgb24, e.g. image/png | image/jpeg), or raw
    RGB uint8 pixels, row-major H x W x 3 (Content-Type application/x-rgb24 with integer X-Width / X-Height headers;
    the body must be exactly X-Width * X-Height * 3 bytes).  Prompts etc. are query parameters
    (``?prompts=car&prompts=tree&threshold=0.5``)."""
    if STATE.get("pipeline") is None or STATE.get("config") is None:
        raise HTTPException(status_code=503, detail="model is still starting")
    ctype = request.headers.get("content-type", "").split(";")[0].strip().lower()
    if ctype == "application/x-rgb24":
        try:
            w, h = int(request.headers["x-width"]), int(request.headers["x-height"])
        except (KeyError, ValueError):
            raise HTTPException(status_code=400, detail="application/x-rgb24 needs integer X-Width and X-Height") from None
        max_side = STATE["config"].max_side
        if max(w, h) > max_side:  # refuse before reading a large body
            raise HTTPException(status_code=400, detail=f"image side {max(w, h)} exceeds SAM3_MAX_SIDE={max_side}")
        decode, tag = _raw_rgb_decoder(w, h), f"rgb{w}x{h}"
    else:
        decode, tag = _decode_bytes, "enc"
    raw = await request.body()
    if not raw:
        raise HTTPException(status_code=400, detail="empty body")
    return await run_in_threadpool(_predict_core, raw, prompts, threshold, max_instances, mask_format, decode, tag)