update

Author: zhangjian

Cargo.toml

@@ -16,8 +16,10 @@ anyhow = { version = "1" }
 slsl = { version = "0.0.7", features = ["rayon"], optional = true }
 aksr = { version = "0.0.7" }
 ab_glyph = { version = "0.2.32" }
-image = { version = "0.25" }
-imageproc = { version = "0.25" }
+image = { version = "0.25", default-features = false, features = [
+  "jpeg", "png", "gif", "webp", "bmp", "ico", "rayon"
+] }
+imageproc = { version = "0.25", default-features = false, features = ["rayon"] }
 ndarray = { version = "0.16.1", features = ["rayon"] }
 half = { version = "2.7.1", features = ["bytemuck", "num-traits"] }
 indicatif = { version = "0.18" }
@@ -35,7 +37,7 @@ dirs = { version = "6.0.0" }
 chrono = { version = "0.4.40" }
 regex = { version = "1.12.2" }
 tempfile = { version = "3.23.0" }
-fast_image_resize = { version = "=5.1.4", features = ["image", "rayon"] }
+fast_image_resize = { version = "=5.1.4", default-features = false, features = ["image", "rayon"] }
 minifb = { version = "0.28.0", optional = true }
 video-rs = { version = "0.10.3", features = ["ndarray"], optional = true }
 ndarray-npy = { version = "0.9.1", optional = true  }
@@ -46,6 +48,7 @@ ort = { version = "=2.0.0-rc.10", default-features = false, optional = true, fea
 ] }
 hf-hub = { version = "0.4.3", default-features = false, features = ["ureq", "native-tls"], optional = true }
 tokenizers = { version = "0.22.1", optional = true }
+lru = "0.16.2"
 
 
 [dev-dependencies]
@@ -54,9 +57,13 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter", "chrono"] }
 
 
 [profile.dev]
-opt-level = 1
+opt-level = 0
 incremental = true
 
+# Optimize dependencies in dev mode (faster runtime, slightly slower first compile)
+[profile.dev.package."*"]
+opt-level = 2
+
 [profile.release]
 opt-level = 3
 lto = "fat"

examples/sam3/README.md

@@ -14,7 +14,7 @@ cargo run -r -F sam3,cuda --example sam3 -- \
 
 # Multiple prompts on multiple images
 cargo run -r -F sam3,cuda --example sam3 -- \
-    --device cuda --dtype q4f16 --batch 2 \
+    --device cuda --dtype q4f16 --vision-batch 2 \
     --source assets/sam3-demo.jpg --source assets/bus.jpg \
     -p bus -p cat -p shoe -p cellphone -p person -p "short hair"
 

examples/sam3/main.rs

@@ -12,13 +12,7 @@ struct Args {
     device: String,
 
     /// source image paths (can specify multiple)
-    #[argh(
-        option,
-        default = "vec![
-        String::from(\"./assets/sam3-demo.jpg\"),
-        // String::from(\"./assets/bus.jpg\")
-    ]"
-    )]
+    #[argh(option, default = "vec![String::from(\"./assets/sam3-demo.jpg\")]")]
     source: Vec<String>,
 
     /// prompts: "text;pos:x,y,w,h;neg:x,y,w,h" (can specify multiple)
