# 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=) 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=)`` 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)