fix(vision): aspect-ratio image processing for Pixtral and Mistral3

docs/supported-models.md

@@ -79,7 +79,7 @@ Implemented VLM variants include:
 - Llama 3.2 Vision (`mllama`): a Llama-3 text backbone whose layers at `cross_attention_layers` (e.g. `[3, 8, 13, 18, 23, 28, 33, 38]`) are gated cross-attention adapters attending to a tiled ViT tower. The image processor picks an optimal tile arrangement (up to `max_num_tiles`, default 4) from the supported aspect ratios, resizes and pads each image into `560x560` tiles, and emits aspect-ratio ids/masks. The tower adds gated tile + position embeddings, runs a local then a global (gated) transformer, and concatenates a set of intermediate hidden states; a `multi_modal_projector` maps those `vision_output_dim` features into the text hidden size. Unlike the LLaVA-style VLMs, image features are not merged into the token stream: they are held as `cross_attention_states` and consulted through the gated cross-attention layers (with per-head `q_norm`/`k_norm`). Text-only prompts leave the cross-attention layers as pass-throughs.
 - LLaVA and LLaVA-Bunny
 - Aya Vision and PaliGemma
-- Pixtral and Mistral 3 VLM wrappers (Mistral 3 VLM supports both the standard Llama/Mistral text backbone and the Mistral4 MLA+MoE backbone; text and image-plus-text are validated on both, including the Mistral Small 4 119B checkpoint)
+- Pixtral and Mistral 3 VLM wrappers (Mistral 3 VLM supports both the standard Llama/Mistral text backbone and the Mistral4 MLA+MoE backbone; text and image-plus-text are validated on both, including the Mistral Small 4 119B checkpoint). Images keep their aspect ratio: each is downscaled so its longest side fits `size.longest_edge` and rounded up to a whole number of merged patches (`patch_size * spatial_merge_size`, so Pixtral rounds to 16 and Mistral 3 to 28), never upscaled. The `[IMG]` placeholder then expands to one patch token per merged cell laid out row by row, with `[IMG_BREAK]` between rows and `[IMG_END]` after the last, so non-square inputs keep the spatial structure the models were trained on. The encoder builds its 2D-RoPE grid per image from the actual patch shape, and multi-image batches with differing sizes run the tower once per image.
 - Qwen2-VL, Qwen2.5-VL, Qwen3-VL, Qwen3.5-VL, and Qwen3-VL MoE
 - PaddleOCR-VL (`paddleocr_vl`): document-OCR VLM pairing a NaViT dynamic-resolution SigLIP-style vision encoder (Conv2d patch embedding, cached bilinearly interpolated learned position embeddings for repeated page grids, 2D vision RoPE, `cu_seqlens`-packed attention with uniform and length-bucketed batched fast paths, and a spatial-merge projector) with a lightweight ERNIE-4.5 text decoder that uses MRoPE. Best for plain OCR, tables, formulas, and chart understanding.
 - Step-3.7 (`step3p7`): StepFun's multimodal model. A 47-block `perception_encoder` ViT tower (Conv2d patchify, learned absolute position embeddings bilinearly resized for the patch grid, 2D vision RoPE, LayerScale, quick-GELU MLP) feeds two stride-2 conv downsamplers and a linear projector into the text hidden size. Each image runs a base pass (728 px, 52x52 grid collapsed to 169 tokens) and, when larger than 728 px, a tiled-patch pass (504 px windows, each 36x36 grid collapsed to 81 tokens); features are ordered patches-first then base per image and scattered into `<im_patch>` placeholders. The text decoder is the Step-3.5 MoE stack (`text_config.model_type: step3p5`), reused as the backbone. Text-only prompts work through the same checkpoint.

src/commands/generate_vlm.rs

@@ -229,6 +229,15 @@ fn print_preparation_summary(summary: VlmPreparationSummary) {
                 image_blocks, total_image_tokens
             );
         }
