SAM 2.1 Hiera-Small video tracking, fp16, 1024 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-small
(Sam2VideoModel, transformers 5.17), following the graph layout of
square-zero-labs/sam2.1-tiny-video-onnx.
- 1024×1024 input, the resolution SAM 2 was trained at. Image features are 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.
The memory bank is SAM 2's own: the prompted frame, the 6 most recent tracked frames and 16 object pointers. For speed over detail, see the 512-input Hiera-Tiny build.
Speed
Measured through the full tracking loop on ONNX Runtime Web 1.30 with WebGPU on an Apple M1 (8-core GPU), plugged in:
| Hiera-Small, 1024 input, fp16 | Hiera-Tiny, 512 input, fp16 | |
|---|---|---|
| per tracked frame | 1.9 s | 0.3 s |
The export script is packages/sam2/scripts/export.py in the Diffusion Studio repository:
export.py small 1024 7 <out-dir>.
Graphs
F = 64 (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,1024,1024] |
feats0 [1,32,256,256], feats1 [1,64,128,128], 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 1024 input), input_labels [1,1,N] int32 |
low_res_mask [1,1,256,256] (the decoder's own logits), high_res_mask [1,1,1024,1024], 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.