@@ -29,37 +23,61 @@ struct Args {
     #[argh(option, default = "0.5")]
     conf: f32,
 
-    /// image batch size min
+    /// show mask
+    #[argh(switch)]
+    show_mask: bool,
+
+    /// dtype
+    #[argh(option, default = "String::from(\"q4f16\")")]
+    dtype: String,
+
+    /// vision encoder batch min
     #[argh(option, default = "1")]
-    batch_min: usize,
+    vision_batch_min: usize,
 
-    /// image batch size
+    /// vision encoder batch opt
     #[argh(option, default = "1")]
-    batch: usize,
+    vision_batch: usize,
 
-    /// image batch size max
+    /// vision encoder batch max
     #[argh(option, default = "8")]
-    batch_max: usize,
+    vision_batch_max: usize,
 
-    /// text batch size min
+    /// text encoder batch min
     #[argh(option, default = "1")]
     text_batch_min: usize,
 
-    /// text batch size
+    /// text encoder batch opt
     #[argh(option, default = "4")]
     text_batch: usize,
 
-    /// text batch size max
+    /// text encoder batch max
     #[argh(option, default = "16")]
     text_batch_max: usize,
 
-    /// dtype
-    #[argh(option, default = "String::from(\"q4f16\")")]
-    dtype: String,
+    /// geometry encoder batch min
+    #[argh(option, default = "1")]
+    geo_batch_min: usize,
 
-    /// show mask
-    #[argh(switch)]
-    show_mask: bool,
+    /// geometry encoder batch opt
+    #[argh(option, default = "8")]
+    geo_batch: usize,
+
+    /// geometry encoder batch max
+    #[argh(option, default = "16")]
+    geo_batch_max: usize,
+
+    /// decoder batch min
+    #[argh(option, default = "1")]
+    decoder_batch_min: usize,
+
+    /// decoder batch opt
+    #[argh(option, default = "1")]
+    decoder_batch: usize,
+
+    /// decoder batch max
+    #[argh(option, default = "8")]
+    decoder_batch_max: usize,
 }
 
 fn main() -> Result<()> {
@@ -80,16 +98,28 @@ fn main() -> Result<()> {
         .collect::<std::result::Result<Vec<_>, _>>()
         .map_err(|e| anyhow::anyhow!("{}", e))?;
 
-    // Build model
+    // Build model with per-encoder batch sizes
     let config = Config::sam3_image_predictor()
         .with_dtype_all(args.dtype.parse()?)
         .with_class_confs(&[args.conf])
-        .with_batch_size_all_min_opt_max(args.batch_min, args.batch, args.batch_max)
+        // Per-encoder batch sizes for TensorRT (min, opt, max)
+        .with_visual_encoder_batch_min_opt_max(
+            args.vision_batch_min,
+            args.vision_batch,
+            args.vision_batch_max,
+        )
         .with_textual_encoder_batch_min_opt_max(
             args.text_batch_min,
             args.text_batch,
             args.text_batch_max,
         )
+        .with_encoder_batch_min_opt_max(args.geo_batch_min, args.geo_batch, args.geo_batch_max)
+        .with_decoder_batch_min_opt_max(
+            args.decoder_batch_min,
+            args.decoder_batch,
+            args.decoder_batch_max,
+        )
+        // Device configuration
         // => 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()?)
@@ -109,7 +139,7 @@ fn main() -> Result<()> {
 
     // DataLoader with batch iteration
     let dataloader = DataLoader::from_paths(&args.source)?
-        .with_batch(args.batch)
+        .with_batch(args.vision_batch)
         .with_progress_bar(true)
         .build()?;
 
@@ -125,10 +155,10 @@ fn main() -> Result<()> {
         }
     }
 
-    // Cache info
-    let (cached_count, mem_mb) = model.cache_stats();
-    println!("Cache: {} texts, ~{}MB", cached_count, mem_mb);
+    // Cache stats
+    model.cache_stats();
 
+    // Performance stats
     usls::perf(false);
 
     Ok(())

src/core/config.rs

@@ -236,46 +236,6 @@ 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());
@@ -860,3 +820,30 @@ impl_ort_config_methods!(Config, coord_decoder);
 impl_ort_config_methods!(Config, visual_projection);
 impl_ort_config_methods!(Config, textual_projection);
 impl_processor_config_methods!(Config, processor);
+
+/// Macro to generate `with_*_batch_min_opt_max` methods for Config
+macro_rules! impl_batch_min_opt_max {
+    ($ty:ty, $field:ident) => {
+        impl $ty {
+            paste::paste! {
+                pub fn [<with_ $field _batch_min_opt_max>](mut self, min: usize, opt: usize, max: usize) -> Self {
+                    self.$field = self.$field.with_ixx(0, 0, (min, opt, max).into());
+                    self
+                }
+            }
+        }
+    };
+}
+
+// Batch size configuration methods
+impl_batch_min_opt_max!(Config, model);
+impl_batch_min_opt_max!(Config, visual);
+impl_batch_min_opt_max!(Config, textual);
+impl_batch_min_opt_max!(Config, encoder);
+impl_batch_min_opt_max!(Config, decoder);
+impl_batch_min_opt_max!(Config, visual_encoder);
+impl_batch_min_opt_max!(Config, textual_encoder);
+impl_batch_min_opt_max!(Config, visual_decoder);
+impl_batch_min_opt_max!(Config, textual_decoder);
+impl_batch_min_opt_max!(Config, visual_projection);
+impl_batch_min_opt_max!(Config, textual_projection);

src/core/dtype.rs

@@ -66,6 +66,21 @@ impl std::str::FromStr for DType {
     }
 }
 
