feat(sam3-image): add TensorRT support (#166)

Author: jamjamjon

examples/sam3/README.md

@@ -1,40 +1,51 @@
 
-### Quick Start
+## Usage
 
 ```bash
-# Text prompt
-cargo run -r -F sam3 -F cuda --example sam3 -- --device cuda --dtype q4f16 --source ./assets/sam3-demo.jpg -p shoe
-cargo run -r -F sam3 -F cuda --example sam3 -- --device cuda --dtype bnb4 --source ./assets/sam3-demo.jpg -p "person in red vest"
-cargo run -r -F sam3 -F cuda --example sam3 -- --device cuda --dtype q8 --source ./assets/sam3-demo.jpg -p "boy in blue vest"
-
-# Visual prompt: a single bbox
-cargo run -r -F sam3 -F cuda --example sam3 -- --device cuda --source ./assets/sam3-demo.jpg -p "visual;pos:480,290,110,360"
+# Basic text prompt
+cargo run -r -F sam3,cuda --example sam3 -- \
+    --device cuda --dtype q4f16 --source assets/sam3-demo.jpg \
+    -p shoe 
+
+# Multiple prompts
+cargo run -r -F sam3,cuda --example sam3 -- \
+    --device cuda --dtype fp16 --source assets/sam3-demo.jpg \
+    -p "person in red vest" --show-mask
+
+# Multiple prompts on multiple images
+cargo run -r -F sam3,cuda --example sam3 -- \
+    --device cuda --dtype q4f16 --batch 2 \
+    --source assets/sam3-demo.jpg --source assets/bus.jpg \
+    -p bus -p cat -p shoe -p cellphone -p person -p "short hair"
+
+# Visual prompt (bbox)
+cargo run -r -F sam3,cuda --example sam3 -- \
+    --device cuda --dtype fp16 --source assets/sam3-demo.jpg \
+    -p "visual;pos:480,290,110,360"
 
 # Visual prompt: multi-boxes prompting(with positive and negative boxes)
-cargo run -r -F sam3 -F cuda --example sam3 -- --device cuda  --source ./assets/sam3-demo.jpg -p "visual;pos:480,290,110,360;neg:370,280,115,375"
+cargo run -r -F sam3 -F cuda --example sam3 -- \
+    --device cuda --source ./assets/sam3-demo.jpg \
+    -p "visual;pos:480,290,110,360;neg:370,280,115,375"
 
 # Text + negative box
-cargo run -r -F sam3 -F cuda --example sam3 -- --device cuda --dtype fp16 --source ./assets/000000136466.jpg -p "handle"
-cargo run -r -F sam3 -F cuda --example sam3 -- --device cuda --dtype fp16 --source ./assets/000000136466.jpg -p "handle;neg:40,183,278,21"
+cargo run -r -F sam3,cuda --example sam3 -- \
+    --device cuda --source assets/000000136466.jpg \
+    -p "handle;neg:40,183,278,21"
 
-# Multiple prompts (Queries)
-cargo run -r -F sam3 -F cuda --example sam3 -- --device cuda --dtype fp16 --source ./assets/sam3-demo.jpg --source ./assets/bus.jpg -p shoe -p face -p person
 ```
 
-
-
-### Prompt Format
+## Prompt Format
 
 ```
-"text;pos:x,y,w,h;neg:x,y,w,h"
+text;pos:x,y,w,h;neg:x,y,w,h
 ```
 
-- `text`: Text description
-- `pos:x,y,w,h`: Positive box (find similar)
-- `neg:x,y,w,h`: Negative box (exclude region)
+- `text` - Text description
+- `pos:x,y,w,h` - Positive box (include region)
+- `neg:x,y,w,h` - Negative box (exclude region)
 
+## Results
 
-### Results
-
-![](https://github.com/jamjamjon/assets/releases/download/sam3/demo.jpg)
 ![](https://github.com/jamjamjon/assets/releases/download/sam3/demo2.jpg)
+![](https://github.com/jamjamjon/assets/releases/download/sam3/demo.jpg)

examples/sam3/main.rs

@@ -29,18 +29,30 @@ struct Args {
     #[argh(option, default = "0.5")]
     conf: f32,
 
-    /// batch size min (default: 1)
+    /// image batch size min
     #[argh(option, default = "1")]
     batch_min: usize,
 
-    /// batch size (default: 1)
+    /// image batch size
     #[argh(option, default = "1")]
     batch: usize,
 
-    /// batch size max (default: 4)
-    #[argh(option, default = "4")]
+    /// image batch size max
+    #[argh(option, default = "8")]
     batch_max: usize,
 
+    /// text batch size min
+    #[argh(option, default = "1")]
+    text_batch_min: usize,
+
+    /// text batch size
+    #[argh(option, default = "4")]
+    text_batch: usize,
+
+    /// text batch size max
+    #[argh(option, default = "16")]
+    text_batch_max: usize,
+
     /// dtype
     #[argh(option, default = "String::from(\"q4f16\")")]
     dtype: String,
@@ -55,7 +67,6 @@ fn main() -> Result<()> {
         .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
         .with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
         .init();
-
     let args: Args = argh::from_env();
 
     // Parse prompts
@@ -71,11 +82,20 @@ fn main() -> Result<()> {
 
     // Build model
     let config = Config::sam3_image_predictor()
-        .with_batch_size_all_min_opt_max(args.batch_min, args.batch, args.batch_max)
-        .with_device_all(args.device.parse()?)
         .with_dtype_all(args.dtype.parse()?)
         .with_class_confs(&[args.conf])
-        .with_num_dry_run_all(1)
+        .with_batch_size_all_min_opt_max(args.batch_min, args.batch, args.batch_max)
+        .with_textual_encoder_batch_min_opt_max(
+            args.text_batch_min,
+            args.text_batch,
+            args.text_batch_max,
+        )
+        // => If your GPU memory is insufficient, you can place some modules on CPU
+        // .with_visual_encoder_device(args.device.parse()?)
+        // .with_textual_encoder_device(args.device.parse()?)
+        // .with_encoder_device(args.device.parse()?)  // geometry-encoder
+        // .with_decoder_device(args.device.parse()?)
+        .with_device_all(args.device.parse()?)
         .commit()?;
     let mut model = SAM3::new(config)?;
 
@@ -105,6 +125,11 @@ fn main() -> Result<()> {
         }
     }
 
+    // Cache info
+    let (cached_count, mem_mb) = model.cache_stats();
+    println!("Cache: {} texts, ~{}MB", cached_count, mem_mb);
+
     usls::perf(false);
+
     Ok(())
 }

src/core/config.rs

@@ -236,6 +236,46 @@ impl Config {
         self
     }
 
+    pub fn with_textual_encoder_batch_min_opt_max(
+        mut self,
+        min: usize,
+        opt: usize,
+        max: usize,
+    ) -> Self {
+        self.textual_encoder = self.textual_encoder.with_ixx(0, 0, (min, opt, max).into());
+        self
+    }
+
+    pub fn with_visual_encoder_batch_min_opt_max(
+        mut self,
+        min: usize,
+        opt: usize,
+        max: usize,
+    ) -> Self {
+        self.visual_encoder = self.visual_encoder.with_ixx(0, 0, (min, opt, max).into());
+        self
+    }
+
+    pub fn with_encoder_encoder_batch_min_opt_max(
+        mut self,
+        min: usize,
+        opt: usize,
+        max: usize,
+    ) -> Self {
+        self.encoder = self.encoder.with_ixx(0, 0, (min, opt, max).into());
+        self
+    }
+
+    pub fn with_decoder_batch_min_opt_max(mut self, min: usize, opt: usize, max: usize) -> Self {
+        self.decoder = self.decoder.with_ixx(0, 0, (min, opt, max).into());
+        self
+    }
+
+    pub fn with_model_batch_min_opt_max(mut self, min: usize, opt: usize, max: usize) -> Self {
+        self.model = self.model.with_ixx(0, 0, (min, opt, max).into());
+        self
+    }
+
     pub fn with_batch_size_all(mut self, batch_size: usize) -> Self {
         self.visual = self.visual.with_ixx(0, 0, batch_size.into());
         self.textual = self.textual.with_ixx(0, 0, batch_size.into());

src/models/sam3/README.md

@@ -5,6 +5,101 @@ A powerful multimodal segmentation model supporting text, bounding box, and comb
 ## References
 
 - Official: [facebookresearch/sam3](https://github.com/facebookresearch/sam3)
+- ONNX Models: [Download](https://github.com/jamjamjon/assets/releases/tag/sam3)
+
+## ONNX Model Specifications
+
+All models support dynamic batch processing.
+
+### Vision Encoder
+```
+Inputs:
+  images                [batch, 3, 1008, 1008]    FLOAT
+
+Outputs:
+  fpn_feat_0            [batch, 256, 288, 288]    FLOAT
+  fpn_feat_1            [batch, 256, 144, 144]    FLOAT
+  fpn_feat_2            [batch, 256, 72, 72]      FLOAT
+  fpn_feat_3            [batch, 256, 36, 36]      FLOAT
+  fpn_pos_0             [batch, 256, 288, 288]    FLOAT
+  fpn_pos_1             [batch, 256, 144, 144]    FLOAT
+  fpn_pos_2             [batch, 256, 72, 72]      FLOAT
+  fpn_pos_3             [batch, 256, 36, 36]      FLOAT
+```
+
+### Text Encoder
+```
+Inputs:
+  input_ids             [batch, 32]               INT64
+  attention_mask        [batch, 32]               INT64
+
+Outputs:
+  text_features         [batch, 32, 256]          FLOAT
+  text_mask             [batch, 32]               BOOL
+```
+
+### Geometry Encoder
+```
+Inputs:
+  input_boxes           [batch, num_boxes, 4]     FLOAT
+  input_boxes_labels    [batch, num_boxes]        INT64
+  fpn_feat              [batch, 256, 72, 72]      FLOAT
+  fpn_pos               [batch, 256, 72, 72]      FLOAT
+
+Outputs:
+  geometry_features     [batch, num_boxes+1, 256] FLOAT
+  geometry_mask         [batch, num_boxes+1]      BOOL
+```
+
+### Decoder
+```
+Inputs:
+  fpn_feat_0            [batch, 256, 288, 288]    FLOAT
+  fpn_feat_1            [batch, 256, 144, 144]    FLOAT
+  fpn_feat_2            [batch, 256, 72, 72]      FLOAT
+  fpn_pos_2             [batch, 256, 72, 72]      FLOAT
+  prompt_features       [batch, prompt_len, 256]  FLOAT
+  prompt_mask           [batch, prompt_len]       BOOL
+
+Outputs:
+  pred_masks            [batch, 200, 288, 288]    FLOAT
+  pred_boxes            [batch, 200, 4]           FLOAT
+  pred_logits           [batch, 200]              FLOAT
+  presence_logits       [batch, 1]                FLOAT
+```
+
+## TensorRT Conversion
+
+```bash
+# Vision Encoder
+trtexec --onnx=sam3_vision_encoder.onnx \
+  --minShapes=images:1x3x1008x1008 \
+  --optShapes=images:4x3x1008x1008 \
+  --maxShapes=images:8x3x1008x1008 \
+  --saveEngine=sam3_vision_encoder.engine
+
+# Text Encoder
+trtexec --onnx=sam3_text_encoder.onnx \
+  --minShapes=input_ids:1x32,attention_mask:1x32 \
+  --optShapes=input_ids:4x32,attention_mask:4x32 \
+  --maxShapes=input_ids:8x32,attention_mask:8x32 \
+  --saveEngine=sam3_text_encoder.engine
+
+# Geometry Encoder
+trtexec --onnx=sam3_geometry_encoder.onnx \
+  --minShapes=input_boxes:1x1x4,input_boxes_labels:1x1,fpn_feat:1x256x72x72,fpn_pos:1x256x72x72 \
+  --optShapes=input_boxes:1x5x4,input_boxes_labels:1x5,fpn_feat:1x256x72x72,fpn_pos:1x256x72x72 \
+  --maxShapes=input_boxes:8x20x4,input_boxes_labels:8x20,fpn_feat:8x256x72x72,fpn_pos:8x256x72x72 \
+  --saveEngine=sam3_geometry_encoder.engine
+
+# Decoder
+trtexec --onnx=sam3_decoder.onnx \
+  --minShapes=fpn_feat_0:1x256x288x288,fpn_feat_1:1x256x144x144,fpn_feat_2:1x256x72x72,fpn_pos_2:1x256x72x72,prompt_features:1x1x256,prompt_mask:1x1 \
+  --optShapes=fpn_feat_0:1x256x288x288,fpn_feat_1:1x256x144x144,fpn_feat_2:1x256x72x72,fpn_pos_2:1x256x72x72,prompt_features:1x33x256,prompt_mask:1x33 \
+  --maxShapes=fpn_feat_0:8x256x288x288,fpn_feat_1:8x256x144x144,fpn_feat_2:8x256x72x72,fpn_pos_2:8x256x72x72,prompt_features:8x60x256,prompt_mask:8x60 \
+  --saveEngine=sam3_decoder.engine
+```
+
 
 ## Example
 

src/models/sam3/config.rs

@@ -11,11 +11,16 @@ impl Config {
     pub fn sam3() -> Self {
         Self::default()
             .with_name("sam3")
-            .with_batch_size_all_min_opt_max(1, 1, 4)
-            .with_visual_encoder_ixx(0, 1, 3.into())
-            .with_visual_encoder_ixx(0, 2, 1008.into())
-            .with_visual_encoder_ixx(0, 3, 1008.into())
-            .with_textual_encoder_ixx(0, 1, 32.into())
+            .with_num_dry_run_all(1)
+            .with_batch_size_all_min_opt_max(1, 1, 10)
+            .with_visual_encoder_ixx(0, 1, 3.into()) // vision-encoder: channels
+            .with_visual_encoder_ixx(0, 2, 1008.into()) // vision-encoder: height
+            .with_visual_encoder_ixx(0, 3, 1008.into()) // vision-encoder: width
+            .with_textual_encoder_ixx(0, 1, 32.into()) // text-encoder: text sequence length
+            .with_encoder_ixx(0, 1, (1, 2, 8).into()) // geometry-encoder: input_boxes, num_boxes
+            .with_encoder_ixx(1, 1, (1, 2, 8).into()) // geometry-encoder: input_boxes_labels, num_boxes
+            .with_decoder_ixx(4, 1, (1, 34, 60).into()) // decoder: prompt_features prompt_len
+            .with_decoder_ixx(5, 1, (1, 34, 60).into()) // decoder: prompt_mask prompt_len
             .with_resize_mode(crate::ResizeMode::FitExact)
             .with_resize_filter("Bilinear")
             .with_image_mean(&[0.5, 0.5, 0.5])