diff --git a/crates/multimodal/src/registry/kimi_k25.rs b/crates/multimodal/src/registry/kimi_k25.rs index 327784d55..994e5cf06 100644 --- a/crates/multimodal/src/registry/kimi_k25.rs +++ b/crates/multimodal/src/registry/kimi_k25.rs @@ -28,9 +28,11 @@ impl ModelProcessorSpec for KimiK25VisionSpec { } fn matches(&self, metadata: &ModelMetadata) -> bool { - // Kimi-K3 reuses K2.5's MoonViT vision stack and `<|media_pad|>` - // placeholder (media_placeholder_token_id 163605), so it shares this - // spec. + // Kimi-K3 uses the same `<|media_pad|>` placeholder + // (media_placeholder_token_id 163605), patchification layout and + // prompt-replacement shape, so it shares this spec. Note that the two + // do *not* share a pixel pipeline — see `vision::processors::kimi_k3` + // for the patch budget and alpha-compositing differences. let id = metadata.model_id.to_ascii_lowercase(); id.contains("kimi") && (id.contains("k2") || id.contains("k3")) || metadata diff --git a/crates/multimodal/src/vision/preprocessor_config.rs b/crates/multimodal/src/vision/preprocessor_config.rs index 9485e6072..25b7d624a 100644 --- a/crates/multimodal/src/vision/preprocessor_config.rs +++ b/crates/multimodal/src/vision/preprocessor_config.rs @@ -278,9 +278,14 @@ impl PreProcessorConfig { .and_then(|v| v.as_u64()) .map(|v| v as usize); } - // Also extract Kimi-specific limits into the extra map - // so processors can read them via get_extra() - for key in ["in_patch_limit", "patch_limit_on_one_side"] { + // Also extract Kimi-specific limits and the K3 transparency settings + // into the extra map so processors can read them via get_extra() + for key in [ + "in_patch_limit", + "patch_limit_on_one_side", + "transparent_bg_config", + "transparent_bg_fill_stage", + ] { if !config.extra.contains_key(key) { if let Some(v) = media_cfg.get(key) { config.extra.insert(key.to_string(), v.clone()); @@ -448,6 +453,9 @@ impl PreProcessorConfig { #[cfg(test)] mod tests { use super::*; + use crate::vision::transforms::{ + TransparentBgConfig, TransparentBgFillStage, TransparentBgPattern, + }; #[test] fn test_parse_clip_config() { @@ -602,4 +610,47 @@ mod tests { assert_eq!(config.get_patch_size(0), 14); assert_eq!(config.merge_size, Some(2)); } + + #[test] + fn test_parse_kimi_k3_transparency_settings() { + // Verbatim excerpt from moonshotai/Kimi-K3's preprocessor_config.json. + let json = r#"{ + "media_proc_cfg": { + "in_patch_limit": 65536, + "patch_size": 14, + "merge_kernel_size": 2, + "patch_limit_on_one_side": 512, + "transparent_bg_config": { + "pattern": "chessboard", + "chessboard_square_size": 8, + "chessboard_square_on_top_left": true, + "chessboard_white_value": 255, + "chessboard_gray_value": 180 + }, + "transparent_bg_fill_stage": "after_resize" + } + }"#; + + let config = PreProcessorConfig::from_json(json).unwrap(); + + assert_eq!(config.get_extra::("in_patch_limit"), Some(65536)); + assert_eq!( + config.get_extra::("patch_limit_on_one_side"), + Some(512) + ); + + let bg = config + .get_extra::("transparent_bg_config") + .expect("transparent_bg_config lifted out of media_proc_cfg"); + assert_eq!(bg.pattern, TransparentBgPattern::Chessboard); + assert_eq!(bg.chessboard_square_size, 8); + assert!(bg.chessboard_square_on_top_left); + assert_eq!(bg.chessboard_white_value, 255); + assert_eq!(bg.chessboard_gray_value, 180); + + assert_eq!( + config.get_extra::("transparent_bg_fill_stage"), + Some(TransparentBgFillStage::AfterResize) + ); + } } diff --git a/crates/multimodal/src/vision/processor.rs b/crates/multimodal/src/vision/processor.rs index 1569bc244..0ae8f0533 100644 --- a/crates/multimodal/src/vision/processor.rs +++ b/crates/multimodal/src/vision/processor.rs @@ -325,7 +325,9 @@ impl VisionProcessorRegistry { Box::new(super::processors::Llama4VisionProcessor::new()), ); - // Register Kimi-K2.5 Vision (also used by Kimi-K3, same MoonViT stack) + // Register the Kimi MoonViT family. K3 shares the stack but ships a + // larger patch budget and chessboard alpha compositing, so it gets its + // own processor rather than reusing K2.5's. registry.register( "kimi-k2", Box::new(super::processors::KimiK25Processor::new()), @@ -336,11 +338,11 @@ impl VisionProcessorRegistry { ); registry.register( "kimi-k3", - Box::new(super::processors::KimiK25Processor::new()), + Box::new(super::processors::KimiK3Processor::new()), ); registry.register( "kimi_k3", - Box::new(super::processors::KimiK25Processor::new()), + Box::new(super::processors::KimiK3Processor::new()), ); registry @@ -394,6 +396,38 @@ mod tests { assert_eq!(processor.model_name(), "inkling"); } + #[test] + fn test_registry_separates_kimi_k25_and_k3() { + // K3's pixel pipeline differs from K2.5's (larger patch budget, alpha + // compositing), so resolving a K3 id to the K2.5 processor is a bug. + let registry = VisionProcessorRegistry::with_defaults(); + + for id in ["moonshotai/Kimi-K3", "moonshotai/Kimi_K3-Instruct"] { + let processor = registry.find(id, None).expect("K3 vision processor"); + assert_eq!(processor.model_name(), "kimi-k3", "{id}"); + } + for id in ["moonshotai/Kimi-K2.5", "moonshotai/Kimi_K2-VL"] { + let processor = registry.find(id, None).expect("K2.5 vision processor"); + assert_eq!(processor.model_name(), "kimi-k2.5", "{id}"); + } + + // model_type fallback must split the same way. + assert_eq!( + registry + .find("internal/checkpoint-final", Some("kimi_k3")) + .expect("K3 by model_type") + .model_name(), + "kimi-k3" + ); + assert_eq!( + registry + .find("internal/checkpoint-final", Some("kimi_k2")) + .expect("K2.5 by model_type") + .model_name(), + "kimi-k2.5" + ); + } + #[test] fn test_registry_find() { let mut registry = VisionProcessorRegistry::new(); diff --git a/crates/multimodal/src/vision/processors/kimi_k25.rs b/crates/multimodal/src/vision/processors/kimi_k25.rs index bcdc58ad4..110e6b24d 100644 --- a/crates/multimodal/src/vision/processors/kimi_k25.rs +++ b/crates/multimodal/src/vision/processors/kimi_k25.rs @@ -1,25 +1,17 @@ //! Kimi-K2.5 (MoonViT) image processor. //! -//! Matches the HuggingFace `KimiK25VisionProcessor` preprocessing pipeline: -//! -//! 1. Compute scale to fit within patch limits (never upscale) -//! 2. Resize with BICUBIC interpolation -//! 3. Zero-pad to make dimensions divisible by factor (patch_size * merge_size) -//! 4. Normalize with [0.5, 0.5, 0.5] mean/std -//! 5. Extract patches as [N, C, patch_size, patch_size] -//! -//! Kimi resizes then zero-pads to make dimensions divisible by the alignment -//! factor (patch_size * merge_size). The model was trained with zero-padded -//! images, so using direct resize-to-aligned would degrade image quality. +//! Matches the HuggingFace `KimiK25VisionProcessor` preprocessing pipeline; the +//! pipeline itself lives in [`super::moonvit`], which Kimi-K3 shares. K2.5 +//! ships no `transparent_bg_config`, so alpha is dropped rather than +//! composited — the reference's `image.convert("RGB")` behavior. -use image::{DynamicImage, GenericImageView}; -use ndarray::Array3; +use image::DynamicImage; +use super::moonvit::{self, MoonVitParams}; use crate::vision::{ preprocessor_config::PreProcessorConfig, - processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, - scratch, - transforms::{self, TransformError}, + processor::{PreprocessedEncoderInputs, VisionPreProcessor}, + transforms::TransformError, }; pub const KIMI_K25_MEAN: [f64; 3] = [0.5, 0.5, 0.5]; @@ -32,21 +24,9 @@ pub const DEFAULT_IN_PATCH_LIMIT: usize = 16384; /// Maximum patches along one spatial dimension pub const DEFAULT_PATCH_LIMIT_ON_ONE_SIDE: usize = 512; -/// Kimi-K2.5 resize configuration for a single image. -struct ResizeConfig { - new_width: usize, - new_height: usize, - pad_width: usize, - pad_height: usize, - num_tokens: usize, -} - #[derive(Debug, Clone)] pub struct KimiK25Processor { - patch_size: usize, - merge_size: usize, - in_patch_limit: usize, - patch_limit_on_one_side: usize, + params: MoonVitParams, } impl Default for KimiK25Processor { @@ -58,177 +38,27 @@ impl Default for KimiK25Processor { impl KimiK25Processor { pub fn new() -> Self { Self { - patch_size: DEFAULT_PATCH_SIZE, - merge_size: DEFAULT_MERGE_SIZE, - in_patch_limit: DEFAULT_IN_PATCH_LIMIT, - patch_limit_on_one_side: DEFAULT_PATCH_LIMIT_ON_ONE_SIDE, + params: MoonVitParams { + patch_size: DEFAULT_PATCH_SIZE, + merge_size: DEFAULT_MERGE_SIZE, + in_patch_limit: DEFAULT_IN_PATCH_LIMIT, + patch_limit_on_one_side: DEFAULT_PATCH_LIMIT_ON_ONE_SIDE, + }, } } pub fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { Self { - patch_size: config.get_patch_size(DEFAULT_PATCH_SIZE), - merge_size: config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), - in_patch_limit: config - .get_extra::("in_patch_limit") - .unwrap_or(DEFAULT_IN_PATCH_LIMIT), - patch_limit_on_one_side: config - .get_extra::("patch_limit_on_one_side") - .unwrap_or(DEFAULT_PATCH_LIMIT_ON_ONE_SIDE), + params: Self::new().params.resolved(config), } } pub fn patch_size(&self) -> usize { - self.patch_size + self.params.patch_size } pub fn merge_size(&self) -> usize { - self.merge_size - } - - #[inline] - fn factor(&self) -> usize { - self.patch_size * self.merge_size - } - - /// Compute resize dimensions and padding, matching HF `navit_resize_image`. - /// - /// Never upscales (scale capped at 1.0). Pads with zeros to align to factor. - fn compute_resize_config(&self, width: usize, height: usize) -> ResizeConfig { - let ps = self.patch_size; - let patches_w = (width / ps).max(1) as f64; - let patches_h = (height / ps).max(1) as f64; - - let s1 = (self.in_patch_limit as f64 / (patches_w * patches_h)).sqrt(); - let s2 = (self.patch_limit_on_one_side * ps) as f64 / width as f64; - let s3 = (self.patch_limit_on_one_side * ps) as f64 / height as f64; - let scale = f64::min(1.0, f64::min(s1, f64::min(s2, s3))); - - let new_w = ((width as f64 * scale) as usize).max(1); - let new_h = ((height as f64 * scale) as usize).max(1); - let new_w = new_w.min(self.patch_limit_on_one_side * ps); - let new_h = new_h.min(self.patch_limit_on_one_side * ps); - - let factor = self.factor(); - let pad_width = (factor - new_w % factor) % factor; - let pad_height = (factor - new_h % factor) % factor; - - let token_height = (new_h + pad_height) / factor; - let token_width = (new_w + pad_width) / factor; - let num_tokens = token_height * token_width; - - ResizeConfig { - new_width: new_w, - new_height: new_h, - pad_width, - pad_height, - num_tokens, - } - } - - /// Fused resize + zero-pad + normalize into a single [C, H_padded, W_padded] tensor. - /// - /// Avoids intermediate allocations by: - /// 1. Allocating the final padded canvas directly - /// 2. Pre-filling with normalized black (bias value) - /// 3. Deinterleaving + normalizing the image region in one pass - fn resize_pad_and_normalize( - image: &DynamicImage, - cfg: &ResizeConfig, - mean: &[f64; 3], - std: &[f64; 3], - ) -> Array3 { - let canvas_h = cfg.new_height + cfg.pad_height; - let canvas_w = cfg.new_width + cfg.pad_width; - - // Resize using SIMD-accelerated BICUBIC (fast_image_resize) - let resized = transforms::resize( - image, - cfg.new_width as u32, - cfg.new_height as u32, - image::imageops::FilterType::CatmullRom, - ); - - let (img_w, img_h, raw) = transforms::rgb_bytes(&resized); - let canvas_pixels = canvas_h * canvas_w; - - // Precompute fused scale/bias: pixel/255 → normalized - // output[c][i] = raw[i*3+c] / 255.0 * (1/std[c]) + (-mean[c]/std[c]) - let scale: [f32; 3] = std::array::from_fn(|c| 1.0 / (255.0 * std[c] as f32)); - let bias: [f32; 3] = std::array::from_fn(|c| -(mean[c] as f32) / (std[c] as f32)); - - // Pooled: this per-image CHW buffer (tens of MB) is recycled by the - // caller after patch extraction, keeping its pages mapped and hot. - let mut data = scratch::take_f32(3 * canvas_pixels); - let (r_plane, rest) = data.split_at_mut(canvas_pixels); - let (g_plane, b_plane) = rest.split_at_mut(canvas_pixels); - - // Pre-fill with normalized black: (0/255 - mean) / std = bias - r_plane.fill(bias[0]); - g_plane.fill(bias[1]); - b_plane.fill(bias[2]); - - // Overwrite image region row-by-row using vectorized deinterleave - let rw = img_w.min(canvas_w); - let rh = img_h.min(canvas_h); - for y in 0..rh { - let src_row = &raw[y * img_w * 3..y * img_w * 3 + rw * 3]; - let dst_offset = y * canvas_w; - transforms::deinterleave_rgb_to_planes( - src_row, - &mut r_plane[dst_offset..dst_offset + rw], - &mut g_plane[dst_offset..dst_offset + rw], - &mut b_plane[dst_offset..dst_offset + rw], - scale, - bias, - ); - } - - #[expect( - clippy::expect_used, - reason = "data has exactly 3*canvas_h*canvas_w elements by construction" - )] - Array3::from_shape_vec((3, canvas_h, canvas_w), data) - .expect("shape matches pre-allocated buffer") - } - - /// Extract [C, patch_size, patch_size] patches from a contiguous [C, H, W] tensor. - /// - /// Uses row-based `copy_from_slice` instead of per-element indexing so the - /// compiler can auto-vectorize the inner copy. - /// Append this image's patches directly into `out` (no per-image intermediate - /// Vec): `out` is the pooled batch buffer pre-sized for the whole request. - fn extract_patches_into(tensor: &Array3, patch_size: usize, out: &mut Vec) { - let channels = tensor.shape()[0]; - let height = tensor.shape()[1]; - let width = tensor.shape()[2]; - - let grid_h = height / patch_size; - let grid_w = width / patch_size; - - // Get contiguous slice for direct row addressing - let flat = tensor.as_standard_layout(); - #[expect( - clippy::expect_used, - reason = "as_standard_layout guarantees contiguous C-order memory" - )] - let data = flat - .as_slice() - .expect("as_standard_layout guarantees contiguous memory"); - - for gh in 0..grid_h { - for gw in 0..grid_w { - let h_start = gh * patch_size; - let w_start = gw * patch_size; - for c in 0..channels { - let plane_offset = c * height * width; - for ph in 0..patch_size { - let row_start = plane_offset + (h_start + ph) * width + w_start; - out.extend_from_slice(&data[row_start..row_start + patch_size]); - } - } - } - } + self.params.merge_size } } @@ -246,86 +76,14 @@ impl VisionPreProcessor for KimiK25Processor { images: &[DynamicImage], config: &PreProcessorConfig, ) -> Result { - if images.is_empty() { - return Err(TransformError::EmptyBatch); - } - - let item_sizes: Vec<(u32, u32)> = images.iter().map(|img| img.dimensions()).collect(); - let mean = config.get_image_mean(); - let std = config.get_image_std(); - - // Pre-size the pooled batch buffer exactly (patch_features per patch = - // 3 * patch_size^2; this is the data plane's hottest allocation). - let patch_features = 3 * self.patch_size * self.patch_size; - let mut estimated_total = 0usize; - for image in images { - let (w, h) = image.dimensions(); - let cfg = self.compute_resize_config(w as usize, h as usize); - let grid_h = (cfg.new_height + cfg.pad_height) / self.patch_size; - let grid_w = (cfg.new_width + cfg.pad_width) / self.patch_size; - estimated_total += grid_h * grid_w * patch_features; - } - let mut all_patches: Vec = scratch::take_f32_cap(estimated_total); - let mut patches_per_image: Vec = Vec::with_capacity(images.len()); - let mut grid_thw_data = Vec::with_capacity(images.len() * 3); - let mut feature_token_counts = Vec::with_capacity(images.len()); - - for image in images { - let (w, h) = image.dimensions(); - let cfg = self.compute_resize_config(w as usize, h as usize); - - // Fused resize + pad + normalize in one pass (avoids 2 extra allocations) - let tensor = Self::resize_pad_and_normalize(image, &cfg, &mean, &std); - - let padded_h = cfg.new_height + cfg.pad_height; - let padded_w = cfg.new_width + cfg.pad_width; - let grid_h = padded_h / self.patch_size; - let grid_w = padded_w / self.patch_size; - let grid_t = 1usize; - - grid_thw_data.push(grid_t as i64); - grid_thw_data.push(grid_h as i64); - grid_thw_data.push(grid_w as i64); - - let num_patches = grid_h * grid_w; - feature_token_counts.push(cfg.num_tokens); - - // Patchify directly into the pooled batch buffer, then recycle the - // CHW tensor's storage (standard layout, offset 0) for the next image. - Self::extract_patches_into(&tensor, self.patch_size, &mut all_patches); - let (storage, _offset) = tensor.into_raw_vec_and_offset(); - scratch::give_f32(storage); - patches_per_image.push(num_patches as i64); - } - - let total_patches: usize = patches_per_image.iter().map(|&n| n as usize).sum(); - let encoder_input = ndarray::Array4::from_shape_vec( - (total_patches, 3, self.patch_size, self.patch_size), - all_patches, - ) - .map_err(|e| { - TransformError::ShapeError(format!( - "Failed to create encoder_input [{total_patches}, 3, {}, {}]: {e}", - self.patch_size, self.patch_size - )) - })?; - - let result = - PreprocessedEncoderInputs::new(encoder_input, feature_token_counts, item_sizes) - .with_extra( - "grid_thws", - ModelSpecificValue::int_2d(grid_thw_data, images.len(), 3), - ) - .with_extra( - "patches_per_image", - ModelSpecificValue::int_1d(patches_per_image), - ); - - Ok(result) + // K2.5 ships no `transparent_bg_config`, so alpha is dropped. + moonvit::preprocess(self.params.resolved(config), images, config, None) } - fn calculate_num_tokens(&self, width: u32, height: u32, _config: &PreProcessorConfig) -> usize { - self.compute_resize_config(width as usize, height as usize) + fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { + self.params + .resolved(config) + .compute_resize_config(width as usize, height as usize) .num_tokens } @@ -343,7 +101,7 @@ mod tests { use image::{Rgb, RgbImage}; use super::*; - use crate::vision::preprocessor_config::PatchSize; + use crate::vision::{preprocessor_config::PatchSize, processor::ModelSpecificValue}; fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { DynamicImage::from(RgbImage::from_pixel(width, height, color)) @@ -354,7 +112,8 @@ mod tests { let p = KimiK25Processor::new(); assert_eq!(p.patch_size(), 14); assert_eq!(p.merge_size(), 2); - assert_eq!(p.factor(), 28); + assert_eq!(p.params.factor(), 28); + assert_eq!(p.params.in_patch_limit, DEFAULT_IN_PATCH_LIMIT); } #[test] @@ -373,7 +132,7 @@ mod tests { fn test_resize_config_no_upscale() { let p = KimiK25Processor::new(); // Small image should NOT be upscaled (scale capped at 1.0) - let cfg = p.compute_resize_config(100, 100); + let cfg = p.params.compute_resize_config(100, 100); assert!(cfg.new_width <= 100); assert!(cfg.new_height <= 100); // Padded dimensions must be factor-aligned @@ -385,7 +144,7 @@ mod tests { fn test_resize_config_large_image_downscaled() { let p = KimiK25Processor::new(); // Large image should be downscaled - let cfg = p.compute_resize_config(4000, 3000); + let cfg = p.params.compute_resize_config(4000, 3000); // Resized dimensions should be smaller than original assert!(cfg.new_width < 4000); assert!(cfg.new_height < 3000); @@ -403,7 +162,7 @@ mod tests { // pad to (600+4=) → let's compute: // factor=28, 400 % 28 = 400 - 14*28 = 400-392 = 8, pad_h = 28-8 = 20 // 600 % 28 = 600 - 21*28 = 600-588 = 12, pad_w = 28-12 = 16 - let cfg = p.compute_resize_config(600, 400); + let cfg = p.params.compute_resize_config(600, 400); assert_eq!(cfg.new_width, 600); assert_eq!(cfg.new_height, 400); assert_eq!(cfg.pad_height, 20); @@ -539,34 +298,65 @@ mod tests { } #[test] - fn test_preprocess_empty_batch_returns_error() { + fn test_k25_drops_alpha_before_resizing() { + // K2.5 ships no transparent_bg_config, so alpha is dropped and the + // stored RGB kept. The reference does that at load time, before any + // resize; resizing RGBA first would premultiply and discard the colour + // under transparent pixels, so transparent red must survive as red. let p = KimiK25Processor::new(); - let config = PreProcessorConfig::default(); - let result = p.preprocess(&[], &config); - assert!(result.is_err()); - } - - #[test] - fn test_from_preprocessor_config_reads_limits() { let config = PreProcessorConfig { - patch_size: Some(PatchSize { - height: Some(14), - width: Some(14), - }), - merge_size: Some(2), - extra: [ - ("in_patch_limit".to_string(), serde_json::json!(8192)), - ( - "patch_limit_on_one_side".to_string(), - serde_json::json!(256), - ), - ] + image_mean: Some(KIMI_K25_MEAN.to_vec()), + image_std: Some(KIMI_K25_STD.to_vec()), + extra: [( + "patch_limit_on_one_side".to_string(), + serde_json::json!(2), // caps the long side at 2 * 14 px + )] .into_iter() .collect(), ..Default::default() }; - let p = KimiK25Processor::from_preprocessor_config(&config); - assert_eq!(p.in_patch_limit, 8192); - assert_eq!(p.patch_limit_on_one_side, 256); + let hidden_red = DynamicImage::from(image::RgbaImage::from_pixel( + 112, + 112, + image::Rgba([255, 0, 0, 0]), + )); + + let result = p.preprocess(&[hidden_red], &config).unwrap(); + // [patches, 3, patch_size, patch_size]. 112px caps to 2 * 14 per side, + // so every patch is content and none of it is padding. + let &[patches, 3, ph, pw] = result.encoder_input.shape() else { + panic!("unexpected shape {:?}", result.encoder_input.shape()); + }; + assert_eq!(patches, 4); + for p in 0..patches { + for y in 0..ph { + for x in 0..pw { + let px = [0, 1, 2].map(|c| result.encoder_input[[p, c, y, x]]); + // mean = std = 0.5, so 255 -> +1.0 and 0 -> -1.0. + assert!( + (px[0] - 1.0).abs() < 1e-3 && px[1] < -0.99 && px[2] < -0.99, + "red under alpha=0 must survive as red at patch {p} ({x},{y}), got {px:?}" + ); + } + } + } + } + + #[test] + fn test_in_patch_limit_resolved_from_config() { + // A model shipping a larger budget must not be capped by the K2.5 default. + let p = KimiK25Processor::new(); + let mut config = PreProcessorConfig::default(); + config + .extra + .insert("in_patch_limit".to_string(), serde_json::json!(65536)); + + let default_tokens = p.calculate_num_tokens(4000, 3000, &PreProcessorConfig::default()); + let raised_tokens = p.calculate_num_tokens(4000, 3000, &config); + assert!( + raised_tokens > default_tokens, + "raising in_patch_limit should raise the token count \ + ({raised_tokens} vs {default_tokens})" + ); } } diff --git a/crates/multimodal/src/vision/processors/kimi_k3.rs b/crates/multimodal/src/vision/processors/kimi_k3.rs new file mode 100644 index 000000000..c6feb4577 --- /dev/null +++ b/crates/multimodal/src/vision/processors/kimi_k3.rs @@ -0,0 +1,448 @@ +//! Kimi-K3 (MoonViT) image processor. +//! +//! K3 runs the same MoonViT stack as K2.5 — see [`super::moonvit`] — but its +//! `preprocessor_config.json` differs in two ways that change the pixels the +//! encoder sees: +//! +//! * `in_patch_limit` is 65536, four times K2.5's 16384, so K3 keeps +//! substantially more resolution before the downscale kicks in. +//! * It ships a `transparent_bg_config` (chessboard, 8px squares, 255/180) with +//! `transparent_bg_fill_stage: "after_resize"`. Transparent pixels are +//! composited over that board instead of having their alpha dropped, which +//! is what `.convert("RGB")` — and a bare `to_rgb8()` — would do. A fully +//! transparent pixel normally stores RGB `(0,0,0)`, so dropping alpha feeds +//! the encoder solid black where the reference feeds a light checkerboard. +//! +//! Defaults here mirror the shipped config; anything the runtime loads from the +//! model's own `preprocessor_config.json` wins at call time. + +use image::DynamicImage; +use serde_json::Value; + +use super::moonvit::{self, MoonVitParams}; +use crate::vision::{ + preprocessor_config::PreProcessorConfig, + processor::{PreprocessedEncoderInputs, VisionPreProcessor}, + transforms::{ + TransformError, TransparentBg, TransparentBgConfig, TransparentBgFillStage, + TransparentBgPattern, + }, +}; + +pub const KIMI_K3_MEAN: [f64; 3] = [0.5, 0.5, 0.5]; +pub const KIMI_K3_STD: [f64; 3] = [0.5, 0.5, 0.5]; + +pub const DEFAULT_PATCH_SIZE: usize = 14; +pub const DEFAULT_MERGE_SIZE: usize = 2; +/// Maximum total patches before merge (`in_patch_limit`) — 4x K2.5's budget. +pub const DEFAULT_IN_PATCH_LIMIT: usize = 65536; +/// Maximum patches along one spatial dimension +pub const DEFAULT_PATCH_LIMIT_ON_ONE_SIDE: usize = 512; + +/// The `transparent_bg_config` shipped with `moonshotai/Kimi-K3`, verbatim from +/// the checkpoint's `preprocessor_config.json`. These values describe K3 only — +/// note the stage is `"after_resize"` where the reference's fallback for a +/// missing key is `"before_resize"`. +fn default_transparent_bg() -> TransparentBg { + TransparentBg { + config: TransparentBgConfig { + pattern: TransparentBgPattern::Chessboard, + chessboard_square_size: 8, + chessboard_square_on_top_left: true, + chessboard_white_value: 255, + chessboard_gray_value: 180, + }, + stage: TransparentBgFillStage::AfterResize, + } +} + +#[derive(Debug, Clone)] +pub struct KimiK3Processor { + params: MoonVitParams, + transparent_bg: Option, +} + +impl Default for KimiK3Processor { + fn default() -> Self { + Self::new() + } +} + +impl KimiK3Processor { + pub fn new() -> Self { + Self { + params: MoonVitParams { + patch_size: DEFAULT_PATCH_SIZE, + merge_size: DEFAULT_MERGE_SIZE, + in_patch_limit: DEFAULT_IN_PATCH_LIMIT, + patch_limit_on_one_side: DEFAULT_PATCH_LIMIT_ON_ONE_SIDE, + }, + transparent_bg: Some(default_transparent_bg()), + } + } + + pub fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { + let base = Self::new(); + Self { + params: base.params.resolved(config), + transparent_bg: base.resolved_transparent_bg(config), + } + } + + pub fn patch_size(&self) -> usize { + self.params.patch_size + } + + pub fn merge_size(&self) -> usize { + self.params.merge_size + } + + /// Overlay the model's own transparency settings, if it ships any. + /// + /// The registry hands out one shared instance, so this is the only point at + /// which a checkpoint's real config can take effect. An absent or malformed + /// key keeps K3's shipped board rather than silently disabling compositing, + /// because `preprocessor_config.json` is optional in this runtime (see + /// `grpc::multimodal::config`). That diverges from the reference, which + /// reads a missing key as "drop alpha", so a checkpoint wanting that path + /// must say `"transparent_bg_config": null` explicitly. + fn resolved_transparent_bg(&self, config: &PreProcessorConfig) -> Option { + if config + .extra + .get("transparent_bg_config") + .is_some_and(Value::is_null) + { + return None; + } + let default = self.transparent_bg?; + Some(TransparentBg { + config: config + .get_extra::("transparent_bg_config") + .unwrap_or(default.config), + stage: config + .get_extra::("transparent_bg_fill_stage") + .unwrap_or(default.stage), + }) + } +} + +impl VisionPreProcessor for KimiK3Processor { + fn default_mean(&self) -> [f64; 3] { + KIMI_K3_MEAN + } + + fn default_std(&self) -> [f64; 3] { + KIMI_K3_STD + } + + fn preprocess( + &self, + images: &[DynamicImage], + config: &PreProcessorConfig, + ) -> Result { + moonvit::preprocess( + self.params.resolved(config), + images, + config, + self.resolved_transparent_bg(config), + ) + } + + fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { + self.params + .resolved(config) + .compute_resize_config(width as usize, height as usize) + .num_tokens + } + + fn model_name(&self) -> &'static str { + "kimi-k3" + } + + fn get_processed_size(&self, _config: &PreProcessorConfig) -> Option<(u32, u32)> { + None + } +} + +#[cfg(test)] +mod tests { + use image::{Rgb, RgbImage, Rgba, RgbaImage}; + use serde_json::json; + + use super::*; + use crate::vision::processors::kimi_k25::{ + KimiK25Processor, DEFAULT_IN_PATCH_LIMIT as K25_LIM, + }; + + fn norm_config() -> PreProcessorConfig { + PreProcessorConfig { + image_mean: Some(KIMI_K3_MEAN.to_vec()), + image_std: Some(KIMI_K3_STD.to_vec()), + ..Default::default() + } + } + + /// Undo the (x/255 - 0.5) / 0.5 normalization to get the byte back. + fn denormalize(v: f32) -> f32 { + (v * 0.5 + 0.5) * 255.0 + } + + #[test] + fn test_defaults_differ_from_k25() { + let p = KimiK3Processor::new(); + assert_eq!(p.patch_size(), 14); + assert_eq!(p.merge_size(), 2); + assert_eq!(p.params.in_patch_limit, 65536); + assert_ne!( + p.params.in_patch_limit, K25_LIM, + "K3's patch budget must not inherit K2.5's" + ); + assert_eq!(p.model_name(), "kimi-k3"); + } + + #[test] + fn test_larger_patch_budget_keeps_more_resolution() { + let k3 = KimiK3Processor::new(); + let k25 = KimiK25Processor::new(); + let config = PreProcessorConfig::default(); + + // 4000x3000 is well past both budgets, so the difference shows up + // directly in how far each model downscales. + let k3_tokens = k3.calculate_num_tokens(4000, 3000, &config); + let k25_tokens = k25.calculate_num_tokens(4000, 3000, &config); + assert!( + k3_tokens > k25_tokens, + "K3 should keep more tokens than K2.5 ({k3_tokens} vs {k25_tokens})" + ); + } + + #[test] + fn test_transparent_pixels_composite_over_chessboard() { + let p = KimiK3Processor::new(); + // 56x56 is factor-aligned (28*2), so there is no padding to confuse the + // check, and 8px squares tile it exactly. + let image = DynamicImage::from(RgbaImage::from_pixel(56, 56, Rgba([0, 0, 0, 0]))); + let result = p.preprocess(&[image], &norm_config()).unwrap(); + + let values: Vec = result + .encoder_input_flat() + .iter() + .map(|&v| denormalize(v)) + .collect(); + + // Every pixel is fully transparent, so the output is the bare board: + // only the two configured grey levels, and both must appear. + assert!( + values.iter().any(|&v| (v - 255.0).abs() < 0.5), + "expected chessboard white (255)" + ); + assert!( + values.iter().any(|&v| (v - 180.0).abs() < 0.5), + "expected chessboard grey (180)" + ); + assert!( + values + .iter() + .all(|&v| (v - 255.0).abs() < 0.5 || (v - 180.0).abs() < 0.5), + "transparent input must not produce anything but board values" + ); + // The bug this guards: dropping alpha would leave normalized -1.0. + assert!( + !result.encoder_input_flat().iter().any(|&v| v < -0.9), + "transparent pixels must not read as solid black" + ); + } + + #[test] + fn test_chessboard_phase_matches_reference() { + // The reference greys a square when `(y//s + x//s) % 2 == 1` for + // chessboard_square_on_top_left=true, so (0,0) is *white* and the + // neighbouring square is grey. An inverted board would still pass a + // "both values present" check, hence this one. + let p = KimiK3Processor::new(); + let image = DynamicImage::from(RgbaImage::from_pixel(56, 56, Rgba([0, 0, 0, 0]))); + let result = p.preprocess(&[image], &norm_config()).unwrap(); + + // Patch 0 is the top-left 14x14 block, channel-first: element 0 is + // R at (0,0), and element 8 is R at (8,0) — the next square over. + let flat = result.encoder_input_flat(); + assert!((denormalize(flat[0]) - 255.0).abs() < 0.5, "(0,0) is white"); + assert!((denormalize(flat[8]) - 180.0).abs() < 0.5, "(8,0) is grey"); + } + + #[test] + fn test_opaque_pixels_are_untouched() { + let p = KimiK3Processor::new(); + let opaque = DynamicImage::from(RgbaImage::from_pixel(56, 56, Rgba([255, 255, 255, 255]))); + let result = p.preprocess(&[opaque], &norm_config()).unwrap(); + assert!( + result + .encoder_input_flat() + .iter() + .all(|&v| (v - 1.0).abs() < 1e-3), + "fully opaque white must stay white" + ); + } + + #[test] + fn test_semi_transparent_blends_toward_background() { + let p = KimiK3Processor::new(); + // Half-opaque black over the board: 0.5*0 + 0.5*bg → 90 or 127. + let image = DynamicImage::from(RgbaImage::from_pixel(56, 56, Rgba([0, 0, 0, 128]))); + let result = p.preprocess(&[image], &norm_config()).unwrap(); + let values: Vec = result + .encoder_input_flat() + .iter() + .map(|&v| denormalize(v)) + .collect(); + let alpha = 128.0 / 255.0; + let expect_grey = (1.0 - alpha) * 180.0; + let expect_white = (1.0 - alpha) * 255.0; + assert!( + values + .iter() + .all(|&v| (v - expect_grey).abs() < 1.5 || (v - expect_white).abs() < 1.5), + "semi-transparent black should land between the board and black" + ); + } + + #[test] + fn test_rgb_input_matches_k25_pipeline() { + // Compositing must be a no-op for images with no alpha channel, so an + // RGB image goes through byte-identically to K2.5. + let config = norm_config(); + let image = DynamicImage::from(RgbImage::from_pixel(196, 140, Rgb([37, 211, 102]))); + + let k3 = KimiK3Processor::new() + .preprocess(std::slice::from_ref(&image), &config) + .unwrap(); + let k25 = KimiK25Processor::new() + .preprocess(&[image], &config) + .unwrap(); + + assert_eq!(k3.encoder_input.shape(), k25.encoder_input.shape()); + assert_eq!(k3.encoder_input_flat(), k25.encoder_input_flat()); + } + + #[test] + fn test_transparent_bg_config_overridden_by_model_config() { + let mut config = norm_config(); + config.extra.insert( + "transparent_bg_config".to_string(), + json!({ "pattern": "white" }), + ); + let p = KimiK3Processor::new(); + let image = DynamicImage::from(RgbaImage::from_pixel(56, 56, Rgba([0, 0, 0, 0]))); + let result = p.preprocess(&[image], &config).unwrap(); + assert!( + result + .encoder_input_flat() + .iter() + .all(|&v| (v - 1.0).abs() < 1e-3), + "a white background config must flatten transparency to pure white" + ); + } + + #[test] + fn test_fill_stage_before_resize_changes_output() { + // The stage is not cosmetic: the board is drawn at the resolution of + // whatever it is painted onto, so a downscaled image differs. + let p = KimiK3Processor::new(); + let mut config = norm_config(); + // Force a downscale so the two stages see different resolutions. + config + .extra + .insert("patch_limit_on_one_side".to_string(), json!(4)); + + let image = DynamicImage::from(RgbaImage::from_pixel(560, 560, Rgba([0, 0, 0, 0]))); + let after = p.preprocess(std::slice::from_ref(&image), &config).unwrap(); + + config.extra.insert( + "transparent_bg_fill_stage".to_string(), + json!("before_resize"), + ); + let before = p.preprocess(&[image], &config).unwrap(); + + assert_eq!(after.encoder_input.shape(), before.encoder_input.shape()); + assert_ne!( + after.encoder_input_flat(), + before.encoder_input_flat(), + "before_resize must not silently behave like after_resize" + ); + } + + #[test] + fn test_from_preprocessor_config_reads_limits() { + let mut config = PreProcessorConfig::default(); + config + .extra + .insert("in_patch_limit".to_string(), json!(1024)); + config + .extra + .insert("patch_limit_on_one_side".to_string(), json!(64)); + let p = KimiK3Processor::from_preprocessor_config(&config); + assert_eq!(p.params.in_patch_limit, 1024); + assert_eq!(p.params.patch_limit_on_one_side, 64); + } + + #[test] + fn test_explicit_null_config_disables_compositing() { + // Absence keeps K3's shipped board, so an explicit null is the only way + // to reach the reference's alpha-dropping path. + let mut config = norm_config(); + config + .extra + .insert("transparent_bg_config".to_string(), Value::Null); + + let p = KimiK3Processor::new(); + assert_eq!(p.resolved_transparent_bg(&config), None); + + let image = DynamicImage::from(RgbaImage::from_pixel(56, 56, Rgba([0, 0, 0, 0]))); + let result = p.preprocess(&[image], &config).unwrap(); + assert!( + result + .encoder_input_flat() + .iter() + .all(|&v| (v + 1.0).abs() < 1e-6), + "an explicit null must fall back to dropping alpha" + ); + } + + #[test] + fn test_after_resize_ignores_colour_hidden_under_alpha() { + // The reference resizes RGBA with PIL before compositing, and PIL + // premultiplies (verified against Pillow 11.2.1), so colour under fully + // transparent pixels cannot bleed into neighbours: swapping it must not + // change the tensor. + let p = KimiK3Processor::new(); + let mut config = norm_config(); + // Cap the long side at 2 * 14 px so a real downscale happens. + config + .extra + .insert("patch_limit_on_one_side".to_string(), json!(2)); + + let outputs = [[0, 0, 0], [255, 0, 255]].map(|hidden| { + let mut img = RgbaImage::new(112, 112); + for (_, y, px) in img.enumerate_pixels_mut() { + *px = if y < 56 { + Rgba([255, 255, 255, 255]) + } else { + Rgba([hidden[0], hidden[1], hidden[2], 0]) + }; + } + p.preprocess(&[DynamicImage::from(img)], &config).unwrap() + }); + + assert_eq!( + outputs[0].encoder_input_flat(), + outputs[1].encoder_input_flat(), + "a straight-alpha resize would let the hidden magenta bleed through" + ); + } + + #[test] + fn test_empty_batch_errors() { + let p = KimiK3Processor::new(); + assert!(p.preprocess(&[], &PreProcessorConfig::default()).is_err()); + } +} diff --git a/crates/multimodal/src/vision/processors/mod.rs b/crates/multimodal/src/vision/processors/mod.rs index 5ff53bbb3..4a1bdfe4c 100644 --- a/crates/multimodal/src/vision/processors/mod.rs +++ b/crates/multimodal/src/vision/processors/mod.rs @@ -13,6 +13,8 @@ //! - **Qwen3-Omni** (`qwen3_omni_vision`): Qwen3 vision preprocessing with Omni video limits and timing metadata //! - **Inkling** (`inkling`): Optional aspect-preserving resize with CLIP normalization and padded patch columns //! - **Kimi-K2.5** (`kimi_k25`): MoonViT resize and zero-padding to patch alignment +//! - **Kimi-K3** (`kimi_k3`): Same MoonViT stack with a 4x patch budget and +//! chessboard alpha compositing //! - **Phi3-Vision** (`phi3_vision`): Dynamic HD transform with 336x336 tiles //! - **Phi4-Vision** (`phi4_vision`): Dynamic HD transform with 448x448 tiles and SiGLIP encoder //! - **LLaMA 4 Vision** (`llama4_vision`): Tile-based processing with 336x336 tiles and global tile @@ -20,8 +22,10 @@ pub mod inkling; pub mod kimi_k25; +pub mod kimi_k3; pub mod llama4_vision; pub mod llava; +pub mod moonvit; pub mod phi3_vision; pub mod phi4_vision; pub mod pixtral; @@ -32,6 +36,7 @@ pub mod qwen_vl_base; pub use inkling::InklingImageProcessor; pub use kimi_k25::KimiK25Processor; +pub use kimi_k3::KimiK3Processor; pub use llama4_vision::Llama4VisionProcessor; pub use llava::{ImageAspectRatio, LlavaNextProcessor, LlavaProcessor}; pub use phi3_vision::Phi3VisionProcessor; diff --git a/crates/multimodal/src/vision/processors/moonvit.rs b/crates/multimodal/src/vision/processors/moonvit.rs new file mode 100644 index 000000000..44405d35a --- /dev/null +++ b/crates/multimodal/src/vision/processors/moonvit.rs @@ -0,0 +1,339 @@ +//! Shared MoonViT preprocessing core for the Kimi vision family. +//! +//! Kimi-K2.5 and Kimi-K3 run the same MoonViT stack: +//! +//! 1. Compute scale to fit within patch limits (never upscale) +//! 2. Resize with BICUBIC interpolation +//! 3. Zero-pad to make dimensions divisible by factor (patch_size * merge_size) +//! 4. Normalize with the configured mean/std +//! 5. Extract patches as [N, C, patch_size, patch_size] +//! +//! Kimi resizes then zero-pads to make dimensions divisible by the alignment +//! factor. The model was trained with zero-padded images, so using direct +//! resize-to-aligned would degrade image quality. +//! +//! The reference `navit_resize_image`, `navit_patchify`, and `normalize` are +//! byte-identical between the two releases. The models diverge only in their +//! configured patch budget and in whether transparent pixels are composited +//! over a background before patchify — see [`MoonVitParams`] and the +//! `transparent_bg` argument to [`preprocess`]. + +use std::borrow::Cow; + +use image::{DynamicImage, GenericImageView}; +use ndarray::Array3; + +use crate::vision::{ + preprocessor_config::PreProcessorConfig, + processor::{ModelSpecificValue, PreprocessedEncoderInputs}, + scratch, + transforms::{self, TransformError, TransparentBg, TransparentBgFillStage}, +}; + +/// MoonViT resize/patchify parameters for a single model. +#[derive(Debug, Clone, Copy)] +pub struct MoonVitParams { + pub patch_size: usize, + pub merge_size: usize, + /// Maximum total patches before merge (`in_patch_limit`). + pub in_patch_limit: usize, + /// Maximum patches along one spatial dimension. + pub patch_limit_on_one_side: usize, +} + +impl MoonVitParams { + #[inline] + pub fn factor(&self) -> usize { + self.patch_size * self.merge_size + } + + /// Overlay whatever the model's `preprocessor_config.json` supplies. + /// + /// The registry hands out one shared processor instance built from + /// compiled-in defaults, so this is the only point where a model's real + /// limits can be applied. K2.5 and K3 ship different `in_patch_limit` + /// values (16384 vs 65536), which is exactly the kind of divergence that + /// silently caps resolution if the defaults win. + pub fn resolved(self, config: &PreProcessorConfig) -> Self { + Self { + // A zero patch or merge size would divide by zero downstream. + patch_size: config.get_patch_size(self.patch_size).max(1), + merge_size: config.merge_size.unwrap_or(self.merge_size).max(1), + in_patch_limit: config + .get_extra::("in_patch_limit") + .unwrap_or(self.in_patch_limit) + .max(1), + patch_limit_on_one_side: config + .get_extra::("patch_limit_on_one_side") + .unwrap_or(self.patch_limit_on_one_side) + .max(1), + } + } + + /// Compute resize dimensions and padding, matching HF `navit_resize_image`. + /// + /// Never upscales (scale capped at 1.0). Pads with zeros to align to factor. + pub fn compute_resize_config(&self, width: usize, height: usize) -> ResizeConfig { + let ps = self.patch_size; + let patches_w = (width / ps).max(1) as f64; + let patches_h = (height / ps).max(1) as f64; + + let s1 = (self.in_patch_limit as f64 / (patches_w * patches_h)).sqrt(); + let s2 = (self.patch_limit_on_one_side * ps) as f64 / width as f64; + let s3 = (self.patch_limit_on_one_side * ps) as f64 / height as f64; + let scale = f64::min(1.0, f64::min(s1, f64::min(s2, s3))); + + let new_w = ((width as f64 * scale) as usize).max(1); + let new_h = ((height as f64 * scale) as usize).max(1); + let new_w = new_w.min(self.patch_limit_on_one_side * ps); + let new_h = new_h.min(self.patch_limit_on_one_side * ps); + + let factor = self.factor(); + let pad_width = (factor - new_w % factor) % factor; + let pad_height = (factor - new_h % factor) % factor; + + let token_height = (new_h + pad_height) / factor; + let token_width = (new_w + pad_width) / factor; + let num_tokens = token_height * token_width; + + ResizeConfig { + new_width: new_w, + new_height: new_h, + pad_width, + pad_height, + num_tokens, + } + } +} + +/// MoonViT resize configuration for a single image. +pub struct ResizeConfig { + pub new_width: usize, + pub new_height: usize, + pub pad_width: usize, + pub pad_height: usize, + pub num_tokens: usize, +} + +/// Fused resize + zero-pad + normalize into a single [C, H_padded, W_padded] tensor. +/// +/// Avoids intermediate allocations by: +/// 1. Allocating the final padded canvas directly +/// 2. Pre-filling with normalized black (bias value) +/// 3. Deinterleaving + normalizing the image region in one pass +/// +/// Alpha handling follows the reference. With a `transparent_bg` the image is +/// composited at the configured stage; K3 ships `"after_resize"`, and the +/// ordering is load-bearing because a chessboard is generated at the resolution +/// of whatever it is painted onto. With `None` alpha is dropped *before* the +/// resize, matching the reference's `.convert("RGB")` at load time. +fn resize_pad_and_normalize( + image: &DynamicImage, + cfg: &ResizeConfig, + mean: &[f64; 3], + std: &[f64; 3], + transparent_bg: Option, +) -> Array3 { + let canvas_h = cfg.new_height + cfg.pad_height; + let canvas_w = cfg.new_width + cfg.pad_width; + + // Nothing to composite over if the image has no alpha to begin with. + let bg = transparent_bg.filter(|_| image.color().has_alpha()); + + // Reduce to RGB up front unless the background is painted after the resize, + // in which case alpha has to survive the convolution. + let pre_flattened = match bg { + Some(b) if b.stage == TransparentBgFillStage::BeforeResize => Some( + DynamicImage::ImageRgb8(transforms::fill_transparent_bg(image, b.config)), + ), + Some(_) => None, + None if image.color().has_alpha() => Some(DynamicImage::ImageRgb8(image.to_rgb8())), + None => None, + }; + let source = pre_flattened.as_ref().unwrap_or(image); + + // SIMD-accelerated BICUBIC (fast_image_resize). Surviving alpha is + // premultiplied, as `PIL.Image.resize` does for RGBA. + let after_resize = bg.is_some_and(|b| b.stage == TransparentBgFillStage::AfterResize); + let resized = transforms::resize( + source, + cfg.new_width as u32, + cfg.new_height as u32, + image::imageops::FilterType::CatmullRom, + ); + + let post_filled = bg + .filter(|_| after_resize) + .map(|b| transforms::fill_transparent_bg(&resized, b.config)); + let (img_w, img_h, raw): (usize, usize, Cow<'_, [u8]>) = match &post_filled { + Some(rgb) => ( + rgb.width() as usize, + rgb.height() as usize, + Cow::Borrowed(rgb.as_raw().as_slice()), + ), + None => transforms::rgb_bytes(&resized), + }; + let canvas_pixels = canvas_h * canvas_w; + + // Precompute fused scale/bias: pixel/255 → normalized + // output[c][i] = raw[i*3+c] / 255.0 * (1/std[c]) + (-mean[c]/std[c]) + let scale: [f32; 3] = std::array::from_fn(|c| 1.0 / (255.0 * std[c] as f32)); + let bias: [f32; 3] = std::array::from_fn(|c| -(mean[c] as f32) / (std[c] as f32)); + + // Pooled: this per-image CHW buffer (tens of MB) is recycled by the + // caller after patch extraction, keeping its pages mapped and hot. + let mut data = scratch::take_f32(3 * canvas_pixels); + let (r_plane, rest) = data.split_at_mut(canvas_pixels); + let (g_plane, b_plane) = rest.split_at_mut(canvas_pixels); + + // Pre-fill with normalized black: (0/255 - mean) / std = bias + r_plane.fill(bias[0]); + g_plane.fill(bias[1]); + b_plane.fill(bias[2]); + + // Overwrite image region row-by-row using vectorized deinterleave + let rw = img_w.min(canvas_w); + let rh = img_h.min(canvas_h); + for y in 0..rh { + let src_row = &raw[y * img_w * 3..y * img_w * 3 + rw * 3]; + let dst_offset = y * canvas_w; + transforms::deinterleave_rgb_to_planes( + src_row, + &mut r_plane[dst_offset..dst_offset + rw], + &mut g_plane[dst_offset..dst_offset + rw], + &mut b_plane[dst_offset..dst_offset + rw], + scale, + bias, + ); + } + + #[expect( + clippy::expect_used, + reason = "data has exactly 3*canvas_h*canvas_w elements by construction" + )] + Array3::from_shape_vec((3, canvas_h, canvas_w), data) + .expect("shape matches pre-allocated buffer") +} + +/// Extract [C, patch_size, patch_size] patches from a contiguous [C, H, W] tensor. +/// +/// Uses row-based `copy_from_slice` instead of per-element indexing so the +/// compiler can auto-vectorize the inner copy. +/// Append this image's patches directly into `out` (no per-image intermediate +/// Vec): `out` is the pooled batch buffer pre-sized for the whole request. +fn extract_patches_into(tensor: &Array3, patch_size: usize, out: &mut Vec) { + let channels = tensor.shape()[0]; + let height = tensor.shape()[1]; + let width = tensor.shape()[2]; + + let grid_h = height / patch_size; + let grid_w = width / patch_size; + + // Get contiguous slice for direct row addressing + let flat = tensor.as_standard_layout(); + #[expect( + clippy::expect_used, + reason = "as_standard_layout guarantees contiguous C-order memory" + )] + let data = flat + .as_slice() + .expect("as_standard_layout guarantees contiguous memory"); + + for gh in 0..grid_h { + for gw in 0..grid_w { + let h_start = gh * patch_size; + let w_start = gw * patch_size; + for c in 0..channels { + let plane_offset = c * height * width; + for ph in 0..patch_size { + let row_start = plane_offset + (h_start + ph) * width + w_start; + out.extend_from_slice(&data[row_start..row_start + patch_size]); + } + } + } + } +} + +/// Run the full MoonViT pipeline over a batch of images. +pub fn preprocess( + params: MoonVitParams, + images: &[DynamicImage], + config: &PreProcessorConfig, + transparent_bg: Option, +) -> Result { + if images.is_empty() { + return Err(TransformError::EmptyBatch); + } + + let item_sizes: Vec<(u32, u32)> = images.iter().map(|img| img.dimensions()).collect(); + let mean = config.get_image_mean(); + let std = config.get_image_std(); + + // Pre-size the pooled batch buffer exactly (patch_features per patch = + // 3 * patch_size^2; this is the data plane's hottest allocation). + let patch_features = 3 * params.patch_size * params.patch_size; + let mut estimated_total = 0usize; + for image in images { + let (w, h) = image.dimensions(); + let cfg = params.compute_resize_config(w as usize, h as usize); + let grid_h = (cfg.new_height + cfg.pad_height) / params.patch_size; + let grid_w = (cfg.new_width + cfg.pad_width) / params.patch_size; + estimated_total += grid_h * grid_w * patch_features; + } + let mut all_patches: Vec = scratch::take_f32_cap(estimated_total); + let mut patches_per_image: Vec = Vec::with_capacity(images.len()); + let mut grid_thw_data = Vec::with_capacity(images.len() * 3); + let mut feature_token_counts = Vec::with_capacity(images.len()); + + for image in images { + let (w, h) = image.dimensions(); + let cfg = params.compute_resize_config(w as usize, h as usize); + + // Fused resize + pad + normalize in one pass (avoids 2 extra allocations) + let tensor = resize_pad_and_normalize(image, &cfg, &mean, &std, transparent_bg); + + let padded_h = cfg.new_height + cfg.pad_height; + let padded_w = cfg.new_width + cfg.pad_width; + let grid_h = padded_h / params.patch_size; + let grid_w = padded_w / params.patch_size; + let grid_t = 1usize; + + grid_thw_data.push(grid_t as i64); + grid_thw_data.push(grid_h as i64); + grid_thw_data.push(grid_w as i64); + + let num_patches = grid_h * grid_w; + feature_token_counts.push(cfg.num_tokens); + + // Patchify directly into the pooled batch buffer, then recycle the + // CHW tensor's storage (standard layout, offset 0) for the next image. + extract_patches_into(&tensor, params.patch_size, &mut all_patches); + let (storage, _offset) = tensor.into_raw_vec_and_offset(); + scratch::give_f32(storage); + patches_per_image.push(num_patches as i64); + } + + let total_patches: usize = patches_per_image.iter().map(|&n| n as usize).sum(); + let encoder_input = ndarray::Array4::from_shape_vec( + (total_patches, 3, params.patch_size, params.patch_size), + all_patches, + ) + .map_err(|e| { + TransformError::ShapeError(format!( + "Failed to create encoder_input [{total_patches}, 3, {}, {}]: {e}", + params.patch_size, params.patch_size + )) + })?; + + Ok( + PreprocessedEncoderInputs::new(encoder_input, feature_token_counts, item_sizes) + .with_extra( + "grid_thws", + ModelSpecificValue::int_2d(grid_thw_data, images.len(), 3), + ) + .with_extra( + "patches_per_image", + ModelSpecificValue::int_1d(patches_per_image), + ), + ) +} diff --git a/crates/multimodal/src/vision/transforms.rs b/crates/multimodal/src/vision/transforms.rs index 10c89ac5e..06db0628a 100644 --- a/crates/multimodal/src/vision/transforms.rs +++ b/crates/multimodal/src/vision/transforms.rs @@ -38,6 +38,123 @@ pub fn rgb_bytes(image: &DynamicImage) -> (usize, usize, std::borrow::Cow<'_, [u } } +/// Background pattern composited underneath a transparent image. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum TransparentBgPattern { + White, + Black, + Gray, + /// Alternating light/dark squares — the familiar image-editor rendering of + /// transparency, which some vision models are trained to read as such. + Chessboard, +} + +/// How to flatten an image that carries an alpha channel. +/// +/// Field names and defaults mirror the reference `TransparentBgConfig` so a +/// model's `preprocessor_config.json` deserializes straight into this. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize)] +#[serde(default)] +pub struct TransparentBgConfig { + pub pattern: TransparentBgPattern, + pub chessboard_square_size: u32, + pub chessboard_square_on_top_left: bool, + pub chessboard_white_value: u8, + pub chessboard_gray_value: u8, +} + +impl Default for TransparentBgConfig { + fn default() -> Self { + Self { + pattern: TransparentBgPattern::Black, + chessboard_square_size: 16, + chessboard_square_on_top_left: true, + chessboard_white_value: 255, + chessboard_gray_value: 200, + } + } +} + +impl TransparentBgConfig { + /// Background grey level for pixel `(x, y)`. + fn background_at(self, x: u32, y: u32) -> u8 { + match self.pattern { + TransparentBgPattern::White => 255, + TransparentBgPattern::Black => 0, + TransparentBgPattern::Gray => 128, + TransparentBgPattern::Chessboard => { + // A square size of 0 would divide by zero here; the reference + // raises on the same input, so clamp instead of panicking. + let size = self.chessboard_square_size.max(1); + let gray_cell = u32::from(self.chessboard_square_on_top_left); + if (y / size + x / size) % 2 == gray_cell { + self.chessboard_gray_value + } else { + self.chessboard_white_value + } + } + } + } +} + +/// Where in the pipeline an alpha-carrying image gets flattened. +/// +/// The distinction is load-bearing for [`TransparentBgPattern::Chessboard`]: +/// the board is generated at the resolution of the image it is composited +/// onto, so flattening before vs. after the resize yields different square +/// sizes relative to the content. `BeforeResize` is the default only because +/// that is the reference's fallback when the key is absent. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TransparentBgFillStage { + #[default] + BeforeResize, + AfterResize, +} + +/// A model's complete alpha-flattening behavior. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct TransparentBg { + pub config: TransparentBgConfig, + pub stage: TransparentBgFillStage, +} + +/// Alpha-composite `image` over the configured background and return RGB. +/// +/// Images with no alpha channel convert straight to RGB, matching the +/// reference's early return. Note that dropping alpha (what `to_rgb8` does on +/// its own) is *not* equivalent: a fully transparent pixel usually stores RGB +/// `(0,0,0)`, so it would read as solid black instead of as background. +pub fn fill_transparent_bg(image: &DynamicImage, config: TransparentBgConfig) -> RgbImage { + if !image.color().has_alpha() { + return image.to_rgb8(); + } + + let rgba = image.to_rgba8(); + let (width, height) = rgba.dimensions(); + let mut out = Vec::with_capacity(width as usize * height as usize * 3); + + for (y, row) in rgba.as_raw().chunks_exact(width as usize * 4).enumerate() { + for (x, px) in row.chunks_exact(4).enumerate() { + let bg = f32::from(config.background_at(x as u32, y as u32)); + let alpha = f32::from(px[3]) / 255.0; + let inv = 1.0 - alpha; + // numpy's `.astype(np.uint8)` truncates rather than rounds, and so + // does an `as` cast; keep them consistent. + out.push((alpha * f32::from(px[0]) + inv * bg) as u8); + out.push((alpha * f32::from(px[1]) + inv * bg) as u8); + out.push((alpha * f32::from(px[2]) + inv * bg) as u8); + } + } + + #[expect( + clippy::expect_used, + reason = "out holds exactly width*height*3 bytes by construction" + )] + RgbImage::from_raw(width, height, out).expect("buffer matches image dimensions") +} + /// Deinterleave interleaved RGB bytes into separate R, G, B f32 planes with /// per-channel `scale` and `bias`: `plane[c][i] = rgb[i*3 + c] * scale[c] + bias[c]`. /// @@ -242,6 +359,10 @@ thread_local! { /// * `width` - Target width /// * `height` - Target height /// * `filter` - Interpolation filter (Nearest, Triangle/Bilinear, CatmullRom/Bicubic, Lanczos3) +/// +/// Alpha-carrying input is premultiplied for the convolution and un-multiplied +/// afterwards (`fast_image_resize` defaults `mul_div_alpha` to `true`), which is +/// what `PIL.Image.resize` does for `RGBA`. pub fn resize(image: &DynamicImage, width: u32, height: u32, filter: FilterType) -> DynamicImage { let pixel_type = match image.pixel_type() { Some(pt) => pt, @@ -1115,6 +1236,8 @@ pub fn bicubic_resize(tensor: &Array3, target_h: usize, target_w: usize) -> #[cfg(test)] mod tests { + use image::Rgba; + use super::*; fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { @@ -1350,4 +1473,97 @@ mod tests { assert_eq!(rgb[1], 64); assert_eq!(rgb[2], 255); } + + fn chessboard(square_size: u32, on_top_left: bool) -> TransparentBgConfig { + TransparentBgConfig { + pattern: TransparentBgPattern::Chessboard, + chessboard_square_size: square_size, + chessboard_square_on_top_left: on_top_left, + chessboard_white_value: 255, + chessboard_gray_value: 180, + } + } + + /// Mirrors the reference `_create_chessboard_background`: + /// `bg[y, x] = gray if (y//s + x//s) % 2 == (1 if on_top_left else 0)`. + fn reference_board(config: TransparentBgConfig, x: u32, y: u32) -> u8 { + let s = config.chessboard_square_size; + let gray_cell = u32::from(config.chessboard_square_on_top_left); + if (y / s + x / s) % 2 == gray_cell { + config.chessboard_gray_value + } else { + config.chessboard_white_value + } + } + + #[test] + fn chessboard_background_matches_reference() { + for on_top_left in [true, false] { + let config = chessboard(8, on_top_left); + let transparent = + DynamicImage::from(image::RgbaImage::from_pixel(24, 24, Rgba([0, 0, 0, 0]))); + let out = fill_transparent_bg(&transparent, config); + for y in 0..24 { + for x in 0..24 { + let expected = reference_board(config, x, y); + assert_eq!( + out.get_pixel(x, y), + &Rgb([expected, expected, expected]), + "({x},{y}) with on_top_left={on_top_left}" + ); + } + } + } + } + + #[test] + fn fill_transparent_bg_blends_partial_alpha() { + // The reference computes `a*img + (1-a)*bg` in float32, then truncates + // via `.astype(np.uint8)`. Check that arithmetic exactly. + let mut img = image::RgbaImage::new(2, 1); + img.put_pixel(0, 0, Rgba([200, 100, 50, 64])); + img.put_pixel(1, 0, Rgba([200, 100, 50, 192])); + let config = TransparentBgConfig { + pattern: TransparentBgPattern::Gray, // flat 128, so no board phase + ..Default::default() + }; + + let out = fill_transparent_bg(&DynamicImage::from(img), config); + for (x, alpha) in [(0u32, 64.0f32), (1, 192.0)] { + let a = alpha / 255.0; + for (c, src) in [200.0f32, 100.0, 50.0].into_iter().enumerate() { + let expected = (a * src + (1.0 - a) * 128.0) as u8; + assert_eq!(out.get_pixel(x, 0)[c], expected, "x={x} channel={c}"); + } + } + } + + #[test] + fn fill_transparent_bg_is_identity_without_alpha() { + let rgb = create_test_image(4, 4, Rgb([12, 34, 56])); + let out = fill_transparent_bg(&rgb, chessboard(2, true)); + assert!(out.pixels().all(|p| p == &Rgb([12, 34, 56]))); + } + + #[test] + fn fill_transparent_bg_survives_zero_square_size() { + // The reference raises on square_size=0; clamp instead of dividing by zero. + let transparent = + DynamicImage::from(image::RgbaImage::from_pixel(4, 4, Rgba([0, 0, 0, 0]))); + let out = fill_transparent_bg(&transparent, chessboard(0, true)); + assert_eq!(out.dimensions(), (4, 4)); + } + + #[test] + fn transparent_bg_config_fills_missing_fields() { + // A model may ship only `pattern`; the rest must fall back rather than + // failing the whole preprocessor_config parse. + let config: TransparentBgConfig = + serde_json::from_str(r#"{"pattern": "chessboard"}"#).unwrap(); + assert_eq!(config.pattern, TransparentBgPattern::Chessboard); + assert_eq!( + config.chessboard_square_size, + TransparentBgConfig::default().chessboard_square_size + ); + } }