+        VlmPreparationSummary::Pixtral {
+            image_blocks,
+            total_image_tokens,
+        } => {
+            println!(
+                "Pixtral/Mistral3: expanded {} image placeholder(s) into row-structured blocks ({} total image tokens)",
+                image_blocks, total_image_tokens
+            );
+        }
         VlmPreparationSummary::GraniteVision {
             image_blocks,
             total_image_tokens,

src/lib.rs

@@ -76,8 +76,8 @@ pub use mlxcel_core::generate::{
 pub use mlxcel_core::speculative::SpeculativeGenerator;
 pub use multimodal::{
     internvl_prompt, kimi_vl_prompt, minicpmo_prompt, moondream2_prompt, moondream3_prompt,
-    phi3v_prompt, phi4_siglip_prompt, phi4mm_prompt, qwen_vl, smolvlm_prompt, video, vlm_prompt,
-    vlm_runtime, youtu_vl_prompt,
+    phi3v_prompt, phi4_siglip_prompt, phi4mm_prompt, pixtral_prompt, qwen_vl, smolvlm_prompt,
+    video, vlm_prompt, vlm_runtime, youtu_vl_prompt,
 };
 
 // Re-export the compute-backend seam (issue #338, reframed to a session engine

src/loaded_model_capabilities.rs

@@ -72,6 +72,10 @@ pub enum VlmRuntimeRef<'a> {
     Ernie45MoeVl(&'a vision::ernie4_5_moe_vl::Ernie45MoeVlModel),
     Qwen3OmniMoe(&'a vision::qwen3_omni_moe::Qwen3OmniMoeModel),
     HunyuanVl(&'a vision::hunyuan_vl::HunyuanVlModel),
+    /// Pixtral / Mistral3 dynamic aspect-ratio runtime. Shares the generic
+    /// `VisionModule` storage but preserves each image's aspect ratio and emits
+    /// `[IMG] / [IMG_BREAK] / [IMG_END]` row structure (see `pixtral_layout`).
+    Pixtral(&'a vision::VisionModule),
     Standard(&'a vision::VisionModule),
 }
 
@@ -120,7 +124,7 @@ pub(crate) fn vision_module_from_runtime(
     runtime: VlmRuntimeRef<'_>,
 ) -> Option<&vision::VisionModule> {
     match runtime {
-        VlmRuntimeRef::Standard(vision) => Some(vision),
+        VlmRuntimeRef::Standard(vision) | VlmRuntimeRef::Pixtral(vision) => Some(vision),
         _ => None,
     }
 }
@@ -187,7 +191,13 @@ impl LoadedModel {
             Self::Lfm2VL(model) => Some(VlmRuntimeRef::Lfm2Vl(model)),
             Self::Gemma3VLM(vlm) => Some(VlmRuntimeRef::Standard(&vlm.vision)),
             Self::Llama4VLM(vlm) => Some(VlmRuntimeRef::Standard(&vlm.vision)),
-            Self::LlavaVLM(vlm) => Some(VlmRuntimeRef::Standard(&vlm.vision)),
+            Self::LlavaVLM(vlm) => Some(if vlm.vision.pixtral_layout.is_some() {
+                // Pixtral / Mistral3 carry a row-structured dynamic layout;
+                // route them to the aspect-ratio-preserving runtime path.
+                VlmRuntimeRef::Pixtral(&vlm.vision)
+            } else {
+                VlmRuntimeRef::Standard(&vlm.vision)
+            }),
             Self::GraniteVisionVLM(model) => Some(VlmRuntimeRef::GraniteVision(model)),
             Self::Granite4VisionVLM(model) => Some(VlmRuntimeRef::Granite4Vision(model)),
             Self::DeepSeekOcrVLM(model) => Some(VlmRuntimeRef::DeepSeekOcr(model)),

src/loaded_model_tests.rs

@@ -65,6 +65,7 @@ fn test_vision_module() -> VisionModule {
         suffix_tokens: vec![13, 14],
         block_prefix_tokens: Vec::new(),
         block_suffix_tokens: Vec::new(),
+        pixtral_layout: None,
     }
 }
 

src/loading/vlm_fastvlm.rs

@@ -200,6 +200,7 @@ pub(crate) fn load_fastvlm_vlm(model_path: &Path) -> Result<LoadedModel> {
         suffix_tokens: Vec::new(),
         block_prefix_tokens: Vec::new(),
         block_suffix_tokens: Vec::new(),
+        pixtral_layout: None,
     };
 
     let vlm = vision::VisionLanguageModel {

src/loading/vlm_gemma.rs

@@ -190,6 +190,7 @@ pub(crate) fn load_gemma3_vlm(model_path: &Path) -> Result<LoadedModel> {
         // the extra \n\n wrapping when expanding BOI tokens from the template.
         block_prefix_tokens: vec![108],
         block_suffix_tokens: vec![108],
+        pixtral_layout: None,
     };
 
     let vlm = vision::VisionLanguageModel {

src/loading/vlm_llava.rs

@@ -350,6 +350,7 @@ pub(crate) fn load_llava_vlm(model_path: &Path) -> Result<LoadedModel> {
         suffix_tokens: Vec::new(),
         block_prefix_tokens: Vec::new(),
         block_suffix_tokens: Vec::new(),
+        pixtral_layout: None,
     };
 
     let vlm = vision::VisionLanguageModel {
@@ -476,6 +477,7 @@ pub(crate) fn load_llava_bunny_vlm(model_path: &Path) -> Result<LoadedModel> {
         suffix_tokens: Vec::new(),
         block_prefix_tokens: Vec::new(),
         block_suffix_tokens: Vec::new(),
+        pixtral_layout: None,
     };
 
     let vlm = vision::VisionLanguageModel {

src/loading/vlm_pixtral.rs

@@ -24,16 +24,111 @@
 use anyhow::Result;
 use mlxcel_core::weights::WeightMap;
 use serde_json::Value;
-use std::path::Path;
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
 
 use crate::LoadedModel;
 use crate::models;
 use crate::vision;
+use crate::vision::processors::pixtral::{PixtralLayout, PixtralProcessor};
 
 use super::llava::infer_llama_config_from_weights;
 use super::{load_vlm_weights_common, read_sanitized_vlm_config, strip_language_model_prefix};
 use crate::model_metadata::is_mistral4_config;
 
+fn read_json_file(path: PathBuf) -> Option<Value> {
+    std::fs::read_to_string(path)
+        .ok()
+        .and_then(|s| serde_json::from_str(&s).ok())
+}
+
+/// Resolve the `[IMG] / [IMG_BREAK] / [IMG_END]` token ids for the Pixtral row
+/// layout. Token strings come from `processor_config.json`; the id mapping comes
+/// from `tokenizer_config.json`'s `added_tokens_decoder`. Falls back to the
+/// well-known ids (10 / 12 / 13) shared by pixtral-12b and mistral-small-3.1
+/// when a file or entry is missing.
+fn resolve_pixtral_row_tokens(model_path: &Path, image_token_id: i32) -> (i32, i32, i32) {
+    let proc_cfg = read_json_file(model_path.join("processor_config.json"));
+    let token_str = |key: &str, default: &str| -> String {
+        proc_cfg
+            .as_ref()
+            .and_then(|c| c.get(key))
+            .and_then(|v| v.as_str())
+            .unwrap_or(default)
+            .to_string()
+    };
+    let img_str = token_str("image_token", "[IMG]");
+    let brk_str = token_str("image_break_token", "[IMG_BREAK]");
+    let end_str = token_str("image_end_token", "[IMG_END]");
+
+    // added_tokens_decoder maps id -> { "content": "[IMG_BREAK]", ... }; invert
+    // to content -> id so we can look each special string up by name.
+    let mut content_to_id: HashMap<String, i32> = HashMap::new();
+    if let Some(tok_cfg) = read_json_file(model_path.join("tokenizer_config.json"))
+        && let Some(map) = tok_cfg
+            .get("added_tokens_decoder")
+            .and_then(|v| v.as_object())
+    {
+        for (id_str, entry) in map {
+            if let (Ok(id), Some(content)) = (
+                id_str.parse::<i32>(),
+                entry.get("content").and_then(|c| c.as_str()),
+            ) {
+                content_to_id.insert(content.to_string(), id);
+            }
+        }
+    }
+
+    let img_id = content_to_id
+        .get(&img_str)
+        .copied()
+        .unwrap_or(image_token_id);
+    let brk_id = content_to_id.get(&brk_str).copied().unwrap_or(12);
+    let end_id = content_to_id.get(&end_str).copied().unwrap_or(13);
+    (img_id, brk_id, end_id)
+}
+
+/// Build the dynamic aspect-ratio [`PixtralLayout`] shared by Pixtral and
+/// Mistral3. `spatial_merge_size` is 1 for Pixtral and 2 for Mistral3.
+fn build_pixtral_layout(
+    model_path: &Path,
+    pixtral_config: &vision::encoders::pixtral::PixtralVisionConfig,
+    image_token_id: i32,
+    spatial_merge_size: usize,
+) -> PixtralLayout {
+    // `size.longest_edge` from the (pre)processor config; both checkpoints set
+    // it equal to vision_config.image_size, which is the fallback.
+    let longest_edge = read_json_file(model_path.join("preprocessor_config.json"))
+        .or_else(|| read_json_file(model_path.join("processor_config.json")))
+        .and_then(|c| {
+            c.get("size")
+                .and_then(|s| s.get("longest_edge"))
+                .and_then(|v| v.as_u64())
+        })
+        .map(|v| v as usize)
+        .unwrap_or(pixtral_config.image_size);
+
+    // The encoder's 2D-RoPE table is sized `image_size / patch_size` per side, so
+    // a resized image may never have more patches per side than that. Clamp the
+    // effective longest edge to `image_size` so a checkpoint whose `longest_edge`
+    // exceeds `image_size` cannot drive the patch grid past the RoPE table (which
+    // would wrap position ids or panic). The shipped checkpoints set the two
+    // equal, so this clamp is a no-op for them.
+    let longest_edge = longest_edge.min(pixtral_config.image_size);
+
+    let processor =
+        PixtralProcessor::new(pixtral_config.patch_size, spatial_merge_size, longest_edge);
+    let (image_token_id, image_break_token_id, image_end_token_id) =
+        resolve_pixtral_row_tokens(model_path, image_token_id);
+
+    PixtralLayout {
+        processor,
+        image_token_id,
+        image_break_token_id,
+        image_end_token_id,
+    }
+}
+
 struct PixtralFamilyContext {
     full_config: Value,
     weights: WeightMap,
@@ -199,7 +294,6 @@ fn build_pixtral_family_context(model_path: &Path) -> Result<PixtralFamilyContex
 pub(crate) fn load_pixtral_vlm(model_path: &Path) -> Result<LoadedModel> {
     use vision::connectors::mlp::MLPProjector;
     use vision::encoders::pixtral::PixtralVisionModel;
-    use vision::processors::siglip::SigLipProcessor;
 
     let context = build_pixtral_family_context(model_path)?;
 
@@ -219,7 +313,15 @@ pub(crate) fn load_pixtral_vlm(model_path: &Path) -> Result<LoadedModel> {
     )
     .map_err(|e| anyhow::anyhow!("Failed to load MLP projector: {}", e))?;
 
-    let processor = SigLipProcessor::new_rescale_only(context.pixtral_config.image_size);
+    // Pixtral has no spatial merge (spatial_merge_size = 1); the row layout and
+    // resize round-up both key off patch_size alone here.
+    let layout = build_pixtral_layout(
+        model_path,
+        &context.pixtral_config,
+        context.image_token_id,
+        1,
+    );
+    let processor = layout.processor.clone();
     let num_patches =
         (context.pixtral_config.image_size / context.pixtral_config.patch_size).pow(2);
     let mm_tokens_per_image = context
@@ -232,7 +334,7 @@ pub(crate) fn load_pixtral_vlm(model_path: &Path) -> Result<LoadedModel> {
         encoder: Box::new(vision_encoder),
         connector: Box::new(connector),
         processor: Box::new(processor),
-        image_token_id: context.image_token_id,
+        image_token_id: layout.image_token_id,
         pad_token_id: context.pad_token_id,
         hidden_size: if context.hidden_size > 0 {
             context.hidden_size
@@ -248,6 +350,7 @@ pub(crate) fn load_pixtral_vlm(model_path: &Path) -> Result<LoadedModel> {
         suffix_tokens: Vec::new(),
         block_prefix_tokens: Vec::new(),
         block_suffix_tokens: Vec::new(),
+        pixtral_layout: Some(layout),
     };
 
     let vlm = vision::VisionLanguageModel {
@@ -262,7 +365,6 @@ pub(crate) fn load_pixtral_vlm(model_path: &Path) -> Result<LoadedModel> {
 pub(crate) fn load_mistral3_vlm(model_path: &Path) -> Result<LoadedModel> {
     use vision::connectors::mistral3::Mistral3Projector;
     use vision::encoders::pixtral::PixtralVisionModel;
-    use vision::processors::siglip::SigLipProcessor;
 
     let context = build_pixtral_family_context(model_path)?;
 
@@ -292,7 +394,13 @@ pub(crate) fn load_mistral3_vlm(model_path: &Path) -> Result<LoadedModel> {
     )
     .map_err(|e| anyhow::anyhow!("Failed to load Mistral3 projector: {}", e))?;
 
-    let processor = SigLipProcessor::new_rescale_only(context.pixtral_config.image_size);
+    let layout = build_pixtral_layout(
+        model_path,
+        &context.pixtral_config,
+        context.image_token_id,
+        spatial_merge_size as usize,
+    );
+    let processor = layout.processor.clone();
     let num_patches_per_side =
         context.pixtral_config.image_size / context.pixtral_config.patch_size;
     let merged_per_side = num_patches_per_side / spatial_merge_size as usize;
@@ -307,7 +415,7 @@ pub(crate) fn load_mistral3_vlm(model_path: &Path) -> Result<LoadedModel> {
         encoder: Box::new(vision_encoder),
         connector: Box::new(connector),
         processor: Box::new(processor),
-        image_token_id: context.image_token_id,
+        image_token_id: layout.image_token_id,
         pad_token_id: context.pad_token_id,
         hidden_size: if context.hidden_size > 0 {
             context.hidden_size
@@ -323,6 +431,7 @@ pub(crate) fn load_mistral3_vlm(model_path: &Path) -> Result<LoadedModel> {
         suffix_tokens: Vec::new(),
         block_prefix_tokens: Vec::new(),
         block_suffix_tokens: Vec::new(),
+        pixtral_layout: Some(layout),
     };
 
     let vlm = vision::VisionLanguageModel {

src/loading/vlm_siglip.rs

@@ -259,6 +259,7 @@ pub(crate) fn load_aya_vision_vlm(model_path: &Path) -> Result<LoadedModel> {
         suffix_tokens: Vec::new(),
         block_prefix_tokens: Vec::new(),
         block_suffix_tokens: Vec::new(),
+        pixtral_layout: None,
     };
 
     let vlm = vision::VisionLanguageModel {
@@ -341,6 +342,7 @@ pub(crate) fn load_paligemma_vlm(model_path: &Path) -> Result<LoadedModel> {
         suffix_tokens: vec![108],    // newline after text prompt
         block_prefix_tokens: Vec::new(),
         block_suffix_tokens: Vec::new(),
+        pixtral_layout: None,
     };
 
     let vlm = vision::VisionLanguageModel {

src/loading/vlm_special.rs

@@ -2291,6 +2291,7 @@ pub(crate) fn load_llama4_vlm(model_path: &Path) -> Result<LoadedModel> {
         suffix_tokens: Vec::new(),
         block_prefix_tokens: Vec::new(),
         block_suffix_tokens: Vec::new(),
+        pixtral_layout: None,
     };
 
     let vlm = vision::VisionLanguageModel {

src/multimodal/mod.rs

@@ -49,6 +49,7 @@ pub mod moondream3_prompt;
 pub mod phi3v_prompt;
 pub mod phi4_siglip_prompt;
 pub mod phi4mm_prompt;
+pub mod pixtral_prompt;
 pub mod qwen_vl;
 pub mod smolvlm_prompt;
 pub mod step3p7_prompt;