diffusionstudio/sam2.1-base-plus-video-onnx-fp16

SAM 2.1 Hiera-Base+ 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-base-plus (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-Base+, 1024 input, fp16	Hiera-Tiny, 512 input, fp16
per tracked frame	2.6 s	0.3 s

The export script is packages/sam2/scripts/export.py in the Diffusion Studio repository: export.py base-plus 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.