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/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..4f1e79dc1 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,69 @@ 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, + // The reference K2.5 processor has no transparency handling: + // alpha is dropped, not composited. + transparent_bg: None, + 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, + // The reference K2.5 processor has no transparency handling: + // alpha is dropped, not composited. + transparent_bg: None, + 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 +111,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 +132,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 +355,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..f49b48ea5 --- /dev/null +++ b/crates/multimodal/src/vision/processors/kimi_k3.rs @@ -0,0 +1,350 @@ +//! 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`]. +//! +//! 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::{borrow::Cow, 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, + // 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() + } + + 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.resolved(config).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, Rgba, RgbaImage}; + + 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() + } + } + + /// 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(); + 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()); + } + + #[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/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..27e0d4aaf --- /dev/null +++ b/crates/multimodal/src/vision/processors/moonvit.rs @@ -0,0 +1,390 @@ +//! 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, TransparentBg, TransparentBgFillStage}, +}; + +/// 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], + /// 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, +} + +/// 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 + } + + 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. + 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, + } + } + + /// 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: + /// 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], + transparent_bg: Option, + ) -> Array3 { + let canvas_h = cfg.new_height + cfg.pad_height; + let canvas_w = cfg.new_width + cfg.pad_width; + + 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; + + // 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, self.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 / 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 + } +} 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); + } + } }