sam3-p150 / code /PYTHON.md
changh95's picture
p150 ETH-dispatch compliance (2026-10-05): default ETH dispatch, 1 CQ, 12x10 in Python API and server; numbers re-measured
3eb7396 verified
|
Raw History Blame Contribute Delete
19.6 kB

tt-sam3: SAM3 on Tenstorrent Blackhole from Python

tt_sam3 runs Meta's SAM3 (Segment Anything with Concepts) text-prompted instance segmentation on one Tenstorrent Blackhole chip. You give an image and one or more noun phrases. You get back masks, boxes and scores for each instance.

Install

Install the package into a Python environment that already has tt-metal's ttnn (and its torch). Run the commands from the model repo root:

pip install -e code/                    # the Python API only
pip install -e "code/[server,test]"     # also the HTTP server (fastapi, uvicorn, ...) and the tests (pytest, httpx)

The package does not declare ttnn or torch, so pip does not replace the tt-metal build of torch.

you want to run extras you need
from tt_sam3 import Sam3 none
the HTTP server (python -m uvicorn tt_sam3.server.app:app) server
the host tests (no device, no weights) server and test: test_server_host.py and test_api_host.py start the server app with fastapi's TestClient
the device tests test, and a Blackhole chip

Run only the host tests (about 1 minute; they do not open a device and do not need the weights):

TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_sam3/tests/test_*_host.py

The weights facebook/sam3 are gated (Meta SAM License). Request access on the Hub and run hf auth login once. The first from_pretrained downloads sam3.pt (3.45 GB) to the HF cache.

Paths in the examples. The examples use the demo image media/source_1.png with a path relative to the model repo root. Run them from the repo root, or give an absolute path. examples/quickstart.py finds the demo image relative to its own file, so you can run it from any directory.

Quickstart

from tt_sam3 import Sam3

# run from the model repo root (the demo image path is relative to it)
with Sam3.from_pretrained(device_id=0, warmup_variants={"prompts": ["car"]}) as model:
    result = model("media/source_1.png", "car")     # the first call is as fast as later calls

print(result)                    # Sam3Result(prompt='car', num_instances=3, scores=[0.973, 0.817, 0.797], ...)
result.masks                     # (N, H, W) bool, original image size
result.boxes                     # (N, 4) float32 [x0, y0, x1, y1], original image pixels
result.scores                    # (N,) float32, high to low

examples/quickstart.py runs this snippet on the demo image and writes an overlay PNG and a JSON summary.

More examples

import numpy as np
from PIL import Image
from tt_sam3 import Sam3

model = Sam3.from_pretrained()                       # device 0, weights from the HF cache, warmed up
# "street.jpg", "a.jpg", ...: your own image files (paths relative to the current directory)

# Several prompts on one image: the image part runs once, each prompt adds about 20 ms.
car, person, tree = model(Image.open("street.jpg"), ["car", "person", "tree"])

# numpy / torch inputs: (H, W, 3) uint8, (3, H, W), grayscale (H, W), or floats in [0, 1].
frame = np.asarray(Image.open("street.jpg"))
r = model(frame, "traffic light", threshold=0.3, max_instances=10)

# Same image again with a new prompt: the image part is reused (about 20 ms instead of about 80 ms).
r2 = model(frame, "bicycle")

# Many images, same prompts: the next image is decoded while the chip works on the current one.
results = model.predict_many(["a.jpg", "b.jpg", "c.jpg"], "dog")

# Only scores and boxes (no upsampling of masks to the image size), plus the raw 288x288 mask logits.
r = model(frame, "car", return_masks=False, return_mask_logits=True)

model.close()                                        # or use "with ... as model:"

API reference

Sam3.from_pretrained(repo_id="facebook/sam3", *, revision=None, checkpoint_path=None, device_id=0, device=None, dispatch="auto", warmup=2, warmup_variants=None, num_threads=None, cache_image=True, verbose=False) -> Sam3

Loads the weights, opens the chip, converts and uploads the weights, then runs the warm-up. The first warm-up compiles the kernels and captures the device traces (minutes with a cold kernel cache, cached in TT_METAL_CACHE). Then it prepares the request variants of warmup_variants (see "Warm-up" below). When it returns, the model is ready and the first real call is as fast as a later call.

