SAM 3 text-prompted segmentation for WebGPU: Standard (fp16) and Compact (8-bit weights)
7eee2d7 verified | """Exports SAM 3 (facebook/sam3) to ONNX for in-browser inference with ONNX Runtime Web on WebGPU. | |
| Three graphs (see `sam3_graphs.py`): `vision` (once per image), `text` and `detector` (once per | |
| phrase), each at a fixed shape: a 1008 × 1008 image, a 32-token phrase, batch 1. | |
| Files, in `<out>/onnx/`, named by graph and precision: | |
| - `*_fp16.onnx`: fp16 weights and activations, inputs and outputs float32. Standard is the three | |
| of these. | |
| - `vision_q8f16.onnx`, `text_q8f16.onnx`: as fp16, but every MatMul weight stored at 8 bits | |
| (`MatMulNBits`, symmetric, block size 128) and expanded on the GPU. Compact is these two and | |
| `detector_fp16.onnx`, which is small and left at fp16. | |
| Weight-only quantisation keeps every node on WebGPU: the usual dynamic "int8" export runs | |
| `MatMulInteger` and `DynamicQuantizeLinear`, which ONNX Runtime Web's WebGPU backend lacks. | |
| Usage: | |
| python export_sam3.py --weights <facebook/sam3 checkout> --out <dir> | |
| Requires torch, transformers (with `Sam3Model`), onnx, onnxscript, onnxconverter-common and | |
| onnxruntime. Runs on the CPU. The fp32 graphs go to `<out>/fp32/`, with external data, and | |
| `--skip-fp32` reuses them. | |
| """ | |
| import argparse | |
| from pathlib import Path | |
| import onnx | |
| import torch | |
| from onnxconverter_common import float16 | |
| from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer | |
| from transformers import Sam3Model | |
| from sam3_graphs import IMAGE_SIZE, TEXT_TOKENS, Sam3Detector, Sam3Text, Sam3Vision, position_encoding_72 | |
| OPSET = 18 | |
| GRAPHS = ("vision", "text", "detector") | |
| QUANTISED = ("vision", "text") | |
| def export_fp32(model, out: Path) -> None: | |
| out.mkdir(parents=True, exist_ok=True) | |
| pixels = torch.zeros(1, 3, IMAGE_SIZE, IMAGE_SIZE) | |
| ids = torch.zeros(1, TEXT_TOKENS, dtype=torch.long) | |
| mask = torch.ones(1, TEXT_TOKENS, dtype=torch.long) | |
| vision = Sam3Vision(model).eval() | |
| with torch.no_grad(): | |
| levels = vision(pixels) | |
| text = Sam3Text(model).eval() | |
| with torch.no_grad(): | |
| features = text(ids, mask) | |
| detector = Sam3Detector(model, position_encoding_72(model)).eval() | |
| jobs = { | |
| "vision": (vision, (pixels,), ["pixel_values"], ["features_288", "features_144", "features_72"]), | |
| "text": (text, (ids, mask), ["input_ids", "attention_mask"], ["text_features"]), | |
| "detector": ( | |
| detector, | |
| (*levels, features, mask), | |
| ["features_288", "features_144", "features_72", "text_features", "text_mask"], | |
| ["pred_logits", "pred_boxes", "presence_logits", "pred_masks", "semantic_seg"], | |
| ), | |
| } | |
| for name, (module, args, inputs, outputs) in jobs.items(): | |
| print(f"exporting {name}") | |
| with torch.no_grad(): | |
| torch.onnx.export( | |
| module, | |
| args, | |
| str(out / f"{name}.onnx"), | |
| input_names=inputs, | |
| output_names=outputs, | |
| opset_version=OPSET, | |
| dynamo=True, | |
| external_data=True, | |
| optimize=True, | |
| ) | |
| def fix_casts(model: onnx.ModelProto) -> int: | |
| """Points each Cast to float32 whose output the converter declared float16 at float16. | |
| `convert_float_to_float16` retypes a Cast's output but leaves its `to`, so a Cast from a | |
| non-float source (the text mask, cast from bool for pooling) still says float32 and the | |
| graph doesn't load. Returns how many it fixed. | |
| """ | |
| declared = {vi.name: vi.type.tensor_type.elem_type for vi in model.graph.value_info} | |
| fixed = 0 | |
| for node in model.graph.node: | |
| if node.op_type != "Cast" or declared.get(node.output[0]) != onnx.TensorProto.FLOAT16: | |
| continue | |
| for attribute in node.attribute: | |
| if attribute.name == "to" and attribute.i == onnx.TensorProto.FLOAT: | |
| attribute.i = onnx.TensorProto.FLOAT16 | |
| fixed += 1 | |
| return fixed | |
| def to_fp16(src: Path, dst: Path) -> None: | |
| dst.mkdir(parents=True, exist_ok=True) | |
| for name in GRAPHS: | |
| model = onnx.load(str(src / f"{name}.onnx")) | |
| # Stale annotations from the export would fight the converter's own inference. | |
| del model.graph.value_info[:] | |
| half = float16.convert_float_to_float16(model, keep_io_types=True, disable_shape_infer=False) | |
| print(f"fp16 {name}: {fix_casts(half)} casts fixed") | |
| onnx.save(half, str(dst / f"{name}_fp16.onnx")) | |
| def to_q8f16(onnx_dir: Path, bits: int, block_size: int) -> None: | |
| for name in QUANTISED: | |
| print(f"{bits}-bit {name}") | |
| quantiser = MatMulNBitsQuantizer( | |
| onnx.load(str(onnx_dir / f"{name}_fp16.onnx")), | |
| bits=bits, | |
| block_size=block_size, | |
| is_symmetric=True, | |
| ) | |
| quantiser.process() | |
| onnx.save(quantiser.model.model, str(onnx_dir / f"{name}_q{bits}f16.onnx")) | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| parser.add_argument("--weights", type=Path, required=True) | |
| parser.add_argument("--out", type=Path, required=True) | |
| parser.add_argument("--bits", type=int, default=8) | |
| parser.add_argument("--block-size", type=int, default=128) | |
| parser.add_argument("--skip-fp32", action="store_true", help="reuse an earlier fp32 export") | |
| args = parser.parse_args() | |
| if not args.skip_fp32: | |
| model = Sam3Model.from_pretrained(args.weights, attn_implementation="eager").eval() | |
| export_fp32(model, args.out / "fp32") | |
| del model | |
| to_fp16(args.out / "fp32", args.out / "onnx") | |
| to_q8f16(args.out / "onnx", args.bits, args.block_size) | |
| if __name__ == "__main__": | |
| main() | |