+impl DType {
+    pub fn size_in_bytes(self) -> usize {
+        match self {
+            Self::Auto => 4, // default to f32
+            Self::Int4 | Self::Uint4 | Self::Bnb4 | Self::Q4 | Self::Fp4e2m1 => 1, // sub-byte, min 1
+            Self::Int8 | Self::Uint8 | Self::Q8 => 1,
+            Self::Fp8e4m3fn | Self::Fp8e4m3fnuz | Self::Fp8e5m2 | Self::Fp8e5m2fnuz => 1,
+            Self::Int16 | Self::Uint16 | Self::Fp16 | Self::Bf16 | Self::Q4f16 => 2,
+            Self::Int32 | Self::Uint32 | Self::Fp32 => 4,
+            Self::Int64 | Self::Uint64 | Self::Fp64 | Self::Complex64 => 8,
+            Self::Complex128 => 16,
+        }
+    }
+}
+
 impl std::fmt::Display for DType {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         let x = match self {

src/core/engine.rs

@@ -40,7 +40,8 @@ impl From<TensorElementType> for DType {
             TensorElementType::Float8E5M2FNUZ => Self::Fp8e5m2fnuz,
             TensorElementType::Complex64 => Self::Complex64,
             TensorElementType::Complex128 => Self::Complex128,
-            _ => todo!(),
+            TensorElementType::Bool => Self::Uint8,
+            TensorElementType::String | TensorElementType::Undefined => Self::Auto,
         }
     }
 }

src/models/sam3/README.md

@@ -70,29 +70,35 @@ Outputs:
 
 ## TensorRT Conversion
 
+### Vision Encoder
 ```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
+### Text Encoder
+```bash
 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
+### Geometry Encoder
+```bash
 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
+### Decoder
+```bash
 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 \

src/models/sam3/config.rs

@@ -3,24 +3,25 @@ use crate::Config;
 /// Model configuration for `SAM3`
 impl Config {
     /// SAM3 base configuration
-    ///
-    /// - Input size: 1008x1008 (FitExact, no aspect ratio preserved)
-    /// - Normalization: mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]
-    /// - Tokenizer: CLIP BPE (max_length=32)
-    /// - Confidence threshold: 0.5
     pub fn sam3() -> Self {
         Self::default()
             .with_name("sam3")
             .with_num_dry_run_all(1)
-            .with_batch_size_all_min_opt_max(1, 1, 10)
+            // Batch sizes: vision=1, text=4, geometry=8, decoder=1
+            .with_visual_encoder_batch_min_opt_max(1, 1, 8)
+            .with_textual_encoder_batch_min_opt_max(1, 4, 16)
+            .with_encoder_batch_min_opt_max(1, 8, 16) // geometry encoder
+            .with_decoder_batch_min_opt_max(1, 1, 8) // decoder (memory intensive)
+            // Shape configurations
             .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_textual_encoder_ixx(0, 1, 32.into()) // text-encoder: sequence length
+            .with_encoder_ixx(0, 1, (1, 2, 8).into()) // geometry-encoder: num_boxes
+            .with_encoder_ixx(1, 1, (1, 2, 8).into()) // geometry-encoder: num_boxes
+            .with_decoder_ixx(4, 1, (1, 34, 60).into()) // decoder: prompt_len
+            .with_decoder_ixx(5, 1, (1, 34, 60).into()) // decoder: prompt_len
+            // Preprocessing
             .with_resize_mode(crate::ResizeMode::FitExact)
             .with_resize_filter("Bilinear")
             .with_image_mean(&[0.5, 0.5, 0.5])