From 75cb1a1e2aadca7605be4b3c25e8ba58632166a5 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 13 Jul 2026 20:26:47 +0900 Subject: [PATCH 1/2] fix(vision): aspect-ratio image processing for Pixtral and Mistral3 Pixtral and Mistral3 image inputs ran through the fixed-square SigLIP rescale-only path, which force-resized every image to a square, applied no CLIP normalization, and emitted a fixed image-token block with no row structure. These models were trained on aspect-ratio-preserving inputs whose image tokens are laid out row by row with [IMG_BREAK] between rows and [IMG_END] at the end, so the old path distorted non-square images and dropped the spatial structure the models rely on. A new PixtralProcessor (src/vision/processors/pixtral.rs) downscales each image so its longest side fits size.longest_edge, rounds each side up to a whole multiple of patch_size * spatial_merge_size (16 for Pixtral, 28 for Mistral3), never upscales, then rescales by 1/255 and CLIP-normalizes. Rounding to the merged-patch factor rather than patch_size alone keeps the pre-merge patch grid even so the Mistral3 PatchMerger never floor-drops a partial merged row; the reference mlx-vlm processor reaches the same factor by passing patch_size * spatial_merge_size to its resize helper. insert_pixtral_image_tokens (src/multimodal/pixtral_prompt.rs) expands each [IMG] placeholder into tokens_h rows of tokens_w [IMG] tokens with [IMG_BREAK] between rows and [IMG_END] after the last, so the [IMG] count equals the post-merge feature count exactly. A dedicated VlmRuntimeRef::Pixtral path (selected when the vision module carries a pixtral_layout) runs the encoder and connector once per image, deriving the 2D-RoPE grid from each image's actual pixel shape and passing (patches_h, patches_w) to the connector's new grid-aware forward_with_grid. Running per image handles multi-image batches whose images have different sizes and cannot share one dense pixel tensor, and the Mistral3 PatchMerger now unfolds the real non-square grid instead of assuming a square. Validated on real checkpoints (pixtral-12b-4bit, mistral-small-3.1-24b-4bit) with a wide 1024x256 image and a 512x512 image: the [IMG] count matches the feature count exactly (Pixtral 1024 and 1024, Mistral3 370 and 361) and both models describe the shapes in the correct spatial order. 22 unit tests cover the resize math (downscale-only, round-up-to-factor, never upscale), the 3x5 row layout, and the placeholder-count equals feature-count invariant across aspect ratios. Refs #779 --- docs/supported-models.md | 2 +- src/commands/generate_vlm.rs | 9 + src/lib.rs | 4 +- src/loaded_model_capabilities.rs | 14 +- src/loaded_model_tests.rs | 1 + src/loading/vlm_fastvlm.rs | 1 + src/loading/vlm_gemma.rs | 1 + src/loading/vlm_llava.rs | 2 + src/loading/vlm_pixtral.rs | 115 ++++++++++++- src/loading/vlm_siglip.rs | 2 + src/loading/vlm_special.rs | 1 + src/multimodal/mod.rs | 1 + src/multimodal/pixtral_prompt.rs | 150 ++++++++++++++++ src/multimodal/pixtral_prompt_tests.rs | 140 +++++++++++++++ src/multimodal/vlm_runtime.rs | 93 +++++++++- src/vision/connectors/mistral3.rs | 41 ++++- src/vision/connectors/mod.rs | 18 ++ src/vision/mod.rs | 7 + src/vision/processors/mod.rs | 1 + src/vision/processors/pixtral.rs | 229 +++++++++++++++++++++++++ src/vision/processors/pixtral_tests.rs | 160 +++++++++++++++++ 21 files changed, 971 insertions(+), 21 deletions(-) create mode 100644 src/multimodal/pixtral_prompt.rs create mode 100644 src/multimodal/pixtral_prompt_tests.rs create mode 100644 src/vision/processors/pixtral.rs create mode 100644 src/vision/processors/pixtral_tests.rs diff --git a/docs/supported-models.md b/docs/supported-models.md index 895430806..e1940e450 100644 --- a/docs/supported-models.md +++ b/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 `` 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. diff --git a/src/commands/generate_vlm.rs b/src/commands/generate_vlm.rs index 605a0fb3e..7cb9278a3 100644 --- a/src/commands/generate_vlm.rs +++ b/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, diff --git a/src/lib.rs b/src/lib.rs index d66d3b917..dc2531c31 100644 --- a/src/lib.rs +++ b/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 diff --git a/src/loaded_model_capabilities.rs b/src/loaded_model_capabilities.rs index 9ae986875..cffc1d8fe 100644 --- a/src/loaded_model_capabilities.rs +++ b/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)), diff --git a/src/loaded_model_tests.rs b/src/loaded_model_tests.rs index dcc4f9cd9..d7bc2a4d5 100644 --- a/src/loaded_model_tests.rs +++ b/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, } } diff --git a/src/loading/vlm_fastvlm.rs b/src/loading/vlm_fastvlm.rs index a4f4e491d..cfc8f7913 100644 --- a/src/loading/vlm_fastvlm.rs +++ b/src/loading/vlm_fastvlm.rs @@ -200,6 +200,7 @@ pub(crate) fn load_fastvlm_vlm(model_path: &Path) -> Result { suffix_tokens: Vec::new(), block_prefix_tokens: Vec::new(), block_suffix_tokens: Vec::new(), + pixtral_layout: None, }; let vlm = vision::VisionLanguageModel { diff --git a/src/loading/vlm_gemma.rs b/src/loading/vlm_gemma.rs index 7f8846f07..a53081f45 100644 --- a/src/loading/vlm_gemma.rs +++ b/src/loading/vlm_gemma.rs @@ -190,6 +190,7 @@ pub(crate) fn load_gemma3_vlm(model_path: &Path) -> Result { // 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 { diff --git a/src/loading/vlm_llava.rs b/src/loading/vlm_llava.rs index ad513e622..22135b567 100644 --- a/src/loading/vlm_llava.rs +++ b/src/loading/vlm_llava.rs @@ -350,6 +350,7 @@ pub(crate) fn load_llava_vlm(model_path: &Path) -> Result { 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 { suffix_tokens: Vec::new(), block_prefix_tokens: Vec::new(), block_suffix_tokens: Vec::new(), + pixtral_layout: None, }; let vlm = vision::VisionLanguageModel { diff --git a/src/loading/vlm_pixtral.rs b/src/loading/vlm_pixtral.rs index 6b1c03288..28dcf0d39 100644 --- a/src/loading/vlm_pixtral.rs +++ b/src/loading/vlm_pixtral.rs @@ -24,16 +24,103 @@ 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 { + 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 = 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::(), + 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); + + 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 +286,6 @@ fn build_pixtral_family_context(model_path: &Path) -> Result Result { 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 +305,15 @@ pub(crate) fn load_pixtral_vlm(model_path: &Path) -> Result { ) .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 +326,7 @@ pub(crate) fn load_pixtral_vlm(model_path: &Path) -> Result { 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 +342,7 @@ pub(crate) fn load_pixtral_vlm(model_path: &Path) -> Result { suffix_tokens: Vec::new(), block_prefix_tokens: Vec::new(), block_suffix_tokens: Vec::new(), + pixtral_layout: Some(layout), }; let vlm = vision::VisionLanguageModel { @@ -262,7 +357,6 @@ pub(crate) fn load_pixtral_vlm(model_path: &Path) -> Result { pub(crate) fn load_mistral3_vlm(model_path: &Path) -> Result { 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 +386,13 @@ pub(crate) fn load_mistral3_vlm(model_path: &Path) -> Result { ) .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 +407,7 @@ pub(crate) fn load_mistral3_vlm(model_path: &Path) -> Result { 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 +423,7 @@ pub(crate) fn load_mistral3_vlm(model_path: &Path) -> Result { suffix_tokens: Vec::new(), block_prefix_tokens: Vec::new(), block_suffix_tokens: Vec::new(), + pixtral_layout: Some(layout), }; let vlm = vision::VisionLanguageModel { diff --git a/src/loading/vlm_siglip.rs b/src/loading/vlm_siglip.rs index 622053fec..45a9cd23c 100644 --- a/src/loading/vlm_siglip.rs +++ b/src/loading/vlm_siglip.rs @@ -259,6 +259,7 @@ pub(crate) fn load_aya_vision_vlm(model_path: &Path) -> Result { 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 { 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 { diff --git a/src/loading/vlm_special.rs b/src/loading/vlm_special.rs index adb46be93..08e9cb994 100644 --- a/src/loading/vlm_special.rs +++ b/src/loading/vlm_special.rs @@ -2291,6 +2291,7 @@ pub(crate) fn load_llama4_vlm(model_path: &Path) -> Result { suffix_tokens: Vec::new(), block_prefix_tokens: Vec::new(), block_suffix_tokens: Vec::new(), + pixtral_layout: None, }; let vlm = vision::VisionLanguageModel { diff --git a/src/multimodal/mod.rs b/src/multimodal/mod.rs index a3e4d8641..7c9cb7408 100644 --- a/src/multimodal/mod.rs +++ b/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; diff --git a/src/multimodal/pixtral_prompt.rs b/src/multimodal/pixtral_prompt.rs new file mode 100644 index 000000000..6851f1f9f --- /dev/null +++ b/src/multimodal/pixtral_prompt.rs @@ -0,0 +1,150 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Pixtral / Mistral3 row-structured image-token expansion. +//! +//! Port of the token layout in +//! https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/pixtral/processing_pixtral.py +//! (`PixtralProcessor.__call__`). Each `[IMG]` placeholder expands to a grid of +//! `[IMG]` tokens laid out row by row: `[IMG_BREAK]` after every row except the +//! last, `[IMG_END]` after the whole image. Only the `[IMG]` tokens are replaced +//! by vision features on merge; `[IMG_BREAK]` and `[IMG_END]` keep their learned +//! text embeddings and carry the spatial structure the model was trained on. +//! +//! The `[IMG]` count per image is exactly `tokens_h * tokens_w`, matching the +//! post-merge feature count from the processor and connector. That equality is +//! the core correctness invariant: a mismatch feeds the model garbage. + +/// Summary of a Pixtral/Mistral3 image-token expansion. +pub struct InsertedPixtralTokens { + /// Number of image blocks emitted (one per image). + pub image_blocks: usize, + /// Total `[IMG]` tokens across all images (== total vision features). + pub total_image_tokens: i32, +} + +/// Append one image's row-structured block: `cols` `[IMG]` per row for `rows` +/// rows, `[IMG_BREAK]` between rows, `[IMG_END]` at the end. +fn push_image_block( + out: &mut Vec, + rows: usize, + cols: usize, + image_token_id: i32, + image_break_token_id: i32, + image_end_token_id: i32, +) { + for row in 0..rows { + for _ in 0..cols { + out.push(image_token_id); + } + if row + 1 < rows { + out.push(image_break_token_id); + } + } + out.push(image_end_token_id); +} + +/// Total tokens (including break/end structure) one block occupies. +fn block_len(rows: usize, cols: usize) -> usize { + rows * cols + rows.saturating_sub(1) + 1 +} + +/// Expand each `[IMG]` placeholder in `prompt_tokens` into its row-structured +/// block using the per-image `(tokens_h, tokens_w)` grids (in order). +/// +/// Two placeholder shapes are supported, mirroring +/// [`crate::deepseekocr_prompt::insert_deepseekocr_image_tokens`]: +/// - one `[IMG]` per image (the chat template's rendered form): each is +/// expanded in place with its own grid; +/// - no `[IMG]` at all: one block per image is spliced right after the leading +/// BOS token. +/// +/// Returns `None` when there are no images, the prompt is empty, or the +/// placeholder count is neither `grids.len()` nor zero. +pub fn insert_pixtral_image_tokens( + prompt_tokens: &mut Vec, + grids: &[(usize, usize)], + image_token_id: i32, + image_break_token_id: i32, + image_end_token_id: i32, +) -> Option { + if prompt_tokens.is_empty() || grids.is_empty() { + return None; + } + + let placeholders = prompt_tokens + .iter() + .filter(|&&t| t == image_token_id) + .count(); + + let total_image_tokens: i32 = grids.iter().map(|&(h, w)| (h * w) as i32).sum(); + let block_budget: usize = grids.iter().map(|&(h, w)| block_len(h, w)).sum(); + + // One placeholder per image: expand each in place with its own grid. + if placeholders == grids.len() { + let mut expanded = Vec::with_capacity(prompt_tokens.len() + block_budget); + let mut idx = 0usize; + for &tok in prompt_tokens.iter() { + if tok == image_token_id { + let (rows, cols) = grids[idx]; + push_image_block( + &mut expanded, + rows, + cols, + image_token_id, + image_break_token_id, + image_end_token_id, + ); + idx += 1; + } else { + expanded.push(tok); + } + } + *prompt_tokens = expanded; + return Some(InsertedPixtralTokens { + image_blocks: grids.len(), + total_image_tokens, + }); + } + + // No placeholder: splice one block per image after the leading BOS token. + if placeholders == 0 { + let bos = prompt_tokens[0]; + let rest = prompt_tokens[1..].to_vec(); + let mut out = Vec::with_capacity(prompt_tokens.len() + block_budget); + out.push(bos); + for &(rows, cols) in grids { + push_image_block( + &mut out, + rows, + cols, + image_token_id, + image_break_token_id, + image_end_token_id, + ); + } + out.extend(rest); + *prompt_tokens = out; + return Some(InsertedPixtralTokens { + image_blocks: grids.len(), + total_image_tokens, + }); + } + + None +} + +#[cfg(test)] +#[path = "pixtral_prompt_tests.rs"] +mod tests; diff --git a/src/multimodal/pixtral_prompt_tests.rs b/src/multimodal/pixtral_prompt_tests.rs new file mode 100644 index 000000000..2d9758105 --- /dev/null +++ b/src/multimodal/pixtral_prompt_tests.rs @@ -0,0 +1,140 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::*; + +// Match the checkpoint ids: [IMG]=10, [IMG_BREAK]=12, [IMG_END]=13. +const IMG: i32 = 10; +const BRK: i32 = 12; +const END: i32 = 13; + +/// Build the expected row-structured block for a `rows x cols` grid. +fn expected_block(rows: usize, cols: usize) -> Vec { + let mut v = Vec::new(); + for row in 0..rows { + v.extend(std::iter::repeat_n(IMG, cols)); + if row + 1 < rows { + v.push(BRK); + } + } + v.push(END); + v +} + +#[test] +fn three_by_five_grid_layout_is_exact() { + // 3 rows x 5 cols: 5 IMG, BREAK, 5 IMG, BREAK, 5 IMG, END. + let block = expected_block(3, 5); + let expected = vec![ + IMG, IMG, IMG, IMG, IMG, BRK, // row 0 + break + IMG, IMG, IMG, IMG, IMG, BRK, // row 1 + break + IMG, IMG, IMG, IMG, IMG, END, // row 2 + end (no break after last row) + ]; + assert_eq!(block, expected); + assert_eq!(block.iter().filter(|&&t| t == IMG).count(), 15); + assert_eq!(block.iter().filter(|&&t| t == BRK).count(), 2); + assert_eq!(block.iter().filter(|&&t| t == END).count(), 1); +} + +#[test] +fn single_placeholder_expands_in_place() { + let mut tokens = vec![1, IMG, 2, 3]; + let stats = insert_pixtral_image_tokens(&mut tokens, &[(3, 5)], IMG, BRK, END).unwrap(); + + let mut expected = vec![1]; + expected.extend(expected_block(3, 5)); + expected.extend([2, 3]); + assert_eq!(tokens, expected); + assert_eq!(stats.image_blocks, 1); + assert_eq!(stats.total_image_tokens, 15); + // The [IMG] count in the final prompt equals the reported feature count. + assert_eq!( + tokens.iter().filter(|&&t| t == IMG).count() as i32, + stats.total_image_tokens + ); +} + +#[test] +fn no_placeholder_splices_after_bos() { + // BOS then plain text; block goes right after BOS. + let mut tokens = vec![1, 7, 8, 9]; + let stats = insert_pixtral_image_tokens(&mut tokens, &[(2, 2)], IMG, BRK, END).unwrap(); + + let mut expected = vec![1]; + expected.extend(expected_block(2, 2)); + expected.extend([7, 8, 9]); + assert_eq!(tokens, expected); + assert_eq!(stats.total_image_tokens, 4); +} + +#[test] +fn multiple_images_each_get_own_grid() { + let mut tokens = vec![1, IMG, 5, IMG, 6]; + let grids = [(2usize, 3usize), (1usize, 4usize)]; + let stats = insert_pixtral_image_tokens(&mut tokens, &grids, IMG, BRK, END).unwrap(); + + let mut expected = vec![1]; + expected.extend(expected_block(2, 3)); + expected.push(5); + expected.extend(expected_block(1, 4)); + expected.push(6); + assert_eq!(tokens, expected); + assert_eq!(stats.image_blocks, 2); + let expected_total: i32 = grids.iter().map(|&(h, w)| (h * w) as i32).sum(); + assert_eq!(stats.total_image_tokens, expected_total); + assert_eq!( + tokens.iter().filter(|&&t| t == IMG).count() as i32, + stats.total_image_tokens + ); +} + +#[test] +fn single_row_has_no_break() { + // A 1-row image emits only IMG*cols followed by END, no break token. + let block = expected_block(1, 4); + assert_eq!(block, vec![IMG, IMG, IMG, IMG, END]); + assert!(!block.contains(&BRK)); +} + +#[test] +fn placeholder_count_mismatch_returns_none() { + // Two placeholders but only one grid -> ambiguous, refuse. + let mut tokens = vec![1, IMG, IMG, 2]; + assert!(insert_pixtral_image_tokens(&mut tokens, &[(2, 2)], IMG, BRK, END).is_none()); +} + +#[test] +fn empty_inputs_return_none() { + let mut empty: Vec = Vec::new(); + assert!(insert_pixtral_image_tokens(&mut empty, &[(2, 2)], IMG, BRK, END).is_none()); + + let mut tokens = vec![1, IMG, 2]; + assert!(insert_pixtral_image_tokens(&mut tokens, &[], IMG, BRK, END).is_none()); +} + +#[test] +fn feature_count_matches_across_aspect_ratios() { + // The [IMG] count must equal tokens_h * tokens_w for a range of shapes. + for &(rows, cols) in &[(1usize, 1usize), (64, 16), (16, 64), (55, 28), (36, 22)] { + let mut tokens = vec![1, IMG, 2]; + let stats = + insert_pixtral_image_tokens(&mut tokens, &[(rows, cols)], IMG, BRK, END).unwrap(); + assert_eq!(stats.total_image_tokens, (rows * cols) as i32); + assert_eq!( + tokens.iter().filter(|&&t| t == IMG).count(), + rows * cols, + "grid {rows}x{cols}" + ); + } +} diff --git a/src/multimodal/vlm_runtime.rs b/src/multimodal/vlm_runtime.rs index 680efe344..20b12e32d 100644 --- a/src/multimodal/vlm_runtime.rs +++ b/src/multimodal/vlm_runtime.rs @@ -35,6 +35,7 @@ use crate::moondream3_prompt::{Moondream3PromptMode, prepare_moondream3_prompt_t use crate::phi3v_prompt::prepare_phi3v_prompt_tokens; use crate::phi4_siglip_prompt::prepare_phi4_siglip_prompt_tokens; use crate::phi4mm_prompt::prepare_phi4mm_prompt_tokens; +use crate::pixtral_prompt::insert_pixtral_image_tokens; use crate::qwen_vl::{insert_qwen_vl_image_tokens, insert_qwen3_omni_image_tokens}; use crate::smolvlm_prompt::insert_smolvlm_image_tokens; use crate::vision::KimiVLModel; @@ -44,7 +45,7 @@ use crate::vision::merge::InputEmbeddings; use crate::vision::processors::ImageProcessor; use crate::vlm_prompt::{ImageTokenBlockInfo, ImageTokenBlockStats, apply_image_token_blocks}; use crate::youtu_vl_prompt::insert_youtu_vl_image_tokens; -use crate::{LoadedModel, VlmRuntimeRef}; +use crate::{LanguageModel, LoadedModel, VlmRuntimeRef}; const MOLMO_V1_BOS_TOKEN_ID: i32 = 151643; @@ -199,6 +200,14 @@ pub enum VlmPreparationSummary { image_blocks: usize, total_image_tokens: i32, }, + /// Pixtral / Mistral3 expanded each `[IMG]` placeholder into a + /// row-structured `[IMG]* [IMG_BREAK] ... [IMG_END]` block sized to the + /// image's aspect ratio. `total_image_tokens` counts only `[IMG]` tokens + /// (== vision features). + Pixtral { + image_blocks: usize, + total_image_tokens: i32, + }, ImageBlocks(ImageTokenBlockStats), } @@ -1442,6 +1451,88 @@ where preparation, })) } + VlmRuntimeRef::Pixtral(vision_module) => { + // Dynamic aspect-ratio path for Pixtral / Mistral3. Each image is + // resized preserving aspect ratio, so images in a multi-image batch + // may have different sizes and cannot be packed into one dense pixel + // tensor. We therefore run the encoder + connector once per image and + // concatenate the projected features in prompt order; the encoder's + // 2D-RoPE grid is derived per call from the actual pixel shape, and + // the connector receives that image's `(patches_h, patches_w)` so the + // Mistral3 PatchMerger unfolds the correct non-square grid. See + // lablup/mlxcel#779. + let layout = vision_module.pixtral_layout.as_ref().ok_or_else(|| { + anyhow::anyhow!("Pixtral runtime selected but pixtral_layout is missing") + })?; + let _ = active_caches; + let _ = image_cache_keys; + + let per_image: Vec<_> = images + .iter() + .map(|img| layout.processor.preprocess_one(img)) + .collect(); + + // Row-structured token expansion: one `[IMG]`-block per image sized to + // its post-merge `(tokens_h, tokens_w)` grid. The `[IMG]` count per + // image equals `tokens_h * tokens_w`, matching the feature count. + let token_grids: Vec<(usize, usize)> = + per_image.iter().map(|p| (p.tokens_h, p.tokens_w)).collect(); + let preparation = insert_pixtral_image_tokens( + prompt_tokens, + &token_grids, + layout.image_token_id, + layout.image_break_token_id, + layout.image_end_token_id, + ) + .map(|stats| VlmPreparationSummary::Pixtral { + image_blocks: stats.image_blocks, + total_image_tokens: stats.total_image_tokens, + }); + + // Text embeddings must be built AFTER expansion so the `[IMG]` + // positions line up with the projected features on merge. + let input_ids_arr = prompt_ids_array(prompt_tokens); + let inputs_embeds = model.embed_tokens(&input_ids_arr).ok_or_else(|| { + anyhow::anyhow!("Text model must support embed_tokens for Pixtral VLM") + })?; + let embed_dtype = mlxcel_core::array_dtype(&inputs_embeds); + + let mut image_features: Option> = None; + for p in &per_image { + // Processor emits channels-first [1, C, H, W]; the encoder wants + // channels-last [1, H, W, C] in the text embedding dtype. + let pv = mlxcel_core::transpose_axes(&p.pixel_values, &[0, 2, 3, 1]); + let pv = mlxcel_core::astype(&pv, embed_dtype); + let encoded = vision_module.encoder.forward(&pv); + let projected = vision_module.connector.forward_with_grid( + &encoded.hidden_states, + p.patches_h as i32, + p.patches_w as i32, + ); + image_features = Some(match image_features { + None => projected, + Some(acc) => mlxcel_core::concatenate(&acc, &projected, 1), + }); + } + + let embeddings = match image_features { + Some(features) => crate::vision::merge::merge_llava( + layout.image_token_id, + &features, + &inputs_embeds, + &input_ids_arr, + ), + None => InputEmbeddings { + inputs_embeds, + attention_mask_4d: None, + }, + }; + + Ok(Some(PreparedVlmEmbeddings { + embeddings, + preparation, + })) + } VlmRuntimeRef::Standard(vision_module) => { let preparation = model .image_token_block_info() diff --git a/src/vision/connectors/mistral3.rs b/src/vision/connectors/mistral3.rs index b6030b3a2..229da6bb8 100644 --- a/src/vision/connectors/mistral3.rs +++ b/src/vision/connectors/mistral3.rs @@ -92,17 +92,25 @@ impl Mistral3Projector { /// Matches Python's unfold (im2col) channels-first ordering: /// [c0_s0, c0_s1, c0_s2, c0_s3, c1_s0, c1_s1, ...] /// - /// Input: [N, D] (squeezed from [1, N, D]) + /// `grid_h` / `grid_w` are the pre-merge patch rows/cols for this image + /// (`target_h / patch_size`, `target_w / patch_size`); for a fixed-square + /// input they both equal `patch_h`, for a dynamic aspect-ratio input they + /// differ. Mirrors + /// https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/mistral3/mistral3.py + /// (`Mistral3PatchMerger.__call__`), which reshapes each image's tokens to + /// its own `(h, w)` grid before the unfold. + /// + /// Input: [N, D] (squeezed from [1, N, D]), N == grid_h * grid_w /// 1. Reshape: [H*W, D] → [H, W, D] /// 2. Reshape: [H/2, 2, W/2, 2, D] /// 3. Transpose to channels-first: [H/2, W/2, D, 2, 2] /// 4. Reshape: [H/2 * W/2, D*4] /// 5. merging_layer: [H/2 * W/2, D] - fn patch_merge(&self, x: &MlxArray) -> UniquePtr { + fn patch_merge(&self, x: &MlxArray, grid_h: i32, grid_w: i32) -> UniquePtr { let shape = mlxcel_core::array_shape(x); let d = shape[shape.len() - 1]; - let h = self.patch_h; - let w = h; // square patches + let h = grid_h; + let w = grid_w; let s = self.spatial_merge_size; let h2 = h / s; let w2 = w / s; @@ -120,8 +128,9 @@ impl Mistral3Projector { } } -impl MultiModalConnector for Mistral3Projector { - fn forward(&self, vision_features: &MlxArray) -> UniquePtr { +impl Mistral3Projector { + /// RMSNorm → PatchMerger(grid_h, grid_w) → Linear → GELU → Linear. + fn project(&self, vision_features: &MlxArray, grid_h: i32, grid_w: i32) -> UniquePtr { // Squeeze batch dim: [1, N, D] → [N, D] let shape = mlxcel_core::array_shape(vision_features); let x = if shape.len() == 3 && shape[0] == 1 { @@ -132,8 +141,8 @@ impl MultiModalConnector for Mistral3Projector { // RMSNorm let x = self.norm.forward(&x); - // PatchMerger - let x = self.patch_merge(&x); + // PatchMerger over this image's actual patch grid + let x = self.patch_merge(&x, grid_h, grid_w); // MLP: Linear → GELU → Linear let x = self.linear_1.forward(&x); let x = mlxcel_core::gelu(&x); @@ -143,3 +152,19 @@ impl MultiModalConnector for Mistral3Projector { mlxcel_core::expand_dims(&x, 0) } } + +impl MultiModalConnector for Mistral3Projector { + fn forward(&self, vision_features: &MlxArray) -> UniquePtr { + // Fixed-square fallback: assume the configured square patch grid. + self.project(vision_features, self.patch_h, self.patch_h) + } + + fn forward_with_grid( + &self, + vision_features: &MlxArray, + grid_h: i32, + grid_w: i32, + ) -> UniquePtr { + self.project(vision_features, grid_h, grid_w) + } +} diff --git a/src/vision/connectors/mod.rs b/src/vision/connectors/mod.rs index 56c4ad285..c496c6d01 100644 --- a/src/vision/connectors/mod.rs +++ b/src/vision/connectors/mod.rs @@ -34,4 +34,22 @@ use mlxcel_core::{MlxArray, UniquePtr}; /// Trait for multi-modal connectors that project vision features to text space pub trait MultiModalConnector { fn forward(&self, vision_features: &MlxArray) -> UniquePtr; + + /// Project vision features when the caller knows the pre-merge patch grid + /// `(grid_h, grid_w)` for a single image. + /// + /// The default ignores the grid and calls [`Self::forward`]; connectors + /// whose spatial merge assumes a square grid (Mistral3's PatchMerger) + /// override this so dynamic aspect-ratio inputs reshape correctly. + /// + /// Used by: Pixtral / Mistral3 dynamic aspect-ratio runtime path. + fn forward_with_grid( + &self, + vision_features: &MlxArray, + grid_h: i32, + grid_w: i32, + ) -> UniquePtr { + let _ = (grid_h, grid_w); + self.forward(vision_features) + } } diff --git a/src/vision/mod.rs b/src/vision/mod.rs index a6865fc2b..025c4b0a1 100644 --- a/src/vision/mod.rs +++ b/src/vision/mod.rs @@ -192,6 +192,13 @@ pub struct VisionModule { /// Tokens to insert after each image block during expansion. /// Gemma3 VLM: `[108]` (\n\n token) to match Python processor behavior. pub block_suffix_tokens: Vec, + /// Dynamic aspect-ratio row-structured layout for Pixtral / Mistral3. + /// + /// When `Some`, the model is routed to the Pixtral runtime path + /// (`VlmRuntimeRef::Pixtral`) which preserves each image's aspect ratio and + /// emits `[IMG] / [IMG_BREAK] / [IMG_END]` row structure instead of the + /// fixed-square LLaVA merge. `None` for every other `VisionModule` family. + pub pixtral_layout: Option, } impl VisionModule { diff --git a/src/vision/processors/mod.rs b/src/vision/processors/mod.rs index 38f7c4ab3..410e22467 100644 --- a/src/vision/processors/mod.rs +++ b/src/vision/processors/mod.rs @@ -40,6 +40,7 @@ pub mod paddleocr_vl; pub mod phi3_v; pub mod phi4_siglip; pub mod phi4mm; +pub mod pixtral; pub mod qwen2_vl; pub mod siglip; pub mod smolvlm; diff --git a/src/vision/processors/pixtral.rs b/src/vision/processors/pixtral.rs new file mode 100644 index 000000000..014d6c411 --- /dev/null +++ b/src/vision/processors/pixtral.rs @@ -0,0 +1,229 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Pixtral / Mistral3 dynamic aspect-ratio image preprocessor. +//! +//! Port of the resize and normalization rules from +//! https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/pixtral/image_processing_pixtral.py +//! (`get_resize_output_image_size`, `PixtralImageProcessor`) and the row-token +//! layout in +//! https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/pixtral/processing_pixtral.py +//! (`PixtralProcessor.__call__`). +//! +//! Unlike the fixed-square SigLIP path these models were trained on, Pixtral and +//! Mistral3 preserve aspect ratio: each image is downscaled so its longest side +//! fits `longest_edge`, then rounded up to a whole number of merged patches. The +//! per-image patch grid drives the encoder's 2D-RoPE position grid and the +//! row-structured `[IMG] / [IMG_BREAK] / [IMG_END]` token expansion, so all three +//! must agree on the same `(rows, cols)`. +//! +//! Round-up factor decision: the upstream resize helper rounds to `patch_size`, +//! but `PixtralProcessor.__call__` calls the image processor with +//! `patch_size = patch_size * spatial_merge_size`, and the token layout divides +//! by that same product (`num_h = h // (patch_size * spatial_merge_size)`). For +//! Pixtral (`spatial_merge_size = 1`) the product equals `patch_size`; for +//! Mistral3 (`spatial_merge_size = 2`) it is `patch_size * 2`. Rounding to +//! `patch_size` alone would let a partial merged row survive the floor division +//! and desync the token count from the emitted features, so this processor +//! rounds each side up to `patch_size * spatial_merge_size` for both families. +//! +//! Used by: Pixtral VLM, Mistral3 VLM + +use super::ImageProcessor; +use image::imageops::FilterType; +use mlxcel_core::{MlxArray, UniquePtr}; + +// The reference `image_mean` / `image_std` carry more decimal digits than f32 +// represents; keep them verbatim from the checkpoint's preprocessor_config so +// the values are diffable against the upstream processor. +/// CLIP normalization mean, shared by Pixtral and Mistral3. +#[allow(clippy::excessive_precision)] +pub const PIXTRAL_IMAGE_MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73]; +/// CLIP normalization std, shared by Pixtral and Mistral3. +#[allow(clippy::excessive_precision)] +pub const PIXTRAL_IMAGE_STD: [f32; 3] = [0.268_629_54, 0.261_302_58, 0.275_777_11]; + +/// One preprocessed image plus the patch/token geometry the encoder, connector +/// and token expansion all key off. +pub struct PixtralPreprocessed { + /// Pixel tensor in channels-first `[1, C, H, W]` layout (f32). + pub pixel_values: UniquePtr, + /// Pre-merge patch rows (`target_h / patch_size`). Drives the encoder grid. + pub patches_h: usize, + /// Pre-merge patch cols (`target_w / patch_size`). + pub patches_w: usize, + /// Post-merge token rows (`target_h / (patch_size * spatial_merge_size)`). + pub tokens_h: usize, + /// Post-merge token cols (`target_w / (patch_size * spatial_merge_size)`). + pub tokens_w: usize, +} + +/// Pixtral / Mistral3 dynamic aspect-ratio processor. +#[derive(Debug, Clone)] +pub struct PixtralProcessor { + pub patch_size: usize, + pub spatial_merge_size: usize, + pub longest_edge: usize, + pub mean: [f32; 3], + pub std: [f32; 3], +} + +impl PixtralProcessor { + pub fn new(patch_size: usize, spatial_merge_size: usize, longest_edge: usize) -> Self { + Self { + patch_size, + spatial_merge_size: spatial_merge_size.max(1), + longest_edge, + mean: PIXTRAL_IMAGE_MEAN, + std: PIXTRAL_IMAGE_STD, + } + } + + /// The whole-patch round-up factor (`patch_size * spatial_merge_size`). + #[inline] + pub fn factor(&self) -> usize { + self.patch_size * self.spatial_merge_size + } + + /// Resized `(target_h, target_w)` in pixels for an original `(h, w)`. + /// + /// Downscale-only: the longest side is scaled to fit `longest_edge`, never + /// upscaled, then each side is rounded up to a whole multiple of + /// [`Self::factor`]. Mirrors `get_resize_output_image_size` with + /// `patch_size = patch_size * spatial_merge_size`. + pub fn target_size(&self, h: usize, w: usize) -> (usize, usize) { + let factor = self.factor(); + let max_edge = self.longest_edge as f64; + let ratio = (h as f64 / max_edge).max(w as f64 / max_edge); + + let (mut hh, mut ww) = (h, w); + if ratio > 1.0 { + hh = (h as f64 / ratio).floor() as usize; + ww = (w as f64 / ratio).floor() as usize; + } + + // (dim - 1) // factor + 1 rounds up to whole merged patches; guard the + // 1px case so the subtraction cannot underflow. + let num_h = hh.saturating_sub(1) / factor + 1; + let num_w = ww.saturating_sub(1) / factor + 1; + (num_h * factor, num_w * factor) + } + + /// Pre-merge patch grid `(patches_h, patches_w)` for a resized image. This + /// is what the encoder derives from the pixel shape and what the Mistral3 + /// PatchMerger unfolds. + pub fn patch_grid(&self, target_h: usize, target_w: usize) -> (usize, usize) { + (target_h / self.patch_size, target_w / self.patch_size) + } + + /// Post-merge token grid `(tokens_h, tokens_w)`; one `[IMG]` token per cell. + pub fn token_grid(&self, target_h: usize, target_w: usize) -> (usize, usize) { + let factor = self.factor(); + (target_h / factor, target_w / factor) + } + + /// Resize (bicubic), rescale by 1/255 and CLIP-normalize a single image, + /// returning the pixel tensor and its patch/token geometry. + pub fn preprocess_one(&self, img: &image::DynamicImage) -> PixtralPreprocessed { + let (orig_h, orig_w) = (img.height() as usize, img.width() as usize); + let (target_h, target_w) = self.target_size(orig_h, orig_w); + let (patches_h, patches_w) = self.patch_grid(target_h, target_w); + let (tokens_h, tokens_w) = self.token_grid(target_h, target_w); + + let pixel_values = self.pixels_for_size(img, target_h, target_w); + + PixtralPreprocessed { + pixel_values, + patches_h, + patches_w, + tokens_h, + tokens_w, + } + } + + /// Resize/normalize `img` to an exact `(target_h, target_w)` and emit a + /// `[1, C, H, W]` f32 tensor. The reference resamples with PIL BICUBIC + /// (`resample = 3`); `CatmullRom` is the `image` crate's cubic filter. + fn pixels_for_size( + &self, + img: &image::DynamicImage, + target_h: usize, + target_w: usize, + ) -> UniquePtr { + let channels = 3usize; + let resized = img.resize_exact(target_w as u32, target_h as u32, FilterType::CatmullRom); + let rgb = resized.to_rgb8(); + + let mut data = vec![0.0f32; channels * target_h * target_w]; + for y in 0..target_h { + for x in 0..target_w { + let pixel = rgb.get_pixel(x as u32, y as u32); + for c in 0..channels { + let val = pixel[c] as f32 / 255.0; + let normalized = (val - self.mean[c]) / self.std[c]; + // [C, H, W] with a leading batch of 1. + let idx = c * target_h * target_w + y * target_w + x; + data[idx] = normalized; + } + } + } + + mlxcel_core::from_slice_f32( + &data, + &[1, channels as i32, target_h as i32, target_w as i32], + ) + } +} + +impl ImageProcessor for PixtralProcessor { + /// Trait entry point. Dynamic aspect-ratio inference runs through the + /// dedicated Pixtral runtime path ([`crate::vlm_runtime`]) which calls + /// [`Self::preprocess_one`] per image and packs features individually, so + /// this batched stack is only exercised for the common single-image case + /// (and same-size batches). Each image is still resized to its own + /// aspect-preserving size; a mixed-size batch cannot be stacked into one + /// dense tensor and must use the per-image runtime path. + fn preprocess(&self, images: &[image::DynamicImage]) -> UniquePtr { + let first = self.preprocess_one(&images[0]); + if images.len() == 1 { + return first.pixel_values; + } + + let mut stacked = mlxcel_core::copy(&first.pixel_values); + for img in &images[1..] { + let next = self.preprocess_one(img); + stacked = mlxcel_core::concatenate(&stacked, &next.pixel_values, 0); + } + stacked + } +} + +/// Row-structured layout descriptor carried on a Pixtral/Mistral3 +/// [`crate::vision::VisionModule`]. Its presence routes the model to the +/// dynamic aspect-ratio runtime path instead of the fixed-square LLaVA merge. +#[derive(Debug, Clone)] +pub struct PixtralLayout { + /// Processor holding the resize/normalize policy and patch geometry. + pub processor: PixtralProcessor, + /// `[IMG]` placeholder id (replaced by vision features on merge). + pub image_token_id: i32, + /// `[IMG_BREAK]` id emitted between patch rows (kept as a text embedding). + pub image_break_token_id: i32, + /// `[IMG_END]` id emitted after the last patch row. + pub image_end_token_id: i32, +} + +#[cfg(test)] +#[path = "pixtral_tests.rs"] +mod tests; diff --git a/src/vision/processors/pixtral_tests.rs b/src/vision/processors/pixtral_tests.rs new file mode 100644 index 000000000..0f9a55ba4 --- /dev/null +++ b/src/vision/processors/pixtral_tests.rs @@ -0,0 +1,160 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::*; +use crate::vision::processors::ImageProcessor; + +fn pixtral() -> PixtralProcessor { + // pixtral-12b-4bit: patch 16, spatial_merge 1, longest_edge 1024. + PixtralProcessor::new(16, 1, 1024) +} + +fn mistral3() -> PixtralProcessor { + // mistral-small-3.1-24b-4bit: patch 14, spatial_merge 2, longest_edge 1540. + PixtralProcessor::new(14, 2, 1540) +} + +fn solid_image(w: u32, h: u32) -> image::DynamicImage { + image::DynamicImage::ImageRgb8(image::RgbImage::from_pixel( + w, + h, + image::Rgb([128, 64, 200]), + )) +} + +#[test] +fn factor_is_patch_times_merge() { + assert_eq!(pixtral().factor(), 16); + assert_eq!(mistral3().factor(), 28); +} + +#[test] +fn small_image_is_never_upscaled() { + // 32x48 is far below longest_edge; ratio <= 1 so no scaling, only the + // round-up to whole patches applies. + let (th, tw) = pixtral().target_size(32, 48); + assert_eq!((th, tw), (32, 48)); + + // A tiny Mistral3 image rounds up to the 28px factor, never to longest_edge. + let (th, tw) = mistral3().target_size(20, 10); + assert_eq!((th, tw), (28, 28)); +} + +#[test] +fn square_within_edge_keeps_aspect() { + let (th, tw) = pixtral().target_size(512, 512); + assert_eq!((th, tw), (512, 512)); + // Regression parity: a square image still produces a square patch grid, so + // the only behavioral change vs. the old fixed-square path is the added + // row-structure tokens, not a different feature count. + assert_eq!(pixtral().token_grid(th, tw), (32, 32)); +} + +#[test] +fn wide_non_square_preserves_ratio() { + // 1024x256 fits the longest edge exactly; the wide aspect ratio survives. + let (th, tw) = pixtral().target_size(1024, 256); + assert_eq!((th, tw), (1024, 256)); + assert_eq!(pixtral().token_grid(th, tw), (64, 16)); + assert_eq!(pixtral().patch_grid(th, tw), (64, 16)); +} + +#[test] +fn oversized_image_is_downscaled_by_longest_edge() { + // 2048x1024 -> ratio 2.0 -> 1024x512, both multiples of 16. + let (th, tw) = pixtral().target_size(2048, 1024); + assert_eq!((th, tw), (1024, 512)); + assert_eq!(pixtral().token_grid(th, tw), (64, 32)); +} + +#[test] +fn dimensions_round_up_to_factor_and_stay_within_edge() { + let procs = [pixtral(), mistral3()]; + let sizes = [ + (1000usize, 500usize), + (777, 333), + (1600, 900), + (3000, 100), + (1, 1), + (1540, 1541), + ]; + for p in &procs { + let factor = p.factor(); + for &(h, w) in &sizes { + let (th, tw) = p.target_size(h, w); + // Whole multiples of the merged-patch factor. + assert_eq!(th % factor, 0, "target_h {th} not multiple of {factor}"); + assert_eq!(tw % factor, 0, "target_w {tw} not multiple of {factor}"); + // Never exceeds the longest edge (both configs have longest_edge + // itself a multiple of the factor). + assert!(th <= p.longest_edge, "target_h {th} > {}", p.longest_edge); + assert!(tw <= p.longest_edge, "target_w {tw} > {}", p.longest_edge); + // Pre-merge patch grid divides evenly by the merge size so no row is + // dropped by the connector's floor division. + let (ph, pw) = p.patch_grid(th, tw); + assert_eq!(ph % p.spatial_merge_size, 0); + assert_eq!(pw % p.spatial_merge_size, 0); + // token grid == patch grid / merge on both axes. + let (tkh, tkw) = p.token_grid(th, tw); + assert_eq!(tkh, ph / p.spatial_merge_size); + assert_eq!(tkw, pw / p.spatial_merge_size); + } + } +} + +#[test] +fn mistral3_partial_merged_row_is_not_dropped() { + // A height that is a whole number of patches but NOT a whole number of + // merged patches would lose a row if we rounded to patch_size alone. The + // factor round-up (28) prevents that: patch grid stays even. + let m = mistral3(); + let (th, tw) = m.target_size(15 * 14, 3 * 14); // 210 x 42 + let (ph, pw) = m.patch_grid(th, tw); + assert_eq!(ph % 2, 0); + assert_eq!(pw % 2, 0); + let (tkh, tkw) = m.token_grid(th, tw); + assert_eq!(tkh, ph / 2); + assert_eq!(tkw, pw / 2); +} + +#[test] +fn clip_normalization_constants() { + let p = pixtral(); + assert_eq!(p.mean, PIXTRAL_IMAGE_MEAN); + assert_eq!(p.std, PIXTRAL_IMAGE_STD); +} + +#[test] +fn preprocess_one_reports_matching_geometry_and_shape() { + let p = mistral3(); + let img = solid_image(1000, 600); // width x height + let pre = p.preprocess_one(&img); + + let (th, tw) = p.target_size(600, 1000); // (h, w) + assert_eq!(p.patch_grid(th, tw), (pre.patches_h, pre.patches_w)); + assert_eq!(p.token_grid(th, tw), (pre.tokens_h, pre.tokens_w)); + + let shape = mlxcel_core::array_shape(&pre.pixel_values); + assert_eq!(shape, vec![1, 3, th as i32, tw as i32]); +} + +#[test] +fn trait_preprocess_single_image_matches_preprocess_one() { + let p = pixtral(); + let img = solid_image(300, 700); + let via_trait = ImageProcessor::preprocess(&p, std::slice::from_ref(&img)); + let shape = mlxcel_core::array_shape(&via_trait); + let (th, tw) = p.target_size(700, 300); + assert_eq!(shape, vec![1, 3, th as i32, tw as i32]); +} From 57581ee14641febe9767637b12bfa3232ded7fcb Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Mon, 13 Jul 2026 20:59:58 +0900 Subject: [PATCH 2/2] fix(pixtral): reject an image-token desync and clamp the resize to the RoPE table The Pixtral runtime ignored the None that insert_pixtral_image_tokens returns on a placeholder/image-count mismatch, leaving the prompt unexpanded and then running the encoder and merge_llava anyway. A caller can trigger the mismatch by placing the literal [IMG] marker in the prompt text, which desyncs the [IMG] positions from the projected features: a wrong result in the common case, and an out-of-range scatter for a tiny single-feature image. Turn the None into a request error before any encoder or merge runs, and assert the rejected expansion leaves the prompt unmutated. Clamp the processor's effective longest_edge to image_size. The encoder's 2D-RoPE table is sized image_size / patch_size per side, and the new aspect-ratio processor reads longest_edge from a separate config field; a checkpoint with longest_edge > image_size could otherwise drive the patch grid past the table and wrap position ids or panic. The shipped checkpoints set the two equal, so the clamp is a no-op for them. --- src/loading/vlm_pixtral.rs | 8 ++++++++ src/multimodal/pixtral_prompt_tests.rs | 10 +++++++++- src/multimodal/vlm_runtime.rs | 25 +++++++++++++++++++++++-- 3 files changed, 40 insertions(+), 3 deletions(-) diff --git a/src/loading/vlm_pixtral.rs b/src/loading/vlm_pixtral.rs index 28dcf0d39..4160b8f6a 100644 --- a/src/loading/vlm_pixtral.rs +++ b/src/loading/vlm_pixtral.rs @@ -108,6 +108,14 @@ fn build_pixtral_layout( .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) = diff --git a/src/multimodal/pixtral_prompt_tests.rs b/src/multimodal/pixtral_prompt_tests.rs index 2d9758105..e5485a8de 100644 --- a/src/multimodal/pixtral_prompt_tests.rs +++ b/src/multimodal/pixtral_prompt_tests.rs @@ -109,9 +109,17 @@ fn single_row_has_no_break() { #[test] fn placeholder_count_mismatch_returns_none() { - // Two placeholders but only one grid -> ambiguous, refuse. + // Two placeholders but only one grid -> ambiguous, refuse. This is the + // path a caller reaches by planting a literal [IMG] marker in the prompt; + // the runtime turns this None into a request error before any encoder or + // merge runs, so the prompt must be left UNEXPANDED (not partially mutated). let mut tokens = vec![1, IMG, IMG, 2]; + let before = tokens.clone(); assert!(insert_pixtral_image_tokens(&mut tokens, &[(2, 2)], IMG, BRK, END).is_none()); + assert_eq!( + tokens, before, + "a rejected expansion must not mutate the prompt" + ); } #[test] diff --git a/src/multimodal/vlm_runtime.rs b/src/multimodal/vlm_runtime.rs index 20b12e32d..d83c016bd 100644 --- a/src/multimodal/vlm_runtime.rs +++ b/src/multimodal/vlm_runtime.rs @@ -1477,14 +1477,35 @@ where // image equals `tokens_h * tokens_w`, matching the feature count. let token_grids: Vec<(usize, usize)> = per_image.iter().map(|p| (p.tokens_h, p.tokens_w)).collect(); - let preparation = insert_pixtral_image_tokens( + let placeholder_count = prompt_tokens + .iter() + .filter(|&&t| t == layout.image_token_id) + .count(); + // Expand before any encoder work. A `None` here means the `[IMG]` + // placeholder count matched neither the image count nor zero, which + // desyncs the prompt from the features. Proceeding would feed + // `merge_llava` an unexpanded prompt whose `[IMG]` positions no + // longer line up with the projected features, producing a wrong + // result and, for a tiny image whose grid is a single feature, an + // out-of-range scatter. Fail the request cleanly instead. A caller + // can provoke this by placing the literal `[IMG]` marker in the + // prompt text, so it is a request error, not an internal invariant. + let stats = insert_pixtral_image_tokens( prompt_tokens, &token_grids, layout.image_token_id, layout.image_break_token_id, layout.image_end_token_id, ) - .map(|stats| VlmPreparationSummary::Pixtral { + .ok_or_else(|| { + anyhow::anyhow!( + "Pixtral image-token expansion failed: the number of [IMG] placeholders \ + ({placeholder_count}) matches neither the image count ({}) nor zero. \ + Remove any literal [IMG] markers from the prompt text.", + token_grids.len() + ) + })?; + let preparation = Some(VlmPreparationSummary::Pixtral { image_blocks: stats.image_blocks, total_image_tokens: stats.total_image_tokens, });