SAM + NVIDIA Triton Inference Server

June 5, 2026 · View on GitHub

Production-grade deployment of Meta's Segment Anything Models (SAM1, SAM2, SAM3) using NVIDIA Triton Inference Server.

Quick Links

CLAUDE.md - Comprehensive architecture and deployment details

Features

Multi-Model Support

SAM1 - Original Segment Anything Model (ViT-H, ViT-L, ViT-B)
SAM2.1 - 40% faster inference with Hiera backbone (tiny, small, base_plus, large)
SAM3 - SAM3 Tracker with multi-scale embeddings

Triton Benefits

Enterprise-grade: Industry-standard inference protocol
Performance: GPU-accelerated ONNX Runtime with dynamic batching support
Scalability: Native multi-GPU support with load balancing
Observability: Built-in Prometheus metrics
Flexibility: Hot-reload models without downtime

Quick Start

Prerequisites

NVIDIA GPU with CUDA support
Docker with NVIDIA Container Toolkit
Pixi - Modern Python package manager

Using Podman? The docker and docker compose commands below have direct podman / podman compose equivalents. Ensure the NVIDIA Container Toolkit is configured for Podman (CDI) and pass GPUs with --device nvidia.com/gpu=all. For example: podman run --rm --device nvidia.com/gpu=all nvcr.io/nvidia/tritonserver:25.01-py3 nvidia-smi.

Verify GPU Access

# Test Docker can access GPUs
docker run --rm --gpus all nvcr.io/nvidia/tritonserver:25.01-py3 nvidia-smi

Setup Options

# 1. Install Pixi (if not already installed)
curl -fsSL https://pixi.sh/install.sh | bash

# 2. Install dependencies
pixi install

# 3. Choose your setup:

# Option A: Setup ALL models (SAM1 + SAM2 + SAM3) - ~5GB download
pixi run setup-all

# Option B: Setup only SAM2 (recommended for quick start)
pixi run setup-sam2

# Option C: Setup individual models
pixi run setup-sam1    # SAM1 only (~2.5GB)
pixi run setup-sam2    # SAM2 only (~350MB)
pixi run setup-sam3    # SAM3 only (~1GB, pre-exported ONNX)

# 4. Start Triton server
docker compose up -d

# 5. Verify deployment
curl http://localhost:8000/v2/models

Test Inference

# Test SAM2
pixi run test-sam2

# Test SAM3
pixi run test-sam3

# Speculative request stress test
pixi run test-speculative

Model Comparison

Model	Input Size	Embedding Shape	Best For
SAM1	1024x1024	(1, 256, 64, 64)	Legacy compatibility, proven accuracy
SAM2	1024x1024	(1, 256, 64, 64)	Production default, video support
SAM3	1008x1008	3 multi-scale	Latest features, text prompts

SAM2 Model Sizes

Model	Parameters	Memory	Speed	Use Case
tiny	39M	2GB	91 FPS	Edge devices, real-time
small	46M	2.5GB	85 FPS	Balanced
base_plus	81M	4GB	64 FPS	Production default
large	224M	8GB	40 FPS	Maximum quality

Architecture Overview

┌─────────────────────────────────────────────────────────────────────────┐
│                       NVIDIA Triton Server                               │
│                                                                          │
│  ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────┐          │
│  │  SAM1 Encoder   │  │  SAM2 Encoder   │  │  SAM3 Encoder   │          │
│  │  (1024x1024)    │  │  (1024x1024)    │  │  (1008x1008)    │          │
│  │  → (256,64,64)  │  │  → (256,64,64)  │  │  → 3 embeddings │          │
│  └────────┬────────┘  └────────┬────────┘  └────────┬────────┘          │
│           │                    │                    │                    │
│  ┌────────▼────────┐  ┌────────▼────────┐  ┌────────▼────────┐          │
│  │  SAM1 Decoder   │  │  SAM2 Decoder   │  │  SAM3 Decoder   │          │
│  │  + prompts      │  │  + prompts      │  │  + prompts      │          │
│  │  → masks        │  │  → masks        │  │  → masks        │          │
│  └─────────────────┘  └─────────────────┘  └─────────────────┘          │
│                                                                          │
│  Ports: HTTP (8000) | gRPC (8001) | Metrics (8002)                      │
└─────────────────────────────────────────────────────────────────────────┘

Two-Stage Workflow

All SAM models use a two-stage inference pipeline optimized for interactive segmentation:

Stage 1: Image Encoding (Expensive)

Run once per image
Generates reusable embeddings
~200-800ms depending on model

Stage 2: Mask Prediction (Fast)

Run many times per image
Uses cached embeddings + point prompts
~10-30ms per prediction

Project Structure