argument meaning
repo_id Hugging Face repo that holds sam3.pt.
revision Weights revision. Default: TT_WEIGHTS_REVISION, else the pinned PINNED_REVISION (3c879f39…), which all published numbers use.
checkpoint_path Local sam3.pt, or a directory that holds it, instead of the Hub. SAM3_WEIGHTS_DIR does the same.
device_id Chip to open. Use 0 when TT_VISIBLE_DEVICES pins one chip.
device An already opened ttnn device. Open it with ttnn.open_device(**tt_sam3.device_open_kwargs()); for ETH dispatch, also set TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE=11,9 (12x10 grid) before tt-metal starts. close() does not close a device that you gave.
dispatch "auto" (default): ETH dispatch, 1 command queue and a 12x10 compute grid when the tt-metal tree has the Blackhole ETH-dispatch patch. This is the p150 target configuration: on a p150 the stock dispatch uses one Tensix column and leaves only 11x10. Without the patch, the model uses the stock dispatch and gives a RuntimeWarning. "eth" selects ETH dispatch directly. "worker" is an explicit opt-in to the stock (Tensix-column) dispatch: 11x10 on a p150, 12x10 only on a Galaxy chip. The outputs are the same in every mode. The model uses only command queue 0.
warmup Number of trace warm-up forwards (default 2). With 0, there is no warm-up at all (also no warmup_variants): the first call compiles and captures.
warmup_variants The request variants to prepare (see "Warm-up"). None (default): DEFAULT_WARMUP_VARIANTS. A dict replaces the defaults of the keys that it names. False: only the trace warm-up.
num_threads Host torch threads. Default: SAM3_NUM_THREADS, else 17, capped at the CPU count.
verbose False (default): a quiet console (see "Console output"). True: tt-metal's INFO log and the tt_sam3 steps.
cache_image Keep the device results of the last image. When the same image comes again (same pixels or same file bytes), the image part is skipped. The output bits are the same. This needs xxhash.

model.warmup(*, prompt_counts=(), image_sizes=(), max_instances=(), prompts=()) -> dict

Prepares more request variants after from_pretrained (same keys as warmup_variants). It is idempotent: a variant that is already prepared is skipped. It returns {variant: seconds} for the variants that ran. model.warmed_variants and model.info["warmup"] show what is prepared and the total warm-up time.

model.warmup(prompts=["bicycle", "traffic light"])    # about 75 ms for each new prompt
model.warmup(image_sizes=[(720, 1280)])               # about 0.1 s

model(image, prompts, *, threshold=0.5, max_instances=50, return_masks=True, return_mask_logits=False)

argument meaning
image A file path, encoded bytes (PNG, JPEG, ...), a PIL.Image, a numpy array or a torch tensor. Arrays can be (H, W), (H, W, 3|4) or (3|4, H, W). uint8 is used as is. Floats in [0, 1] are converted like the upstream Sam3Processor (x * 255.999, truncated). Alpha is dropped. Any size is accepted. The model squashes the image to 1008x1008 (no aspect-ratio padding, as upstream).
prompts One noun phrase (str) or a list of noun phrases. Each prompt is an independent query.
threshold Keep the instances with score > threshold. The upstream default is 0.5.
max_instances Maximum instances per prompt (1 to 200), highest scores first.
return_masks False skips the mask upsample (masks is None).
return_mask_logits Also return the raw (N, 288, 288) mask logits.

Returns one Sam3Result for a str prompt, or a list of Sam3Result (same order) for a list of prompts.

Sam3Result

field type and shape meaning
prompt str The prompt (stripped).
scores np.ndarray (N,) float32 sigmoid(logit) * presence, sorted high to low. This is the upstream Sam3Processor score.
boxes np.ndarray (N, 4) float32 [x0, y0, x1, y1] in original image pixels, clamped to the image.
masks np.ndarray (N, H, W) bool or None Instance masks at the original image size. A pixel is True when the bilinear upsample of the logits is > 0 (sigmoid > 0.5), as upstream.
presence float Probability that the concept is in the image at all.
width, height int Original image size.
mask_logits np.ndarray (N, 288, 288) float32 or None Raw mask logits in the 1008x1008 frame (only with return_mask_logits=True).

len(result) is N. result["masks"] works like the upstream Sam3Processor state dict. result.to_dict() gives a JSON-friendly summary (no masks). Upstream Sam3Processor gives masks as (N, 1, H, W) torch tensors in query order. Here the masks are (N, H, W) numpy arrays, sorted by score. Use torch.from_numpy(result.masks)[:, None] to get the upstream layout.

model.predict_many(images, prompts, **kwargs) -> list

Calls model(image, prompts, **kwargs) for each image and returns the results in the same order. The next image is read and decoded on a host thread while the chip runs the current image. The device work and the outputs are the same as separate calls.

model.close(), with Sam3.from_pretrained() as model:

close() releases the device traces and buffers, stops the host worker threads and closes the device (only if from_pretrained opened it). It is safe to call more than once. A model that is still open when Python exits is closed at exit. One process can use a chip for only one model at a time.

