obraia/sam2.1-tiny-video-onnx-fp16

SAM 2.1 tiny, video pipeline, ONNX fp16

fp16 version of the four graphs of facebook/sam2.1-hiera-tiny's promptable video / memory pipeline, for in-browser propagation with onnxruntime-web on WebGPU (shader-f16). Made for the subject masking of Gausify Orbit, a video-to-Gaussian-Splatting app.

Converted from the fp32 export jax-image-tools/sam21-tiny-video-onnx (commit 3cb090d9d0b4c897281fdcac7765c0f3a888c49f). Inputs and outputs keep the same names, shapes and fp32 types, so code written for that export works unchanged. The memory bank is not a graph; constants.json (copied from the source repo) has its geometry.

file	size
vision_encoder.onnx	55.2 MB
mask_decoder.onnx	9.1 MB
memory_attention.onnx	16.4 MB
memory_encoder.onnx	2.8 MB
constants.json	memory-bank geometry (from the source repo)
export_fp16.py	conversion script

Changes from the source export

1. Vision encoder: the two If nodes of the Hiera position-embedding interpolation always take their else branch (an Identity) at the fixed 1024×1024 input. They are replaced by that branch. Outputs are identical (max |Δ| = 0 on a random input). This also removes a wrong declared shape in the unused then branch, which onnxruntime-web 1.30 rejects.
2. fp16: converted with onnxruntime.transformers.float16.convert_float_to_float16 (keep_io_types=True). Range, Resize, Sin, Cos and CumSum (the positional-encoding math) stay in fp32, along with the converter's default block list.

Script: export_fp16.py in this repo (needs onnx, onnxruntime 1.19+ and numpy): python3 export_fp16.py <source dir> <out dir>. It reproduces these files byte for byte.

Accuracy and speed

Measured in headless Chrome on an Apple M4 (WebGPU), against the fp32 graphs driven by the same memory-bank code:

• Masks: IoU ≥ 0.993 per frame against fp32 on two orbit videos (48 and 60 frames). IoU against ground-truth masks on a synthetic orbit was unchanged. Mean IoU 0.996 per frame against the fp32 PyTorch propagate_in_video_iterator.
• Speed: ~0.78 s/frame for full propagation with a full memory bank, versus ~1.15 s/frame in fp32.

License

Apache-2.0, inherited from facebook/sam2.1-hiera-tiny (SAM 2 by Meta). This repository redistributes a modified version: see "Changes from the source export" above.