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.

Downloads last month
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for diffusionstudio/sam2.1-tiny-video-onnx-fp16

Quantized
(10)
this model