Upload 6 files

README.md

@@ -0,0 +1,285 @@
+---
+license: other
+license_name: sam-3-license
+license_link: https://ai.meta.com/resources/models-and-libraries/sam-license/
+base_model: facebook/sam3
+pipeline_tag: mask-generation
+tags:
+  - sam3
+  - sam-3
+  - onnx
+  - image-segmentation
+  - text-promptable
+  - open-vocabulary
+  - concept-segmentation
+library_name: transformers
+---
+
+# SAM 3 — Text-Promptable Image Segmentation (ONNX)
+
+ONNX export of Meta's **[SAM 3](https://huggingface.co/facebook/sam3)** image model — the text-promptable concept segmentation variant. Run open-vocabulary image segmentation from a text prompt (`"seed"`, `"cat"`, `"yellow school bus"`) in any environment that has an ONNX runtime: Python, C++, Rust, browsers via WebAssembly/WebGPU.
+
+To my knowledge this is the **first public ONNX export of SAM 3 image with text prompts**. Other community exports (e.g. [`onnx-community/sam3-tracker-ONNX`](https://huggingface.co/onnx-community/sam3-tracker-ONNX)) cover the tracker variant which only accepts point/box prompts.
+
+## Why this exists
+
+[`facebook/sam3`](https://huggingface.co/facebook/sam3) is published in PyTorch only. Running it in browsers, mobile, or any non-Python environment requires ONNX. As of this export's publication:
+
+- `optimum-onnx` does not have native SAM 3 support (the CLI fails with `Trying to export a sam3_video model, that is a custom or unsupported architecture`).
+- Meta's [`facebookresearch/sam3`](https://github.com/facebookresearch/sam3) repository ships no ONNX export tooling.
+- The existing community ONNX export (`vietanhdev/segment-anything-3-onnx-models`) targets Python `onnxruntime` only and isn't compatible with `transformers.js` / `onnxruntime-web`.
+- SegmentLens (sam3.ai) paywalls text-prompt SAM 3 behind server-side cloud inference rather than shipping it in-browser.
+
+This export was produced by hand-wrapping the three sub-modules of `Sam3Model` and calling `torch.onnx.export` directly on each. Validated end-to-end against the original PyTorch model — bit-equivalent detection count and box locations on a held-out test image.
+
+## What's in this repo
+
+```
+sam3-text-onnx/
+├── vision_encoder.onnx          # graph (6.2 MB)
+├── vision_encoder.onnx.data     # weights (1.84 GB)
+├── text_encoder.onnx            # graph (3.0 MB)
+├── text_encoder.onnx.data       # weights (1.35 GB)
+├── decoder.onnx                 # graph + weights inline (96 MB)
+├── export_sam3_vision.py        # script that produced vision_encoder.onnx
+├── export_sam3_text.py          # script that produced text_encoder.onnx
+├── export_sam3_decoder.py       # script that produced decoder.onnx
+└── validate_sam3_e2e.py         # end-to-end Python validation harness
+```
+
+Total: **~3.3 GB at fp32**. Quantization is recommended for browser deployment — see [Quantization & next steps](#quantization--next-steps).
+
+## Architecture
+
+SAM 3 is structured so the vision encoder runs **once per image** and produces multi-scale FPN features. The text encoder runs **once per prompt**. The decoder consumes both and produces masks/boxes/scores. This means changing the prompt while keeping the same image only re-runs the cheap text + decoder path:
+
+```
+image ──► vision_encoder.onnx ────────────────┐
+                                              ▼
+text prompt ──► text_encoder.onnx ──► decoder.onnx ──► pred_masks
+                                                       pred_boxes (xyxy normalized)
+                                                       pred_logits (sigmoid → scores)
+```
+
+## Usage — Python with onnxruntime
+
+```python
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+from transformers import AutoTokenizer
+from transformers.models.sam3.image_processing_sam3 import Sam3ImageProcessor
+
+MODEL_ID = "facebook/sam3"  # for the preprocessors only; weights come from ONNX
+
+# Preprocessors (still come from HF — they're tiny and have no ONNX equivalent)
+image_processor = Sam3ImageProcessor.from_pretrained(MODEL_ID)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# Load the three ONNX components
+vision_sess = ort.InferenceSession("vision_encoder.onnx", providers=["CPUExecutionProvider"])
+text_sess = ort.InferenceSession("text_encoder.onnx", providers=["CPUExecutionProvider"])
+decoder_sess = ort.InferenceSession("decoder.onnx", providers=["CPUExecutionProvider"])
+
+# Preprocess
+image = Image.open("your-image.png").convert("RGB")
+pixel_values = image_processor(images=image, return_tensors="np")["pixel_values"]
+
+encoded = tokenizer("seed", return_tensors="np", padding="max_length", max_length=32, truncation=True)
+input_ids = encoded["input_ids"].astype(np.int64)
+attention_mask = encoded["attention_mask"].astype(np.int64)
+
+# 1. Vision encoder — produces 4 FPN feature maps + position encodings
+v_out = vision_sess.run(None, {"pixel_values": pixel_values})
+fpn_hidden_states = v_out[0:4]       # spatial scales 288, 144, 72, 36 at 1008x1008 input
+fpn_position_encoding = v_out[4:8]
+
+# 2. Text encoder — projects "seed" to a 256-dim feature
+text_features = text_sess.run(None, {
+    "input_ids": input_ids,
+    "attention_mask": attention_mask,
+})[0]
+
+# 3. Decoder — uses first 3 FPN levels + only the last position encoding
+pred_masks, pred_boxes, pred_logits = decoder_sess.run(None, {
+    "fpn_hidden_state_0": fpn_hidden_states[0],
+    "fpn_hidden_state_1": fpn_hidden_states[1],
+    "fpn_hidden_state_2": fpn_hidden_states[2],
+    "fpn_position_encoding_2": fpn_position_encoding[2],
+    "text_features": text_features,
+    "attention_mask": attention_mask,
+})
+
+# Convert logits to scores in [0, 1]
+scores = 1.0 / (1.0 + np.exp(-pred_logits))
+
+# Filter by confidence
+keep = scores[0] > 0.5
+print(f"Detections: {keep.sum()}")
+print(f"Boxes (xyxy normalized): {pred_boxes[0, keep]}")
+print(f"Scores: {scores[0, keep]}")
+```
+
+A complete runnable version is in `validate_sam3_e2e.py`.
+
+## Usage — browser with onnxruntime-web
+
+The same recipe in JavaScript. Note: `transformers.js` does not currently have a `Sam3Model` JS class for the image variant, so you call `onnxruntime-web` directly. You still need `transformers.js` for the tokenizer.
+
+```js
+import { AutoTokenizer } from "@huggingface/transformers";
+import * as ort from "onnxruntime-web";
+
+const tokenizer = await AutoTokenizer.from_pretrained("facebook/sam3");
+
+const visionSess = await ort.InferenceSession.create(
+  "https://huggingface.co/danilobukvic/sam3-text-onnx/resolve/main/vision_encoder.onnx",
+  { executionProviders: ["webgpu"] }
+);
+const textSess = await ort.InferenceSession.create(
+  "https://huggingface.co/danilobukvic/sam3-text-onnx/resolve/main/text_encoder.onnx",
+  { executionProviders: ["webgpu"] }
+);
+const decoderSess = await ort.InferenceSession.create(
+  "https://huggingface.co/danilobukvic/sam3-text-onnx/resolve/main/decoder.onnx",
+  { executionProviders: ["webgpu"] }
+);
+
+// Preprocess image to [1, 3, 1008, 1008] Float32Array with ImageNet normalization
+// (see Sam3ImageProcessor for exact mean/std values — you'll need to port this)
+const pixelValues = preprocessImage(imageBitmap);
+
+// Tokenize prompt to length 32
+const { input_ids, attention_mask } = await tokenizer("seed", {
+  padding: "max_length",
+  max_length: 32,
+  truncation: true,
+});
+
+// Run the pipeline (same as Python above) ...
+```
+
+## Input/output contracts
+
+### vision_encoder.onnx
+
+| | Name | Shape | Type | Notes |
+|---|---|---|---|---|
+| **In** | `pixel_values` | `[batch, 3, 1008, 1008]` | float32 | ImageNet-normalized. Size is fixed: SAM 3 precomputes positional embeddings for 1008×1008. |
+| **Out** | `fpn_hidden_state_0..3` | `[batch, 256, H, W]` | float32 | Spatial scales: 288, 144, 72, 36 (for 1008 input) |
+| **Out** | `fpn_position_encoding_0..3` | `[batch, 256, H, W]` | float32 | Matching position encodings |
+
+### text_encoder.onnx
+
+| | Name | Shape | Type | Notes |
+|---|---|---|---|---|
+| **In** | `input_ids` | `[batch, 32]` | int64 | CLIP-style tokens; SAM 3 uses `max_position_embeddings=32` (shorter than standard CLIP's 77) |
+| **In** | `attention_mask` | `[batch, 32]` | int64 | 1 for real tokens, 0 for padding |
+| **Out** | `text_features` | `[batch, 32, 256]` | float32 | Projected from CLIP's 1024-dim to SAM 3's 256-dim DETR space |
+
+### decoder.onnx
+
+| | Name | Shape | Type | Notes |
+|---|---|---|---|---|
+| **In** | `fpn_hidden_state_0` | `[batch, 256, 288, 288]` | float32 | From vision encoder (largest scale) |
+| **In** | `fpn_hidden_state_1` | `[batch, 256, 144, 144]` | float32 | From vision encoder |
+| **In** | `fpn_hidden_state_2` | `[batch, 256, 72, 72]` | float32 | From vision encoder (smallest scale) |
+| **In** | `fpn_position_encoding_2` | `[batch, 256, 72, 72]` | float32 | Only the smallest scale's PE is used (others were optimized out by the tracer) |
+| **In** | `text_features` | `[batch, 32, 256]` | float32 | From text encoder |
+| **In** | `attention_mask` | `[batch, 32]` | int64 | Same mask passed to text encoder |
+| **Out** | `pred_masks` | `[batch, 200, 288, 288]` | float32 | 200 query slots, each a low-resolution mask |
+| **Out** | `pred_boxes` | `[batch, 200, 4]` | float32 | Boxes in xyxy format, **normalized to [0, 1]** |
+| **Out** | `pred_logits` | `[batch, 200]` | float32 | Apply sigmoid for [0, 1] confidence scores |
+
+## Performance
+
+Measured on CPU (Intel laptop), single image:
+
+| Component | Time |
+|---|---|
+| vision_encoder | ~150 s |
+| text_encoder | ~7 s |
+| decoder | ~8 s |
+| **Total per image** | **~170 s** |
+| Re-run with different prompt (same image) | **~15 s** (vision cached) |
+
+GPU (CUDA execution provider, RTX-class) should be roughly 10-15× faster. WebGPU in modern Chrome should fall between the two.
+
+The architecture is designed for prompt iteration: cache the vision encoder output, re-run text + decoder per prompt.
+
+## Validation
+
+Compared end-to-end against the original PyTorch `Sam3Model.forward()` on a microscope-style seed image with prompt `"seed"`:
+
+| Metric | PyTorch SAM 3 | This ONNX export |
+|---|---|---|
+| Detections > 0.5 | 12 | 12 |
+| Top scores | 0.66–0.93 | 0.66–0.93 |
+| Box locations | clustered at seeds | clustered at seeds |
+
+Numerically equivalent within tracing noise. See `validate_sam3_e2e.py` to reproduce.
+
+## Quantization & next steps
+
+The fp32 export totals ~3.3 GB which is impractical for browser deployment. The natural follow-up is per-component quantization, mirroring what `onnx-community/sam3-tracker-ONNX` does for the tracker variant:
+
+```js
+dtype: {
+  vision_encoder: "q4",                    // ~470 MB
+  prompt_encoder_mask_decoder: "fp32",     // kept at full precision
+}
+```
+
+Estimated sizes after quantization:
+
+| Component | fp32 | fp16 | q8 | q4 |
+|---|---|---|---|---|
+| vision_encoder | 1.84 GB | ~900 MB | ~470 MB | ~240 MB |
+| text_encoder | 1.35 GB | ~675 MB | ~340 MB | ~170 MB |
+| decoder | 96 MB | ~48 MB | ~25 MB | ~13 MB |
+| **Total** | **3.3 GB** | **~1.6 GB** | **~830 MB** | **~420 MB** |
+
+I haven't published quantized versions yet. If you build them, PRs welcome.
+
+## Known caveats and TODOs
+
+- **Fixed input size**: The vision encoder's positional embeddings are precomputed for 1008×1008 input. Other sizes will produce shape mismatches. Use the SAM 3 image processor to handle resizing/padding.
+- **Geometry prompts not supported**: This export covers the text-only path. SAM 3's optional box/point prompts are not wired in — would need a separate decoder variant.
+- **Dynamic batch is dynamic but untested**: The export uses symbolic batch dim but I've only validated batch=1. Higher batch sizes should work but no guarantees.
+- **Tracer warnings during export**: A few `TracerWarning: Converting a tensor to a Python boolean...` were emitted. These bake config flags (e.g. `is_causal`, attention backend selection) into the graph at export time. Fine for inference with the same model config, but means the ONNX isn't reusable across configs.
+- **No `transformers.js` integration**: At time of writing, `transformers.js` only has `Sam3TrackerModel` (point/box) but not `Sam3Model` (text-promptable image). Until that lands, use `onnxruntime-web` directly.
+
+## How this was built
+
+Three Python scripts, one per sub-module, each calling `torch.onnx.export` on a thin wrapper around the SAM 3 component:
+
+1. `export_sam3_vision.py` — wraps `Sam3VisionModel`, flattens `Sam3VisionEncoderOutput` into a tuple of tensors. Uses the default (dynamo) exporter.
+2. `export_sam3_text.py` — wraps `CLIPTextModelWithProjection` + the `text_projection` Linear layer.
+3. `export_sam3_decoder.py` — wraps `detr_encoder` + `detr_decoder` + `mask_decoder` + `dot_product_scoring` + box head as a single pipeline. **Uses `dynamo=False`** (legacy tracer) — the dynamo exporter crashes on the attention reshape-after-transpose pattern.
+
+All three are in this repo. See `validate_sam3_e2e.py` for the end-to-end pipeline.
+
+## License
+
+This export inherits the license of the original `facebook/sam3` weights: the **[SAM 3 License Agreement](https://ai.meta.com/resources/models-and-libraries/sam-license/)**. Read it before using these weights — it includes restrictions on commercial use.
+
+## Citation
+
+The underlying model is Meta's SAM 3:
+
+```bibtex
+@article{ravi2025sam3,
+  title   = {SAM 3: Segment Anything with Concepts},
+  author  = {Ravi, Nikhila and others},
+  journal = {arXiv preprint arXiv:2511.16719},
+  year    = {2025}
+}
+```
+
+If this ONNX export is useful to you, a star on the repo or a mention is appreciated.
+
+## Acknowledgments
+
+- **Meta AI** for releasing SAM 3 and its weights
+- **HuggingFace** for the `transformers` integration that made the architecture introspectable
+- **PyTorch team** for `torch.onnx.export` (especially the legacy `dynamo=False` path, which is what got the decoder across the line)

decoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb6a7bbfe1a13d4d2141900f4c32c50636841def66ca89de50b6b3fa5fde4bf8
+size 95951921

export_sam3_decoder.py

@@ -0,0 +1,232 @@
+"""Export SAM3's decoder pipeline to ONNX.
+
+Phase 3c — the hard one. Bundles geometry_encoder (skipped for now since
+text-only prompts don't use it), detr_encoder, detr_decoder, mask_decoder,
+and dot_product_scoring into a single ONNX file that takes pre-computed
+vision FPN features + projected text features as input.
+
+The point: avoid re-running the heavy vision/text encoders. The browser will
+run vision_encoder.onnx and text_encoder.onnx once each, cache the outputs,
+then call this decoder.onnx for the actual segmentation per image+prompt.
+
+Inputs (all tensors, no structured types):
+    fpn_hidden_state_0,1,2:       3 FPN levels at spatial scales 288, 144, 72
+    fpn_position_encoding_0,1,2:  matching position encodings
+    text_features:                [B, 32, 256] projected text features
+    attention_mask:               [B, 32] int64 (1=real token, 0=pad)
+
+Outputs:
+    pred_masks:   [B, num_queries, H, W]
+    pred_boxes:   [B, num_queries, 4]   xyxy format
+    pred_logits:  [B, num_queries]      classification scores
+"""
+
+from pathlib import Path
+import torch
+from torch import nn
+from transformers import AutoTokenizer, Sam3Model
+from transformers.models.sam3.image_processing_sam3 import Sam3ImageProcessor
+from transformers.models.sam3.modeling_sam3 import (
+    inverse_sigmoid,
+    box_cxcywh_to_xyxy,
+)
+from PIL import Image
+
+OUTPUT_DIR = Path("sam3-onnx-test")
+OUTPUT_DIR.mkdir(exist_ok=True)
+OUTPUT_FILE = OUTPUT_DIR / "decoder.onnx"
+
+MODEL_ID = "facebook/sam3"
+
+
+class WrappedDecoder(nn.Module):
+    """Bundles the SAM3 decoder pipeline (detr_encoder → detr_decoder → mask_decoder).
+
+    Skips geometry prompts entirely — text-only path. Mirrors the relevant
+    portion of Sam3Model.forward() but with flat tensor I/O for ONNX export.
+    """
+
+    def __init__(self, full_model: Sam3Model):
+        super().__init__()
+        self.detr_encoder = full_model.detr_encoder
+        self.detr_decoder = full_model.detr_decoder
+        self.mask_decoder = full_model.mask_decoder
+        self.dot_product_scoring = full_model.dot_product_scoring
+
+    def forward(
+        self,
+        fpn_hidden_state_0: torch.Tensor,
+        fpn_hidden_state_1: torch.Tensor,
+        fpn_hidden_state_2: torch.Tensor,
+        fpn_position_encoding_0: torch.Tensor,
+        fpn_position_encoding_1: torch.Tensor,
+        fpn_position_encoding_2: torch.Tensor,
+        text_features: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ):
+        fpn_hidden_states = (fpn_hidden_state_0, fpn_hidden_state_1, fpn_hidden_state_2)
+        fpn_position_encoding = (
+            fpn_position_encoding_0,
+            fpn_position_encoding_1,
+            fpn_position_encoding_2,
+        )
+
+        text_mask = attention_mask.bool()
+        combined_prompt_features = text_features
+        combined_prompt_mask = text_mask
+
+        # 1. DETR encoder operates on the smallest (most-pooled) FPN level + text
+        encoder_outputs = self.detr_encoder(
+            vision_features=[fpn_hidden_states[-1]],
+            text_features=combined_prompt_features,
+            vision_pos_embeds=[fpn_position_encoding[-1]],
+            text_mask=combined_prompt_mask,
+        )
+
+        # 2. DETR decoder produces object queries
+        decoder_outputs = self.detr_decoder(
+            vision_features=encoder_outputs.last_hidden_state,
+            text_features=encoder_outputs.text_features,
+            vision_pos_encoding=encoder_outputs.pos_embeds_flattened,
+            text_mask=combined_prompt_mask,
+            spatial_shapes=encoder_outputs.spatial_shapes,
+        )
+
+        # 3. Box predictions: refine reference boxes via decoder's box head
+        all_box_offsets = self.detr_decoder.box_head(decoder_outputs.intermediate_hidden_states)
+        reference_boxes_inv_sig = inverse_sigmoid(decoder_outputs.reference_boxes)
+        all_pred_boxes_cxcywh = (reference_boxes_inv_sig + all_box_offsets).sigmoid()
+        all_pred_boxes = box_cxcywh_to_xyxy(all_pred_boxes_cxcywh)
+
+        # 4. Classification scores: dot product between queries and text
+        all_pred_logits = self.dot_product_scoring(
+            decoder_hidden_states=decoder_outputs.intermediate_hidden_states,
+            text_features=encoder_outputs.text_features,
+            text_mask=combined_prompt_mask,
+        ).squeeze(-1)
+
+        # We only return the FINAL decoder layer's predictions (the typical case)
+        pred_logits = all_pred_logits[-1]
+        pred_boxes = all_pred_boxes[-1]
+        decoder_hidden_states = decoder_outputs.intermediate_hidden_states[-1]
+
+        # 5. Mask decoder produces the actual segmentation masks
+        mask_outputs = self.mask_decoder(
+            decoder_queries=decoder_hidden_states,
+            backbone_features=list(fpn_hidden_states),
+            encoder_hidden_states=encoder_outputs.last_hidden_state,
+            prompt_features=combined_prompt_features,
+            prompt_mask=combined_prompt_mask,
+        )
+
+        return mask_outputs.pred_masks, pred_boxes, pred_logits
+
+
+def main() -> None:
+    print(f"Loading {MODEL_ID} ...")
+    model = Sam3Model.from_pretrained(MODEL_ID)
+    model.eval()
+
+    # Build real inputs end-to-end using the actual vision + text encoders.
+    # We don't want to fabricate fake FPN tensors — they have to match the
+    # exact shape and statistical distribution the decoder was trained on.
+    print("\nBuilding real inputs by running the encoders ...")
+    image_processor = Sam3ImageProcessor.from_pretrained(MODEL_ID)
+    dummy_pil = Image.new("RGB", (640, 480), color=(128, 128, 128))
+    pixel_values = image_processor(images=dummy_pil, return_tensors="pt")["pixel_values"]
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    max_len = int(model.config.text_config.max_position_embeddings)
+    encoded = tokenizer(
+        "seed",
+        return_tensors="pt",
+        padding="max_length",
+        max_length=max_len,
+        truncation=True,
+    )
+
+    with torch.no_grad():
+        vision_out = model.vision_encoder(pixel_values)
+        text_out = model.get_text_features(
+            input_ids=encoded["input_ids"], attention_mask=encoded["attention_mask"]
+        )
+
+    # The decoder uses FPN[:-1] (first 3 of 4 levels)
+    fpn_h = vision_out.fpn_hidden_states[:-1]
+    fpn_p = vision_out.fpn_position_encoding[:-1]
+    text_features = text_out.pooler_output
+    attention_mask = encoded["attention_mask"]
+
+    print(f"  FPN hidden states: {len(fpn_h)} tensors")
+    for i, t in enumerate(fpn_h):
+        print(f"    [{i}] shape={tuple(t.shape)}")
+    print(f"  text_features: {tuple(text_features.shape)}")
+    print(f"  attention_mask: {tuple(attention_mask.shape)}")
+
+    # Smoke test the wrapped decoder
+    print("\nSmoke testing wrapped decoder in PyTorch ...")
+    wrapped = WrappedDecoder(model).eval()
+    with torch.no_grad():
+        pred_masks, pred_boxes, pred_logits = wrapped(
+            fpn_h[0], fpn_h[1], fpn_h[2],
+            fpn_p[0], fpn_p[1], fpn_p[2],
+            text_features,
+            attention_mask,
+        )
+    print(f"  pred_masks: shape={tuple(pred_masks.shape)} dtype={pred_masks.dtype}")
+    print(f"  pred_boxes: shape={tuple(pred_boxes.shape)} dtype={pred_boxes.dtype}")
+    print(f"  pred_logits: shape={tuple(pred_logits.shape)} dtype={pred_logits.dtype}")
+    print(f"    logits mean={pred_logits.mean().item():.4f} std={pred_logits.std().item():.4f}")
+
+    # Export
+    print(f"\nExporting to {OUTPUT_FILE} ...")
+    torch.onnx.export(
+        wrapped,
+        (
+            fpn_h[0], fpn_h[1], fpn_h[2],
+            fpn_p[0], fpn_p[1], fpn_p[2],
+            text_features,
+            attention_mask,
+        ),
+        str(OUTPUT_FILE),
+        input_names=[
+            "fpn_hidden_state_0", "fpn_hidden_state_1", "fpn_hidden_state_2",
+            "fpn_position_encoding_0", "fpn_position_encoding_1", "fpn_position_encoding_2",
+            "text_features",
+            "attention_mask",
+        ],
+        output_names=["pred_masks", "pred_boxes", "pred_logits"],
+        dynamic_axes={
+            "fpn_hidden_state_0": {0: "batch", 2: "h0", 3: "w0"},
+            "fpn_hidden_state_1": {0: "batch", 2: "h1", 3: "w1"},
+            "fpn_hidden_state_2": {0: "batch", 2: "h2", 3: "w2"},
+            "fpn_position_encoding_0": {0: "batch", 2: "h0", 3: "w0"},
+            "fpn_position_encoding_1": {0: "batch", 2: "h1", 3: "w1"},
+            "fpn_position_encoding_2": {0: "batch", 2: "h2", 3: "w2"},
+            "text_features": {0: "batch", 1: "text_seq"},
+            "attention_mask": {0: "batch", 1: "text_seq"},
+            "pred_masks": {0: "batch"},
+            "pred_boxes": {0: "batch"},
+            "pred_logits": {0: "batch"},
+        },
+        opset_version=18,
+        do_constant_folding=True,
+        verbose=False,
+        # SAM3's attention layers use .reshape() on transposed tensors with a
+        # dynamic batch dim, which trips PyTorch's new dynamo exporter (it can't
+        # trace the view-vs-copy decision symbolically). The legacy torch.jit.trace
+        # path handles this pattern fine. Force it.
+        dynamo=False,
+    )
+
+    size_mb = OUTPUT_FILE.stat().st_size / (1024 * 1024)
+    print(f"\n✅ Exported decoder: {OUTPUT_FILE} ({size_mb:.1f} MB graph)")
+
+    print("\nFiles in output dir:")
+    for f in sorted(OUTPUT_DIR.iterdir()):
+        size_mb = f.stat().st_size / (1024 * 1024)
+        print(f"  {f.name}: {size_mb:.1f} MB")
+
+
+if __name__ == "__main__":
+    main()

export_sam3_text.py

@@ -0,0 +1,119 @@
+"""Export SAM3's text encoder to ONNX.
+
+Phase 3b of the SAM3 custom export. The text encoder is a CLIP-style
+text model (CLIPTextModelWithProjection) followed by a linear projection
+that maps CLIP's 1024-dim output to SAM3's 256-dim DETR space.
+
+Input:
+    input_ids:      [batch, seq_len] int64 — tokenized text
+    attention_mask: [batch, seq_len] int64 — 1 for real tokens, 0 for padding
+
+Output:
+    projected_text_features: [batch, seq_len, 256] float32
+
+This matches Sam3Model.get_text_features() in transformers source.
+"""
+
+from pathlib import Path
+import torch
+from torch import nn
+from transformers import AutoTokenizer, Sam3Model
+
+OUTPUT_DIR = Path("sam3-onnx-test")
+OUTPUT_DIR.mkdir(exist_ok=True)
+OUTPUT_FILE = OUTPUT_DIR / "text_encoder.onnx"
+
+MODEL_ID = "facebook/sam3"
+
+
+class WrappedTextEncoder(nn.Module):
+    """Bundles CLIP text encoder + the projection layer into a single module.
+
+    Mirrors what Sam3Model.get_text_features() does, but as a pure nn.Module
+    suitable for torch.onnx.export (no kwargs, no return_dict, flat tensor I/O).
+    """
+
+    def __init__(self, text_encoder: nn.Module, text_projection: nn.Module):
+        super().__init__()
+        self.text_encoder = text_encoder
+        self.text_projection = text_projection
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        text_outputs = self.text_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            return_dict=True,
+        )
+        last_hidden_state = text_outputs.last_hidden_state
+        return self.text_projection(last_hidden_state)
+
+
+def main() -> None:
+    print(f"Loading {MODEL_ID} ...")
+    model = Sam3Model.from_pretrained(MODEL_ID)
+    model.eval()
+
+    text_encoder = model.text_encoder
+    text_projection = model.text_projection
+    print(f"Text encoder type: {type(text_encoder).__name__}")
+    print(f"Text projection: {text_projection}")
+
+    # Build a real dummy input via the tokenizer — this is the canonical input shape
+    print("\nBuilding tokenizer ...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    # SAM3 uses a customized CLIP text model with max_position_embeddings=32
+    # (shorter than standard CLIP's 77 — Meta tuned it for concept prompts).
+    # Read the actual limit from the config so we stay in sync if it ever changes.
+    max_len = int(model.config.text_config.max_position_embeddings)
+    print(f"  max_position_embeddings (from config): {max_len}")
+    encoded = tokenizer(
+        "seed",
+        return_tensors="pt",
+        padding="max_length",
+        max_length=max_len,
+        truncation=True,
+    )
+    dummy_input_ids = encoded["input_ids"]
+    dummy_attention_mask = encoded["attention_mask"]
+    print(f"  input_ids: {tuple(dummy_input_ids.shape)} dtype={dummy_input_ids.dtype}")
+    print(f"  attention_mask: {tuple(dummy_attention_mask.shape)} dtype={dummy_attention_mask.dtype}")
+
+    # Smoke test in PyTorch
+    print("\nSmoke testing wrapped text encoder in PyTorch ...")
+    wrapped = WrappedTextEncoder(text_encoder, text_projection).eval()
+    with torch.no_grad():
+        pt_out = wrapped(dummy_input_ids, dummy_attention_mask)
+    print(f"  Output: shape={tuple(pt_out.shape)} dtype={pt_out.dtype}")
+    print(f"  mean={pt_out.mean().item():.4f} std={pt_out.std().item():.4f}")
+
+    # Export
+    print(f"\nExporting to {OUTPUT_FILE} ...")
+    torch.onnx.export(
+        wrapped,
+        (dummy_input_ids, dummy_attention_mask),
+        str(OUTPUT_FILE),
+        input_names=["input_ids", "attention_mask"],
+        output_names=["text_features"],
+        dynamic_axes={
+            "input_ids": {0: "batch", 1: "sequence_length"},
+            "attention_mask": {0: "batch", 1: "sequence_length"},
+            "text_features": {0: "batch", 1: "sequence_length"},
+        },
+        opset_version=18,
+        do_constant_folding=True,
+        verbose=False,
+    )
+
+    onnx_size = OUTPUT_FILE.stat().st_size / (1024 * 1024)
+    print(f"\n✅ Exported text encoder: {OUTPUT_FILE} ({onnx_size:.1f} MB graph)")
+    print("   (weights may be in a sibling .data file if model > 2GB)")
+
+    # Show all files produced
+    print("\nFiles in output dir:")
+    for f in sorted(OUTPUT_DIR.iterdir()):
+        size_mb = f.stat().st_size / (1024 * 1024)
+        print(f"  {f.name}: {size_mb:.1f} MB")
+
+
+if __name__ == "__main__":
+    main()

export_sam3_vision.py

@@ -0,0 +1,118 @@
+"""Export SAM3's vision encoder to ONNX.
+
+Phase 3a of the SAM3 custom export attempt. We bypass optimum-cli's architecture
+catalog and call torch.onnx.export directly. The vision encoder is a standard
+ViT-style image-to-features model — the simplest piece of SAM3 to export.
+
+If this works, we move on to text encoder + decoder. If it fails, we learn
+exactly what SAM3 op isn't ONNX-compatible.
+
+Usage:
+    python export_sam3_vision.py
+"""
+
+from pathlib import Path
+import torch
+from PIL import Image
+from transformers import Sam3Model
+from transformers.models.sam3.image_processing_sam3 import Sam3ImageProcessor
+
+OUTPUT_DIR = Path("sam3-onnx-test")
+OUTPUT_DIR.mkdir(exist_ok=True)
+OUTPUT_FILE = OUTPUT_DIR / "vision_encoder.onnx"
+
+MODEL_ID = "facebook/sam3"
+
+
+def main() -> None:
+    print(f"Loading {MODEL_ID} ...")
+    model = Sam3Model.from_pretrained(MODEL_ID)
+    model.eval()
+
+    vision_encoder = model.vision_encoder
+    print(f"Vision encoder type: {type(vision_encoder).__name__}")
+
+    # Print key config values so we know what shape the model expects
+    vc = model.config.vision_config
+    print(
+        f"  vision_config.image_size: {getattr(vc, 'image_size', '?')}\n"
+        f"  vision_config.patch_size: {getattr(vc, 'patch_size', '?')}"
+    )
+
+    # Use the official image processor to produce a correctly-shaped tensor.
+    # SAM3 precomputes positional embeddings for a fixed feature-map size,
+    # so we can't just pass random shapes.
+    print("\nBuilding image processor ...")
+    image_processor = Sam3ImageProcessor.from_pretrained(MODEL_ID)
+    dummy_pil = Image.new("RGB", (640, 480), color="white")
+    inputs = image_processor(images=dummy_pil, return_tensors="pt")
+    dummy_pixel_values = inputs["pixel_values"]
+    print(f"  Processed pixel_values: {tuple(dummy_pixel_values.shape)} dtype={dummy_pixel_values.dtype}")
+
+    # Smoke test in PyTorch first — if this fails, no point trying ONNX.
+    print("\nSmoke testing vision encoder in PyTorch ...")
+    with torch.no_grad():
+        out = vision_encoder(dummy_pixel_values)
+    print(
+        f"  PyTorch output type: {type(out).__name__}\n"
+        f"  fpn_hidden_states: {len(out.fpn_hidden_states)} tensors\n"
+        f"    first shape: {out.fpn_hidden_states[0].shape}\n"
+        f"  fpn_position_encoding: {len(out.fpn_position_encoding)} tensors\n"
+        f"    first shape: {out.fpn_position_encoding[0].shape}"
+    )
+
+    # The encoder returns a structured output. For ONNX we need flat tensors.
+    # We wrap it to flatten the output tuples.
+    class WrappedVisionEncoder(torch.nn.Module):
+        def __init__(self, inner: torch.nn.Module):
+            super().__init__()
+            self.inner = inner
+
+        def forward(self, pixel_values: torch.Tensor):
+            out = self.inner(pixel_values)
+            # Flatten fpn_hidden_states and fpn_position_encoding into a flat tuple
+            return (*out.fpn_hidden_states, *out.fpn_position_encoding)
+
+    wrapped = WrappedVisionEncoder(vision_encoder).eval()
+
+    # Confirm the wrapped version works
+    print("Smoke testing wrapped encoder ...")
+    with torch.no_grad():
+        wrapped_out = wrapped(dummy_pixel_values)
+    print(f"  Flat outputs: {len(wrapped_out)} tensors")
+    for i, t in enumerate(wrapped_out):
+        print(f"    [{i}] shape={t.shape}, dtype={t.dtype}")
+
+    # Construct output names
+    n_fpn = len(out.fpn_hidden_states)
+    output_names = (
+        [f"fpn_hidden_state_{i}" for i in range(n_fpn)]
+        + [f"fpn_position_encoding_{i}" for i in range(n_fpn)]
+    )
+
+    # Dynamic axes: batch and spatial dims can vary
+    dynamic_axes = {"pixel_values": {0: "batch", 2: "height", 3: "width"}}
+    for name in output_names:
+        # FPN outputs are [batch, channels, h, w]
+        dynamic_axes[name] = {0: "batch", 2: "height", 3: "width"}
+
+    print(f"\nExporting to {OUTPUT_FILE} ...")
+    torch.onnx.export(
+        wrapped,
+        (dummy_pixel_values,),
+        str(OUTPUT_FILE),
+        input_names=["pixel_values"],
+        output_names=output_names,
+        dynamic_axes=dynamic_axes,
+        opset_version=18,
+        do_constant_folding=True,
+        verbose=False,
+    )
+
+    size_mb = OUTPUT_FILE.stat().st_size / (1024 * 1024)
+    print(f"\n✅ Exported vision encoder: {OUTPUT_FILE} ({size_mb:.1f} MB)")
+    print("\nNext step: validate with onnxruntime in Python.")
+
+
+if __name__ == "__main__":
+    main()

validate_sam3_e2e.py

@@ -0,0 +1,155 @@
+"""End-to-end validation of the full SAM3 ONNX pipeline.
+
+Runs vision_encoder.onnx → text_encoder.onnx → decoder.onnx with a real seed
+image and a "seed" text prompt, and prints the predicted boxes/masks/scores.
+
+If this produces sensible output, we have a complete in-browser-ready SAM3.
+"""
+
+from pathlib import Path
+import sys
+import time
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+from transformers import AutoTokenizer
+from transformers.models.sam3.image_processing_sam3 import Sam3ImageProcessor
+
+OUTPUT_DIR = Path("sam3-onnx-test")
+MODEL_ID = "facebook/sam3"
+
+VISION_ONNX = OUTPUT_DIR / "vision_encoder.onnx"
+TEXT_ONNX = OUTPUT_DIR / "text_encoder.onnx"
+DECODER_ONNX = OUTPUT_DIR / "decoder.onnx"
+
+
+def main() -> None:
+    # Allow optional CLI args: <image_path> <text_prompt>
+    image_path = sys.argv[1] if len(sys.argv) > 1 else None
+    text_prompt = sys.argv[2] if len(sys.argv) > 2 else "seed"
+
+    # --- Load preprocessors (from HF, NOT from the ONNX files) ----------------
+    print("Loading preprocessors ...")
+    image_processor = Sam3ImageProcessor.from_pretrained(MODEL_ID)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+    # --- Prep image -----------------------------------------------------------
+    if image_path:
+        print(f"Loading image: {image_path}")
+        pil_img = Image.open(image_path).convert("RGB")
+    else:
+        print("No image given — using gray 640x480 placeholder")
+        pil_img = Image.new("RGB", (640, 480), color=(128, 128, 128))
+    pixel_values = image_processor(images=pil_img, return_tensors="np")["pixel_values"]
+    print(f"  pixel_values: shape={pixel_values.shape} dtype={pixel_values.dtype}")
+
+    # --- Prep text ------------------------------------------------------------
+    print(f"Tokenizing prompt: {text_prompt!r}")
+    # We hardcode max_len=32 since that's SAM3's config — keeps script standalone
+    encoded = tokenizer(
+        text_prompt,
+        return_tensors="np",
+        padding="max_length",
+        max_length=32,
+        truncation=True,
+    )
+    input_ids = encoded["input_ids"].astype(np.int64)
+    attention_mask = encoded["attention_mask"].astype(np.int64)
+    print(f"  input_ids: shape={input_ids.shape} dtype={input_ids.dtype}")
+
+    # --- Build ONNX sessions --------------------------------------------------
+    providers = ["CPUExecutionProvider"]
+    print(f"\nLoading ONNX sessions on {providers} ...")
+    t0 = time.time()
+    vision_sess = ort.InferenceSession(str(VISION_ONNX), providers=providers)
+    text_sess = ort.InferenceSession(str(TEXT_ONNX), providers=providers)
+    decoder_sess = ort.InferenceSession(str(DECODER_ONNX), providers=providers)
+    print(f"  Loaded in {time.time() - t0:.1f}s")
+
+    # Print actual input/output names so we know what the exporter kept.
+    # The legacy tracer drops unused inputs, so the decoder may not have all
+    # the names we declared in the export script.
+    print("\n  Decoder ONNX inputs:")
+    for inp in decoder_sess.get_inputs():
+        print(f"    {inp.name}: shape={inp.shape} dtype={inp.type}")
+    print("  Decoder ONNX outputs:")
+    for outp in decoder_sess.get_outputs():
+        print(f"    {outp.name}: shape={outp.shape} dtype={outp.type}")
+
+    # --- 1. Vision encoder ----------------------------------------------------
+    print("\n[1/3] Running vision encoder ...")
+    t0 = time.time()
+    v_out = vision_sess.run(None, {"pixel_values": pixel_values})
+    print(f"  Done in {time.time() - t0:.1f}s")
+    # v_out is a list of 8 tensors in our defined order:
+    # fpn_hidden_state_0..3, fpn_position_encoding_0..3
+    fpn_h = v_out[0:4]
+    fpn_p = v_out[4:8]
+    for i, t in enumerate(fpn_h):
+        print(f"  fpn_hidden_state_{i}: shape={t.shape}")
+
+    # --- 2. Text encoder ------------------------------------------------------
+    print("\n[2/3] Running text encoder ...")
+    t0 = time.time()
+    t_out = text_sess.run(
+        None,
+        {"input_ids": input_ids, "attention_mask": attention_mask},
+    )
+    text_features = t_out[0]
+    print(f"  Done in {time.time() - t0:.1f}s")
+    print(f"  text_features: shape={text_features.shape}")
+
+    # --- 3. Decoder pipeline --------------------------------------------------
+    # Decoder uses first 3 of 4 FPN levels (forward() does fpn_hidden_states[:-1])
+    print("\n[3/3] Running decoder ...")
+    t0 = time.time()
+    # Build feed dict dynamically — the legacy tracer may have dropped unused inputs
+    candidate_inputs = {
+        "fpn_hidden_state_0": fpn_h[0],
+        "fpn_hidden_state_1": fpn_h[1],
+        "fpn_hidden_state_2": fpn_h[2],
+        "fpn_position_encoding_0": fpn_p[0],
+        "fpn_position_encoding_1": fpn_p[1],
+        "fpn_position_encoding_2": fpn_p[2],
+        "text_features": text_features,
+        "attention_mask": attention_mask,
+    }
+    expected_input_names = {inp.name for inp in decoder_sess.get_inputs()}
+    feed = {k: v for k, v in candidate_inputs.items() if k in expected_input_names}
+    missing = expected_input_names - feed.keys()
+    if missing:
+        raise RuntimeError(f"Decoder expects inputs we didn't provide: {missing}")
+    dropped = candidate_inputs.keys() - feed.keys()
+    if dropped:
+        print(f"  (note: tracer optimized out unused inputs: {sorted(dropped)})")
+    d_out = decoder_sess.run(None, feed)
+    pred_masks, pred_boxes, pred_logits = d_out
+    print(f"  Done in {time.time() - t0:.1f}s")
+    print(f"  pred_masks: shape={pred_masks.shape} dtype={pred_masks.dtype}")
+    print(f"  pred_boxes: shape={pred_boxes.shape} dtype={pred_boxes.dtype}")
+    print(f"  pred_logits: shape={pred_logits.shape} dtype={pred_logits.dtype}")
+
+    # --- Inspect top detections -----------------------------------------------
+    # Apply sigmoid to logits to get scores in [0, 1]
+    scores = 1.0 / (1.0 + np.exp(-pred_logits))  # sigmoid
+    scores_b0 = scores[0]
+    top_k = 10
+    top_idx = np.argsort(-scores_b0)[:top_k]
+
+    print(f"\nTop {top_k} detections by score:")
+    print(f"  {'idx':>4} {'score':>7} {'box (xyxy normalized)':>30}")
+    for i in top_idx:
+        x1, y1, x2, y2 = pred_boxes[0, i]
+        print(f"  {i:>4} {scores_b0[i]:>7.4f}  ({x1:.3f}, {y1:.3f}, {x2:.3f}, {y2:.3f})")
+
+    # How many detections above a reasonable threshold?
+    for thresh in (0.5, 0.3, 0.1, 0.05):
+        n_above = (scores_b0 > thresh).sum()
+        print(f"  detections with score > {thresh}: {n_above}")
+
+    print("\n✅ End-to-end ONNX pipeline ran without error.")
+    print("   If scores look reasonable for your test image, we're shipping.")
+
+
+if __name__ == "__main__":
+    main()