triton-sam2/
├── README.md                 # This file
├── CLAUDE.md                 # Detailed architecture docs
├── docker-compose.yml        # Triton server deployment
├── pyproject.toml            # Pixi tasks and dependencies
│
├── triton_sam/               # Python client module
│   ├── client.py             # SAM2TritonClient (sync)
│   ├── speculative_client.py # Async client with cancellation
│   └── tests/
│
├── scripts/
│   ├── download_sam1.sh      # Download SAM1 checkpoints
│   ├── download_sam2.sh      # Download SAM2 checkpoints
│   ├── download_sam3.sh      # Download SAM3 checkpoints
│   ├── download_sam3_onnx.py # Download pre-exported SAM3 ONNX
│   ├── export_sam1_to_onnx.py # Export SAM1 to ONNX
│   ├── export_sam2_to_onnx.py # Export SAM2 to ONNX
│   └── export_sam3_to_onnx.py # Export SAM3 to ONNX
│
├── model_repository/         # Triton model repository
│   ├── sam1_encoder/
│   │   ├── 1/model.onnx
│   │   └── config.pbtxt
│   ├── sam1_decoder/
│   │   ├── 1/model.onnx
│   │   └── config.pbtxt
│   ├── sam2_encoder/
│   │   ├── 1/model.onnx
│   │   └── config.pbtxt
│   ├── sam2_decoder/
│   │   ├── 1/model.onnx
│   │   └── config.pbtxt
│   ├── sam3_encoder/
│   │   ├── 1/vision_encoder.onnx
│   │   └── config.pbtxt
│   └── sam3_decoder/
│       ├── 1/prompt_encoder_mask_decoder.onnx
│       └── config.pbtxt
│
├── checkpoints/              # Downloaded model weights
├── sam1_repo/                # Cloned segment-anything repo
└── sam2_repo/                # Cloned segment-anything-2 repo

Pixi Tasks Reference

Setup Tasks

# Complete setup (all models)
pixi run setup-all      # SAM1 + SAM2 + SAM3

# Individual model setup
pixi run setup-sam1     # Download, clone repo, export SAM1
pixi run setup-sam2     # Download, clone repo, export SAM2 (alias: setup)
pixi run setup-sam3     # Download pre-exported SAM3 ONNX

Download Tasks

# SAM1 checkpoints
pixi run download-sam1-h    # ViT-Huge (2.5GB, recommended)
pixi run download-sam1-l    # ViT-Large
pixi run download-sam1-b    # ViT-Base

# SAM2 checkpoints
pixi run download-tiny      # 38.9M params
pixi run download-small     # 46M params
pixi run download-base      # 80.8M params (recommended)
pixi run download-large     # 224.4M params

# SAM3 ONNX models
pixi run download-sam3-onnx # Pre-exported from HuggingFace

Export Tasks

pixi run export-sam1    # Export SAM1 to ONNX
pixi run export-sam2    # Export SAM2 to ONNX
pixi run export-onnx    # Alias for export-sam2

Test Tasks

pixi run test-sam2          # Basic SAM2 inference test
pixi run test-sam3          # SAM3 inference test
pixi run test-speculative   # Stress test with cancellation

Benchmark Tasks

pixi run benchmark-sam2     # SAM2 performance benchmark
pixi run benchmark-sam3     # SAM3 performance benchmark

Python Client Usage

Basic Client (SAM2TritonClient)

from triton_sam import SAM2TritonClient

# Initialize client (supports sam2 or sam3)
client = SAM2TritonClient("localhost:8000", model_type="sam2")

# Encode image once (cached)
client.set_image("image.jpg")

# Predict masks from point prompts
masks, iou = client.predict(
    point_coords=[[512, 512]],  # (x, y) in original image space
    point_labels=[1]             # 1=foreground, 0=background
)

# Threshold logits at 0 for binary mask
binary_mask = (masks[0, 0] > 0).astype(np.uint8)

Speculative Client (Async with Cancellation)

from triton_sam import SpeculativeSAM2Client, queue_multiple_requests
import asyncio

async def interactive_segmentation():
    client = SpeculativeSAM2Client("localhost:8000")
    client.set_image("image.jpg")

    session_id = "user_session_1"

    # Queue many requests (simulating mouse movement)
    coords_list = [np.array([[x, y]]) for x, y in mouse_positions]
    labels_list = [np.array([1]) for _ in mouse_positions]

    tasks = await queue_multiple_requests(
        client, coords_list, labels_list, session_id
    )

    # Cancel intermediate requests when user stops
    client.cancel_session_requests(session_id)

    # Get final result
    result = await wait_for_latest_result(tasks, client, session_id)

Integration with SAM Service API

This Triton deployment integrates with the SAM Service FastAPI application. See the SAM_service/ directory for the API that provides:

POST /embedded_model - Generate embeddings (supports model_version param)
POST /from_model - Generate masks from embeddings
GET /models - List available models
GET /health - Health check

Monitoring

# Server health
curl http://localhost:8000/v2/health/ready

# List models
curl http://localhost:8000/v2/models

# Prometheus metrics
curl http://localhost:8002/metrics | grep nv_inference

# Docker logs
docker compose logs -f

Troubleshooting

GPU Not Detected

# Verify NVIDIA Container Toolkit
docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi

Models Not Loading

# Check model files exist
ls -la model_repository/*/1/

# Check Triton logs
docker compose logs triton | grep -i error

Out of Memory

Use smaller model (tiny or small for SAM2)
Reduce instance count in config.pbtxt
Check other GPU processes: nvidia-smi

Performance Expectations

Latency (base_plus on modern GPU)

Encoder: ~300ms per image
Decoder: ~15ms per mask
End-to-end (1 image, 10 masks): ~450ms

Memory Requirements

SAM1 ViT-H: ~4GB
SAM2 base_plus: ~4GB
SAM3: ~6GB

License

Released under the Janelia Open-Source Software License.

References

SAM Paper - Original Segment Anything
SAM2 Paper - Segment Anything 2
SAM2 GitHub
NVIDIA Triton