From 2d28c75e4fe46242a88d5f2bf9b8bf6e57c24304 Mon Sep 17 00:00:00 2001 From: key4ng Date: Tue, 28 Jul 2026 14:59:14 -0700 Subject: [PATCH 1/2] refactor(multimodal): extract shared MoonViT base for the Kimi family Kimi-K2.5 and Kimi-K3 both encode images with MoonViT and differ only in their preprocessing parameters, but SMG registered `kimi-k3` against `KimiK25Processor` directly, leaving nowhere to express a K3-specific difference. The reference implementation instead factors the NaViT resize/patchify helpers into a shared `media_utils` module and keeps a thin per-model processor class on top. Mirror that split, using the same shape `qwen_vl_base` already provides for the Qwen VL family: - Move the MoonViT pipeline into `processors::moonvit` as `MoonVitProcessorBase`, parameterized by `MoonVitConfig`. - Reduce `KimiK25Processor` to a thin wrapper that `Deref`s to the base. - Add `KimiK3Processor` as a second wrapper with `model_name: "kimi-k3"`. - Point the `kimi-k3` / `kimi_k3` registrations at it. This commit is a no-op: `KimiK3Processor` reuses K2.5's parameters, so both processors still produce byte-identical output. A test asserts that directly, which anchors the behavior change that follows. Signed-off-by: key4ng --- crates/multimodal/src/vision/processor.rs | 9 +- .../src/vision/processors/kimi_k25.rs | 325 +++-------------- .../src/vision/processors/kimi_k3.rs | 211 +++++++++++ .../multimodal/src/vision/processors/mod.rs | 4 + .../src/vision/processors/moonvit.rs | 338 ++++++++++++++++++ 5 files changed, 613 insertions(+), 274 deletions(-) create mode 100644 crates/multimodal/src/vision/processors/kimi_k3.rs create mode 100644 crates/multimodal/src/vision/processors/moonvit.rs diff --git a/crates/multimodal/src/vision/processor.rs b/crates/multimodal/src/vision/processor.rs index 1569bc244..4705d8c68 100644 --- a/crates/multimodal/src/vision/processor.rs +++ b/crates/multimodal/src/vision/processor.rs @@ -325,7 +325,7 @@ impl VisionProcessorRegistry { Box::new(super::processors::Llama4VisionProcessor::new()), ); - // Register Kimi-K2.5 Vision (also used by Kimi-K3, same MoonViT stack) + // Register Kimi-K2.5 Vision registry.register( "kimi-k2", Box::new(super::processors::KimiK25Processor::new()), @@ -334,13 +334,16 @@ impl VisionProcessorRegistry { "kimi_k2", Box::new(super::processors::KimiK25Processor::new()), ); + + // Register Kimi-K3 Vision. Shares K2.5's MoonViT stack but keeps its own + // processor: the reference ships a separate class whose parameters differ. 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 diff --git a/crates/multimodal/src/vision/processors/kimi_k25.rs b/crates/multimodal/src/vision/processors/kimi_k25.rs index bcdc58ad4..ed3eee97d 100644 --- a/crates/multimodal/src/vision/processors/kimi_k25.rs +++ b/crates/multimodal/src/vision/processors/kimi_k25.rs @@ -1,25 +1,22 @@ //! Kimi-K2.5 (MoonViT) image processor. //! -//! Matches the HuggingFace `KimiK25VisionProcessor` preprocessing pipeline: +//! Matches the HuggingFace `KimiK25VisionProcessor` preprocessing pipeline, +//! which is implemented in [`super::moonvit`] and shared with Kimi-K3. This +//! module supplies only K2.5's defaults. //! -//! 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. +//! K2.5's checkpoint declares no `transparent_bg_config`, and its reference +//! processor has no transparency handling at all, so alpha-carrying images are +//! flattened by dropping the alpha channel — matching `image.convert("RGB")`. + +use std::ops::Deref; -use image::{DynamicImage, GenericImageView}; -use ndarray::Array3; +use image::DynamicImage; +use super::moonvit::{MoonVitConfig, MoonVitProcessorBase}; 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 +29,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, + inner: MoonVitProcessorBase, } impl Default for KimiK25Processor { @@ -56,189 +41,63 @@ impl Default for KimiK25Processor { } impl KimiK25Processor { + /// Create a processor with Kimi-K2.5's published defaults. 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, + inner: MoonVitProcessorBase::new(MoonVitConfig { + 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, + mean: KIMI_K25_MEAN, + std: KIMI_K25_STD, + model_name: "kimi-k2.5", + }), } } 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), + inner: MoonVitProcessorBase::new(MoonVitConfig { + 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), + mean: KIMI_K25_MEAN, + std: KIMI_K25_STD, + model_name: "kimi-k2.5", + }), } } pub fn patch_size(&self) -> usize { - self.patch_size + self.inner.patch_size() } pub fn merge_size(&self) -> usize { - self.merge_size + self.inner.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") - } +impl Deref for KimiK25Processor { + type Target = MoonVitProcessorBase; - /// 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]); - } - } - } - } + fn deref(&self) -> &Self::Target { + &self.inner } } impl VisionPreProcessor for KimiK25Processor { fn default_mean(&self) -> [f64; 3] { - KIMI_K25_MEAN + self.inner.default_mean() } fn default_std(&self) -> [f64; 3] { - KIMI_K25_STD + self.inner.default_std() } fn preprocess( @@ -246,95 +105,19 @@ 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) + self.inner.preprocess(images, config) } - fn calculate_num_tokens(&self, width: u32, height: u32, _config: &PreProcessorConfig) -> usize { - self.compute_resize_config(width as usize, height as usize) - .num_tokens + fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { + self.inner.calculate_num_tokens(width, height, config) } fn model_name(&self) -> &'static str { - "kimi-k2.5" + self.inner.model_name() } - fn get_processed_size(&self, _config: &PreProcessorConfig) -> Option<(u32, u32)> { - None + fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> { + self.inner.get_processed_size(config) } } @@ -343,7 +126,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)) @@ -566,7 +349,7 @@ mod tests { ..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); + assert_eq!(p.in_patch_limit(), 8192); + assert_eq!(p.patch_limit_on_one_side(), 256); } } 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..bc40a0f80 --- /dev/null +++ b/crates/multimodal/src/vision/processors/kimi_k3.rs @@ -0,0 +1,211 @@ +//! Kimi-K3 (MoonViT) image processor. +//! +//! K3 encodes images with the same MoonViT stack as K2.5 — implemented in +//! [`super::moonvit`] — but its reference processor is a separate class with its +//! own parameters, so it gets its own thin wrapper here rather than sharing +//! K2.5's. This mirrors how [`super::qwen2_vl`] and [`super::qwen3_vl`] sit on +//! [`super::qwen_vl_base`]. +//! +//! Known remaining difference, deliberately left for a follow-up change: K3's +//! checkpoint raises `in_patch_limit` to 65536, so large images should keep +//! roughly 4x the visual tokens they get from K2.5's budget. Reading the +//! budget from the checkpoint touches token accounting on both models and is +//! tracked separately from this module's transparency handling. + +use std::ops::Deref; + +use image::DynamicImage; + +use super::{ + kimi_k25::{ + DEFAULT_IN_PATCH_LIMIT, DEFAULT_MERGE_SIZE, DEFAULT_PATCH_LIMIT_ON_ONE_SIDE, + DEFAULT_PATCH_SIZE, KIMI_K25_MEAN, KIMI_K25_STD, + }, + moonvit::{MoonVitConfig, MoonVitProcessorBase}, +}; +use crate::vision::{ + preprocessor_config::PreProcessorConfig, + processor::{PreprocessedEncoderInputs, VisionPreProcessor}, + transforms::TransformError, +}; + +/// K3 ships the same MoonViT normalization constants as K2.5. +pub const KIMI_K3_MEAN: [f64; 3] = KIMI_K25_MEAN; +pub const KIMI_K3_STD: [f64; 3] = KIMI_K25_STD; + +#[derive(Debug, Clone)] +pub struct KimiK3Processor { + inner: MoonVitProcessorBase, +} + +impl Default for KimiK3Processor { + fn default() -> Self { + Self::new() + } +} + +impl KimiK3Processor { + /// Create a processor with Kimi-K3's defaults. + pub fn new() -> Self { + Self { + inner: MoonVitProcessorBase::new(Self::base_config( + DEFAULT_PATCH_SIZE, + DEFAULT_MERGE_SIZE, + DEFAULT_IN_PATCH_LIMIT, + DEFAULT_PATCH_LIMIT_ON_ONE_SIDE, + )), + } + } + + pub fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { + Self { + inner: MoonVitProcessorBase::new(Self::base_config( + config.get_patch_size(DEFAULT_PATCH_SIZE), + config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), + config + .get_extra::("in_patch_limit") + .unwrap_or(DEFAULT_IN_PATCH_LIMIT), + config + .get_extra::("patch_limit_on_one_side") + .unwrap_or(DEFAULT_PATCH_LIMIT_ON_ONE_SIDE), + )), + } + } + + fn base_config( + patch_size: usize, + merge_size: usize, + in_patch_limit: usize, + patch_limit_on_one_side: usize, + ) -> MoonVitConfig { + MoonVitConfig { + patch_size, + merge_size, + in_patch_limit, + patch_limit_on_one_side, + mean: KIMI_K3_MEAN, + std: KIMI_K3_STD, + model_name: "kimi-k3", + } + } + + pub fn patch_size(&self) -> usize { + self.inner.patch_size() + } + + pub fn merge_size(&self) -> usize { + self.inner.merge_size() + } +} + +impl Deref for KimiK3Processor { + type Target = MoonVitProcessorBase; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +impl VisionPreProcessor for KimiK3Processor { + fn default_mean(&self) -> [f64; 3] { + self.inner.default_mean() + } + + fn default_std(&self) -> [f64; 3] { + self.inner.default_std() + } + + fn preprocess( + &self, + images: &[DynamicImage], + config: &PreProcessorConfig, + ) -> Result { + self.inner.preprocess(images, config) + } + + fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { + self.inner.calculate_num_tokens(width, height, config) + } + + fn model_name(&self) -> &'static str { + self.inner.model_name() + } + + fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> { + self.inner.get_processed_size(config) + } +} + +#[cfg(test)] +mod tests { + use image::{Rgb, RgbImage}; + + use super::*; + use crate::vision::processors::KimiK25Processor; + + fn kimi_config() -> PreProcessorConfig { + PreProcessorConfig { + image_mean: Some(KIMI_K3_MEAN.to_vec()), + image_std: Some(KIMI_K3_STD.to_vec()), + ..Default::default() + } + } + + #[test] + fn defaults_match_the_moonvit_stack() { + let p = KimiK3Processor::new(); + assert_eq!(p.patch_size(), 14); + assert_eq!(p.merge_size(), 2); + assert_eq!(p.factor(), 28); + assert_eq!(p.default_mean(), KIMI_K3_MEAN); + assert_eq!(p.default_std(), KIMI_K3_STD); + } + + #[test] + fn model_name_is_distinct_from_k25() { + assert_eq!(KimiK3Processor::new().model_name(), "kimi-k3"); + assert_ne!( + KimiK3Processor::new().model_name(), + KimiK25Processor::new().model_name() + ); + } + + #[test] + fn from_preprocessor_config_reads_limits() { + let config = PreProcessorConfig { + extra: [ + ("in_patch_limit".to_string(), serde_json::json!(65536)), + ( + "patch_limit_on_one_side".to_string(), + serde_json::json!(512), + ), + ] + .into_iter() + .collect(), + ..Default::default() + }; + let p = KimiK3Processor::from_preprocessor_config(&config); + assert_eq!(p.in_patch_limit(), 65536); + assert_eq!(p.patch_limit_on_one_side(), 512); + } + + /// Opaque images must go through K3 exactly as they go through K2.5: the + /// two processors share one pipeline, and only alpha-carrying inputs are + /// allowed to diverge. + #[test] + fn opaque_images_match_k25_byte_for_byte() { + let image = DynamicImage::from(RgbImage::from_pixel(600, 400, Rgb([37, 128, 220]))); + let config = kimi_config(); + + let k3 = KimiK3Processor::new() + .preprocess(std::slice::from_ref(&image), &config) + .unwrap(); + let k25 = KimiK25Processor::new() + .preprocess(std::slice::from_ref(&image), &config) + .unwrap(); + + assert_eq!(k3.encoder_input.shape(), k25.encoder_input.shape()); + assert_eq!(k3.feature_token_counts, k25.feature_token_counts); + assert_eq!(k3.encoder_input_flat(), k25.encoder_input_flat()); + } +} diff --git a/crates/multimodal/src/vision/processors/mod.rs b/crates/multimodal/src/vision/processors/mod.rs index 5ff53bbb3..68ddbf1e9 100644 --- a/crates/multimodal/src/vision/processors/mod.rs +++ b/crates/multimodal/src/vision/processors/mod.rs @@ -13,6 +13,7 @@ //! - **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 as K2.5 with K3's own parameters //! - **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 +21,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 +35,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..6252fb60a --- /dev/null +++ b/crates/multimodal/src/vision/processors/moonvit.rs @@ -0,0 +1,338 @@ +//! Shared base implementation for MoonViT-based image processors. +//! +//! Kimi-K2.5 and Kimi-K3 both encode images with MoonViT and differ only in +//! their preprocessing parameters. This mirrors how the reference +//! implementation factors the NaViT resize/patchify helpers into a shared +//! `media_utils` module and keeps a thin per-model processor class on top, and +//! how [`super::qwen_vl_base`] serves the Qwen VL family here. +//! +//! # Processing 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 the checkpoint's mean/std +//! 5. Extract patches as [N, C, patch_size, patch_size] +//! +//! MoonViT resizes then zero-pads to make dimensions divisible by the alignment +//! factor (patch_size * merge_size). The models were trained with zero-padded +//! images, so using direct resize-to-aligned would degrade image quality. + +use image::{DynamicImage, GenericImageView}; +use ndarray::Array3; + +use crate::vision::{ + preprocessor_config::PreProcessorConfig, + processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, + scratch, + transforms::{self, TransformError}, +}; + +/// Parameters distinguishing one MoonViT processor variant from another. +#[derive(Debug, Clone)] +pub struct MoonVitConfig { + /// Vision encoder patch size. + pub patch_size: usize, + /// Merge size for token reduction. + 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, + /// Default normalization mean, used when the checkpoint declares none. + pub mean: [f64; 3], + /// Default normalization std, used when the checkpoint declares none. + pub std: [f64; 3], + /// Model name for identification. + pub model_name: &'static str, +} + +/// MoonViT resize configuration for a single image. +pub(super) struct ResizeConfig { + pub(super) new_width: usize, + pub(super) new_height: usize, + pub(super) pad_width: usize, + pub(super) pad_height: usize, + pub(super) num_tokens: usize, +} + +/// Shared MoonViT preprocessing pipeline, parameterized by [`MoonVitConfig`]. +#[derive(Debug, Clone)] +pub struct MoonVitProcessorBase { + config: MoonVitConfig, +} + +impl MoonVitProcessorBase { + pub fn new(config: MoonVitConfig) -> Self { + Self { config } + } + + pub fn patch_size(&self) -> usize { + self.config.patch_size + } + + pub fn merge_size(&self) -> usize { + self.config.merge_size + } + + pub fn in_patch_limit(&self) -> usize { + self.config.in_patch_limit + } + + pub fn patch_limit_on_one_side(&self) -> usize { + self.config.patch_limit_on_one_side + } + + /// Dimension alignment factor: `patch_size * merge_size`. + #[inline] + pub fn factor(&self) -> usize { + self.config.patch_size * self.config.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. + pub(super) fn compute_resize_config(&self, width: usize, height: usize) -> ResizeConfig { + let ps = self.config.patch_size; + let patches_w = (width / ps).max(1) as f64; + let patches_h = (height / ps).max(1) as f64; + + let s1 = (self.config.in_patch_limit as f64 / (patches_w * patches_h)).sqrt(); + let s2 = (self.config.patch_limit_on_one_side * ps) as f64 / width as f64; + let s3 = (self.config.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.config.patch_limit_on_one_side * ps); + let new_h = new_h.min(self.config.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]); + } + } + } + } + } +} + +impl VisionPreProcessor for MoonVitProcessorBase { + fn default_mean(&self) -> [f64; 3] { + self.config.mean + } + + fn default_std(&self) -> [f64; 3] { + self.config.std + } + + fn preprocess( + &self, + images: &[DynamicImage], + config: &PreProcessorConfig, + ) -> Result { + if images.is_empty() { + return Err(TransformError::EmptyBatch); + } + + let patch_size = self.config.patch_size; + 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 * patch_size * 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) / patch_size; + let grid_w = (cfg.new_width + cfg.pad_width) / 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 / patch_size; + let grid_w = padded_w / 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, 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, patch_size, patch_size), + all_patches, + ) + .map_err(|e| { + TransformError::ShapeError(format!( + "Failed to create encoder_input [{total_patches}, 3, {patch_size}, \ + {patch_size}]: {e}" + )) + })?; + + 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) + } + + fn calculate_num_tokens(&self, width: u32, height: u32, _config: &PreProcessorConfig) -> usize { + self.compute_resize_config(width as usize, height as usize) + .num_tokens + } + + fn model_name(&self) -> &'static str { + self.config.model_name + } + + fn get_processed_size(&self, _config: &PreProcessorConfig) -> Option<(u32, u32)> { + None + } +} From d926d91eee679b53fd8bf1eece20f5d7fed0016a Mon Sep 17 00:00:00 2001 From: key4ng Date: Tue, 28 Jul 2026 15:20:09 -0700 Subject: [PATCH 2/2] fix(multimodal): composite transparent images for Kimi-K3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit K3's reference processor composites alpha-carrying images over a background before the NaViT resize and patchify; SMG dropped the alpha channel instead. That is not the same operation. A fully transparent PNG pixel usually stores RGB (0,0,0) under its alpha, so dropping alpha normalizes it to -1.0 — solid black — where the model was trained to see the chessboard the checkpoint describes. Logos, icons, screenshots with rounded corners, and anything else served as RGBA reached the encoder with its background inverted. Add the compositing step to the MoonViT base and let K3 drive it: - `transforms`: `TransparentBgConfig` / `TransparentBgPattern` / `TransparentBgFillStage` deserialize straight from a checkpoint, and `fill_transparent_bg` blends `a*img + (1-a)*bg` with the reference's truncating uint8 cast. - `transforms::resize_straight_alpha`: our SIMD resizer premultiplies by alpha by default, which drags RGB toward zero wherever alpha is partial. Pillow convolves channels independently, so compositing after a resize needs the non-premultiplied variant. - `PreProcessorConfig::transparent_bg()` lifts the two keys out of `media_proc_cfg` and returns them together, since the fill stage is only meaningful alongside a config. - `KimiK3Processor` resolves them per request rather than hardcoding them, matching the reference, which falls back to dropping alpha when `transparent_bg_config` is absent. K2.5 declares `transparent_bg: None` because its reference class has no transparency handling at all. The fill stage is load-bearing, not cosmetic: a chessboard is generated at the resolution of the image it lands on, so `before_resize` and `after_resize` produce different square sizes relative to the content. K3's checkpoint ships `after_resize`. Opaque images are untouched — they skip the compositing branch entirely, and a test asserts a declared background does not perturb them. Out of scope, tracked separately: K3's checkpoint also raises `in_patch_limit` to 65536, which changes token accounting. Signed-off-by: key4ng --- .../src/vision/preprocessor_config.rs | 87 +++++- .../src/vision/processors/kimi_k25.rs | 6 + .../src/vision/processors/kimi_k3.rs | 145 +++++++++- .../src/vision/processors/moonvit.rs | 70 ++++- crates/multimodal/src/vision/transforms.rs | 260 +++++++++++++++++- 5 files changed, 552 insertions(+), 16 deletions(-) diff --git a/crates/multimodal/src/vision/preprocessor_config.rs b/crates/multimodal/src/vision/preprocessor_config.rs index 9485e6072..7f340b668 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 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()); @@ -434,6 +439,20 @@ impl PreProcessorConfig { .and_then(|v| serde_json::from_value(v.clone()).ok()) } + /// The checkpoint's alpha-flattening behavior, if it declares any. + /// + /// `None` means the checkpoint asks for no compositing, matching the + /// reference's behavior when `transparent_bg_config` is absent. The fill + /// stage is only meaningful alongside a config, so it is read here rather + /// than exposed on its own. + pub fn transparent_bg(&self) -> Option { + let config = self.get_extra::("transparent_bg_config")?; + let stage = self + .get_extra::("transparent_bg_fill_stage") + .unwrap_or_default(); + Some(transforms::TransparentBg { config, stage }) + } + // Common default values pub const CLIP_MEAN: [f64; 3] = [0.48145466, 0.4578275, 0.40821073]; pub const CLIP_STD: [f64; 3] = [0.26862954, 0.26130258, 0.27577711]; @@ -601,5 +620,67 @@ mod tests { assert_eq!(config.get_patch_size(0), 14); assert_eq!(config.merge_size, Some(2)); + + // K2.5 declares no transparency handling. + assert_eq!(config.transparent_bg(), None); + } + + #[test] + fn test_parse_kimi_k3_transparency_settings() { + // 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": 16, + "chessboard_square_on_top_left": true, + "chessboard_white_value": 255, + "chessboard_gray_value": 200 + }, + "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 + .transparent_bg() + .expect("transparent_bg_config lifted out of media_proc_cfg"); + assert_eq!( + bg.config.pattern, + transforms::TransparentBgPattern::Chessboard + ); + assert_eq!(bg.config.chessboard_square_size, 16); + assert!(bg.config.chessboard_square_on_top_left); + assert_eq!(bg.config.chessboard_white_value, 255); + assert_eq!(bg.config.chessboard_gray_value, 200); + assert_eq!(bg.stage, transforms::TransparentBgFillStage::AfterResize); + } + + #[test] + fn test_transparency_fill_stage_defaults_without_explicit_key() { + // The reference falls back to "before_resize" when the key is absent. + let json = r#"{ + "media_proc_cfg": { + "transparent_bg_config": { "pattern": "white" } + } + }"#; + + let bg = PreProcessorConfig::from_json(json) + .unwrap() + .transparent_bg() + .expect("config present"); + assert_eq!(bg.config.pattern, transforms::TransparentBgPattern::White); + assert_eq!(bg.stage, transforms::TransparentBgFillStage::BeforeResize); } } diff --git a/crates/multimodal/src/vision/processors/kimi_k25.rs b/crates/multimodal/src/vision/processors/kimi_k25.rs index ed3eee97d..4f1e79dc1 100644 --- a/crates/multimodal/src/vision/processors/kimi_k25.rs +++ b/crates/multimodal/src/vision/processors/kimi_k25.rs @@ -51,6 +51,9 @@ impl KimiK25Processor { patch_limit_on_one_side: DEFAULT_PATCH_LIMIT_ON_ONE_SIDE, mean: KIMI_K25_MEAN, std: KIMI_K25_STD, + // The reference K2.5 processor has no transparency handling: + // alpha is dropped, not composited. + transparent_bg: None, model_name: "kimi-k2.5", }), } @@ -69,6 +72,9 @@ impl KimiK25Processor { .unwrap_or(DEFAULT_PATCH_LIMIT_ON_ONE_SIDE), mean: KIMI_K25_MEAN, std: KIMI_K25_STD, + // The reference K2.5 processor has no transparency handling: + // alpha is dropped, not composited. + transparent_bg: None, model_name: "kimi-k2.5", }), } diff --git a/crates/multimodal/src/vision/processors/kimi_k3.rs b/crates/multimodal/src/vision/processors/kimi_k3.rs index bc40a0f80..f49b48ea5 100644 --- a/crates/multimodal/src/vision/processors/kimi_k3.rs +++ b/crates/multimodal/src/vision/processors/kimi_k3.rs @@ -6,13 +6,21 @@ //! K2.5's. This mirrors how [`super::qwen2_vl`] and [`super::qwen3_vl`] sit on //! [`super::qwen_vl_base`]. //! +//! The difference that matters for output: K3 composites alpha-carrying images +//! over a background the checkpoint describes, where K2.5 just drops the alpha +//! channel. Dropping it is not equivalent — a fully transparent pixel usually +//! stores RGB `(0,0,0)`, so the model sees solid black where it was trained to +//! see a chessboard. The settings are read from the request's +//! `PreProcessorConfig` rather than hardcoded, matching the reference, which +//! falls back to dropping alpha when `transparent_bg_config` is absent. +//! //! Known remaining difference, deliberately left for a follow-up change: K3's //! checkpoint raises `in_patch_limit` to 65536, so large images should keep //! roughly 4x the visual tokens they get from K2.5's budget. Reading the //! budget from the checkpoint touches token accounting on both models and is //! tracked separately from this module's transparency handling. -use std::ops::Deref; +use std::{borrow::Cow, ops::Deref}; use image::DynamicImage; @@ -85,10 +93,27 @@ impl KimiK3Processor { patch_limit_on_one_side, mean: KIMI_K3_MEAN, std: KIMI_K3_STD, + // Resolved per request from the checkpoint, never hardcoded here: + // the reference reads `transparent_bg_config` out of the model's + // config and falls back to dropping alpha when the key is absent. + transparent_bg: None, model_name: "kimi-k3", } } + /// Apply the checkpoint's transparency settings, if it declares any. + /// + /// The registry hands out a default-constructed processor, so a request's + /// `PreProcessorConfig` is the first point at which these values are + /// known. Mirrors [`super::qwen2_vl`]'s call-time re-resolution, narrowed + /// to transparency: nothing else here depends on the checkpoint. + fn resolved(&self, config: &PreProcessorConfig) -> Cow<'_, MoonVitProcessorBase> { + match config.transparent_bg() { + Some(bg) => Cow::Owned(self.inner.with_transparent_bg(Some(bg))), + None => Cow::Borrowed(&self.inner), + } + } + pub fn patch_size(&self) -> usize { self.inner.patch_size() } @@ -120,7 +145,7 @@ impl VisionPreProcessor for KimiK3Processor { images: &[DynamicImage], config: &PreProcessorConfig, ) -> Result { - self.inner.preprocess(images, config) + self.resolved(config).preprocess(images, config) } fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { @@ -138,7 +163,7 @@ impl VisionPreProcessor for KimiK3Processor { #[cfg(test)] mod tests { - use image::{Rgb, RgbImage}; + use image::{Rgb, RgbImage, Rgba, RgbaImage}; use super::*; use crate::vision::processors::KimiK25Processor; @@ -151,6 +176,28 @@ mod tests { } } + /// A config carrying the transparency keys as a checkpoint would ship them. + fn kimi_config_with_bg(bg: serde_json::Value, stage: &str) -> PreProcessorConfig { + PreProcessorConfig { + extra: [ + ("transparent_bg_config".to_string(), bg), + ( + "transparent_bg_fill_stage".to_string(), + serde_json::json!(stage), + ), + ] + .into_iter() + .collect(), + ..kimi_config() + } + } + + /// With mean = std = 0.5, a channel byte `v` normalizes to `v/127.5 - 1`: + /// black is -1.0 and white is +1.0. + fn normalized(byte: u8) -> f32 { + f32::from(byte) / 127.5 - 1.0 + } + #[test] fn defaults_match_the_moonvit_stack() { let p = KimiK3Processor::new(); @@ -208,4 +255,96 @@ mod tests { assert_eq!(k3.feature_token_counts, k25.feature_token_counts); assert_eq!(k3.encoder_input_flat(), k25.encoder_input_flat()); } + + #[test] + fn transparent_pixels_composite_over_the_configured_background() { + // A fully transparent PNG stores RGB (0,0,0) under the alpha channel, + // so dropping alpha normalizes it to -1.0 — solid black. Compositing + // over the checkpoint's white background must instead give +1.0. + let image = DynamicImage::from(RgbaImage::from_pixel(56, 56, Rgba([0, 0, 0, 0]))); + let config = kimi_config_with_bg(serde_json::json!({"pattern": "white"}), "after_resize"); + + let out = KimiK3Processor::new() + .preprocess(std::slice::from_ref(&image), &config) + .unwrap(); + + assert!( + out.encoder_input_flat() + .iter() + .all(|v| (v - normalized(255)).abs() < 1e-6), + "transparent pixels must read as the background, not as black" + ); + } + + #[test] + fn transparency_config_does_not_reach_k25() { + // K2.5's reference processor has no transparency handling at all, so + // the same config must leave it dropping alpha. + let image = DynamicImage::from(RgbaImage::from_pixel(56, 56, Rgba([0, 0, 0, 0]))); + let config = kimi_config_with_bg(serde_json::json!({"pattern": "white"}), "after_resize"); + + let out = KimiK25Processor::new() + .preprocess(std::slice::from_ref(&image), &config) + .unwrap(); + + assert!( + out.encoder_input_flat() + .iter() + .all(|v| (v - normalized(0)).abs() < 1e-6), + "K2.5 must ignore a background it does not declare" + ); + } + + #[test] + fn opaque_images_ignore_the_transparency_config() { + // Declaring a background must not perturb images without alpha, which + // is the overwhelming majority of traffic. + let image = DynamicImage::from(RgbImage::from_pixel(600, 400, Rgb([37, 128, 220]))); + let bg = serde_json::json!({"pattern": "chessboard"}); + + let plain = KimiK3Processor::new() + .preprocess(std::slice::from_ref(&image), &kimi_config()) + .unwrap(); + let with_bg = KimiK3Processor::new() + .preprocess( + std::slice::from_ref(&image), + &kimi_config_with_bg(bg, "after_resize"), + ) + .unwrap(); + + assert_eq!(plain.encoder_input_flat(), with_bg.encoder_input_flat()); + } + + #[test] + fn fill_stage_paints_the_board_at_its_own_resolution() { + // The chessboard is generated at the resolution of the image it lands + // on, so flattening before vs. after the resize changes the square + // size relative to the content. Force a 4x downscale and check the two + // stages disagree — the key is not cosmetic. + let image = DynamicImage::from(RgbaImage::from_pixel(112, 112, Rgba([0, 0, 0, 0]))); + let bg = serde_json::json!({ + "pattern": "chessboard", + "chessboard_square_size": 8, + }); + + let outputs = ["before_resize", "after_resize"].map(|stage| { + let mut config = kimi_config_with_bg(bg.clone(), stage); + config + .extra + .insert("patch_limit_on_one_side".to_string(), serde_json::json!(2)); + KimiK3Processor::from_preprocessor_config(&config) + .preprocess(std::slice::from_ref(&image), &config) + .unwrap() + }); + + assert_eq!( + outputs[0].encoder_input.shape(), + outputs[1].encoder_input.shape(), + "only the pixel values should differ" + ); + assert_ne!( + outputs[0].encoder_input_flat(), + outputs[1].encoder_input_flat() + ); + } } diff --git a/crates/multimodal/src/vision/processors/moonvit.rs b/crates/multimodal/src/vision/processors/moonvit.rs index 6252fb60a..27e0d4aaf 100644 --- a/crates/multimodal/src/vision/processors/moonvit.rs +++ b/crates/multimodal/src/vision/processors/moonvit.rs @@ -25,7 +25,7 @@ use crate::vision::{ preprocessor_config::PreProcessorConfig, processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, scratch, - transforms::{self, TransformError}, + transforms::{self, TransformError, TransparentBg, TransparentBgFillStage}, }; /// Parameters distinguishing one MoonViT processor variant from another. @@ -43,6 +43,10 @@ pub struct MoonVitConfig { pub mean: [f64; 3], /// Default normalization std, used when the checkpoint declares none. pub std: [f64; 3], + /// How to flatten alpha-carrying images, when the model honors a + /// background at all. `None` drops alpha, which is what MoonViT variants + /// without transparency handling do. + pub transparent_bg: Option, /// Model name for identification. pub model_name: &'static str, } @@ -89,6 +93,23 @@ impl MoonVitProcessorBase { self.config.patch_size * self.config.merge_size } + pub fn transparent_bg(&self) -> Option { + self.config.transparent_bg + } + + /// Clone with a different alpha-flattening behavior. + /// + /// Transparency settings arrive with the checkpoint rather than at + /// construction, so the per-model wrapper applies them per request. + pub fn with_transparent_bg(&self, transparent_bg: Option) -> Self { + Self { + config: MoonVitConfig { + transparent_bg, + ..self.config.clone() + }, + } + } + /// Compute resize dimensions and padding, matching HF `navit_resize_image`. /// /// Never upscales (scale capped at 1.0). Pads with zeros to align to factor. @@ -124,6 +145,41 @@ impl MoonVitProcessorBase { } } + /// Resize to `cfg`'s dimensions, flattening alpha along the way. + /// + /// Opaque images take the plain resize path unchanged. When the checkpoint + /// declares a background and the image actually carries alpha, the fill + /// stage decides which side of the resize composites, because a chessboard + /// is generated at the resolution of the image it lands on. + fn resize_and_flatten_alpha( + image: &DynamicImage, + cfg: &ResizeConfig, + transparent_bg: Option, + ) -> DynamicImage { + let width = cfg.new_width as u32; + let height = cfg.new_height as u32; + let filter = image::imageops::FilterType::CatmullRom; + + // Resize using SIMD-accelerated BICUBIC (fast_image_resize) + let Some(bg) = transparent_bg.filter(|_| image.color().has_alpha()) else { + return transforms::resize(image, width, height, filter); + }; + + match bg.stage { + TransparentBgFillStage::BeforeResize => { + let flattened = + DynamicImage::from(transforms::fill_transparent_bg(image, bg.config)); + transforms::resize(&flattened, width, height, filter) + } + TransparentBgFillStage::AfterResize => { + // Straight (non-premultiplied) alpha: RGB under transparent + // pixels must survive the resize for compositing to see it. + let resized = transforms::resize_straight_alpha(image, width, height, filter); + DynamicImage::from(transforms::fill_transparent_bg(&resized, bg.config)) + } + } + } + /// Fused resize + zero-pad + normalize into a single [C, H_padded, W_padded] tensor. /// /// Avoids intermediate allocations by: @@ -135,17 +191,12 @@ impl MoonVitProcessorBase { 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; - // 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 resized = Self::resize_and_flatten_alpha(image, cfg, transparent_bg); let (img_w, img_h, raw) = transforms::rgb_bytes(&resized); let canvas_pixels = canvas_h * canvas_w; @@ -274,7 +325,8 @@ impl VisionPreProcessor for MoonVitProcessorBase { 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 tensor = + Self::resize_pad_and_normalize(image, &cfg, &mean, &std, self.transparent_bg()); let padded_h = cfg.new_height + cfg.pad_height; let padded_w = cfg.new_width + cfg.pad_width; diff --git a/crates/multimodal/src/vision/transforms.rs b/crates/multimodal/src/vision/transforms.rs index 10c89ac5e..4e65c4951 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]`. /// @@ -243,12 +360,39 @@ thread_local! { /// * `height` - Target height /// * `filter` - Interpolation filter (Nearest, Triangle/Bilinear, CatmullRom/Bicubic, Lanczos3) pub fn resize(image: &DynamicImage, width: u32, height: u32, filter: FilterType) -> DynamicImage { + resize_inner(image, width, height, filter, true) +} + +/// Resize without premultiplying by alpha, matching `PIL.Image.resize`. +/// +/// Pillow convolves every channel independently, alpha included. Our SIMD +/// resizer premultiplies by default (`ResizeOptions::mul_div_alpha`), which +/// changes RGB wherever alpha is partial and would diverge from any reference +/// pipeline that composites *after* the resize. +pub fn resize_straight_alpha( + image: &DynamicImage, + width: u32, + height: u32, + filter: FilterType, +) -> DynamicImage { + resize_inner(image, width, height, filter, false) +} + +fn resize_inner( + image: &DynamicImage, + width: u32, + height: u32, + filter: FilterType, + premultiply_alpha: bool, +) -> DynamicImage { let pixel_type = match image.pixel_type() { Some(pt) => pt, None => return image.resize_exact(width, height, filter), }; let mut dst = FirImage::new(width, height, pixel_type); - let options = ResizeOptions::new().resize_alg(to_fir_algorithm(filter)); + let options = ResizeOptions::new() + .resize_alg(to_fir_algorithm(filter)) + .use_alpha(premultiply_alpha); let ok = RESIZER.with(|r| r.borrow_mut().resize(image, &mut dst, &options).is_ok()); if !ok { return image.resize_exact(width, height, filter); @@ -1115,6 +1259,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 +1496,116 @@ 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 + ); + } + + #[test] + fn straight_alpha_resize_keeps_rgba_and_ignores_premultiply() { + // Compositing after the resize needs the alpha channel intact, and RGB + // under transparent pixels must not be scaled toward zero. + let mut img = image::RgbaImage::new(4, 4); + for y in 0..4 { + for x in 0..4 { + img.put_pixel(x, y, Rgba([255, 128, 0, 0])); + } + } + let out = resize_straight_alpha(&DynamicImage::from(img), 2, 2, FilterType::CatmullRom); + assert!(out.color().has_alpha(), "alpha must survive the resize"); + let rgba = out.to_rgba8(); + for px in rgba.pixels() { + assert_eq!(px[0], 255, "premultiplication would zero this out"); + assert_eq!(px[3], 0); + } + } }