jax-image-tools/sam3-video-onnx

ONNX export of facebook/sam3's promptable video / memory pipeline, for in-browser propagation with onnxruntime-web on WebGPU.

The published promptable-SAM exports stop at encoder + decoder, which segments one frame but cannot propagate a prompt across frames — that needs the memory bank. This repo carries the graphs for it, including the object_pointer output that image-only exports drop (their checkpoints discard object_pointer_proj).

Files

file	size
vision_encoder.onnx	1733.6 MB
mask_decoder.onnx	17.0 MB
memory_attention.onnx	33.0 MB
memory_encoder.onnx	5.3 MB
constants.json	bank geometry for the JS side

Graphs

vision_encoder     pixel_values (1,3,1008,1008) -> FPN feats + position encodings
mask_decoder       feats + point prompts -> pred_masks, high_res_masks, iou_scores,
                                            object_pointer, object_score_logits
memory_attention   current feats + memory bank -> memory-conditioned feats
memory_encoder     feats + predicted mask -> memory tokens + position encoding

The memory bank is not a graph. Upstream keeps it as dicts keyed by integer frame index and every branch is on a Python int or bool, never a tensor value, so it needs no tracing and belongs in JS. constants.json carries the geometry.

memory_attention takes the spatial blocks and the object pointers as separate, dynamic inputs, so it is exact for every bank state propagation produces — no padding on early frames.

Fidelity

Gated end to end against fp32 PyTorch propagate_in_video_iterator, with every neural call routed through these graphs:

Worst per-frame IoU 1.000000 against fp32 PyTorch propagate_in_video_iterator, worst max|Δ| on mask logits 3.6e-02, on a 4-frame clip with slow drift approximating z-stack slice-to-slice change, with the vision encoder, mask decoder, memory attention and memory encoder all routed through these graphs and the bank assembly left to upstream.

These graphs are fp32. The vision encoder is ~1.7 GB and is a reference export, not a browser-ready one — it also encodes at ~2.5 s/frame on WebGPU (measured on an M1 Max at q4 via the image tracker). For in-browser stack propagation see jax-image-tools/edgetam-video-onnx, which is 62 MB total and ~40 ms/frame.

Provenance

Exported with browser-onnx-tools (export/export_sam_video_onnx.py), gated by export/validate_sam_video.py.

License

These weights are NOT Apache-2.0. They are ONNX conversions of facebook/sam3, which Meta ships under the custom SAM License (terms). Converting a format does not relicense a model: everything here remains subject to those terms, including the acknowledgement requirement for publications, the use-case restrictions, and the export-control and sanctions conditions. Read the licence before use or redistribution.

The upstream checkpoint is gated — you must accept Meta's terms on the model page to obtain it. These converted graphs are provided for convenience and do not waive that.