Helpers

  • tt_sam3.load_image(image) -> np.ndarray: the (H, W, 3) uint8 RGB array that the model uses for any supported input.
  • tt_sam3.device_open_kwargs(device_id=0, dispatch="auto") -> dict: the ttnn.open_device arguments that from_pretrained uses (l1_small_size, trace_region_size, and for ETH dispatch dispatch_core_config and num_command_queues=1).
  • model.info: the checkpoint path, the revision, the dispatch, the compute grid and the host paths in use.

Warm-up

After from_pretrained returns, the first real call of a prepared variant is as fast as a later call (within about 2 %, see the table). The warm-up uses real calls on generated photo-like images of the real sizes and dummy noun phrases. These calls go through the full path: the device traces, the host text-encoder threads, PNG and JPEG decode, the input kernels, the kept-instance readback, the mask upsample and its buffers. The warm-up keeps no result. The only exception is the text features of the prompts that you give in prompts, which go into the normal text cache. A call computes and caches the same values. The image cache is empty after the warm-up. The outputs do not change: the digests of all 20 x 9 measured calls were bit-identical before and after the change.

Variants

key default what it prepares
prompt_counts (1, 2, 3, 4) Calls with 1..k prompts (the server's SAM3_MAX_PROMPTS default is 4).
image_sizes ((1080, 1920), (1008, 1008)) (height, width) inputs. 1008x1008 uses its own input kernel (no resample).
max_instances (1, 2, 4, 8, 16, 32, 64, 200) The device readback bucket for the number of kept instances (powers of two up to 64; more than 64 = full mask readback).
prompts () Your noun phrases. Their text features are computed now.

Cost of the defaults: about 3.2 s more in from_pretrained (15.4..16.3 s -> 18.7..19.5 s with a warm kernel cache). No more device memory, because all traces and readback buffers exist after the trace warm-up. Each prompt in prompts adds about 75 ms.

What the defaults give, and what needs your prompts

  • The defaults prepare the host first-use costs (decoders, the first use of the text encoder) and the request shapes (1..4 prompts, the two image sizes, all kept-instance buckets). Alone, they make call 1 faster for calls with several prompts (2 prompts: 375 -> 238 ms, 4 prompts: 414 -> 363 ms), but not for one prompt (210 -> 207 ms).
  • For a fixed prompt set, call 1 still runs the text encoder for each of your prompts (+70..+175 ms with the defaults). Give your prompts in warmup_variants={"prompts": [...]} to remove this: call 1 / steady goes from 1.5..2.0x to 0.96..1.02x.
  • The text encoder has SAM3_TEXT_WORKERS (default 4) threads. A new prompt goes to the first idle thread, so one new prompt per call always uses the same thread, and 4 new prompts use 4 threads in parallel. The warm-up runs its text encodes on its own threads and stops them at the end, so the text threads of your calls are the same as without the warm-up. The time of a call with new prompts does not change.

What the warm-up cannot prepare

  • A new prompt. The prompts are open vocabulary. The host text encoder runs once for each prompt that the model has not seen yet, and the result goes into the text cache (256 entries). This costs about 75 to 150 ms on the call that uses the prompt first (measured: +65..+150 ms for one new prompt, about +225 ms for 4 new prompts, which encode in parallel). A later call with a new prompt pays the same cost, so this is not a first-call cost. To remove it from the first call, give your prompts in warmup_variants={"prompts": [...]} or call model.warmup(prompts=[...]).
  • Other image sizes. Any size works without a first-use cost above about 1 ms (only the resample tables of the new size are computed). In the measurement, the JPEG set has 9 different sizes and none of them was slower on first use.

Measured

Chip 11 of a Galaxy Blackhole host (ETH dispatch, 1 CQ, 12x10 grid; shared host), one fresh process per row, perf/warmup/first_call_bench.py. Every call uses a NEW real image (18 photos, never the warm-up input, never the previous image). "steady" is the median of calls 3..20. The image set is 1920x1080 PNG unless stated (sq: 1008x1008 PNG, jpg: JPEG files of 9 sizes, 333x500 to 1600x900). Columns: from_pretrained, then call 1 / 2 / 10 / steady in ms.

variant before after, defaults after, prompts listed
1 prompt ("car") 15.9 s, 210 / 132 / 128 / 130 18.8 s, 207 / 128 / 128 / 129 19.6 s, 127 / 128 / 128 / 128
2 prompts 15.7 s, 375 / 157 / 165 / 152 19.1 s, 238 / 149 / 149 / 149 19.5 s, 147 / 154 / 149 / 149
4 prompts 16.2 s, 414 / 197 / 196 / 188 19.2 s, 363 / 197 / 198 / 188 19.7 s, 187 / 197 / 193 / 188
1 prompt, numpy input 15.6 s, 154 / 82 / 84 / 81 19.4 s, 165 / 83 / 83 / 82 20.2 s, 81 / 82 / 84 / 81
1 prompt, return_masks=False, return_mask_logits=True 16.0 s, 191 / 129 / 129 / 128 18.9 s, 208 / 129 / 128 / 128 19.8 s, 128 / 128 / 127 / 128
1 prompt, threshold=0.05, max_instances=200 21.4 s, 198 / 131 / 129 / 133 19.0 s, 207 / 129 / 128 / 133 19.9 s, 128 / 131 / 126 / 134
1 prompt, sq 16.1 s, 177 / 106 / 107 / 108 19.1 s, 193 / 106 / 106 / 106 20.2 s, 105 / 105 / 106 / 106
1 prompt, jpg 16.0 s, 157 / 85 / 82 / 83 19.2 s, 166 / 84 / 83 / 84 19.9 s, 86 / 84 / 83 / 84
a new prompt on every call 15.9 s, 217 / 228 / 219 / 209 18.8 s, 208 / 233 / 223 / 222 -

Re-measured 2026-10-05 in the p150 configuration (chip 9, ETH dispatch, 1 CQ, 12x10, as reported by model.info; a new set of 24 1920x1080 PNGs made from the photos in the tt-metal tree, because the earlier set was lost in a host crash; perf/warmup/results/eth_1cq/):

variant defaults prompts listed
1 prompt ("car") 19.6 s, 222 / 128 / 126 / 125 20.7 s, 123 / 123 / 134 / 127
1 prompt, numpy input 19.3 s, 152 / 80 / 79 / 80 19.2 s, 79 / 80 / 80 / 80

With the prompts listed, call 1 is 0.96x / 0.99x of the steady median. The opt-in dispatch="worker" (12x10 on this Galaxy chip) gives the same output digests (DIGESTS IDENTICAL) and 189 / 123 / 124 / 125 ms with the defaults.

"Prompts listed" = warmup_variants={"prompts": ["car", "person", "tree", "road"]} (the other keys at their defaults). The first-call cost that remains with the defaults is the text encoder for prompts that are not listed (see above). The output digests of every call were bit-identical to "before". The JSON of every run is in perf/warmup/results/ (before, after, afterp).

Steady state when every call has NEW prompts (the host text encoder is most of the call), the same images and prompts, old (54d1f2a, no variant warm-up) and new interleaved in fresh processes (perf/warmup/results/ab_54d1f2a/):

workload 54d1f2a steady (ms) now steady (ms)
1 new prompt per call, numpy input, 12 pairs 175.4 (mean) 176.7 (mean); difference per pair -16.5..+13.6, median +1.5
same, warmup_variants=False, 6 rounds - 172.9 (mean)
1 new prompt per call, JPEG files, 3 pairs 181.8 185.4
4 new prompts per call, 2 pairs 354.7 360.9
1 fixed prompt, numpy input (the published path), 4 pairs 81.3 81.1

The differences are inside the run-to-run noise of this shared host (about +-10 ms per run). An earlier version of the warm-up used every text thread, and a single new prompt then moved between 4 threads: 171.8 vs 160.7 ms, slower in 8 of 8 pairs (perf/warmup/results/ab_4df6bea/). The text pool and the warm-up's own text threads (see above) remove this.

Console output

from_pretrained(verbose=False) (the default) keeps the console quiet:

  • It sets TT_LOGGER_LEVEL=error (tt-metal's C++ log) and LOGURU_LEVEL=WARNING (ttnn's Python log). It does this only when you have not set them and only before ttnn is imported (they are read once, at import).
  • It hides two known harmless start-up warnings: sam3's CUDA autocast decorators on a CPU-only torch, and numba's TBB version check.
  • Errors and all other warnings (also the RuntimeWarnings of tt_sam3) are shown.

With verbose=True, nothing is changed and the tt_sam3 logger prints its steps (load, device open, warm-up) to stderr. To see tt-metal's log while verbose=False, set TT_LOGGER_LEVEL=info before you start Python.

Notes

  • Same results as the HTTP server: model(...) runs the same device traces, host preprocessing, instance selection and mask upsample as POST /predict. The server rounds scores to 4 and boxes to 2 decimals in its JSON; the Python API gives the unrounded float32 values.
  • Thread safety: calls are serialized with a lock (one chip). Calls from several threads are safe, but they do not run in parallel.
  • Text prompts only. The model does not use box or point prompts (enable_inst_interactivity=False).
  • Timing on a Galaxy Blackhole chip (ETH dispatch, 1 CQ, 12x10 grid: the p150 configuration; warm, demo 1920x1080 image, one prompt): about 80 ms for a new image, about 20 ms for each more prompt or for a repeated image. See SERVING.md and OPT_REPORT.md in the model repo for the measured numbers.
  • Environment knobs of the server (TT_FUSED, SAM3_*) also apply here. The defaults are the validated configuration. Do not change them unless you validate the change again.