SAM 2.1 Hiera-Tiny video tracking, fp16, 512 input (ONNX)
The full SAM 2.1 video tracker, including the memory encoder, memory attention and
object pointers, as five fixed-shape ONNX graphs for ONNX Runtime on WebGPU. It is used by
Diffusion Studio's object mask tool. The graphs were
exported from the transformers port of facebook/sam2.1-hiera-tiny
(Sam2VideoModel, transformers 5.17), following the graph layout of
square-zero-labs/sam2.1-tiny-video-onnx.
It is built for speed on integrated GPUs:
- 512×512 input instead of 1024. Image features are 32×32 instead of 64×64.
- fp16 weights and compute, with float32 inputs and outputs (a cast at each graph boundary). Input-independent position encodings are computed in float32 at export time.
Halving the input makes memory attention about 16 times less work and the image encoder about 4 times less. The memory bank is SAM 2's own: the prompted frame, the 6 most recent tracked frames and 16 object pointers.
The trade-off is detail. Small objects and thin structures are resolved less finely, and a small object's identity is weaker: after a scene cut, the tracker is more likely to latch onto a similar-looking object than to report the object as gone.
Speed
Measured through the full tracking loop on ONNX Runtime Web 1.30 with WebGPU on an Apple M1 (8-core GPU), plugged in:
| 1024 input, fp16 | 512 input, fp16 | |
|---|---|---|
| per tracked frame | 1.3 s | 0.25 s |
The export script is packages/sam2/scripts/export.py in the Diffusion Studio repository.
It takes the input size and the number of memories as parameters.
Graphs
F = 32 (input / 16), R = 7·F² + 64 (frame memories, then 16 object pointers as 4 tokens each).
| file | inputs | outputs |
|---|---|---|
onnx/vision_encoder.onnx |
pixel_values [1,3,512,512] |
feats0 [1,32,128,128], feats1 [1,64,64,64], feats2 [1,256,F,F], feats2_no_mem (prompted-frame variant), vision_pos_embed [1,256,F,F] |
onnx/mask_decoder.onnx |
feats0, feats1, feats2_cond, input_points [1,1,N,2] (pixels of the 512 input), input_labels [1,1,N] int32 |
low_res_mask [1,1,256,256] (upsampled logits), high_res_mask [1,1,512,512], iou [1,1], object_score_logits [1,1,1], object_pointer [1,1,256] |
onnx/memory_encoder.onnx |
feats2, high_res_mask, object_score_logits [1,1], binarize (1 for prompted frames) |
memory_tokens [F²,1,64], memory_pos [F²,1,64] |
onnx/memory_attention.onnx |
current_vision_features [F²,1,256], current_vision_position_embeddings [F²,1,256], memory [R,1,64], memory_pos [R,1,64] |
conditioned_feats [1,256,F,F] |
onnx/pointer_tpos.onnx |
normalized_diffs [16] |
pointer_pos [16,64] |
The mask decoder chooses its output as the video predictor does. It considers several
candidate masks when there is at most one real point (every tracked frame, and a single
click), and otherwise uses a single mask with the stability fallback. Masks and object
pointers are suppressed in-graph when object_score_logits ≤ 0.
constants.json carries the input size, the number of memories, the normalization
constants and the 7×64 temporal positional encoding table. The prompted frame's memory uses
row 6, and the memory k frames back uses row k − 1.