diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 1c6c20e36..bba8fd40a 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -33,7 +33,6 @@ /crates/kv_index @slin1237 /crates/mcp @key4ng @CatherineSue @slin1237 @zhoug9127 /crates/mesh @tonyluj @llfl @slin1237 -/crates/multimodal @slin1237 @CatherineSue /crates/protocols @CatherineSue @key4ng /crates/reasoning_parser @CatherineSue /crates/tokenizer @slin1237 @CatherineSue diff --git a/.github/labeler.yml b/.github/labeler.yml index b7623f4de..2885d9976 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -66,10 +66,6 @@ data-connector: - changed-files: - any-glob-to-any-file: 'crates/data_connector/**/*' -multimodal: - - changed-files: - - any-glob-to-any-file: 'crates/multimodal/**/*' - protocols: - changed-files: - any-glob-to-any-file: 'crates/protocols/**/*' diff --git a/.github/workflows/nightly-benchmark.yml b/.github/workflows/nightly-benchmark.yml index b1e77f7ac..dfc0592f0 100644 --- a/.github/workflows/nightly-benchmark.yml +++ b/.github/workflows/nightly-benchmark.yml @@ -48,7 +48,7 @@ jobs: path: | bindings/python/dist/*.whl clients/python/smg_client/types/_generated.py - key: nightly-wheel-${{ runner.os }}-${{ hashFiles('Cargo.lock', '**/Cargo.toml', 'crates/auth/src/**', 'bindings/python/src/**', 'clients/openapi-gen/src/**', 'crates/data_connector/src/**', 'crates/grpc_client/src/**', 'crates/kv_index/src/**', 'crates/mcp/src/**', 'crates/mesh/src/**', 'model_gateway/src/**', 'crates/multimodal/src/**', 'crates/protocols/src/**', 'crates/reasoning_parser/src/**', 'crates/tokenizer/src/**', 'crates/tool_parser/src/**', 'crates/wasm/src/**', 'crates/workflow/src/**') }} + key: nightly-wheel-${{ runner.os }}-${{ hashFiles('Cargo.lock', '**/Cargo.toml', 'crates/auth/src/**', 'bindings/python/src/**', 'clients/openapi-gen/src/**', 'crates/data_connector/src/**', 'crates/grpc_client/src/**', 'crates/kv_index/src/**', 'crates/mcp/src/**', 'crates/mesh/src/**', 'model_gateway/src/**', 'crates/protocols/src/**', 'crates/reasoning_parser/src/**', 'crates/tokenizer/src/**', 'crates/tool_parser/src/**', 'crates/wasm/src/**', 'crates/workflow/src/**') }} restore-keys: | nightly-wheel-${{ runner.os }}- diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 9a24ee660..1146cf2bb 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -272,7 +272,7 @@ jobs: - name: Setup Rust uses: ./.github/actions/setup-rust - - name: Verify default multimodal build does not require OpenCV + - name: Verify default multimodal dependency build does not require OpenCV run: | source "$HOME/.cargo/env" cargo check -p llm-multimodal @@ -300,12 +300,6 @@ jobs: rustup toolchain install nightly --profile minimal cargo +nightly fmt -- --check - - name: Generate vision golden fixtures - run: | - python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu - python -m pip install transformers pillow numpy scipy - python crates/multimodal/scripts/generate_vision_golden.py - - name: Run Rust tests timeout-minutes: 30 run: | @@ -413,6 +407,7 @@ jobs: with: filters: | common: + - 'Cargo.toml' - 'model_gateway/**' - 'crates/protocols/**' - 'bindings/**' @@ -434,7 +429,6 @@ jobs: - 'crates/tool_parser/**' chat-completions: - 'crates/reasoning_parser/**' - - 'crates/multimodal/**' - 'crates/grpc_client/**' - 'grpc_servicer/**' - 'e2e_test/chat_completions/**' diff --git a/.github/workflows/release-crates.yml b/.github/workflows/release-crates.yml index 16287e39a..948cb8054 100644 --- a/.github/workflows/release-crates.yml +++ b/.github/workflows/release-crates.yml @@ -59,8 +59,6 @@ jobs: path: crates/mcp - crate: smg-grpc-client path: crates/grpc_client - - crate: llm-multimodal - path: crates/multimodal - crate: smg-mesh path: crates/mesh steps: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 465548021..78d6d7c2c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -32,7 +32,6 @@ repos: src/proto/.*\.proto$| src/mesh/proto/.*\.proto$| src/tokenizer/chat_template\.rs| - crates/multimodal/tests/fixtures/.*| target/.* )$ diff --git a/Cargo.toml b/Cargo.toml index 7bf1f0eb1..6ed7d1339 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["model_gateway", "crates/protocols", "crates/reasoning_parser", "crates/tool_parser", "crates/workflow", "crates/tokenizer", "crates/auth", "crates/mcp", "crates/kv_index", "crates/data_connector", "crates/multimodal", "crates/mm_rdma", "crates/wasm", "crates/mesh", "crates/grpc_client", "bindings/python", "bindings/golang", "clients/rust", "clients/openapi-gen", "crates/mock_worker"] +members = ["model_gateway", "crates/protocols", "crates/reasoning_parser", "crates/tool_parser", "crates/workflow", "crates/tokenizer", "crates/auth", "crates/mcp", "crates/kv_index", "crates/data_connector", "crates/mm_rdma", "crates/wasm", "crates/mesh", "crates/grpc_client", "bindings/python", "bindings/golang", "clients/rust", "clients/openapi-gen", "crates/mock_worker"] resolver = "2" [workspace.dependencies] @@ -14,7 +14,7 @@ smg-auth = { version = "1.2.2", path = "crates/auth" } smg-mcp = { version = "2.3.2", path = "crates/mcp" } kv-index = { version = "1.3.2", path = "crates/kv_index" } smg-data-connector = { version = "2.3.2", path = "crates/data_connector", package = "data-connector" } -llm-multimodal = { version = "1.7.1", path = "crates/multimodal" } +llm-multimodal = { version = "1.7.1", git = "https://github.com/smg-project/llm-multimodal.git", rev = "7df38e53f99aefaebe86934f010aa8084ec99b2f" } smg-wasm = { version = "1.1.3", path = "crates/wasm", package = "smg-wasm" } smg-mesh = { version = "1.4.2", path = "crates/mesh", package = "smg-mesh" } smg-grpc-client = { version = "1.8.0", path = "crates/grpc_client" } diff --git a/crates/grpc_client/proto/tokenspeed_scheduler.proto b/crates/grpc_client/proto/tokenspeed_scheduler.proto index f37d9204b..fc673267a 100644 --- a/crates/grpc_client/proto/tokenspeed_scheduler.proto +++ b/crates/grpc_client/proto/tokenspeed_scheduler.proto @@ -10,7 +10,7 @@ import "common.proto"; // from the cross-engine admin messages in smg.grpc.common (flush/profile). // Trimmed to text+multimodal generation (no embed, no LoRA, no hidden-state // forwarding). Multimodal carries preprocessed tensors only — media fetch + -// per-model preprocess happen in the gateway (see crates/multimodal). EPD: a +// per-model preprocess happen in the gateway (see the llm-multimodal crate). EPD: a // prefill request may carry `EncodeBootstrapInfo` so its multimodal embeddings // arrive from an encode worker over Mooncake (see tokenspeed_encoder.proto). service TokenSpeedScheduler { @@ -205,7 +205,7 @@ message MultimodalItem { } // Multimodal inputs for vision/audio models. Tensors are produced by the -// gateway's per-model preprocessor (crates/multimodal); the servicer +// gateway's per-model preprocessor (llm-multimodal); the servicer // only reconstructs them and hands them to the engine — no preprocess. message MultimodalInputs { // Per-modality itemized payload. Each item owns its tensor payload, diff --git a/crates/multimodal/Cargo.toml b/crates/multimodal/Cargo.toml deleted file mode 100644 index 4d4955e67..000000000 --- a/crates/multimodal/Cargo.toml +++ /dev/null @@ -1,66 +0,0 @@ -[package] -name = "llm-multimodal" -version = "1.7.1" -edition = "2021" -description = "Multimodal processing for vision and other modalities" -license = "Apache-2.0" -repository = "https://github.com/lightseekorg/smg" -authors = [ - "Simo Lin ", - "Chang Su ", -] -keywords = ["vision", "multimodal", "image-processing", "llm"] -categories = ["multimedia::images", "science"] - -[lib] -name = "llm_multimodal" - -[features] -default = [] -opencv-video = ["dep:opencv"] - -[dependencies] -base64 = "0.22" -hf-hub = { version = "0.5.0", default-features = false, features = ["tokio", "rustls-tls"] } -bytes = { version = "1.12.0", features = ["serde"] } -fast_image_resize = { version = "6.0.0", features = ["image"] } -image = { version = "0.25.10", default-features = false, features = ["png", "jpeg", "gif", "bmp", "ico", "tiff", "webp"] } -libloading = "0.8" -ndarray = "0.17" -once_cell = "1.21.4" -rayon = "1.12" -rustfft = "6.4" -symphonia = { version = "0.6", default-features = false, features = ["all"] } -opencv = { version = "0.99.0", default-features = false, features = ["clang-runtime", "imgproc", "videoio"], optional = true } -reqwest = { workspace = true, features = ["stream"] } -serde = { workspace = true, features = ["derive"] } -serde_bytes = "0.11" -serde_json.workspace = true -tempfile = "3.27" -thiserror.workspace = true -tokio = { workspace = true, features = ["sync", "fs", "rt-multi-thread", "process", "time"] } -tracing.workspace = true -url = "2.5.8" - -# Workspace crates -llm-tokenizer.workspace = true -anyhow.workspace = true -blake3.workspace = true - -[dev-dependencies] -criterion = { version = "0.8", features = ["html_reports"] } -futures.workspace = true -http.workspace = true -npyz = { version = "0.9", features = ["npz"] } -tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } - -[build-dependencies] -cc = "1" -pkg-config = "0.3" - -[[bench]] -name = "image_preprocess" -harness = false - -[lints] -workspace = true diff --git a/crates/multimodal/benches/image_preprocess.rs b/crates/multimodal/benches/image_preprocess.rs deleted file mode 100644 index 1bf47c7ce..000000000 --- a/crates/multimodal/benches/image_preprocess.rs +++ /dev/null @@ -1,367 +0,0 @@ -//! Benchmark: SMG image preprocessing vs HF processor baseline. -//! -//! Measures the time for model-specific image preprocessing (resize, normalize, -//! patchify) at various image sizes. Compare results with the companion Python -//! script `scripts/bench_image_preprocess.py` which benchmarks HF transformers. -//! -//! Run: cargo bench -p llm-multimodal --bench image_preprocess - -#![allow(clippy::unwrap_used, clippy::expect_used)] - -use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; -use image::{imageops::FilterType, DynamicImage, RgbImage}; -use llm_multimodal::vision::{ - preprocessor_config::PreProcessorConfig, - processors::{Llama4VisionProcessor, Qwen2VLProcessor, Qwen3VLProcessor}, - transforms, VisionPreProcessor, -}; - -/// Create a synthetic RGB image with some variation (not all zeros). -fn make_test_image(width: u32, height: u32) -> DynamicImage { - let img = RgbImage::from_fn(width, height, |x, y| { - image::Rgb([(x % 256) as u8, (y % 256) as u8, ((x + y) % 256) as u8]) - }); - DynamicImage::ImageRgb8(img) -} - -fn load_preprocessor_config(model_path: &str) -> Option { - let config_path = format!("{model_path}/preprocessor_config.json"); - let json = std::fs::read_to_string(&config_path).ok()?; - PreProcessorConfig::from_json(&json).ok() -} - -// ── Full pipeline benchmarks ───────────────────────────────────── - -fn bench_qwen3_vl(c: &mut Criterion) { - let processor = Qwen3VLProcessor::new(); - let config = - load_preprocessor_config("/raid/models/Qwen/Qwen3-VL-8B-Instruct").unwrap_or_else(|| { - PreProcessorConfig::from_json( - r#"{"do_resize": true, "size": {"shortest_edge": 3136, "longest_edge": 12845056}}"#, - ) - .unwrap() - }); - - let sizes: &[(u32, u32)] = &[ - (224, 224), - (640, 480), - (1024, 768), - (1920, 1080), - (3840, 2160), - ]; - - let mut group = c.benchmark_group("qwen3_vl_preprocess"); - for &(w, h) in sizes { - let image = make_test_image(w, h); - let images = [image]; - group.bench_with_input( - BenchmarkId::new("single", format!("{w}x{h}")), - &images, - |b, imgs| { - b.iter(|| processor.preprocess(imgs, &config).unwrap()); - }, - ); - } - group.finish(); - - // Batch benchmarks - let mut group = c.benchmark_group("qwen3_vl_batch"); - for batch_size in [3, 5, 10] { - let images: Vec = (0..batch_size) - .map(|i| make_test_image(640 + i * 10, 480 + i * 10)) - .collect(); - group.bench_with_input( - BenchmarkId::new("640x480", format!("batch{batch_size}")), - &images, - |b, imgs| { - b.iter(|| processor.preprocess(imgs, &config).unwrap()); - }, - ); - } - group.finish(); - - // Extreme: very small and very large - let mut group = c.benchmark_group("qwen3_vl_extreme"); - let extremes: &[(u32, u32, &str)] = &[ - (32, 32, "tiny_32x32"), - (50, 50, "small_50x50"), - (100, 2000, "tall_100x2000"), - (2000, 100, "wide_2000x100"), - (4096, 4096, "huge_4096x4096"), - ]; - for &(w, h, label) in extremes { - let image = make_test_image(w, h); - let images = [image]; - group.bench_with_input(BenchmarkId::new("single", label), &images, |b, imgs| { - b.iter(|| processor.preprocess(imgs, &config).unwrap()); - }); - } - group.finish(); -} - -fn bench_qwen2_vl(c: &mut Criterion) { - let processor = Qwen2VLProcessor::new(); - let config = - load_preprocessor_config("/raid/models/Qwen/Qwen2-VL-2B-Instruct").unwrap_or_else(|| { - PreProcessorConfig::from_json( - r#"{"do_resize": true, "size": {"shortest_edge": 3136, "longest_edge": 12845056}}"#, - ) - .unwrap() - }); - - let sizes: &[(u32, u32)] = &[(224, 224), (640, 480), (1024, 768), (1920, 1080)]; - - let mut group = c.benchmark_group("qwen2_vl_preprocess"); - for &(w, h) in sizes { - let image = make_test_image(w, h); - let images = [image]; - group.bench_with_input( - BenchmarkId::new("single", format!("{w}x{h}")), - &images, - |b, imgs| { - b.iter(|| processor.preprocess(imgs, &config).unwrap()); - }, - ); - } - group.finish(); -} - -fn bench_llama4(c: &mut Criterion) { - let processor = Llama4VisionProcessor::new(); - let config = - load_preprocessor_config("/raid/models/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8") - .unwrap_or_else(|| { - PreProcessorConfig::from_json( - r#"{"do_resize": true, "size": {"height": 336, "width": 336}}"#, - ) - .unwrap() - }); - - let sizes: &[(u32, u32)] = &[ - (224, 224), - (336, 336), - (640, 480), - (1024, 768), - (1920, 1080), - ]; - - let mut group = c.benchmark_group("llama4_preprocess"); - for &(w, h) in sizes { - let image = make_test_image(w, h); - let images = [image]; - group.bench_with_input( - BenchmarkId::new("single", format!("{w}x{h}")), - &images, - |b, imgs| { - b.iter(|| processor.preprocess(imgs, &config).unwrap()); - }, - ); - } - group.finish(); -} - -// ── Per-step profiling benchmarks ──────────────────────────────── - -fn bench_individual_steps(c: &mut Criterion) { - let sizes: &[(u32, u32)] = &[(640, 480), (1024, 768), (1920, 1080)]; - - // Step 1: Resize only - let mut group = c.benchmark_group("step_resize"); - for &(w, h) in sizes { - let image = make_test_image(w, h); - // Qwen3-VL target: smart_resize result - let processor = Qwen3VLProcessor::new(); - let (th, tw) = processor.smart_resize(h as usize, w as usize).unwrap(); - group.bench_with_input( - BenchmarkId::new("fir_bilinear", format!("{w}x{h}")), - &image, - |b, img| { - b.iter(|| transforms::resize(img, tw as u32, th as u32, FilterType::Triangle)); - }, - ); - } - group.finish(); - - // Step 2: to_tensor only - let mut group = c.benchmark_group("step_to_tensor"); - for &(w, h) in sizes { - // Use a pre-resized image to isolate to_tensor cost - let image = make_test_image(w, h); - group.bench_with_input( - BenchmarkId::new("rgb8", format!("{w}x{h}")), - &image, - |b, img| { - b.iter(|| transforms::to_tensor(img)); - }, - ); - } - group.finish(); - - // Step 3: normalize only - let mut group = c.benchmark_group("step_normalize"); - let mean = [0.5, 0.5, 0.5]; - let std = [0.5, 0.5, 0.5]; - for &(w, h) in sizes { - let image = make_test_image(w, h); - let tensor = transforms::to_tensor(&image); - group.bench_with_input( - BenchmarkId::new("f32", format!("{w}x{h}")), - &tensor, - |b, t| { - b.iter_batched( - || t.clone(), - |mut fresh| transforms::normalize(&mut fresh, &mean, &std), - criterion::BatchSize::SmallInput, - ); - }, - ); - } - group.finish(); -} - -fn bench_llama4_steps(c: &mut Criterion) { - let processor = Llama4VisionProcessor::new(); - let config = - load_preprocessor_config("/raid/models/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8") - .unwrap_or_else(|| { - PreProcessorConfig::from_json( - r#"{"do_resize": true, "size": {"height": 336, "width": 336}}"#, - ) - .unwrap() - }); - - // 1024x768 is the worst case (1.8x slower than HF) - let sizes: &[(u32, u32)] = &[(640, 480), (1024, 768), (1920, 1080)]; - - let mut group = c.benchmark_group("llama4_steps"); - for &(w, h) in sizes { - let image = make_test_image(w, h); - - // Full preprocess - group.bench_with_input( - BenchmarkId::new("full_preprocess", format!("{w}x{h}")), - &image, - |b, img| { - let imgs = [img.clone()]; - b.iter(|| processor.preprocess(&imgs, &config).unwrap()); - }, - ); - - // Just resize - group.bench_with_input( - BenchmarkId::new("resize_only", format!("{w}x{h}")), - &image, - |b, img| { - b.iter(|| transforms::resize(img, 336, 336, FilterType::Triangle)); - }, - ); - - // to_tensor_and_normalize on tile-sized image - let tile_img = make_test_image(336, 336); - group.bench_with_input( - BenchmarkId::new("tensor_normalize_336", format!("{w}x{h}")), - &tile_img, - |b, img| { - b.iter(|| { - transforms::to_tensor_and_normalize(img, &[0.5, 0.5, 0.5], &[0.5, 0.5, 0.5]) - }); - }, - ); - } - group.finish(); -} - -criterion_group!( - benches, - bench_qwen3_vl, - bench_qwen2_vl, - bench_llama4, - bench_llama4_steps, - bench_individual_steps, - bench_fused_to_tensor_normalize, - bench_to_rgb8, - bench_resize_detailed, -); -criterion_main!(benches); - -fn bench_fused_to_tensor_normalize(c: &mut Criterion) { - let sizes: &[(u32, u32)] = &[(640, 480), (1024, 768), (1920, 1080)]; - let mean = [0.5, 0.5, 0.5]; - let std = [0.5, 0.5, 0.5]; - - let mut group = c.benchmark_group("step_to_tensor_normalize_fused"); - for &(w, h) in sizes { - let image = make_test_image(w, h); - group.bench_with_input( - BenchmarkId::new("fused", format!("{w}x{h}")), - &image, - |b, img| { - b.iter(|| transforms::to_tensor_and_normalize(img, &mean, &std)); - }, - ); - } - group.finish(); -} - -fn bench_to_rgb8(c: &mut Criterion) { - let sizes: &[(u32, u32)] = &[(640, 480), (1024, 768), (1920, 1080)]; - - let mut group = c.benchmark_group("step_to_rgb8"); - for &(w, h) in sizes { - let image = make_test_image(w, h); - group.bench_with_input( - BenchmarkId::new("rgb8", format!("{w}x{h}")), - &image, - |b, img| { - b.iter(|| img.to_rgb8()); - }, - ); - } - group.finish(); - - // Also test when image is already RGB8 (should be free) - let mut group = c.benchmark_group("step_to_rgb8_noop"); - for &(w, h) in sizes { - let image = make_test_image(w, h); - let rgb = DynamicImage::ImageRgb8(image.to_rgb8()); - group.bench_with_input( - BenchmarkId::new("already_rgb8", format!("{w}x{h}")), - &rgb, - |b, img| { - b.iter(|| img.to_rgb8()); - }, - ); - } - group.finish(); -} - -fn bench_resize_detailed(c: &mut Criterion) { - let processor = Qwen3VLProcessor::new(); - - // Benchmark: make_test_image + resize + convert back to DynamicImage - let mut group = c.benchmark_group("step_resize_full_pipeline"); - let sizes: &[(u32, u32)] = &[(640, 480), (1024, 768), (1920, 1080)]; - for &(w, h) in sizes { - let image = make_test_image(w, h); - let (th, tw) = processor.smart_resize(h as usize, w as usize).unwrap(); - - // fir resize (our path) - group.bench_with_input( - BenchmarkId::new("fir", format!("{w}x{h}->{tw}x{th}")), - &image, - |b, img| { - b.iter(|| transforms::resize(img, tw as u32, th as u32, FilterType::Triangle)); - }, - ); - - // image crate resize (old path, for comparison) - group.bench_with_input( - BenchmarkId::new("image_crate", format!("{w}x{h}->{tw}x{th}")), - &image, - |b, img| { - b.iter(|| img.resize_exact(tw as u32, th as u32, FilterType::Triangle)); - }, - ); - } - group.finish(); -} diff --git a/crates/multimodal/build.rs b/crates/multimodal/build.rs deleted file mode 100644 index 8ec6387cf..000000000 --- a/crates/multimodal/build.rs +++ /dev/null @@ -1,31 +0,0 @@ -use std::env; - -fn main() -> Result<(), Box> { - println!("cargo:rerun-if-changed=src/opencv_buffer_capture.cpp"); - println!("cargo:rerun-if-env-changed=OPENCV_INCLUDE_PATHS"); - if env::var_os("CARGO_FEATURE_OPENCV_VIDEO").is_none() { - return Ok(()); - } - - let mut build = cc::Build::new(); - build - .cpp(true) - .file("src/opencv_buffer_capture.cpp") - .flag_if_supported("-std=c++17"); - - if let Some(paths) = env::var_os("OPENCV_INCLUDE_PATHS") { - for path in env::split_paths(&paths) { - build.include(path); - } - } else { - let opencv = pkg_config::Config::new() - .cargo_metadata(false) - .probe("opencv4")?; - for path in opencv.include_paths { - build.include(path); - } - } - - build.compile("smg_opencv_buffer_capture"); - Ok(()) -} diff --git a/crates/multimodal/scripts/generate_qwen_preprocess_fingerprints.py b/crates/multimodal/scripts/generate_qwen_preprocess_fingerprints.py deleted file mode 100755 index 08ed29409..000000000 --- a/crates/multimodal/scripts/generate_qwen_preprocess_fingerprints.py +++ /dev/null @@ -1,153 +0,0 @@ -#!/usr/bin/env python3 -"""Generate compact HuggingFace reference fingerprints for Qwen preprocessing. - -The processor classes are constructed locally, so this script does not download -model weights or configuration. Its output is checked into the Rust integration -test as an external correctness oracle for resize, normalization, and patchify. -""" - -import json - -import numpy as np -from PIL import Image -from PIL import __version__ as pillow_version -from transformers import Qwen2VLImageProcessor -from transformers import __version__ as transformers_version -from transformers.models.qwen3_vl.video_processing_qwen3_vl import ( - Qwen3VLVideoProcessor, - smart_resize, -) - -CASES = ((37, 23), (259, 194)) - - -def make_image(width: int, height: int, seed: int = 0) -> Image.Image: - y, x = np.indices((height, width), dtype=np.uint32) - pixels = np.stack( - ( - (seed + x * 7 + y * 3) % 256, - (seed + x * 5 + y * 11) % 256, - (seed + x + y * 2) % 256, - ), - axis=-1, - ).astype(np.uint8) - return Image.fromarray(pixels, mode="RGB") - - -def fingerprint_bytes(values: np.ndarray) -> str: - contiguous = np.ascontiguousarray(values) - value = 0xCBF29CE484222325 - for byte in contiguous.tobytes(): - value ^= byte - value = value * 0x100000001B3 & 0xFFFFFFFFFFFFFFFF - return f"{value:016x}" - - -def processor_cases(name: str, processor: Qwen2VLImageProcessor) -> list[dict]: - results = [] - for width, height in CASES: - output = processor( - images=make_image(width, height), - do_normalize=False, - return_tensors="np", - ) - values = output["pixel_values"] - patch_u8 = np.rint(values * 255.0).astype(np.uint8) - results.append( - { - "model": name, - "width": width, - "height": height, - "shape": list(values.shape), - "grid_thw": output["image_grid_thw"][0].tolist(), - "fnv1a_patch_u8": fingerprint_bytes(patch_u8), - } - ) - return results - - -def qwen3_video_case(processor: Qwen3VLVideoProcessor) -> dict: - width, height = 37, 35 - seeds = (3, 101, 177) - temporal_patch_size = 2 - target_height, target_width = smart_resize( - len(seeds), - height, - width, - temporal_factor=temporal_patch_size, - factor=32, - min_pixels=4096, - max_pixels=25165824, - ) - # Keep PIL as the resize oracle for SMG's PIL-compatible kernel, then use - # the real HF video processor for temporal padding and patchification. - frames = [ - make_image(width, height, seed).resize( - (target_width, target_height), - Image.Resampling.BICUBIC, - ) - for seed in seeds - ] - output = processor( - videos=[frames], - do_resize=False, - do_normalize=False, - do_sample_frames=False, - return_tensors="pt", - ) - values = output["pixel_values_videos"].cpu().numpy() - patch_u8 = np.rint(values * 255.0).astype(np.uint8) - return { - "model": "qwen3_vl", - "width": width, - "height": height, - "frame_count": len(seeds), - "shape": list(values.shape), - "grid_thw": output["video_grid_thw"][0].tolist(), - "fnv1a_patch_u8": fingerprint_bytes(patch_u8), - } - - -def main() -> None: - qwen2 = Qwen2VLImageProcessor( - patch_size=14, - merge_size=2, - temporal_patch_size=2, - min_pixels=256 * 28 * 28, - max_pixels=1280 * 28 * 28, - image_mean=[0.48145466, 0.4578275, 0.40821073], - image_std=[0.26862954, 0.26130258, 0.27577711], - resample=Image.Resampling.BICUBIC, - ) - qwen3 = Qwen2VLImageProcessor( - patch_size=16, - merge_size=2, - temporal_patch_size=2, - min_pixels=65536, - max_pixels=16777216, - image_mean=[0.5, 0.5, 0.5], - image_std=[0.5, 0.5, 0.5], - resample=Image.Resampling.BICUBIC, - ) - qwen3_video = Qwen3VLVideoProcessor( - patch_size=16, - merge_size=2, - temporal_patch_size=2, - size={"shortest_edge": 4096, "longest_edge": 25165824}, - image_mean=[0.5, 0.5, 0.5], - image_std=[0.5, 0.5, 0.5], - resample=Image.Resampling.BICUBIC, - do_sample_frames=False, - ) - document = { - "generator": "generate_qwen_preprocess_fingerprints.py", - "transformers": transformers_version, - "pillow": pillow_version, - "cases": processor_cases("qwen2_vl", qwen2) + processor_cases("qwen3_vl", qwen3), - "video_cases": [qwen3_video_case(qwen3_video)], - } - print(json.dumps(document, indent=2, sort_keys=True)) - - -if __name__ == "__main__": - main() diff --git a/crates/multimodal/scripts/generate_vision_golden.py b/crates/multimodal/scripts/generate_vision_golden.py deleted file mode 100755 index a3f29a653..000000000 --- a/crates/multimodal/scripts/generate_vision_golden.py +++ /dev/null @@ -1,689 +0,0 @@ -#!/usr/bin/env python3 -"""Generate golden outputs for vision processor testing. - -This script generates reference outputs from HuggingFace transformers -that are used to verify the Rust image preprocessors produce identical results. - -Usage: - # Generate all golden outputs - python crates/multimodal/scripts/generate_vision_golden.py - - # Generate for specific model - python crates/multimodal/scripts/generate_vision_golden.py --model llava - - # Use specific image - python crates/multimodal/scripts/generate_vision_golden.py \ - --image crates/multimodal/tests/fixtures/images/square.jpg -""" - -import argparse -import json -import os -import sys -from pathlib import Path - -import numpy as np -from PIL import Image - -# Model configurations -MODELS = { - "llava": { - "model_id": "llava-hf/llava-1.5-7b-hf", - "processor_class": "CLIPImageProcessor", - "description": "Standard CLIP processing (no expand-to-square)", - }, - "llava_pad": { - "model_id": "liuhaotian/llava-v1.5-7b", - "processor_class": "CLIPImageProcessor", - "description": "With expand-to-square (image_aspect_ratio=pad)", - }, - "llava_next": { - "model_id": "llava-hf/llava-v1.6-mistral-7b-hf", - "processor_class": "LlavaNextImageProcessor", - "description": "Multi-crop anyres processing", - }, - "qwen2_vl": { - "model_id": "Qwen/Qwen2-VL-7B-Instruct", - "processor_class": "Qwen2VLImageProcessor", - "description": "Dynamic resolution with smart resize", - }, - "qwen3_vl": { - "model_id": "Qwen/Qwen3-VL-8B-Instruct", - "processor_class": "Qwen2VLImageProcessorFast", - "description": "Dynamic resolution with patch_size=16 and [0.5,0.5,0.5] normalization", - }, - "phi3_vision": { - "model_id": "microsoft/Phi-3-vision-128k-instruct", - "processor_class": "Phi3VImageProcessor", - "description": "Dynamic HD transform with 336x336 tiles", - }, - "phi4_vision": { - "model_id": "microsoft/Phi-4-multimodal-instruct", - "processor_class": "Phi4MMImageProcessor", - "description": "Dynamic HD transform with 448x448 tiles and SiGLIP encoder", - }, - "llama4_vision": { - "model_id": "meta-llama/Llama-4-Scout-17B-16E-Instruct", - "processor_class": "Llama4ImageProcessorFast", - "description": "Tile-based processing with 336x336 tiles and global tile", - }, - "pixtral": { - "model_id": "mistralai/Pixtral-12B-2409", - "processor_class": "PixtralImageProcessor", - "description": "Dynamic resolution with CLIP normalization and bicubic resize", - }, -} - -# Default test images -DEFAULT_IMAGES = [ - "crates/multimodal/tests/fixtures/images/square.jpg", - "crates/multimodal/tests/fixtures/images/tall.jpg", - "crates/multimodal/tests/fixtures/images/wide.jpg", - "crates/multimodal/tests/fixtures/images/small.jpg", - "crates/multimodal/tests/fixtures/images/tiny.jpg", - "crates/multimodal/tests/fixtures/images/very_tall.jpg", - "crates/multimodal/tests/fixtures/images/very_wide.jpg", - "crates/multimodal/tests/fixtures/images/large.jpg", - "crates/multimodal/tests/fixtures/images/odd_dims.jpg", - "crates/multimodal/tests/fixtures/images/grayscale.jpg", -] - - -def expand_to_square(image: Image.Image, background_color: tuple) -> Image.Image: - """Expand image to square by padding with background color. - - This matches the LLaVA preprocessing pipeline where images are - first expanded to square before being processed by CLIP. - """ - width, height = image.size - if width == height: - return image - elif width > height: - # Pad vertically - new_image = Image.new("RGB", (width, width), background_color) - paste_y = (width - height) // 2 - new_image.paste(image, (0, paste_y)) - return new_image - else: - # Pad horizontally - new_image = Image.new("RGB", (height, height), background_color) - paste_x = (height - width) // 2 - new_image.paste(image, (paste_x, 0)) - return new_image - - -def generate_golden_llava(image_path: str, output_dir: str) -> dict: - """Generate golden output for LLaVA 1.5 (standard CLIP processing). - - This uses standard CLIP processing WITHOUT expand-to-square. - Matches behavior of llava-hf/* models where image_aspect_ratio is not set. - - LLaVA 1.5 preprocessing pipeline: - 1. Resize so shortest edge = 336 (preserving aspect ratio) - 2. Center crop to 336x336 - 3. Normalize with CLIP mean/std - """ - from transformers import CLIPImageProcessor - - processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14-336") - image = Image.open(image_path).convert("RGB") - original_size = image.size - - # Standard CLIP processing (no expand-to-square) - outputs = processor(images=image, return_tensors="np") - pixel_values = outputs["pixel_values"] - - # Calculate expected token count - # LLaVA 1.5: (336 / 14)^2 = 576 tokens - patch_size = 14 - image_size = 336 - num_tokens = (image_size // patch_size) ** 2 - - return { - "pixel_values": pixel_values, - "original_size": original_size, - "num_tokens": num_tokens, - "processor_config": processor.to_dict(), - } - - -def generate_golden_llava_pad(image_path: str, output_dir: str) -> dict: - """Generate golden output for LLaVA 1.5 with expand-to-square (pad mode). - - This uses expand-to-square preprocessing. - Matches behavior of liuhaotian/llava-* models where image_aspect_ratio = "pad". - - LLaVA 1.5 pad mode preprocessing pipeline: - 1. Expand image to square by padding with mean color - 2. Resize to 336x336 - 3. Normalize with CLIP mean/std - """ - from transformers import CLIPImageProcessor - - processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14-336") - image = Image.open(image_path).convert("RGB") - original_size = image.size - - # LLaVA-specific: expand to square with mean color padding - # CLIP mean values converted to 0-255 range - clip_mean = (0.48145466, 0.4578275, 0.40821073) - mean_color = tuple(int(m * 255) for m in clip_mean) - image = expand_to_square(image, mean_color) - - # Process image with CLIP processor - outputs = processor(images=image, return_tensors="np") - pixel_values = outputs["pixel_values"] - - # Calculate expected token count - # LLaVA 1.5: (336 / 14)^2 = 576 tokens - patch_size = 14 - image_size = 336 - num_tokens = (image_size // patch_size) ** 2 - - return { - "pixel_values": pixel_values, - "original_size": original_size, - "num_tokens": num_tokens, - "processor_config": processor.to_dict(), - } - - -def generate_golden_llava_next(image_path: str, output_dir: str) -> dict: - """Generate golden output for LLaVA-NeXT (anyres).""" - try: - from transformers import LlavaNextImageProcessor - except ImportError: - print("LlavaNextImageProcessor not available, skipping llava_next") - return None - - processor = LlavaNextImageProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf") - image = Image.open(image_path).convert("RGB") - original_size = image.size - - # Process image - outputs = processor(images=image, return_tensors="np") - pixel_values = outputs["pixel_values"] - - # Get additional outputs if available - image_sizes = outputs.get("image_sizes") - - result = { - "pixel_values": pixel_values, - "original_size": original_size, - "processor_config": processor.to_dict(), - } - - if image_sizes is not None: - result["image_sizes"] = np.array(image_sizes) - - return result - - -def generate_golden_qwen2_vl(image_path: str, output_dir: str) -> dict: - """Generate golden output for Qwen2-VL. - - Qwen2-VL uses dynamic resolution with smart resize: - 1. Smart resize to fit within min/max pixel bounds - 2. Align dimensions to (patch_size * merge_size) boundary - 3. Normalize with CLIP mean/std - 4. Returns image_grid_thw for position encoding - - Default parameters: - - patch_size: 14 - - merge_size: 2 - - min_pixels: 256 * 28 * 28 = 200,704 - - max_pixels: 1280 * 28 * 28 = 1,003,520 - - temporal_patch_size: 2 - """ - try: - from transformers import Qwen2VLImageProcessor - except ImportError: - print("Qwen2VLImageProcessor not available, skipping qwen2_vl") - return None - - processor = Qwen2VLImageProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") - image = Image.open(image_path).convert("RGB") - original_size = image.size - - # Process image - outputs = processor(images=image, return_tensors="np") - pixel_values = outputs["pixel_values"] - image_grid_thw = outputs.get("image_grid_thw") - - # Get config values for token calculation - patch_size = processor.patch_size - merge_size = processor.merge_size - temporal_patch_size = getattr(processor, "temporal_patch_size", 2) - min_pixels = processor.min_pixels - max_pixels = processor.max_pixels - - # Calculate number of tokens - # tokens = (T * H * W) / merge_size² - if image_grid_thw is not None: - # image_grid_thw has shape [batch, 3] with [T, H, W] - grid_thw = image_grid_thw[0] # First (and only) image - num_tokens = int(np.prod(grid_thw) / (merge_size**2)) - else: - num_tokens = None - - result = { - "pixel_values": pixel_values, - "original_size": original_size, - "processor_config": processor.to_dict(), - } - - if image_grid_thw is not None: - result["image_grid_thw"] = np.array(image_grid_thw) - - if num_tokens is not None: - result["num_tokens"] = num_tokens - - # Add debug info - result["config_info"] = { - "patch_size": patch_size, - "merge_size": merge_size, - "temporal_patch_size": temporal_patch_size, - "min_pixels": min_pixels, - "max_pixels": max_pixels, - } - - return result - - -def save_golden(model_key: str, image_name: str, data: dict, output_dir: str): - """Save golden output to files.""" - model_dir = Path(output_dir) / model_key - model_dir.mkdir(parents=True, exist_ok=True) - - # Save numpy data - npz_data = {k: v for k, v in data.items() if isinstance(v, np.ndarray)} - npz_data["original_size"] = np.array(data["original_size"]) - if "num_tokens" in data: - npz_data["num_tokens"] = np.array([data["num_tokens"]]) - - npz_path = model_dir / f"golden_{image_name}.npz" - np.savez(npz_path, **npz_data) - print(f" Saved: {npz_path}") - - # Save processor config (only once per model) - config_path = model_dir / "preprocessor_config.json" - if not config_path.exists() and "processor_config" in data: - with open(config_path, "w") as f: - json.dump(data["processor_config"], f, indent=2) - print(f" Saved: {config_path}") - - -def generate_golden_qwen3_vl(image_path: str, output_dir: str) -> dict: - """Generate golden output for Qwen3-VL. - - Qwen3-VL uses dynamic resolution with smart resize similar to Qwen2-VL - but with different parameters: - - patch_size: 16 (vs 14 in Qwen2-VL) - - factor: 32 (vs 28 in Qwen2-VL) - - normalization: [0.5, 0.5, 0.5] mean/std (vs CLIP values in Qwen2-VL) - - Default parameters: - - patch_size: 16 - - merge_size: 2 - - temporal_patch_size: 2 - """ - from transformers import AutoProcessor - - processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-8B-Instruct", trust_remote_code=True) - image = Image.open(image_path).convert("RGB") - original_size = image.size - - # Process image using the image processor directly - outputs = processor.image_processor(images=image, return_tensors="pt") - - # Convert to numpy for saving - pixel_values = outputs["pixel_values"].numpy() - image_grid_thw = outputs.get("image_grid_thw") - if image_grid_thw is not None: - image_grid_thw = image_grid_thw.numpy() - - # Get config values - img_processor = processor.image_processor - patch_size = getattr(img_processor, "patch_size", 16) - merge_size = getattr(img_processor, "merge_size", 2) - temporal_patch_size = getattr(img_processor, "temporal_patch_size", 2) - - # Calculate number of tokens - if image_grid_thw is not None: - grid_thw = image_grid_thw[0] - num_tokens = int(np.prod(grid_thw) / (merge_size**2)) - else: - num_tokens = None - - result = { - "pixel_values": pixel_values, - "original_size": original_size, - "processor_config": img_processor.to_dict(), - } - - if image_grid_thw is not None: - result["image_grid_thw"] = image_grid_thw - - if num_tokens is not None: - result["num_tokens"] = num_tokens - - # Add debug info - result["config_info"] = { - "patch_size": patch_size, - "merge_size": merge_size, - "temporal_patch_size": temporal_patch_size, - } - - return result - - -def generate_golden_phi3_vision(image_path: str, output_dir: str) -> dict: - """Generate golden output for Phi3-Vision. - - Phi3-Vision uses Dynamic HD transform: - 1. If width < height, transpose image - 2. Calculate scale: while scale * ceil(scale/ratio) <= hd_num: scale++ - 3. Resize to new_w = scale * 336, new_h = new_w / ratio - 4. Pad height to multiple of 336 (centered, white padding) - 5. If transposed, transpose back - 6. Normalize with CLIP mean/std - 7. Create global image (336x336 via bicubic) - 8. Reshape into tiles [num_tiles, 3, 336, 336] - 9. Concatenate [global, tiles] and pad to [num_crops+1, 3, 336, 336] - - Default parameters: - - num_crops: 16 - - num_img_tokens: 144 (per tile) - - normalization: CLIP mean/std - """ - from transformers import AutoImageProcessor - - processor = AutoImageProcessor.from_pretrained( - "microsoft/Phi-3-vision-128k-instruct", trust_remote_code=True - ) - image = Image.open(image_path).convert("RGB") - original_size = image.size - - # Process image - outputs = processor(images=image, return_tensors="np") - pixel_values = outputs["pixel_values"] - image_sizes = outputs.get("image_sizes") - num_img_tokens = outputs.get("num_img_tokens") - - result = { - "pixel_values": pixel_values, - "original_size": original_size, - "processor_config": processor.to_dict(), - } - - if image_sizes is not None: - result["image_sizes"] = np.array(image_sizes) - - if num_img_tokens is not None: - result["num_img_tokens"] = np.array(num_img_tokens) - - # Add debug info - result["config_info"] = { - "num_crops": processor.num_crops, - "num_img_tokens": processor.num_img_tokens, - } - - return result - - -def generate_golden_phi4_vision(image_path: str, output_dir: str) -> dict: - """Generate golden output for Phi4-Vision (Phi-4-multimodal). - - Phi4-Vision uses Dynamic HD transform similar to Phi3 but with: - - Base resolution: 448 (vs 336 in Phi3) - - Normalization: [0.5, 0.5, 0.5] mean/std (vs CLIP in Phi3) - - Default dynamic_hd: 36 (vs 16 num_crops in Phi3) - - Uses SiGLIP vision encoder (vs CLIP in Phi3) - - Has per-crop attention masks - - Token count formula: - 256 + 1 + mask_sum + mask_col0_sum + 16 - - Note: Phi4 uses 'input_image_embeds' key instead of 'pixel_values' - """ - from transformers import AutoProcessor - - processor = AutoProcessor.from_pretrained( - "microsoft/Phi-4-multimodal-instruct", trust_remote_code=True - ) - image = Image.open(image_path).convert("RGB") - original_size = image.size - - # Process image using the image processor directly - outputs = processor.image_processor(images=image, return_tensors="np") - - # Phi4 uses 'input_image_embeds' instead of 'pixel_values' - pixel_values = outputs.get("input_image_embeds") - pixel_attention_mask = outputs.get("image_attention_mask") - image_sizes = outputs.get("image_sizes") - num_img_tokens = outputs.get("num_img_tokens") - - result = { - "pixel_values": pixel_values, - "original_size": original_size, - "processor_config": processor.image_processor.to_dict(), - } - - if pixel_attention_mask is not None: - result["pixel_attention_mask"] = np.array(pixel_attention_mask) - - if image_sizes is not None: - result["image_sizes"] = np.array(image_sizes) - - if num_img_tokens is not None: - result["num_img_tokens"] = np.array(num_img_tokens) - - # Add debug info - result["config_info"] = { - "dynamic_hd": getattr(processor.image_processor, "dynamic_hd", 36), - "base_resolution": 448, - } - - return result - - -def generate_golden_llama4_vision(image_path: str, output_dir: str) -> dict: - """Generate golden output for LLaMA 4 Vision. - - LLaMA 4 Vision uses tile-based processing: - 1. Find supported resolutions based on max_patches (default 16) - 2. Get best fit resolution for the image (minimize upscaling) - 3. Resize preserving aspect ratio - 4. Pad with black (0) to target dimensions - 5. Normalize with [0.5, 0.5, 0.5] mean/std - 6. Split into tiles of 336x336 - 7. If multiple tiles, add global tile at the end - - Output: - - pixel_values: [1, num_tiles, 3, 336, 336] - - aspect_ratios: [1, 2] with [h_tiles, w_tiles] - - Token count: num_tiles * (336 / 14)² = num_tiles * 576 - """ - from transformers.models.llama4 import Llama4ImageProcessorFast - - processor = Llama4ImageProcessorFast() - image = Image.open(image_path).convert("RGB") - original_size = image.size - - # Process image - Llama4 only supports PyTorch tensors - outputs = processor(images=image, return_tensors="pt") - # Convert to numpy (need to convert from bfloat16 to float32 first) - pixel_values = outputs["pixel_values"].float().numpy() - aspect_ratios = outputs.get("aspect_ratios") - if aspect_ratios is not None: - aspect_ratios = aspect_ratios.numpy() - - result = { - "pixel_values": pixel_values, - "original_size": original_size, - "processor_config": processor.to_dict(), - } - - if aspect_ratios is not None: - result["aspect_ratios"] = aspect_ratios - - # Calculate num_tokens from aspect_ratios - if aspect_ratios is not None: - h_tiles = int(aspect_ratios[0][0]) - w_tiles = int(aspect_ratios[0][1]) - num_tiles = h_tiles * w_tiles - # Add 1 for global tile if num_tiles > 1 - total_tiles = num_tiles + 1 if num_tiles > 1 else num_tiles - tokens_per_tile = (336 // 14) ** 2 # 576 - num_tokens = total_tiles * tokens_per_tile - result["num_tokens"] = num_tokens - - # Add debug info - result["config_info"] = { - "tile_size": 336, - "max_patches": processor.max_patches, - "resize_to_max_canvas": processor.resize_to_max_canvas, - } - - return result - - -def generate_golden_pixtral(image_path: str, output_dir: str) -> dict: - """Generate golden output for Pixtral/Mistral3 Vision. - - Pixtral uses dynamic resolution processing: - 1. If image exceeds longest_edge (default 1024), scale down proportionally - 2. Resize to dimensions that are multiples of patch_size (default 16) - 3. Use bicubic interpolation for resize - 4. Normalize with CLIP mean/std - - Output: - - pixel_values: [1, 3, H, W] where H, W are multiples of patch_size - - image_sizes: [(H, W)] - - Token count: (H / patch_size) * (W / patch_size) - """ - from transformers import PixtralImageProcessor - - processor = PixtralImageProcessor.from_pretrained("mistral-community/pixtral-12b") - image = Image.open(image_path).convert("RGB") - original_size = image.size - - # Process image - outputs = processor(images=image, return_tensors="np") - pixel_values = outputs["pixel_values"] - image_sizes = outputs.get("image_sizes") - - result = { - "pixel_values": pixel_values, - "original_size": original_size, - "processor_config": processor.to_dict(), - } - - if image_sizes is not None: - result["image_sizes"] = np.array(image_sizes) - - # Calculate num_tokens from image_sizes - if image_sizes is not None: - h, w = image_sizes[0] - patch_size = getattr(processor, "patch_size", {"height": 16, "width": 16}) - if isinstance(patch_size, dict): - patch_h = patch_size.get("height", 16) - patch_w = patch_size.get("width", 16) - else: - patch_h = patch_w = patch_size - num_tokens = (h // patch_h) * (w // patch_w) - result["num_tokens"] = num_tokens - - # Add debug info - result["config_info"] = { - "longest_edge": processor.size.get("longest_edge", 1024), - "patch_size": processor.patch_size, - "image_mean": processor.image_mean, - "image_std": processor.image_std, - } - - return result - - -def generate_for_model(model_key: str, image_paths: list, output_dir: str): - """Generate golden outputs for a specific model.""" - print(f"\nGenerating golden outputs for {model_key}...") - - generator_fn = { - "llava": generate_golden_llava, - "llava_pad": generate_golden_llava_pad, - "llava_next": generate_golden_llava_next, - "qwen2_vl": generate_golden_qwen2_vl, - "qwen3_vl": generate_golden_qwen3_vl, - "phi3_vision": generate_golden_phi3_vision, - "phi4_vision": generate_golden_phi4_vision, - "llama4_vision": generate_golden_llama4_vision, - "pixtral": generate_golden_pixtral, - }.get(model_key) - - if generator_fn is None: - print(f" No generator for {model_key}, skipping") - return - - for image_path in image_paths: - if not os.path.exists(image_path): - print(f" Image not found: {image_path}, skipping") - continue - - image_name = Path(image_path).stem - print(f" Processing {image_name}...") - - try: - data = generator_fn(image_path, output_dir) - if data is not None: - save_golden(model_key, image_name, data, output_dir) - print(f" pixel_values shape: {data['pixel_values'].shape}") - print( - f" pixel_values range: [{data['pixel_values'].min():.4f}, {data['pixel_values'].max():.4f}]" - ) - except Exception as e: - print(f" Error: {e}") - - -def main(): - parser = argparse.ArgumentParser( - description="Generate golden outputs for vision processor testing" - ) - parser.add_argument("--model", "-m", help="Specific model to generate (default: all)") - parser.add_argument("--image", "-i", action="append", help="Specific image path(s)") - parser.add_argument( - "--output-dir", - "-o", - default="crates/multimodal/tests/fixtures/golden", - help="Output directory for golden files", - ) - args = parser.parse_args() - - # Determine which images to use - image_paths = args.image if args.image else DEFAULT_IMAGES - - # Determine which models to generate - if args.model: - if args.model not in MODELS: - print(f"Unknown model: {args.model}") - print(f"Available: {list(MODELS.keys())}") - sys.exit(1) - models_to_generate = [args.model] - else: - models_to_generate = list(MODELS.keys()) - - print(f"Output directory: {args.output_dir}") - print(f"Images: {image_paths}") - print(f"Models: {models_to_generate}") - - # Generate golden outputs - for model_key in models_to_generate: - generate_for_model(model_key, image_paths, args.output_dir) - - print("\nDone!") - - -if __name__ == "__main__": - main() diff --git a/crates/multimodal/src/audio/decode.rs b/crates/multimodal/src/audio/decode.rs deleted file mode 100644 index f7118cd61..000000000 --- a/crates/multimodal/src/audio/decode.rs +++ /dev/null @@ -1,533 +0,0 @@ -//! Audio decode helpers shared by model-specific audio preprocessors. -//! -//! The default path mirrors SMG video decode: use an in-process decoder first, -//! then fall back to an external FFmpeg binary for difficult containers/codecs. - -use std::{ - io::{Cursor, Write}, - mem::size_of, - process::{Output, Stdio}, - sync::OnceLock, - time::{Duration, Instant}, -}; - -use symphonia::{ - core::{ - codecs::audio::AudioDecoderOptions, - errors::Error as SymphoniaError, - formats::{probe::Hint, FormatOptions, TrackType}, - io::MediaSourceStream, - meta::MetadataOptions, - }, - default::{get_codecs, get_probe}, -}; -use tokio::{process::Command, task, time}; -use tracing::debug; - -use crate::error::TransformError; - -const DEFAULT_AUDIO_PROCESS_TIMEOUT: Duration = Duration::from_secs(30); -const DEFAULT_AUDIO_MAX_DECODED_BYTES: usize = 256 * 1024 * 1024; - -static AUDIO_PROCESS_TIMEOUT: OnceLock = OnceLock::new(); -static AUDIO_MAX_DECODED_BYTES: OnceLock = OnceLock::new(); - -#[derive(Debug, Clone, PartialEq)] -pub struct DecodedAudio { - pub samples: Vec, - pub sample_rate: usize, -} - -pub async fn decode_audio_mono_f32(bytes: &[u8]) -> Result { - match audio_decode_backend_override() { - Some("symphonia") => decode_audio_with_symphonia_blocking(bytes).await, - Some("ffmpeg") => decode_audio_with_ffmpeg(bytes).await, - Some(backend) => Err(TransformError::ShapeError(format!( - "unsupported SMG_AUDIO_DECODE_BACKEND={backend}; expected auto, symphonia, or ffmpeg" - ))), - None => match decode_audio_with_symphonia_blocking(bytes).await { - Ok(decoded) => Ok(decoded), - Err(symphonia_error) => { - debug!( - error = %symphonia_error, - "smg_mm_timing audio_decode_auto_symphonia_fallback" - ); - decode_audio_with_ffmpeg(bytes).await.map_err(|ffmpeg_error| { - TransformError::ShapeError(format!( - "Symphonia audio decode failed: {symphonia_error}; ffmpeg fallback failed: {ffmpeg_error}" - )) - }) - } - }, - } -} - -async fn decode_audio_with_symphonia_blocking( - bytes: &[u8], -) -> Result { - let bytes = bytes.to_vec(); - task::spawn_blocking(move || decode_audio_mono_f32_symphonia(&bytes)) - .await - .map_err(|e| TransformError::ShapeError(format!("Symphonia decode task failed: {e}")))? -} - -pub(crate) fn decode_audio_mono_f32_symphonia( - bytes: &[u8], -) -> Result { - decode_audio_mono_f32_symphonia_with_limits( - bytes, - audio_max_decoded_bytes(), - audio_process_timeout(), - ) -} - -fn decode_audio_mono_f32_symphonia_with_limits( - bytes: &[u8], - max_decoded_bytes: usize, - timeout: Duration, -) -> Result { - let started = Instant::now(); - let mut hint = Hint::new(); - if let Some(ext) = audio_extension_hint(bytes) { - hint.with_extension(ext); - } - - let cursor = Cursor::new(bytes.to_vec()); - let media_source = MediaSourceStream::new(Box::new(cursor), Default::default()); - let mut format = get_probe() - .probe( - &hint, - media_source, - FormatOptions::default(), - MetadataOptions::default(), - ) - .map_err(|e| TransformError::ShapeError(format!("Symphonia probe failed: {e}")))?; - - let track = format.default_track(TrackType::Audio).ok_or_else(|| { - TransformError::ShapeError("Symphonia found no supported audio track".to_string()) - })?; - let track_id = track.id; - let audio_params = track - .codec_params - .as_ref() - .and_then(|params| params.audio()) - .ok_or_else(|| { - TransformError::ShapeError( - "Symphonia audio track is missing codec parameters".to_string(), - ) - })?; - let mut decoder = get_codecs() - .make_audio_decoder(audio_params, &AudioDecoderOptions::default()) - .map_err(|e| TransformError::ShapeError(format!("Symphonia decoder failed: {e}")))?; - - let mut sample_rate = audio_params.sample_rate.map(|rate| rate as usize); - let mut mono = Vec::new(); - let mut interleaved = Vec::new(); - loop { - ensure_symphonia_deadline(started, timeout)?; - let Some(packet) = format.next_packet().map_err(|error| { - TransformError::ShapeError(format!("Symphonia packet read failed: {error}")) - })? - else { - break; - }; - if packet.track_id != track_id { - continue; - } - - let audio_buf = match decoder.decode(&packet) { - Ok(decoded) => decoded, - Err(SymphoniaError::DecodeError(_)) => continue, - Err(SymphoniaError::IoError(error)) - if error.kind() == std::io::ErrorKind::UnexpectedEof => - { - break; - } - Err(error) => { - return Err(TransformError::ShapeError(format!( - "Symphonia packet decode failed: {error}" - ))); - } - }; - ensure_symphonia_deadline(started, timeout)?; - - let spec = audio_buf.spec(); - sample_rate = Some(spec.rate() as usize); - let channels = spec.channels().count(); - if channels == 0 { - return Err(TransformError::ShapeError( - "decoded audio has zero channels".to_string(), - )); - } - - let interleaved_samples = audio_buf.samples_interleaved(); - ensure_decoded_sample_limit(0, interleaved_samples, max_decoded_bytes)?; - let additional_samples = interleaved_samples / channels; - ensure_decoded_sample_limit(mono.len(), additional_samples, max_decoded_bytes)?; - mono.try_reserve(additional_samples).map_err(|error| { - TransformError::ShapeError(format!( - "failed to reserve {additional_samples} decoded audio samples: {error}" - )) - })?; - interleaved.resize(interleaved_samples, 0.0); - audio_buf.copy_to_slice_interleaved(&mut interleaved); - for frame in interleaved.chunks_exact(channels) { - mono.push(frame.iter().copied().sum::() / channels as f32); - } - } - - let sample_rate = sample_rate.ok_or_else(|| { - TransformError::ShapeError("decoded audio is missing sample rate".to_string()) - })?; - finish_decoded_audio(mono, sample_rate) -} - -async fn decode_audio_with_ffmpeg(bytes: &[u8]) -> Result { - let input_file = write_temp_audio_file_async(bytes).await?; - let sample_rate = probe_audio_sample_rate(input_file.path()).await?; - // Ask FFmpeg for one sample beyond our limit so a longer stream is - // distinguishable from a valid stream whose size is exactly the limit. - let output_limit = audio_max_decoded_bytes() - .saturating_add(size_of::()) - .to_string(); - - let mut command = Command::new("ffmpeg"); - command - .args(["-hide_banner", "-loglevel", "error", "-nostdin", "-i"]) - .arg(input_file.path()) - .args([ - "-map", - "0:a:0", - "-vn", - "-ac", - "1", - "-fs", - &output_limit, - "-f", - "f32le", - "-sample_fmt", - "flt", - "pipe:1", - ]); - let output = run_audio_command_output(command, "ffmpeg").await?; - if !output.status.success() { - return Err(TransformError::ShapeError(format!( - "ffmpeg failed: {}", - String::from_utf8_lossy(&output.stderr) - ))); - } - ensure_decoded_byte_limit(output.stdout.len(), audio_max_decoded_bytes())?; - if output.stdout.len() % 4 != 0 { - return Err(TransformError::ShapeError(format!( - "ffmpeg f32le output has trailing partial sample: {} bytes", - output.stdout.len() % 4 - ))); - } - - let samples = output - .stdout - .chunks_exact(4) - .map(|bytes| { - let mut sample = [0_u8; size_of::()]; - sample.copy_from_slice(bytes); - f32::from_le_bytes(sample) - }) - .collect(); - finish_decoded_audio(samples, sample_rate) -} - -fn ensure_symphonia_deadline(started: Instant, timeout: Duration) -> Result<(), TransformError> { - if started.elapsed() >= timeout { - return Err(TransformError::ShapeError(format!( - "Symphonia timed out after {:.3} seconds", - timeout.as_secs_f64() - ))); - } - Ok(()) -} - -fn ensure_decoded_sample_limit( - existing_samples: usize, - additional_samples: usize, - max_decoded_bytes: usize, -) -> Result<(), TransformError> { - let total_samples = existing_samples - .checked_add(additional_samples) - .ok_or_else(|| { - TransformError::ShapeError("decoded audio sample count overflow".to_string()) - })?; - let decoded_bytes = total_samples.checked_mul(size_of::()).ok_or_else(|| { - TransformError::ShapeError("decoded audio byte size overflow".to_string()) - })?; - ensure_decoded_byte_limit(decoded_bytes, max_decoded_bytes) -} - -fn ensure_decoded_byte_limit( - decoded_bytes: usize, - max_decoded_bytes: usize, -) -> Result<(), TransformError> { - if decoded_bytes > max_decoded_bytes { - return Err(TransformError::ShapeError(format!( - "decoded audio payload is {decoded_bytes} bytes, exceeding SMG_AUDIO_MAX_DECODED_BYTES={max_decoded_bytes}" - ))); - } - Ok(()) -} - -fn finish_decoded_audio( - samples: Vec, - sample_rate: usize, -) -> Result { - if samples.is_empty() { - return Err(TransformError::ShapeError( - "decoded audio produced no samples".to_string(), - )); - } - Ok(DecodedAudio { - samples, - sample_rate, - }) -} - -async fn probe_audio_sample_rate(input_path: &std::path::Path) -> Result { - let mut command = Command::new("ffprobe"); - command - .args([ - "-v", - "error", - "-select_streams", - "a:0", - "-show_entries", - "stream=sample_rate", - "-of", - "default=noprint_wrappers=1:nokey=1", - ]) - .arg(input_path); - let output = run_audio_command_output(command, "ffprobe").await?; - if !output.status.success() { - return Err(TransformError::ShapeError(format!( - "ffprobe failed: {}", - String::from_utf8_lossy(&output.stderr) - ))); - } - let stdout = String::from_utf8_lossy(&output.stdout); - stdout - .lines() - .find_map(|line| line.trim().parse::().ok()) - .filter(|rate| *rate > 0) - .ok_or_else(|| { - TransformError::ShapeError(format!("failed to parse ffprobe sample rate: {stdout:?}")) - }) -} - -async fn write_temp_audio_file_async( - bytes: &[u8], -) -> Result { - let bytes = bytes.to_vec(); - task::spawn_blocking(move || write_temp_audio_file(&bytes)) - .await - .map_err(|e| TransformError::ShapeError(format!("audio tempfile task failed: {e}")))? -} - -fn write_temp_audio_file(bytes: &[u8]) -> Result { - let mut input_file = tempfile::Builder::new() - .prefix("smg-audio-") - .suffix(audio_temp_suffix(bytes)) - .tempfile() - .map_err(|e| TransformError::ShapeError(format!("audio tempfile failed: {e}")))?; - input_file - .write_all(bytes) - .map_err(|e| TransformError::ShapeError(format!("audio tempfile write failed: {e}")))?; - input_file - .flush() - .map_err(|e| TransformError::ShapeError(format!("audio tempfile flush failed: {e}")))?; - Ok(input_file) -} - -async fn run_audio_command_output( - mut command: Command, - program: &'static str, -) -> Result { - command - .stdin(Stdio::null()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .kill_on_drop(true); - let child = command.spawn().map_err(|e| { - if e.kind() == std::io::ErrorKind::NotFound { - TransformError::ShapeError(format!( - "{program} executable not found; install {program} for audio decode fallback" - )) - } else { - TransformError::ShapeError(format!("{program} spawn failed: {e}")) - } - })?; - - let timeout = audio_process_timeout(); - match time::timeout(timeout, child.wait_with_output()).await { - Ok(Ok(output)) => Ok(output), - Ok(Err(error)) => Err(TransformError::ShapeError(format!( - "{program} wait failed: {error}" - ))), - Err(_) => Err(TransformError::ShapeError(format!( - "{program} timed out after {:.3} seconds", - timeout.as_secs_f64() - ))), - } -} - -fn audio_decode_backend_override() -> Option<&'static str> { - static BACKEND: OnceLock> = OnceLock::new(); - BACKEND - .get_or_init(|| { - std::env::var("SMG_AUDIO_DECODE_BACKEND") - .ok() - .map(|value| value.trim().to_ascii_lowercase()) - .filter(|value| !value.is_empty() && value != "auto") - }) - .as_deref() -} - -fn audio_process_timeout() -> Duration { - *AUDIO_PROCESS_TIMEOUT.get_or_init(|| { - std::env::var("SMG_AUDIO_PROCESS_TIMEOUT_SECS") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|seconds| seconds.is_finite() && *seconds > 0.0) - .map(Duration::from_secs_f64) - .unwrap_or(DEFAULT_AUDIO_PROCESS_TIMEOUT) - }) -} - -fn audio_max_decoded_bytes() -> usize { - *AUDIO_MAX_DECODED_BYTES.get_or_init(|| { - std::env::var("SMG_AUDIO_MAX_DECODED_BYTES") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|bytes| *bytes > 0) - .unwrap_or(DEFAULT_AUDIO_MAX_DECODED_BYTES) - }) -} - -fn audio_temp_suffix(bytes: &[u8]) -> &'static str { - if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && bytes.get(8..12) == Some(b"WAVE") { - return ".wav"; - } - if bytes.starts_with(b"fLaC") { - return ".flac"; - } - if bytes.starts_with(b"ID3") || is_mp3_frame_sync(bytes) { - return ".mp3"; - } - if bytes.starts_with(b"OggS") { - return ".ogg"; - } - if bytes.len() >= 12 && bytes.get(4..8) == Some(b"ftyp") { - return ".m4a"; - } - if bytes.starts_with(&[0x1a, 0x45, 0xdf, 0xa3]) { - return ".webm"; - } - if bytes.len() >= 12 - && bytes.starts_with(b"FORM") - && matches!(bytes.get(8..12), Some(b"AIFF" | b"AIFC")) - { - return ".aiff"; - } - if bytes.len() >= 4 && bytes.starts_with(b"caff") { - return ".caf"; - } - ".audio" -} - -fn audio_extension_hint(bytes: &[u8]) -> Option<&'static str> { - match audio_temp_suffix(bytes) { - ".wav" => Some("wav"), - ".flac" => Some("flac"), - ".mp3" => Some("mp3"), - ".ogg" => Some("ogg"), - ".m4a" => Some("m4a"), - ".webm" => Some("webm"), - ".aiff" => Some("aiff"), - ".caf" => Some("caf"), - _ => None, - } -} - -fn is_mp3_frame_sync(bytes: &[u8]) -> bool { - bytes.len() >= 2 && bytes[0] == 0xff && (bytes[1] & 0xe0) == 0xe0 -} - -#[cfg(test)] -mod tests { - use super::*; - - fn wav_i16_mono(sample_rate: u32, samples: &[i16]) -> Vec { - let data_bytes = samples.len() as u32 * 2; - let mut bytes = Vec::new(); - bytes.extend_from_slice(b"RIFF"); - bytes.extend_from_slice(&(36 + data_bytes).to_le_bytes()); - bytes.extend_from_slice(b"WAVEfmt "); - bytes.extend_from_slice(&16_u32.to_le_bytes()); - bytes.extend_from_slice(&1_u16.to_le_bytes()); - bytes.extend_from_slice(&1_u16.to_le_bytes()); - bytes.extend_from_slice(&sample_rate.to_le_bytes()); - bytes.extend_from_slice(&(sample_rate * 2).to_le_bytes()); - bytes.extend_from_slice(&2_u16.to_le_bytes()); - bytes.extend_from_slice(&16_u16.to_le_bytes()); - bytes.extend_from_slice(b"data"); - bytes.extend_from_slice(&data_bytes.to_le_bytes()); - for sample in samples { - bytes.extend_from_slice(&sample.to_le_bytes()); - } - bytes - } - - #[test] - fn symphonia_decodes_wav_to_mono_f32() { - let wav = wav_i16_mono(16_000, &[0, 16_384, -16_384]); - let decoded = decode_audio_mono_f32_symphonia(&wav).unwrap(); - assert_eq!(decoded.sample_rate, 16_000); - assert_eq!(decoded.samples.len(), 3); - assert!(decoded.samples[0].abs() < 1e-6); - assert!((decoded.samples[1] - 0.5).abs() < 1e-4); - assert!((decoded.samples[2] + 0.5).abs() < 1e-4); - } - - #[test] - fn symphonia_enforces_decoded_byte_limit() { - let wav = wav_i16_mono(16_000, &[0, 1, 2]); - let error = decode_audio_mono_f32_symphonia_with_limits( - &wav, - 2 * size_of::(), - Duration::from_secs(1), - ) - .unwrap_err(); - - assert!(error.to_string().contains("SMG_AUDIO_MAX_DECODED_BYTES=8")); - } - - #[test] - fn symphonia_enforces_decode_deadline() { - let wav = wav_i16_mono(16_000, &[0]); - let error = decode_audio_mono_f32_symphonia_with_limits(&wav, usize::MAX, Duration::ZERO) - .unwrap_err(); - - assert!(error.to_string().contains("Symphonia timed out")); - } - - #[test] - fn empty_decoded_audio_is_rejected() { - let error = finish_decoded_audio(Vec::new(), 16_000).unwrap_err(); - assert!(error.to_string().contains("produced no samples")); - } - - #[test] - fn audio_suffixes_cover_common_containers() { - assert_eq!(audio_temp_suffix(b"fLaC..."), ".flac"); - assert_eq!(audio_temp_suffix(b"ID3..."), ".mp3"); - assert_eq!(audio_temp_suffix(b"OggS..."), ".ogg"); - assert_eq!(audio_temp_suffix(b"\x1a\x45\xdf\xa3..."), ".webm"); - assert_eq!(audio_temp_suffix(b"\0\0\0\x18ftypM4A "), ".m4a"); - } -} diff --git a/crates/multimodal/src/audio/mod.rs b/crates/multimodal/src/audio/mod.rs deleted file mode 100644 index dff8c6d76..000000000 --- a/crates/multimodal/src/audio/mod.rs +++ /dev/null @@ -1,10 +0,0 @@ -//! Audio preprocessing implementations. - -pub mod decode; -pub mod processor; -pub mod processors; -pub(crate) mod transforms; - -pub use decode::{decode_audio_mono_f32, DecodedAudio}; -pub use processor::AudioPreProcessor; -pub use processors::{Qwen3AudioParams, Qwen3AudioProcessor}; diff --git a/crates/multimodal/src/audio/processor.rs b/crates/multimodal/src/audio/processor.rs deleted file mode 100644 index 8e953ca5c..000000000 --- a/crates/multimodal/src/audio/processor.rs +++ /dev/null @@ -1,16 +0,0 @@ -use std::sync::Arc; - -use crate::{encoder_inputs::PreprocessedEncoderInputs, error::TransformError, types::AudioClip}; - -/// Audio preprocessing contract for a model family. -/// -/// The concrete processor for a model is selected by its -/// [`ModelProcessorSpec::audio_processor`](crate::registry::ModelProcessorSpec::audio_processor), -/// which owns audio-processor selection alongside the model's prompt/placeholder -/// logic. -pub trait AudioPreProcessor: Send + Sync { - fn preprocess( - &self, - clips: &[Arc], - ) -> Result; -} diff --git a/crates/multimodal/src/audio/processors/mod.rs b/crates/multimodal/src/audio/processors/mod.rs deleted file mode 100644 index 77d27a3f6..000000000 --- a/crates/multimodal/src/audio/processors/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -//! Model-specific audio preprocessing implementations. - -mod qwen3_audio; - -pub use qwen3_audio::{Qwen3AudioParams, Qwen3AudioProcessor}; diff --git a/crates/multimodal/src/audio/processors/qwen3_audio.rs b/crates/multimodal/src/audio/processors/qwen3_audio.rs deleted file mode 100644 index 0670714a4..000000000 --- a/crates/multimodal/src/audio/processors/qwen3_audio.rs +++ /dev/null @@ -1,643 +0,0 @@ -//! Qwen3 audio preprocessing using a Whisper-compatible log-mel frontend. - -use std::sync::Arc; - -use ndarray::{Array2, Array3}; -use rustfft::{num_complex::Complex32, Fft, FftPlanner}; -use serde_json::Value; - -use crate::{ - audio::{ - transforms::{bandlimited_resample, hann_window, mel_basis}, - AudioPreProcessor, DecodedAudio, - }, - encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}, - error::TransformError, - types::AudioClip, - vision::PreProcessorConfig, -}; - -/// Parameters used by the Qwen3 audio frontend. -#[derive(Debug, Clone, PartialEq)] -pub struct Qwen3AudioParams { - pub sample_rate: usize, - pub n_mels: usize, - pub n_fft: usize, - pub hop_length: usize, - pub n_window: usize, - pub padding_value: f32, - pub max_samples: Option, -} - -impl Default for Qwen3AudioParams { - fn default() -> Self { - Self { - sample_rate: 16_000, - n_mels: 128, - n_fft: 400, - hop_length: 160, - n_window: 50, - padding_value: 0.0, - max_samples: None, - } - } -} - -impl Qwen3AudioParams { - pub fn from_configs(model_config: &Value, preprocessor_config: &PreProcessorConfig) -> Self { - let mut params = Self::default(); - - if let Some(value) = preprocessor_config.get_extra::("sampling_rate") { - params.sample_rate = value; - } - if let Some(value) = preprocessor_config.get_extra::("feature_size") { - params.n_mels = value; - } else if let Some(value) = find_model_usize( - model_config, - &[ - &["thinker_config", "audio_config", "num_mel_bins"], - &["audio_config", "num_mel_bins"], - ], - ) { - params.n_mels = value; - } - if let Some(value) = preprocessor_config.get_extra::("n_fft") { - params.n_fft = value; - } - if let Some(value) = preprocessor_config.get_extra::("hop_length") { - params.hop_length = value; - } - if let Some(value) = preprocessor_config - .get_extra::("n_window") - .or_else(|| { - find_model_usize( - model_config, - &[ - &["thinker_config", "audio_config", "n_window"], - &["audio_config", "n_window"], - ], - ) - }) - { - params.n_window = value; - } - if let Some(value) = preprocessor_config.get_extra::("padding_value") { - params.padding_value = value; - } - // Qwen processors set `padding=true, truncation=false`: n_samples is - // the default padding target, not an input limit. Honor it only when a - // checkpoint or deployment explicitly enables truncation. A custom - // max_samples remains available as an operational hard limit. - params.max_samples = preprocessor_config - .get_extra::("max_samples") - .or_else(|| { - preprocessor_config - .get_extra::("truncation") - .filter(|enabled| *enabled) - .and_then(|_| { - preprocessor_config - .get_extra::("n_samples") - .or_else(|| { - preprocessor_config - .get_extra::("chunk_length") - .and_then(|seconds| seconds.checked_mul(params.sample_rate)) - }) - }) - }); - - params - } -} - -fn find_model_usize(config: &Value, paths: &[&[&str]]) -> Option { - paths.iter().find_map(|path| { - let mut value = config; - for key in *path { - value = value.get(*key)?; - } - value.as_u64().and_then(|value| usize::try_from(value).ok()) - }) -} - -#[derive(Debug, Clone)] -pub struct Qwen3AudioProcessor { - params: Qwen3AudioParams, -} - -impl Default for Qwen3AudioProcessor { - fn default() -> Self { - Self::new() - } -} - -impl Qwen3AudioProcessor { - pub fn new() -> Self { - Self { - params: Qwen3AudioParams::default(), - } - } - - pub fn with_params(params: Qwen3AudioParams) -> Self { - Self { params } - } - - pub fn from_configs(model_config: &Value, preprocessor_config: &PreProcessorConfig) -> Self { - Self::with_params(Qwen3AudioParams::from_configs( - model_config, - preprocessor_config, - )) - } - - pub fn params(&self) -> &Qwen3AudioParams { - &self.params - } - - pub fn preprocess_decoded_clips( - &self, - clips: Vec, - ) -> Result { - self.validate_params()?; - if clips.is_empty() { - return Err(TransformError::EmptyBatch); - } - - let mut waveforms = Vec::with_capacity(clips.len()); - for (clip_index, clip) in clips.into_iter().enumerate() { - if clip.sample_rate == 0 { - return Err(TransformError::ShapeError( - "decoded audio sample rate must be positive".to_string(), - )); - } - if clip.samples.is_empty() { - return Err(TransformError::ShapeError( - "decoded audio contains no samples".to_string(), - )); - } - if clip.samples.iter().any(|sample| !sample.is_finite()) { - return Err(TransformError::ShapeError( - "decoded audio contains a non-finite sample".to_string(), - )); - } - let samples = if clip.sample_rate == self.params.sample_rate { - clip.samples - } else { - bandlimited_resample(&clip.samples, clip.sample_rate, self.params.sample_rate)? - }; - let mut samples = samples; - if let Some(max_samples) = self.params.max_samples { - samples.truncate(max_samples); - } - if samples.is_empty() { - return Err(TransformError::ShapeError( - "decoded audio contains no samples after truncation".to_string(), - )); - } - if samples.len() < self.params.hop_length { - return Err(TransformError::ShapeError(format!( - "Qwen3 audio clip at batch index {clip_index} has {} samples after resampling and truncation; at least {} are required", - samples.len(), - self.params.hop_length - ))); - } - waveforms.push(samples); - } - - let max_samples = waveforms.iter().map(Vec::len).max().unwrap_or(0); - // Whisper's centered STFT yields floor(samples / hop) + 1 frames and - // then drops the final frame. This matches padding=True with the batch - // padded to its longest waveform. - let max_frames = max_samples / self.params.hop_length; - if max_frames == 0 { - return Err(TransformError::ShapeError(format!( - "Qwen3 audio requires at least {} samples after resampling", - self.params.hop_length - ))); - } - - let batch_size = waveforms.len(); - let feature_values = batch_size - .checked_mul(self.params.n_mels) - .and_then(|value| value.checked_mul(max_frames)) - .ok_or_else(|| { - TransformError::ShapeError("Qwen3 audio feature size overflow".to_string()) - })?; - let mut all_features = Vec::with_capacity(feature_values); - let mut attention_mask = Vec::with_capacity(batch_size * max_frames); - let mut feature_lengths = Vec::with_capacity(batch_size); - let mut token_counts = Vec::with_capacity(batch_size); - let mut item_sizes = Vec::with_capacity(batch_size); - let mut planner = FftPlanner::::new(); - let fft = planner.plan_fft_forward(self.params.n_fft); - - for waveform in waveforms { - let original_samples = waveform.len(); - let feature_length = (original_samples / self.params.hop_length).min(max_frames); - let mut padded = waveform; - padded.resize(max_samples, self.params.padding_value); - let features = whisper_log_mel(&padded, max_frames, &self.params, fft.as_ref())?; - all_features.extend(features.into_raw_vec_and_offset().0); - - attention_mask.extend((0..max_frames).map(|frame| i64::from(frame < feature_length))); - feature_lengths.push(feature_length as i64); - token_counts.push(qwen3_audio_output_length( - feature_length, - self.params.n_window, - )); - item_sizes.push((self.params.n_mels as u32, feature_length as u32)); - } - - let encoder_input = - Array3::from_shape_vec((batch_size, self.params.n_mels, max_frames), all_features) - .map_err(|error| { - TransformError::ShapeError(format!( - "failed to create Qwen3 audio input [{batch_size}, {}, {max_frames}]: {error}", - self.params.n_mels - )) - })?; - - Ok( - PreprocessedEncoderInputs::new(encoder_input, token_counts, item_sizes) - .with_extra( - "feature_attention_mask", - ModelSpecificValue::int_2d(attention_mask, batch_size, max_frames), - ) - .with_extra( - "audio_feature_lengths", - ModelSpecificValue::int_1d(feature_lengths), - ), - ) - } - - pub fn preprocess_decoded(&self, decoded: DecodedAudio) -> Result, TransformError> { - let output = self.preprocess_decoded_clips(vec![decoded])?; - output - .encoder_input - .into_dimensionality::() - .map_err(|error| TransformError::ShapeError(error.to_string()))? - .index_axis_move(ndarray::Axis(0), 0) - .into_dimensionality::() - .map_err(|error| TransformError::ShapeError(error.to_string())) - } - - fn validate_params(&self) -> Result<(), TransformError> { - if self.params.sample_rate == 0 - || self.params.n_mels == 0 - || self.params.n_fft == 0 - || self.params.hop_length == 0 - || self.params.n_window == 0 - { - return Err(TransformError::ShapeError( - "Qwen3 audio sample rate, mel bins, FFT size, hop length, and window size must be positive" - .to_string(), - )); - } - if self.params.n_fft < self.params.hop_length { - return Err(TransformError::ShapeError(format!( - "Qwen3 audio n_fft ({}) must be at least hop_length ({})", - self.params.n_fft, self.params.hop_length - ))); - } - if !self.params.padding_value.is_finite() { - return Err(TransformError::ShapeError( - "Qwen3 audio padding_value must be finite".to_string(), - )); - } - if self.params.max_samples == Some(0) { - return Err(TransformError::ShapeError( - "Qwen3 audio max_samples must be positive".to_string(), - )); - } - Ok(()) - } -} - -impl AudioPreProcessor for Qwen3AudioProcessor { - fn preprocess( - &self, - clips: &[Arc], - ) -> Result { - self.preprocess_decoded_clips(clips.iter().map(|clip| clip.decoded().clone()).collect()) - } -} - -fn qwen_audio_cnn_output_length(mut input_length: usize) -> usize { - for _ in 0..3 { - input_length = input_length.div_ceil(2); - } - input_length -} - -/// Output tokens produced by Qwen's chunked audio encoder for a log-mel length. -fn qwen3_audio_output_length(input_length: usize, n_window: usize) -> usize { - debug_assert!(n_window > 0); - let chunk_size = 2 * n_window; - let full_windows = input_length / chunk_size; - let remainder = input_length % chunk_size; - full_windows * qwen_audio_cnn_output_length(chunk_size) - + qwen_audio_cnn_output_length(remainder) -} - -fn whisper_log_mel( - samples: &[f32], - frame_count: usize, - params: &Qwen3AudioParams, - fft: &dyn Fft, -) -> Result, TransformError> { - let center_pad = params.n_fft / 2; - let padded = reflect_pad(samples, center_pad); - let fft_bins = params.n_fft / 2 + 1; - let window = hann_window(params.n_fft); - let mel_filters = mel_basis(params.sample_rate, params.n_fft, params.n_mels); - let mut buffer = vec![Complex32::new(0.0, 0.0); params.n_fft]; - let mut output = vec![0.0_f32; params.n_mels * frame_count]; - - for frame in 0..frame_count { - let start = frame * params.hop_length; - let end = start + params.n_fft; - let frame_samples = padded.get(start..end).ok_or_else(|| { - TransformError::ShapeError(format!( - "Qwen3 audio STFT frame {frame} lies outside padded waveform" - )) - })?; - for index in 0..params.n_fft { - buffer[index] = Complex32::new(frame_samples[index] * window[index], 0.0); - } - fft.process(&mut buffer); - - for mel in 0..params.n_mels { - let filter = &mel_filters[mel * fft_bins..(mel + 1) * fft_bins]; - let mut value = 0.0_f32; - for bin in 0..fft_bins { - let fft_value = buffer[bin]; - let power = fft_value - .re - .mul_add(fft_value.re, fft_value.im * fft_value.im); - value = filter[bin].mul_add(power, value); - } - output[mel * frame_count + frame] = value.max(1e-10).log10(); - } - } - - let peak = output.iter().copied().fold(f32::NEG_INFINITY, f32::max); - let floor = peak - 8.0; - for value in &mut output { - *value = (value.max(floor) + 4.0) / 4.0; - } - - Array2::from_shape_vec((params.n_mels, frame_count), output).map_err(|error| { - TransformError::ShapeError(format!( - "failed to create Qwen log-mel input [{}, {frame_count}]: {error}", - params.n_mels - )) - }) -} - -fn reflect_pad(samples: &[f32], padding: usize) -> Vec { - if padding == 0 { - return samples.to_vec(); - } - if samples.len() == 1 { - return vec![samples[0]; samples.len() + 2 * padding]; - } - - let mut padded = Vec::with_capacity(samples.len() + 2 * padding); - for position in -(padding as isize)..(samples.len() + padding) as isize { - padded.push(samples[reflect_index(position, samples.len())]); - } - padded -} - -fn reflect_index(mut index: isize, len: usize) -> usize { - let last = len as isize - 1; - while index < 0 || index > last { - if index < 0 { - index = -index; - } - if index > last { - index = 2 * last - index; - } - } - index as usize -} - -#[cfg(test)] -mod tests { - use super::*; - - fn decoded(samples: usize) -> DecodedAudio { - DecodedAudio { - samples: vec![0.0; samples], - sample_rate: 16_000, - } - } - - #[test] - fn silence_matches_whisper_normalization() { - let features = Qwen3AudioProcessor::new() - .preprocess_decoded(decoded(1600)) - .unwrap(); - assert_eq!(features.shape(), &[128, 10]); - assert!(features.iter().all(|value| (*value + 1.5).abs() < 1e-6)); - } - - #[test] - fn log_mel_matches_numpy_whisper_reference() { - let samples = (0..1000) - .map(|index| ((index % 23) as f32 - 11.0) / 32.0) - .collect(); - let features = Qwen3AudioProcessor::new() - .preprocess_decoded(DecodedAudio { - samples, - sample_rate: 16_000, - }) - .unwrap(); - assert_eq!(features.shape(), &[128, 6]); - - for ((mel, frame), expected) in [ - ((0, 0), 0.71467185), - ((10, 0), 0.735_666_9), - ((20, 1), 0.36943567), - ((32, 3), 0.79836977), - ((64, 5), -0.582_248_9), - ((100, 5), 0.60166824), - ((127, 5), 0.52557164), - ] { - assert!( - (features[[mel, frame]] - expected).abs() < 2e-4, - "log-mel mismatch at ({mel}, {frame}): {} vs {expected}", - features[[mel, frame]] - ); - } - let sum: f64 = features.iter().map(|&value| f64::from(value)).sum(); - assert!((sum - 125.08224487).abs() < 0.02, "feature sum {sum}"); - } - - #[test] - fn log_mel_boundary_matches_hf_whisper_reference() { - let samples = (0..1600) - .map(|index| ((index % 23) as f32 - 11.0) / 32.0) - .collect(); - let features = Qwen3AudioProcessor::new() - .preprocess_decoded(DecodedAudio { - samples, - sample_rate: 16_000, - }) - .unwrap(); - - assert_eq!(features.shape(), &[128, 10]); - for (mel, expected) in [(64, 0.017_234_564), (100, 0.601_944_7), (127, 0.525_99)] { - assert!( - (features[[mel, 9]] - expected).abs() < 2e-4, - "last-frame log-mel mismatch at ({mel}, 9): {} vs {expected}", - features[[mel, 9]] - ); - } - } - - #[test] - fn batches_variable_lengths_with_feature_mask() { - let output = Qwen3AudioProcessor::new() - .preprocess_decoded_clips(vec![decoded(1000), decoded(801)]) - .unwrap(); - - assert_eq!(output.encoder_input.shape(), &[2, 128, 6]); - assert_eq!(output.feature_token_counts, vec![1, 1]); - assert_eq!(output.item_sizes, vec![(128, 6), (128, 5)]); - assert!(matches!( - output.model_specific.get("audio_feature_lengths"), - Some(ModelSpecificValue::IntTensor { data, shape }) - if data == &vec![6, 5] && shape == &vec![2] - )); - assert!(matches!( - output.model_specific.get("feature_attention_mask"), - Some(ModelSpecificValue::IntTensor { data, shape }) - if data == &vec![1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0] - && shape == &vec![2, 6] - )); - } - - #[test] - fn rejects_sub_hop_clip_in_mixed_batch() { - let error = Qwen3AudioProcessor::new() - .preprocess_decoded_clips(vec![decoded(320), decoded(159)]) - .unwrap_err(); - - match error { - TransformError::ShapeError(message) => assert_eq!( - message, - "Qwen3 audio clip at batch index 1 has 159 samples after resampling and truncation; at least 160 are required" - ), - other => panic!("expected a shape error, got {other}"), - } - } - - #[test] - fn accepts_per_clip_hop_length_boundary() { - let processor = Qwen3AudioProcessor::new(); - let expected_token_counts = vec![ - qwen3_audio_output_length( - 320 / processor.params().hop_length, - processor.params().n_window, - ), - qwen3_audio_output_length( - 160 / processor.params().hop_length, - processor.params().n_window, - ), - ]; - assert_eq!(expected_token_counts, vec![1, 1]); - - let output = processor - .preprocess_decoded_clips(vec![decoded(320), decoded(160)]) - .unwrap(); - - assert_eq!(output.encoder_input.shape(), &[2, 128, 2]); - assert_eq!(output.feature_token_counts, expected_token_counts); - assert_eq!(output.item_sizes, vec![(128, 2), (128, 1)]); - } - - #[test] - fn reads_preprocessor_and_nested_model_config() { - let preprocessor = PreProcessorConfig::from_json( - r#"{"sampling_rate": 16000, "feature_size": 64, "n_fft": 320, "hop_length": 80, "n_samples": 4800000}"#, - ) - .unwrap(); - let params = Qwen3AudioParams::from_configs( - &serde_json::json!({ - "thinker_config": {"audio_config": {"num_mel_bins": 96, "n_window": 64}} - }), - &preprocessor, - ); - assert_eq!(params.sample_rate, 16_000); - assert_eq!(params.n_mels, 64); - assert_eq!(params.n_fft, 320); - assert_eq!(params.hop_length, 80); - assert_eq!(params.n_window, 64); - assert_eq!(params.max_samples, None); - - let truncating_preprocessor = PreProcessorConfig::from_json( - r#"{"sampling_rate": 16000, "n_samples": 4800000, "truncation": true}"#, - ) - .unwrap(); - let params = - Qwen3AudioParams::from_configs(&serde_json::json!({}), &truncating_preprocessor); - assert_eq!(params.max_samples, Some(4_800_000)); - - let params = Qwen3AudioParams::from_configs( - &serde_json::json!({ - "thinker_config": {"audio_config": {"num_mel_bins": 96}} - }), - &PreProcessorConfig::default(), - ); - assert_eq!(params.n_mels, 96); - assert_eq!(params.max_samples, None); - } - - #[test] - fn n_samples_is_padding_target_not_implicit_truncation() { - let preprocessor = PreProcessorConfig::from_json(r#"{"n_samples": 320}"#).unwrap(); - let processor = Qwen3AudioProcessor::from_configs(&serde_json::json!({}), &preprocessor); - let output = processor - .preprocess_decoded_clips(vec![decoded(800)]) - .unwrap(); - - assert_eq!(output.encoder_input.shape(), &[1, 128, 5]); - assert_eq!(output.item_sizes, vec![(128, 5)]); - } - - #[test] - fn truncates_to_explicit_audio_limit() { - let processor = Qwen3AudioProcessor::with_params(Qwen3AudioParams { - max_samples: Some(320), - ..Default::default() - }); - let output = processor - .preprocess_decoded_clips(vec![decoded(800)]) - .unwrap(); - - assert_eq!(output.encoder_input.shape(), &[1, 128, 2]); - assert_eq!(output.feature_token_counts, vec![1]); - assert_eq!(output.item_sizes, vec![(128, 2)]); - } - - #[test] - fn qwen_chunked_encoder_output_lengths_match_reference_formula() { - assert_eq!(qwen3_audio_output_length(0, 50), 0); - assert_eq!(qwen3_audio_output_length(1, 50), 1); - assert_eq!(qwen3_audio_output_length(8, 50), 1); - assert_eq!(qwen3_audio_output_length(9, 50), 2); - assert_eq!(qwen3_audio_output_length(99, 50), 13); - assert_eq!(qwen3_audio_output_length(100, 50), 13); - assert_eq!(qwen3_audio_output_length(101, 50), 14); - assert_eq!(qwen3_audio_output_length(3000, 50), 390); - assert_eq!(qwen3_audio_output_length(17, 4), 3); - } - - #[test] - fn reflect_padding_matches_numpy_convention() { - assert_eq!( - reflect_pad(&[1.0, 2.0, 3.0], 2), - vec![3.0, 2.0, 1.0, 2.0, 3.0, 2.0, 1.0] - ); - assert_eq!(reflect_pad(&[2.0], 2), vec![2.0; 5]); - } -} diff --git a/crates/multimodal/src/audio/transforms.rs b/crates/multimodal/src/audio/transforms.rs deleted file mode 100644 index 25c7dd909..000000000 --- a/crates/multimodal/src/audio/transforms.rs +++ /dev/null @@ -1,242 +0,0 @@ -//! Shared digital-signal-processing transforms for audio frontends. - -use std::{f32::consts::PI as PI_F32, f64::consts::PI}; - -use crate::error::TransformError; - -pub(super) fn hann_window(window_size: usize) -> Vec { - (0..window_size) - .map(|i| (0.5 - 0.5 * (2.0 * PI * i as f64 / window_size as f64).cos()) as f32) - .collect() -} - -fn hz_to_mel(frequency: f64) -> f64 { - let f_sp = 200.0 / 3.0; - let min_log_hz = 1000.0; - let min_log_mel = min_log_hz / f_sp; - let logstep = 6.4_f64.ln() / 27.0; - if frequency >= min_log_hz { - min_log_mel + (frequency / min_log_hz).ln() / logstep - } else { - frequency / f_sp - } -} - -fn mel_to_hz(mel: f64) -> f64 { - let f_sp = 200.0 / 3.0; - let min_log_hz = 1000.0; - let min_log_mel = min_log_hz / f_sp; - let logstep = 6.4_f64.ln() / 27.0; - if mel >= min_log_mel { - min_log_hz * (logstep * (mel - min_log_mel)).exp() - } else { - mel * f_sp - } -} - -/// Build a Slaney-normalized mel filter bank. -pub(super) fn mel_basis(sample_rate: usize, n_fft: usize, n_mels: usize) -> Vec { - let fft_bins = n_fft / 2 + 1; - let mut fft_freqs = Vec::with_capacity(fft_bins); - for bin in 0..fft_bins { - fft_freqs.push(bin as f64 * sample_rate as f64 / n_fft as f64); - } - - let mel_min = hz_to_mel(0.0); - let mel_max = hz_to_mel(sample_rate as f64 / 2.0); - let mut mel_edges = Vec::with_capacity(n_mels + 2); - for i in 0..n_mels + 2 { - let t = i as f64 / (n_mels + 1) as f64; - mel_edges.push(mel_to_hz(mel_min + (mel_max - mel_min) * t)); - } - - let mel_widths: Vec = mel_edges.windows(2).map(|w| w[1] - w[0]).collect(); - let mut weights = vec![0.0_f32; n_mels * fft_bins]; - for mel in 0..n_mels { - let enorm = 2.0 / (mel_edges[mel + 2] - mel_edges[mel]); - for (bin, &freq) in fft_freqs.iter().enumerate() { - let lower = (freq - mel_edges[mel]) / mel_widths[mel]; - let upper = (mel_edges[mel + 2] - freq) / mel_widths[mel + 1]; - weights[mel * fft_bins + bin] = lower.min(upper).max(0.0).mul_add(enorm, 0.0) as f32; - } - } - weights -} - -/// Match torchaudio's default `functional.resample`: band-limited sinc -/// interpolation with a Hann window, filter width 6, and rolloff 0.99. -pub(super) fn bandlimited_resample( - samples: &[f32], - src_sample_rate: usize, - dst_sample_rate: usize, -) -> Result, TransformError> { - const LOWPASS_FILTER_WIDTH: f32 = 6.0; - const ROLLOFF: f32 = 0.99; - - if src_sample_rate == 0 || dst_sample_rate == 0 { - return Err(TransformError::ShapeError( - "audio resampling rates must be positive".to_string(), - )); - } - if samples.is_empty() || src_sample_rate == dst_sample_rate { - return Ok(samples.to_vec()); - } - - let gcd = greatest_common_divisor(src_sample_rate, dst_sample_rate); - let orig_freq = src_sample_rate / gcd; - let new_freq = dst_sample_rate / gcd; - let base_freq = orig_freq.min(new_freq) as f32 * ROLLOFF; - let width = (LOWPASS_FILTER_WIDTH * orig_freq as f32 / base_freq).ceil() as usize; - let kernel_len = width - .checked_mul(2) - .and_then(|value| value.checked_add(orig_freq)) - .ok_or_else(|| TransformError::ShapeError("audio resample kernel is too large".into()))?; - let kernel_values = new_freq - .checked_mul(kernel_len) - .ok_or_else(|| TransformError::ShapeError("audio resample kernel size overflow".into()))?; - let mut kernels = Vec::new(); - kernels.try_reserve_exact(kernel_values).map_err(|error| { - TransformError::ShapeError(format!("failed to allocate audio resample kernel: {error}")) - })?; - - let orig_freq_f32 = orig_freq as f32; - let new_freq_f32 = new_freq as f32; - let scale = base_freq / orig_freq_f32; - for phase in 0..new_freq { - for kernel_index in 0..kernel_len { - let idx = (kernel_index as f32 - width as f32) / orig_freq_f32; - let mut t = (idx - phase as f32 / new_freq_f32) * base_freq; - t = t.clamp(-LOWPASS_FILTER_WIDTH, LOWPASS_FILTER_WIDTH); - let window = (t * PI_F32 / LOWPASS_FILTER_WIDTH / 2.0).cos().powi(2); - let radians = t * PI_F32; - let sinc = if radians == 0.0 { - 1.0 - } else { - radians.sin() / radians - }; - kernels.push(sinc * window * scale); - } - } - - let target_len_u128 = (samples.len() as u128 * new_freq as u128).div_ceil(orig_freq as u128); - let target_len = usize::try_from(target_len_u128).map_err(|_| { - TransformError::ShapeError("resampled audio length exceeds usize".to_string()) - })?; - let mut output = Vec::new(); - output.try_reserve_exact(target_len).map_err(|error| { - TransformError::ShapeError(format!("failed to allocate resampled audio: {error}")) - })?; - - for block in 0..samples.len().div_ceil(orig_freq) { - let input_start = block * orig_freq; - for phase in 0..new_freq { - if output.len() == target_len { - return Ok(output); - } - let kernel = &kernels[phase * kernel_len..(phase + 1) * kernel_len]; - let mut value = 0.0_f32; - for (kernel_index, &coefficient) in kernel.iter().enumerate() { - let padded_index = input_start + kernel_index; - if padded_index >= width { - let sample_index = padded_index - width; - if let Some(&sample) = samples.get(sample_index) { - value = sample.mul_add(coefficient, value); - } - } - } - output.push(value); - } - } - - Ok(output) -} - -fn greatest_common_divisor(mut lhs: usize, mut rhs: usize) -> usize { - while rhs != 0 { - (lhs, rhs) = (rhs, lhs % rhs); - } - lhs -} - -#[cfg(test)] -mod tests { - use super::bandlimited_resample; - use crate::error::TransformError; - - #[test] - fn bandlimited_resample_matches_torchaudio_golden_vector() { - let input = [0.0, 0.25, -0.5, 0.75, -1.0, 0.5, 0.125, -0.25]; - // torchaudio 2.11 functional.resample(input, 8000, 16000), using - // the default Hann-windowed sinc kernel. - let expected = [ - 0.012_859_77, - 0.347_330_45, - 0.230_903_15, - -0.388_803_3, - -0.476_089_15, - 0.325_961_17, - 0.724_112_45, - -0.117_642_37, - -0.975_495_4, - -0.540_876_27, - 0.479_874_4, - 0.698_470_23, - 0.138_904_54, - -0.281_137_47, - -0.257_479_25, - -0.095_786_646, - ]; - - let actual = bandlimited_resample(&input, 8_000, 16_000).unwrap(); - - assert_eq!(actual.len(), expected.len()); - for (index, (actual, expected)) in actual.iter().zip(expected).enumerate() { - assert!( - (actual - expected).abs() < 2e-5, - "resampled sample {index}: {actual} != {expected}" - ); - } - } - - #[test] - fn bandlimited_resample_uses_torchaudio_output_length_formula() { - for (input_len, source_rate, destination_rate) in [ - (17, 16_000, 8_000), - (8, 8_000, 11_025), - (161, 48_000, 44_100), - ] { - let input = vec![0.0; input_len]; - let output = bandlimited_resample(&input, source_rate, destination_rate).unwrap(); - let expected_len = - (input_len as u128 * destination_rate as u128).div_ceil(source_rate as u128); - - assert_eq!( - output.len(), - expected_len as usize, - "unexpected output length for {source_rate} Hz to {destination_rate} Hz" - ); - } - } - - #[test] - fn bandlimited_resample_handles_empty_and_passthrough_inputs() { - assert_eq!( - bandlimited_resample(&[], 8_000, 16_000).unwrap(), - Vec::::new() - ); - - let input = [0.25, -0.5, 1.0]; - assert_eq!(bandlimited_resample(&input, 16_000, 16_000).unwrap(), input); - } - - #[test] - fn bandlimited_resample_rejects_zero_rates() { - for (source_rate, destination_rate) in [(0, 16_000), (16_000, 0)] { - assert!(matches!( - bandlimited_resample(&[0.0], source_rate, destination_rate), - Err(TransformError::ShapeError(message)) - if message == "audio resampling rates must be positive" - )); - } - } -} diff --git a/crates/multimodal/src/encoder_inputs.rs b/crates/multimodal/src/encoder_inputs.rs deleted file mode 100644 index ed4600412..000000000 --- a/crates/multimodal/src/encoder_inputs.rs +++ /dev/null @@ -1,328 +0,0 @@ -//! Shared encoder-input types for all encoder-backed modalities. - -use std::{borrow::Cow, collections::HashMap}; - -use anyhow::{Context, Result as AnyhowResult}; -use ndarray::{Array, ArrayD, Dimension}; - -use crate::types::FieldLayout; - -/// Model-specific auxiliary output values. -#[derive(Debug, Clone)] -pub enum ModelSpecificValue { - /// A tensor with shape information (data as flat vec, shape as dims) - Tensor { data: Vec, shape: Vec }, - - /// A tensor of integers (e.g., aspect_ratio_ids) - IntTensor { data: Vec, shape: Vec }, - - /// A tensor of unsigned integers (e.g., image_grid_thw) - UintTensor { data: Vec, shape: Vec }, - - /// Simple integer value - Int(i64), - - /// Simple float value - Float(f64), - - /// List of integers - IntVec(Vec), - - /// List of unsigned integers - UintVec(Vec), - - /// List of floats - FloatVec(Vec), - - /// List of tuples (e.g., media item sizes) - TupleVec(Vec<(u32, u32)>), - - /// Boolean flag - Bool(bool), -} - -impl ModelSpecificValue { - /// Create a 1D uint tensor from a vector. - pub fn uint_1d(data: Vec) -> Self { - let len = data.len(); - Self::UintTensor { - data, - shape: vec![len], - } - } - - /// Create a 2D uint tensor. - pub fn uint_2d(data: Vec, rows: usize, cols: usize) -> Self { - Self::UintTensor { - data, - shape: vec![rows, cols], - } - } - - /// Create a 1D int tensor from a vector. - pub fn int_1d(data: Vec) -> Self { - let len = data.len(); - Self::IntTensor { - data, - shape: vec![len], - } - } - - /// Create a 2D int tensor. - pub fn int_2d(data: Vec, rows: usize, cols: usize) -> Self { - Self::IntTensor { - data, - shape: vec![rows, cols], - } - } - - /// Interpret this value as per-item flat sizes. - pub fn as_flat_sizes(&self) -> AnyhowResult> { - match self { - Self::IntTensor { data, .. } => data - .iter() - .map(|&v| usize::try_from(v).context("negative flat size")) - .collect(), - Self::UintTensor { data, .. } => Ok(data.iter().map(|&v| v as usize).collect()), - Self::IntVec(values) => values - .iter() - .map(|&v| usize::try_from(v).context("negative flat size")) - .collect(), - Self::UintVec(values) => Ok(values.iter().map(|&v| v as usize).collect()), - _ => Err(anyhow::anyhow!("unsupported flat sizes value type")), - } - } - - /// Slice item-batched metadata along the first dimension. - pub fn slice_first_dim(&self, start: usize, len: usize) -> AnyhowResult { - match self { - Self::Tensor { data, shape } => { - let (data, shape) = slice_tensor_first_dim(data, shape, start, len)?; - Ok(Self::Tensor { data, shape }) - } - Self::IntTensor { data, shape } => { - let (data, shape) = slice_tensor_first_dim(data, shape, start, len)?; - Ok(Self::IntTensor { data, shape }) - } - Self::UintTensor { data, shape } => { - let (data, shape) = slice_tensor_first_dim(data, shape, start, len)?; - Ok(Self::UintTensor { data, shape }) - } - Self::IntVec(values) => Ok(Self::IntVec(slice_1d(values, start, len)?.to_vec())), - Self::UintVec(values) => Ok(Self::UintVec(slice_1d(values, start, len)?.to_vec())), - Self::FloatVec(values) => Ok(Self::FloatVec(slice_1d(values, start, len)?.to_vec())), - Self::TupleVec(values) => Ok(Self::TupleVec(slice_1d(values, start, len)?.to_vec())), - _ => Ok(self.clone()), - } - } -} - -fn slice_tensor_first_dim( - data: &[T], - shape: &[usize], - start: usize, - len: usize, -) -> AnyhowResult<(Vec, Vec)> { - let first_dim = *shape - .first() - .ok_or_else(|| anyhow::anyhow!("cannot slice scalar tensor"))?; - let end = start - .checked_add(len) - .ok_or_else(|| anyhow::anyhow!("tensor slice range overflow"))?; - anyhow::ensure!( - end <= first_dim, - "tensor first-dimension slice {start}..{end} exceeds {first_dim}" - ); - let row_width = shape[1..] - .iter() - .try_fold(1usize, |acc, &dim| acc.checked_mul(dim)) - .ok_or_else(|| anyhow::anyhow!("tensor row width overflow"))?; - let data_start = start - .checked_mul(row_width) - .ok_or_else(|| anyhow::anyhow!("tensor data start overflow"))?; - let data_len = len - .checked_mul(row_width) - .ok_or_else(|| anyhow::anyhow!("tensor data length overflow"))?; - let data_end = data_start - .checked_add(data_len) - .ok_or_else(|| anyhow::anyhow!("tensor data end overflow"))?; - anyhow::ensure!( - data_end <= data.len(), - "tensor slice data range {data_start}..{data_end} exceeds {}", - data.len() - ); - let mut new_shape = shape.to_vec(); - new_shape[0] = len; - Ok((data[data_start..data_end].to_vec(), new_shape)) -} - -fn slice_1d(values: &[T], start: usize, len: usize) -> AnyhowResult<&[T]> { - let end = start - .checked_add(len) - .ok_or_else(|| anyhow::anyhow!("slice range overflow"))?; - values - .get(start..end) - .ok_or_else(|| anyhow::anyhow!("slice range {start}..{end} exceeds {}", values.len())) -} - -/// Preprocessed encoder inputs ready for model consumption. -#[derive(Debug, Clone)] -pub struct PreprocessedEncoderInputs { - /// Primary encoder input as a dynamic-dimensional float32 tensor. - pub encoder_input: ArrayD, - - /// Number of encoder feature tokens per media item in the batch. - pub feature_token_counts: Vec, - - /// Modality-specific item size metadata before preprocessing. - /// - /// The exact tuple order follows each processor/model contract. Auxiliary - /// shape tensors that need a fixed order should be emitted in - /// `model_specific`. - pub item_sizes: Vec<(u32, u32)>, - - /// Model-specific auxiliary outputs. - pub model_specific: HashMap, -} - -impl PreprocessedEncoderInputs { - /// Create encoder inputs backed by a tensor of any dimensionality. - pub fn new( - encoder_input: Array, - feature_token_counts: Vec, - item_sizes: Vec<(u32, u32)>, - ) -> Self { - Self { - encoder_input: encoder_input.into_dyn(), - feature_token_counts, - item_sizes, - model_specific: HashMap::new(), - } - } - - /// Add a model-specific value. - pub fn with_extra(mut self, key: impl Into, value: ModelSpecificValue) -> Self { - self.model_specific.insert(key.into(), value); - self - } - - /// Get the number of media items represented by this preprocessed batch. - pub fn batch_size(&self) -> usize { - self.item_sizes.len() - } - - /// Get the number of dimensions of encoder_input. - pub fn ndim(&self) -> usize { - self.encoder_input.ndim() - } - - /// Get total number of encoder feature tokens across all media items. - pub fn total_feature_tokens(&self) -> usize { - self.feature_token_counts.iter().sum() - } - - /// Get the primary encoder input as a flat f32 slice without copying if possible. - pub fn encoder_input_flat(&self) -> Cow<'_, [f32]> { - match self.encoder_input.as_slice() { - Some(slice) => Cow::Borrowed(slice), - None => Cow::Owned(self.encoder_input.iter().copied().collect()), - } - } - - /// Get the shape of the primary encoder input as a vector. - pub fn encoder_input_shape(&self) -> Vec { - self.encoder_input.shape().to_vec() - } - - /// Extract batched tensor keys from explicit field layout declarations. - pub fn batched_keys(layouts: &HashMap) -> Vec { - layouts - .iter() - .filter(|(_, layout)| matches!(layout, FieldLayout::Batched)) - .map(|(key, _)| key.clone()) - .collect() - } - - /// Extract flat-slicing tensor keys from explicit field layout declarations. - /// - /// Returns a map of tensor name to sizes tensor name. - pub fn flat_keys(layouts: &HashMap) -> HashMap { - layouts - .iter() - .filter_map(|(key, layout)| match layout { - FieldLayout::Flat { sizes_key } => Some((key.clone(), sizes_key.clone())), - FieldLayout::Batched => None, - }) - .collect() - } -} - -#[cfg(test)] -mod tests { - use ndarray::Array4; - - use super::*; - - #[test] - fn encoder_input_accessors_are_modality_neutral() { - let inputs = PreprocessedEncoderInputs::new( - Array4::::zeros((2, 3, 4, 5)), - vec![6, 7], - vec![(4, 5), (8, 9)], - ); - - assert_eq!(inputs.batch_size(), 2); - assert_eq!(inputs.ndim(), 4); - assert_eq!(inputs.total_feature_tokens(), 13); - assert_eq!(inputs.encoder_input_shape(), vec![2, 3, 4, 5]); - } - - #[test] - fn encoder_inputs_accept_model_specific_values() { - let inputs = PreprocessedEncoderInputs::new( - Array4::::zeros((1, 3, 224, 224)), - vec![196], - vec![(224, 224)], - ) - .with_extra( - "image_grid_thw", - ModelSpecificValue::uint_1d(vec![1, 16, 16]), - ) - .with_extra("aspect_ratio_id", ModelSpecificValue::Int(0)); - - assert!(inputs.model_specific.contains_key("image_grid_thw")); - assert!(inputs.model_specific.contains_key("aspect_ratio_id")); - } - - #[test] - fn model_specific_value_tensor_constructors_set_shapes() { - assert!(matches!( - ModelSpecificValue::uint_1d(vec![1, 2, 3]), - ModelSpecificValue::UintTensor { data, shape } - if data == vec![1, 2, 3] && shape == vec![3] - )); - assert!(matches!( - ModelSpecificValue::uint_2d(vec![1, 2, 3, 4], 2, 2), - ModelSpecificValue::UintTensor { data, shape } - if data == vec![1, 2, 3, 4] && shape == vec![2, 2] - )); - assert!(matches!( - ModelSpecificValue::int_1d(vec![1, 2, 3]), - ModelSpecificValue::IntTensor { data, shape } - if data == vec![1, 2, 3] && shape == vec![3] - )); - assert!(matches!( - ModelSpecificValue::int_2d(vec![1, 2, 3, 4], 2, 2), - ModelSpecificValue::IntTensor { data, shape } - if data == vec![1, 2, 3, 4] && shape == vec![2, 2] - )); - } - - #[test] - fn encoder_input_flat_preserves_values() { - let encoder_input = Array4::from_shape_vec((1, 1, 2, 2), vec![1.0, 2.0, 3.0, 4.0]).unwrap(); - let inputs = PreprocessedEncoderInputs::new(encoder_input, vec![4], vec![(2, 2)]); - - assert_eq!(inputs.encoder_input_flat(), vec![1.0, 2.0, 3.0, 4.0]); - } -} diff --git a/crates/multimodal/src/error.rs b/crates/multimodal/src/error.rs deleted file mode 100644 index 37798c26c..000000000 --- a/crates/multimodal/src/error.rs +++ /dev/null @@ -1,68 +0,0 @@ -use std::time::Duration; - -use thiserror::Error; - -pub type MultiModalResult = Result; - -/// Errors that can occur while transforming media into encoder inputs. -#[derive(Debug, Error)] -pub enum TransformError { - #[error("Invalid tensor shape: expected {expected}, got {actual:?}")] - InvalidShape { - expected: String, - actual: Vec, - }, - - #[error("Empty batch: cannot stack zero tensors")] - EmptyBatch, - - #[error("Inconsistent tensor shapes in batch")] - InconsistentShapes, - - #[error("Shape error: {0}")] - ShapeError(String), -} - -#[derive(Debug, Error)] -pub enum MediaConnectorError { - #[error("unsupported media scheme: {0}")] - UnsupportedScheme(String), - #[error("invalid media URL: {0}")] - InvalidUrl(String), - #[error("media domain '{0}' is not in the allow list")] - DisallowedDomain(String), - #[error("local media path is not allowed: {0}")] - DisallowedLocalPath(String), - #[error("HTTP error while fetching media: {0}")] - Http(#[from] reqwest::Error), - #[error("I/O error while reading media: {0}")] - Io(#[from] std::io::Error), - #[error("base64 decode error: {0}")] - Base64Decode(#[from] base64::DecodeError), - #[error("data URL parse error: {0}")] - DataUrl(String), - #[error("{media} input payload exceeds the maximum size of {limit} bytes")] - PayloadTooLarge { media: &'static str, limit: usize }, - #[error("media decode task failed: {0}")] - Blocking(#[from] tokio::task::JoinError), - #[error("image decode error: {0}")] - Image(#[from] image::ImageError), - #[error("audio decode error: {0}")] - AudioDecode(String), - #[error("video decode error: {0}")] - VideoDecode(String), - #[error("media fetch timed out after {0:?}")] - Timeout(Duration), -} - -#[derive(Debug, Error)] -pub enum MultiModalError { - #[error(transparent)] - Media(#[from] MediaConnectorError), - #[error("unsupported content part: {0}")] - UnsupportedContent(&'static str), - #[error("tracker task join error: {0}")] - Join(#[from] tokio::task::JoinError), - #[error("tracker validation error: {0}")] - Validation(String), -} diff --git a/crates/multimodal/src/hasher.rs b/crates/multimodal/src/hasher.rs deleted file mode 100644 index 2019a3f68..000000000 --- a/crates/multimodal/src/hasher.rs +++ /dev/null @@ -1,72 +0,0 @@ -use std::collections::BTreeMap; - -/// Compute a blake3 hex-digest hash for a single image's raw bytes. -pub fn hash_image(raw_bytes: &[u8]) -> String { - blake3::hash(raw_bytes).to_hex().to_string() -} - -/// TODO(yechan): Decide whether video hashes should cover the full encoded -/// payload or a normalized representation of the sampled frames. -/// Compute a blake3 hex-digest hash for a single video's raw bytes. -pub fn hash_video(raw_bytes: &[u8]) -> String { - blake3::hash(raw_bytes).to_hex().to_string() -} - -/// Compute a blake3 hex-digest hash for a single audio payload's raw bytes. -pub fn hash_audio(raw_bytes: &[u8]) -> String { - blake3::hash(raw_bytes).to_hex().to_string() -} - -/// Compute per-image hashes keyed by modality. -/// -/// Returns a `BTreeMap` of per-modality hash lists, -/// e.g. `{"image": ["abc123...", "def456..."]}`. -pub fn hash_images(raw_bytes: &[impl AsRef<[u8]>]) -> BTreeMap> { - let hashes: Vec = raw_bytes.iter().map(|b| hash_image(b.as_ref())).collect(); - let mut map = BTreeMap::new(); - if !hashes.is_empty() { - map.insert("image".to_string(), hashes); - } - map -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_hash_deterministic() { - let data = b"test image bytes"; - assert_eq!(hash_image(data), hash_image(data)); - } - - #[test] - fn test_hash_different_inputs() { - let a = b"image A"; - let b = b"image B"; - assert_ne!(hash_image(a), hash_image(b)); - } - - #[test] - fn test_hash_images_empty() { - let empty: Vec> = vec![]; - let result = hash_images(&empty); - assert!(result.is_empty()); - } - - #[test] - fn test_hash_images_keyed_by_modality() { - let images = vec![b"img1".to_vec(), b"img2".to_vec()]; - let result = hash_images(&images); - assert_eq!(result.len(), 1); - assert!(result.contains_key("image")); - assert_eq!(result["image"].len(), 2); - } - - #[test] - fn test_hash_is_hex() { - let hash = hash_image(b"test"); - assert!(hash.chars().all(|c| c.is_ascii_hexdigit())); - assert_eq!(hash.len(), 64); // blake3 produces 256-bit = 64 hex chars - } -} diff --git a/crates/multimodal/src/hub.rs b/crates/multimodal/src/hub.rs deleted file mode 100644 index 702c86a55..000000000 --- a/crates/multimodal/src/hub.rs +++ /dev/null @@ -1,51 +0,0 @@ -//! HuggingFace Hub integration for downloading model config files. -//! -//! When the model source is a HuggingFace model ID (e.g. `Qwen/Qwen3-VL-8B-Instruct`) -//! rather than a local directory, this module downloads `config.json` and -//! `preprocessor_config.json` / video processor config files to the local HF -//! cache and returns the resolved path. - -use std::path::{Path, PathBuf}; - -use anyhow::Context; -use hf_hub::api::tokio::ApiBuilder; - -const HF_TOKEN_ENV: &str = "HF_TOKEN"; - -/// Resolve a model source to a local directory containing `config.json`. -/// -/// If the source is already a local directory with `config.json`, returns it as-is. -/// Otherwise, treats the source as a HuggingFace model ID and downloads -/// `config.json` plus optional processor configs to the local HF cache. -pub async fn resolve_model_config_dir(source: &str) -> anyhow::Result { - let path = Path::new(source); - if path.join("config.json").exists() { - return Ok(path.to_path_buf()); - } - - let mut builder = ApiBuilder::from_env().with_progress(false); - if let Ok(token) = std::env::var(HF_TOKEN_ENV) { - if !token.is_empty() { - builder = builder.with_token(Some(token)); - } - } - let api = builder - .build() - .context("Failed to build HuggingFace API client")?; - let repo = api.model(source.to_string()); - - let config_path = repo - .get("config.json") - .await - .with_context(|| format!("Failed to download config.json for model '{source}'"))?; - - // Best-effort download of processor configs (not all models have them). - let _ = repo.get("preprocessor_config.json").await; - let _ = repo.get("video_preprocessor_config.json").await; - let _ = repo.get("processor_config.json").await; - - config_path - .parent() - .map(|p| p.to_path_buf()) - .ok_or_else(|| anyhow::anyhow!("Invalid HF cache path for model '{source}'")) -} diff --git a/crates/multimodal/src/jpeg_turbo.rs b/crates/multimodal/src/jpeg_turbo.rs deleted file mode 100644 index b24c53b97..000000000 --- a/crates/multimodal/src/jpeg_turbo.rs +++ /dev/null @@ -1,182 +0,0 @@ -//! Runtime (dlopen) binding to libjpeg-turbo's TurboJPEG API for JPEG decode. -//! -//! PIL/Pillow (and therefore vLLM) decode JPEGs with libjpeg-turbo using its -//! default options: accurate (islow) integer IDCT and "fancy" (bilinear) chroma -//! upsampling. The pure-Rust `image`/`zune-jpeg` decoder differs by a few levels -//! per pixel, which the vision encoder amplifies into a large embedding shift, -//! making TokenSpeed's multimodal accuracy diverge from vLLM. Decoding through -//! libjpeg-turbo with the same defaults makes SMG's pixel values match vLLM's. -//! -//! We load libturbojpeg at RUNTIME via `dlopen` rather than linking it, so the -//! crate (and every consumer — including the Go/Python bindings and CI builds -//! that don't ship libturbojpeg) compiles on any platform with no build script -//! and no link-time dependency. Where the shared library is present (the serving -//! image), decode goes through it for PIL parity; where it's absent, -//! `decode_jpeg_rgb` returns `None` and the caller falls back to the pure-Rust -//! decoder. Default flags (0) select accurate DCT + fancy upsampling, matching -//! Pillow. -//! -//! This module is the crate's only FFI surface, so it locally overrides the -//! workspace-wide `unsafe_code = "deny"` for the C bindings. -#![allow(unsafe_code)] - -use std::{ - os::raw::{c_int, c_uchar, c_ulong, c_void}, - sync::OnceLock, -}; - -use image::{DynamicImage, RgbImage}; -use libloading::{Library, Symbol}; - -type TjHandle = *mut c_void; -const TJPF_RGB: c_int = 0; - -type TjInitDecompress = unsafe extern "C" fn() -> TjHandle; -type TjDecompressHeader3 = unsafe extern "C" fn( - TjHandle, - *const c_uchar, - c_ulong, - *mut c_int, - *mut c_int, - *mut c_int, - *mut c_int, -) -> c_int; -type TjDecompress2 = unsafe extern "C" fn( - TjHandle, - *const c_uchar, - c_ulong, - *mut c_uchar, - c_int, - c_int, - c_int, - c_int, - c_int, -) -> c_int; -type TjDestroy = unsafe extern "C" fn(TjHandle) -> c_int; - -/// Resolved TurboJPEG entry points. Holds the loaded `Library` so the function -/// pointers stay valid for the process lifetime. -struct TurboJpeg { - _lib: Library, - init: TjInitDecompress, - header: TjDecompressHeader3, - decompress: TjDecompress2, - destroy: TjDestroy, -} - -// The function pointers are plain C entry points with no shared mutable state; -// the library handle is kept alive for the process and never mutated. -unsafe impl Send for TurboJpeg {} -unsafe impl Sync for TurboJpeg {} - -fn load_turbojpeg() -> Option { - // Try the runtime soname first (shipped by the runtime package), then the - // dev symlink and common macOS names. - const CANDIDATES: &[&str] = &[ - "libturbojpeg.so.0", - "libturbojpeg.so", - "libturbojpeg.0.dylib", - "libturbojpeg.dylib", - ]; - // SAFETY: loading a system shared library by name; we only resolve the four - // documented TurboJPEG symbols below and keep the handle for their lifetime. - let lib = CANDIDATES - .iter() - .find_map(|name| unsafe { Library::new(name) }.ok())?; - // SAFETY: each symbol is resolved against the just-loaded library with the - // signature documented by the TurboJPEG API. We copy the bare function - // pointers out (dropping the borrowing `Symbol`s) and keep `lib` alive in - // the returned struct, so the pointers remain valid. - let (init, header, decompress, destroy) = unsafe { - let init: Symbol = lib.get(b"tjInitDecompress\0").ok()?; - let header: Symbol = lib.get(b"tjDecompressHeader3\0").ok()?; - let decompress: Symbol = lib.get(b"tjDecompress2\0").ok()?; - let destroy: Symbol = lib.get(b"tjDestroy\0").ok()?; - (*init, *header, *decompress, *destroy) - }; - Some(TurboJpeg { - _lib: lib, - init, - header, - decompress, - destroy, - }) -} - -/// Process-wide cached TurboJPEG binding, or `None` if the library is absent. -fn turbojpeg() -> Option<&'static TurboJpeg> { - static TJ: OnceLock> = OnceLock::new(); - TJ.get_or_init(load_turbojpeg).as_ref() -} - -/// True if `bytes` start with the JPEG SOI marker. -pub fn is_jpeg(bytes: &[u8]) -> bool { - bytes.len() >= 3 && bytes[0] == 0xFF && bytes[1] == 0xD8 && bytes[2] == 0xFF -} - -/// Decode a JPEG to an RGB8 `DynamicImage` via libjpeg-turbo (PIL-compatible -/// defaults). Returns `None` on any failure — including libturbojpeg not being -/// available at runtime — so the caller can fall back to the pure-Rust decoder. -pub fn decode_jpeg_rgb(bytes: &[u8]) -> Option { - if !is_jpeg(bytes) { - return None; - } - let tj = turbojpeg()?; - // SAFETY: - // - `tj.init` returns a handle that is null-checked before use; every early - // return below calls `tj.destroy(handle)` first, so the handle is freed - // exactly once and never used after destruction. - // - `bytes` is a live `&[u8]`; its ptr/len describe a valid immutable region - // for the duration of the calls (libjpeg-turbo only reads it). - // - `buf` is an owned `Vec` sized to exactly `w*h*3` (overflow-checked) - // from the header dimensions, and is the sole alias passed to the decoder; - // with `pitch=0` (=> `w*3`) and `TJPF_RGB` the decoder writes at most - // `w*h*3` bytes, so no out-of-bounds write occurs. The image is built only - // after `rc == 0` confirms a successful, complete write. - unsafe { - let handle = (tj.init)(); - if handle.is_null() { - return None; - } - let (mut w, mut h, mut subsamp, mut colorspace) = (0_i32, 0_i32, 0_i32, 0_i32); - let hdr = (tj.header)( - handle, - bytes.as_ptr(), - bytes.len() as c_ulong, - &mut w, - &mut h, - &mut subsamp, - &mut colorspace, - ); - if hdr != 0 || w <= 0 || h <= 0 { - (tj.destroy)(handle); - return None; - } - let (wu, hu) = (w as usize, h as usize); - // Guard against absurd dimensions before allocating. - let nbytes = match wu.checked_mul(hu).and_then(|p| p.checked_mul(3)) { - Some(n) => n, - None => { - (tj.destroy)(handle); - return None; - } - }; - let mut buf = vec![0_u8; nbytes]; - let rc = (tj.decompress)( - handle, - bytes.as_ptr(), - bytes.len() as c_ulong, - buf.as_mut_ptr(), - w, - 0, // pitch = 0 -> width * pixelsize - h, - TJPF_RGB, - 0, // default flags: accurate IDCT + fancy upsampling (matches Pillow) - ); - (tj.destroy)(handle); - if rc != 0 { - return None; - } - RgbImage::from_raw(w as u32, h as u32, buf).map(DynamicImage::ImageRgb8) - } -} diff --git a/crates/multimodal/src/lib.rs b/crates/multimodal/src/lib.rs deleted file mode 100644 index ff6ce601c..000000000 --- a/crates/multimodal/src/lib.rs +++ /dev/null @@ -1,32 +0,0 @@ -pub mod audio; -pub mod encoder_inputs; -pub mod error; -pub mod hasher; -pub mod hub; -pub mod jpeg_turbo; -pub mod media; -#[cfg(feature = "opencv-video")] -mod opencv_buffer; -pub mod registry; -pub mod tracker; -pub mod types; -pub mod vision; - -pub use audio::AudioPreProcessor; -pub use encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}; -pub use error::{MediaConnectorError, MultiModalError, MultiModalResult, TransformError}; -pub use media::{ - ImageFetchConfig, MediaConnector, MediaConnectorConfig, MediaSource, VideoFetchConfig, -}; -pub use registry::{ModelMetadata, ModelProcessorSpec, ModelRegistry}; -pub use tracker::{AsyncMultiModalTracker, TrackerOutput}; -pub use types::{ - AudioClip, AudioSource, EncoderFieldLayouts, FieldLayout, ImageDetail, ImageFrame, ImageSize, - ImageSource, MediaContentPart, Modality, MultiModalData, MultiModalUUIDs, PlaceholderRange, - PromptReplacement, RgbFrameRef, TokenId, TrackedMedia, VideoClip, VideoSource, -}; -// Re-export vision processing components -pub use vision::{ - LlavaNextProcessor, LlavaProcessor, PreProcessorConfig, VisionPreProcessor, - VisionProcessorRegistry, -}; diff --git a/crates/multimodal/src/media.rs b/crates/multimodal/src/media.rs deleted file mode 100644 index 29684d6ee..000000000 --- a/crates/multimodal/src/media.rs +++ /dev/null @@ -1,2361 +0,0 @@ -#[cfg(feature = "opencv-video")] -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::{ - collections::HashSet, - io::Write, - path::PathBuf, - process::{Output, Stdio}, - sync::{Arc, OnceLock}, - time::{Duration, Instant}, -}; - -use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; -use bytes::{Bytes, BytesMut}; -#[cfg(feature = "opencv-video")] -use opencv::{ - core::{Mat, Vector}, - prelude::*, - videoio, -}; -use reqwest::Client; -use tokio::{fs, io::AsyncReadExt, process::Command, task, time}; -use tracing::info; -use url::Url; - -use crate::audio::decode_audio_mono_f32; - -const DEFAULT_VIDEO_PROCESS_TIMEOUT: Duration = Duration::from_secs(30); -const DEFAULT_IMAGE_MAX_INPUT_BYTES: usize = 256 * 1024 * 1024; -const DEFAULT_VIDEO_MAX_INPUT_BYTES: usize = 256 * 1024 * 1024; -const DEFAULT_VIDEO_MAX_DECODED_BYTES: usize = 1024 * 1024 * 1024; -const DEFAULT_AUDIO_MAX_INPUT_BYTES: usize = 256 * 1024 * 1024; -const _: () = assert!(DEFAULT_VIDEO_MAX_INPUT_BYTES < DEFAULT_VIDEO_MAX_DECODED_BYTES); -static VIDEO_DECODE_BACKEND: OnceLock> = OnceLock::new(); -static LOG_VIDEO_DECODE_TIMING: OnceLock = OnceLock::new(); -static VIDEO_PROCESS_TIMEOUT: OnceLock = OnceLock::new(); -static IMAGE_MAX_INPUT_BYTES: OnceLock = OnceLock::new(); -static VIDEO_MAX_INPUT_BYTES: OnceLock = OnceLock::new(); -static VIDEO_MAX_DECODED_BYTES: OnceLock = OnceLock::new(); -static AUDIO_MAX_INPUT_BYTES: OnceLock = OnceLock::new(); -#[cfg(feature = "opencv-video")] -static ACTIVE_OPENCV_DECODES: AtomicUsize = AtomicUsize::new(0); -#[cfg(feature = "opencv-video")] -static AVAILABLE_OPENCV_CPUS: OnceLock = OnceLock::new(); -#[cfg(feature = "opencv-video")] -const MAX_OPENCV_DECODER_THREADS: usize = 8; -#[cfg(feature = "opencv-video")] -const OPENCV_DECODE_BURST_COALESCE: Duration = Duration::from_millis(5); -#[cfg(feature = "opencv-video")] -const OPENCV_LOW_CONCURRENCY_LIMIT: usize = 8; -#[cfg(feature = "opencv-video")] -const OPENCV_LOW_CONCURRENCY_CPU_MULTIPLIER: usize = 2; -#[cfg(feature = "opencv-video")] -const OPENCV_HIGH_CONCURRENCY_CPU_BUDGET_NUMERATOR: usize = 6; -#[cfg(feature = "opencv-video")] -const OPENCV_HIGH_CONCURRENCY_CPU_BUDGET_DENOMINATOR: usize = 7; - -use super::{ - error::MediaConnectorError, - types::{ - AudioClip, AudioSource, DecodedRgbFrame, DecodedRgbVideo, ImageDetail, ImageFrame, - ImageSource, VideoClip, VideoSource, - }, -}; - -#[derive(Clone)] -pub struct MediaConnectorConfig { - pub allowed_domains: Option>, - pub allowed_local_media_path: Option, - pub fetch_timeout: Duration, -} - -impl Default for MediaConnectorConfig { - fn default() -> Self { - Self { - allowed_domains: None, - allowed_local_media_path: None, - fetch_timeout: Duration::from_secs(10), - } - } -} - -#[derive(Clone, Copy, Debug)] -pub struct ImageFetchConfig { - pub detail: ImageDetail, -} - -impl Default for ImageFetchConfig { - fn default() -> Self { - Self { - detail: ImageDetail::Auto, - } - } -} - -#[derive(Clone, Copy, Debug)] -pub struct VideoFetchConfig { - pub min_frames: usize, - pub max_frames: usize, - pub sample_fps: f32, -} - -impl Default for VideoFetchConfig { - fn default() -> Self { - Self { - min_frames: 4, - max_frames: 768, - sample_fps: 2.0, - } - } -} - -#[derive(Debug, Clone)] -pub enum MediaSource { - Url(String), - DataUrl(String), - InlineBytes(Vec), - File(PathBuf), -} - -#[derive(Clone)] -pub struct MediaConnector { - client: Client, - allowed_domains: Option>, - allowed_local_media_path: Option, - fetch_timeout: Duration, -} - -impl MediaConnector { - pub fn new(client: Client, config: MediaConnectorConfig) -> Result { - let allowed_domains = config.allowed_domains.map(|domains| { - domains - .into_iter() - .map(|d| d.to_ascii_lowercase()) - .collect::>() - }); - - let allowed_local_media_path = if let Some(path) = config.allowed_local_media_path { - Some(std::fs::canonicalize(path)?) - } else { - None - }; - - Ok(Self { - client, - allowed_domains, - allowed_local_media_path, - fetch_timeout: config.fetch_timeout, - }) - } - - pub async fn fetch_image( - &self, - source: MediaSource, - cfg: ImageFetchConfig, - ) -> Result, MediaConnectorError> { - match source { - MediaSource::Url(url) => self.fetch_http_image(url, cfg).await, - MediaSource::DataUrl(data_url) => self.fetch_data_url(data_url, cfg).await, - MediaSource::InlineBytes(bytes) => { - self.decode_image(bytes.into(), cfg.detail, ImageSource::InlineBytes) - .await - } - MediaSource::File(path) => self.fetch_file(path, cfg).await, - } - } - - pub async fn fetch_video( - &self, - source: MediaSource, - cfg: VideoFetchConfig, - ) -> Result, MediaConnectorError> { - match source { - MediaSource::Url(url) => self.fetch_http_video(url, cfg).await, - MediaSource::DataUrl(data_url) => self.fetch_video_data_url(data_url, cfg).await, - MediaSource::InlineBytes(bytes) => { - self.decode_video(bytes.into(), cfg, VideoSource::InlineBytes) - .await - } - MediaSource::File(path) => self.fetch_video_file(path, cfg).await, - } - } - - pub async fn fetch_audio( - &self, - source: MediaSource, - ) -> Result, MediaConnectorError> { - match source { - MediaSource::Url(url) => self.fetch_http_audio(url).await, - MediaSource::DataUrl(data_url) => self.fetch_audio_data_url(data_url).await, - MediaSource::InlineBytes(bytes) => { - self.decode_audio(bytes.into(), AudioSource::InlineBytes) - .await - } - MediaSource::File(path) => self.fetch_audio_file(path).await, - } - } - - async fn fetch_http_image( - &self, - url: String, - cfg: ImageFetchConfig, - ) -> Result, MediaConnectorError> { - let parsed = Url::parse(&url).map_err(|_| MediaConnectorError::InvalidUrl(url.clone()))?; - self.ensure_domain_allowed(&parsed)?; - - let mut req = self.client.get(parsed.as_str()); - if self.fetch_timeout > Duration::ZERO { - req = req.timeout(self.fetch_timeout); - } - - let resp = req.send().await.map_err(|err| { - if err.is_timeout() { - MediaConnectorError::Timeout(self.fetch_timeout) - } else { - MediaConnectorError::Http(err) - } - })?; - - let resp = resp.error_for_status()?; - let bytes = collect_http_body_with_limit(resp, image_max_input_bytes(), "image").await?; - self.decode_image( - bytes, - cfg.detail, - ImageSource::Url { - url: parsed.to_string(), - }, - ) - .await - } - - async fn fetch_data_url( - &self, - data_url: String, - cfg: ImageFetchConfig, - ) -> Result, MediaConnectorError> { - let (metadata, data) = data_url - .split_once(',') - .ok_or_else(|| MediaConnectorError::DataUrl("missing comma in data url".into()))?; - - if !metadata.ends_with(";base64") { - return Err(MediaConnectorError::DataUrl( - "only base64 encoded data URLs are supported".into(), - )); - } - - let data = data.trim(); - let decoded = decode_base64_with_limit(data, image_max_input_bytes(), "image")?; - self.decode_image(decoded.into(), cfg.detail, ImageSource::DataUrl) - .await - } - - async fn fetch_video_data_url( - &self, - data_url: String, - cfg: VideoFetchConfig, - ) -> Result, MediaConnectorError> { - let (metadata, data) = data_url - .split_once(',') - .ok_or_else(|| MediaConnectorError::DataUrl("missing comma in data url".into()))?; - - if !metadata.ends_with(";base64") { - return Err(MediaConnectorError::DataUrl( - "only base64 encoded data URLs are supported".into(), - )); - } - - let data = data.trim(); - let decoded = decode_base64_with_limit(data, video_max_input_bytes(), "video")?; - self.decode_video(decoded.into(), cfg, VideoSource::DataUrl) - .await - } - - async fn fetch_audio_data_url( - &self, - data_url: String, - ) -> Result, MediaConnectorError> { - let (metadata, data) = data_url - .split_once(',') - .ok_or_else(|| MediaConnectorError::DataUrl("missing comma in data url".into()))?; - - if !metadata.ends_with(";base64") { - return Err(MediaConnectorError::DataUrl( - "only base64 encoded data URLs are supported".into(), - )); - } - - let data = data.trim(); - let decoded = decode_base64_with_limit(data, audio_max_input_bytes(), "audio")?; - self.decode_audio(decoded.into(), AudioSource::DataUrl) - .await - } - - async fn fetch_file( - &self, - path: PathBuf, - cfg: ImageFetchConfig, - ) -> Result, MediaConnectorError> { - let allowed_root = self - .allowed_local_media_path - .as_ref() - .ok_or_else(|| MediaConnectorError::DisallowedLocalPath(path.display().to_string()))?; - - let canonical = fs::canonicalize(&path).await?; - if !canonical.starts_with(allowed_root) { - return Err(MediaConnectorError::DisallowedLocalPath( - path.display().to_string(), - )); - } - - let bytes = read_file_with_limit(&canonical, image_max_input_bytes(), "image").await?; - self.decode_image(bytes, cfg.detail, ImageSource::File { path: canonical }) - .await - } - - async fn fetch_http_video( - &self, - url: String, - cfg: VideoFetchConfig, - ) -> Result, MediaConnectorError> { - let parsed = Url::parse(&url).map_err(|_| MediaConnectorError::InvalidUrl(url.clone()))?; - self.ensure_domain_allowed(&parsed)?; - - let mut req = self.client.get(parsed.as_str()); - if self.fetch_timeout > Duration::ZERO { - req = req.timeout(self.fetch_timeout); - } - - let resp = req.send().await.map_err(|err| { - if err.is_timeout() { - MediaConnectorError::Timeout(self.fetch_timeout) - } else { - MediaConnectorError::Http(err) - } - })?; - - let resp = resp.error_for_status()?; - let bytes = collect_http_body_with_limit(resp, video_max_input_bytes(), "video").await?; - self.decode_video( - bytes, - cfg, - VideoSource::Url { - url: parsed.to_string(), - }, - ) - .await - } - - async fn fetch_http_audio(&self, url: String) -> Result, MediaConnectorError> { - let parsed = Url::parse(&url).map_err(|_| MediaConnectorError::InvalidUrl(url.clone()))?; - self.ensure_domain_allowed(&parsed)?; - - let mut req = self.client.get(parsed.as_str()); - if self.fetch_timeout > Duration::ZERO { - req = req.timeout(self.fetch_timeout); - } - - let resp = req.send().await.map_err(|err| { - if err.is_timeout() { - MediaConnectorError::Timeout(self.fetch_timeout) - } else { - MediaConnectorError::Http(err) - } - })?; - - let resp = resp.error_for_status()?; - let bytes = collect_http_body_with_limit(resp, audio_max_input_bytes(), "audio").await?; - self.decode_audio( - bytes, - AudioSource::Url { - url: parsed.to_string(), - }, - ) - .await - } - - async fn fetch_video_file( - &self, - path: PathBuf, - cfg: VideoFetchConfig, - ) -> Result, MediaConnectorError> { - let allowed_root = self - .allowed_local_media_path - .as_ref() - .ok_or_else(|| MediaConnectorError::DisallowedLocalPath(path.display().to_string()))?; - - let canonical = fs::canonicalize(&path).await?; - if !canonical.starts_with(allowed_root) { - return Err(MediaConnectorError::DisallowedLocalPath( - path.display().to_string(), - )); - } - - let bytes = read_file_with_limit(&canonical, video_max_input_bytes(), "video").await?; - self.decode_video(bytes, cfg, VideoSource::File { path: canonical }) - .await - } - - async fn fetch_audio_file(&self, path: PathBuf) -> Result, MediaConnectorError> { - let allowed_root = self - .allowed_local_media_path - .as_ref() - .ok_or_else(|| MediaConnectorError::DisallowedLocalPath(path.display().to_string()))?; - - let canonical = fs::canonicalize(&path).await?; - if !canonical.starts_with(allowed_root) { - return Err(MediaConnectorError::DisallowedLocalPath( - path.display().to_string(), - )); - } - - let bytes = read_file_with_limit(&canonical, audio_max_input_bytes(), "audio").await?; - self.decode_audio(bytes, AudioSource::File { path: canonical }) - .await - } - - fn ensure_domain_allowed(&self, url: &Url) -> Result<(), MediaConnectorError> { - if let Some(allowed) = &self.allowed_domains { - let host = url - .host_str() - .map(|h| h.to_ascii_lowercase()) - .ok_or_else(|| MediaConnectorError::InvalidUrl(url.to_string()))?; - if !allowed.contains(&host) { - return Err(MediaConnectorError::DisallowedDomain(host)); - } - } - Ok(()) - } - - async fn decode_image( - &self, - bytes: Bytes, - detail: ImageDetail, - source: ImageSource, - ) -> Result, MediaConnectorError> { - ensure_input_byte_limit(bytes.len(), image_max_input_bytes(), "image")?; - let hash = crate::hasher::hash_image(&bytes); - - // Decode JPEGs through libjpeg-turbo (PIL-compatible defaults: accurate - // IDCT + fancy upsampling) so pixel values match vLLM bit-for-bit; the - // pure-Rust decoder diverges by a few levels, which the vision encoder - // amplifies into an embedding shift. Non-JPEG inputs and any turbojpeg - // failure fall back to the `image` crate. - let bytes_for_decode = bytes.clone(); - let image = task::spawn_blocking( - move || -> Result { - if let Some(img) = crate::jpeg_turbo::decode_jpeg_rgb(&bytes_for_decode) { - return Ok(img); - } - let cursor = std::io::Cursor::new(bytes_for_decode); - let reader = image::ImageReader::new(cursor).with_guessed_format()?; - Ok(reader.decode()?) - }, - ) - .await - .map_err(MediaConnectorError::Blocking)??; - - Ok(Arc::new(ImageFrame::new( - image, bytes, detail, source, hash, - ))) - } - - async fn decode_audio( - &self, - bytes: Bytes, - source: AudioSource, - ) -> Result, MediaConnectorError> { - ensure_input_byte_limit(bytes.len(), audio_max_input_bytes(), "audio")?; - let hash = crate::hasher::hash_audio(&bytes); - let decoded = decode_audio_mono_f32(&bytes) - .await - .map_err(|e| MediaConnectorError::AudioDecode(e.to_string()))?; - Ok(Arc::new(AudioClip::new(bytes, decoded, source, hash))) - } - - async fn decode_video( - &self, - bytes: Bytes, - cfg: VideoFetchConfig, - source: VideoSource, - ) -> Result, MediaConnectorError> { - ensure_input_byte_limit(bytes.len(), video_max_input_bytes(), "video")?; - if cfg.max_frames == 0 { - return Err(MediaConnectorError::VideoDecode( - "max_frames must be greater than 0".to_string(), - )); - } - if cfg.min_frames == 0 { - return Err(MediaConnectorError::VideoDecode( - "min_frames must be greater than 0".to_string(), - )); - } - if cfg.min_frames > cfg.max_frames { - return Err(MediaConnectorError::VideoDecode( - "min_frames must be less than or equal to max_frames".to_string(), - )); - } - if !cfg.sample_fps.is_finite() || cfg.sample_fps <= 0.0 { - return Err(MediaConnectorError::VideoDecode( - "sample_fps must be finite and greater than 0".to_string(), - )); - } - - let hash = crate::hasher::hash_video(&bytes); - let decoded = decode_video_frames(bytes.clone(), cfg).await?; - - let clip = match decoded { - DecodedVideoFrames::Images { frames, sample_fps } => { - VideoClip::new_with_sample_fps(frames, bytes, source, hash, sample_fps) - } - DecodedVideoFrames::Rgb { video, sample_fps } => { - VideoClip::new_rgb_with_sample_fps(video, bytes, source, hash, sample_fps) - } - }; - Ok(Arc::new(clip)) - } -} - -async fn read_file_with_limit( - path: &std::path::Path, - limit: usize, - media: &'static str, -) -> Result { - let file = fs::File::open(path).await?; - let limit_u64 = u64::try_from(limit).unwrap_or(u64::MAX); - if file.metadata().await?.len() > limit_u64 { - return Err(MediaConnectorError::PayloadTooLarge { media, limit }); - } - - // Read at most one byte beyond the limit. The post-read exact check also - // covers a file growing after the metadata check. - let mut reader = file.take(limit_u64.saturating_add(1)); - let mut bytes = Vec::new(); - reader.read_to_end(&mut bytes).await?; - ensure_input_byte_limit(bytes.len(), limit, media)?; - Ok(Bytes::from(bytes)) -} - -fn ensure_input_byte_limit( - input_bytes: usize, - limit: usize, - media: &'static str, -) -> Result<(), MediaConnectorError> { - checked_payload_length(0, input_bytes, limit, media).map(|_| ()) -} - -async fn collect_http_body_with_limit( - mut response: reqwest::Response, - limit: usize, - media: &'static str, -) -> Result { - if response - .content_length() - .is_some_and(|length| length > limit as u64) - { - return Err(MediaConnectorError::PayloadTooLarge { media, limit }); - } - - let mut body = BytesMut::new(); - while let Some(chunk) = response.chunk().await? { - checked_payload_length(body.len(), chunk.len(), limit, media)?; - body.extend_from_slice(&chunk); - } - Ok(body.freeze()) -} - -fn checked_payload_length( - current: usize, - additional: usize, - limit: usize, - media: &'static str, -) -> Result { - current - .checked_add(additional) - .filter(|length| *length <= limit) - .ok_or(MediaConnectorError::PayloadTooLarge { media, limit }) -} - -fn decode_base64_with_limit( - encoded: &str, - limit: usize, - media: &'static str, -) -> Result, MediaConnectorError> { - // A padded base64 encoding of at most `limit` bytes needs no more than - // ceil(limit / 3) * 4 input bytes. Reject longer strings before the base64 - // decoder allocates; the exact decoded-length check below handles the up - // to two-byte slack at the boundary. - let max_encoded_len = (limit as u128).div_ceil(3) * 4; - if encoded.len() as u128 > max_encoded_len { - return Err(MediaConnectorError::PayloadTooLarge { media, limit }); - } - - let decoded = BASE64_STANDARD.decode(encoded)?; - checked_payload_length(0, decoded.len(), limit, media)?; - Ok(decoded) -} - -fn env_byte_limit(cache: &'static OnceLock, env_var: &str, default: usize) -> usize { - *cache.get_or_init(|| { - std::env::var(env_var) - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|bytes| *bytes > 0) - .unwrap_or(default) - }) -} - -fn image_max_input_bytes() -> usize { - env_byte_limit( - &IMAGE_MAX_INPUT_BYTES, - "SMG_IMAGE_MAX_INPUT_BYTES", - DEFAULT_IMAGE_MAX_INPUT_BYTES, - ) -} - -fn video_max_input_bytes() -> usize { - env_byte_limit( - &VIDEO_MAX_INPUT_BYTES, - "SMG_VIDEO_MAX_INPUT_BYTES", - DEFAULT_VIDEO_MAX_INPUT_BYTES, - ) -} - -fn audio_max_input_bytes() -> usize { - env_byte_limit( - &AUDIO_MAX_INPUT_BYTES, - "SMG_AUDIO_MAX_INPUT_BYTES", - DEFAULT_AUDIO_MAX_INPUT_BYTES, - ) -} - -enum DecodedVideoFrames { - Images { - frames: Vec, - sample_fps: f32, - }, - Rgb { - video: DecodedRgbVideo, - sample_fps: f32, - }, -} - -async fn decode_video_frames( - bytes: Bytes, - cfg: VideoFetchConfig, -) -> Result { - #[cfg(feature = "opencv-video")] - let input_bytes = bytes.len(); - match video_decode_backend_override() { - Some("ffmpeg") => decode_video_bytes_with_ffmpeg(bytes, cfg).await, - Some("opencv") => { - #[cfg(feature = "opencv-video")] - { - let opencv_bytes = bytes.clone(); - let result = task::spawn_blocking(move || { - decode_video_with_opencv_bytes_logged(opencv_bytes, input_bytes, cfg) - }) - .await - .map_err(MediaConnectorError::Blocking)?; - match result { - Ok(frames) => Ok(frames), - Err(error) => { - if log_video_decode_timing_enabled() { - info!( - error = %error, - "smg_mm_timing video_decode_opencv_buffer_fallback" - ); - } - decode_video_bytes_with_tempfile(bytes, cfg) - .await - .map_err(|fallback_error| { - MediaConnectorError::VideoDecode(format!( - "buffered OpenCV decode failed: {error}; tempfile OpenCV fallback failed: {fallback_error}" - )) - }) - } - } - } - #[cfg(not(feature = "opencv-video"))] - { - Err(MediaConnectorError::VideoDecode( - "SMG_VIDEO_DECODE_BACKEND=opencv requires the opencv-video feature".to_string(), - )) - } - } - Some(backend) => Err(MediaConnectorError::VideoDecode(format!( - "unsupported SMG_VIDEO_DECODE_BACKEND={backend}; expected auto, opencv, or ffmpeg" - ))), - None => { - #[cfg(feature = "opencv-video")] - { - let opencv_bytes = bytes.clone(); - let opencv_result = task::spawn_blocking(move || { - decode_video_with_opencv_bytes_logged(opencv_bytes, input_bytes, cfg) - }) - .await - .map_err(MediaConnectorError::Blocking)?; - match opencv_result { - Ok(frames) => Ok(frames), - Err(opencv_error) => { - if log_video_decode_timing_enabled() { - info!( - error = %opencv_error, - "smg_mm_timing video_decode_auto_opencv_fallback" - ); - } - decode_video_bytes_with_tempfile(bytes, cfg) - .await - .map_err(|fallback_error| { - MediaConnectorError::VideoDecode(format!( - "buffered OpenCV decode failed: {opencv_error}; tempfile fallback failed: {fallback_error}" - )) - }) - } - } - } - #[cfg(not(feature = "opencv-video"))] - { - decode_video_bytes_with_ffmpeg(bytes, cfg).await - } - } - } -} - -#[cfg(feature = "opencv-video")] -async fn decode_video_bytes_with_tempfile( - bytes: Bytes, - cfg: VideoFetchConfig, -) -> Result { - let input_bytes = bytes.len(); - let input_file = { - let bytes = bytes.clone(); - task::spawn_blocking(move || write_temp_video_file(&bytes)) - .await - .map_err(MediaConnectorError::Blocking)?? - }; - decode_video_frames_from_path(input_file.path(), input_bytes, cfg).await -} - -async fn decode_video_bytes_with_ffmpeg( - bytes: Bytes, - cfg: VideoFetchConfig, -) -> Result { - let input_bytes = bytes.len(); - let input_file = { - let bytes = bytes.clone(); - task::spawn_blocking(move || write_temp_video_file(&bytes)) - .await - .map_err(MediaConnectorError::Blocking)?? - }; - let input_path = input_file.path().to_path_buf(); - decode_video_with_ffmpeg(&input_path, input_bytes, cfg).await -} - -#[cfg(feature = "opencv-video")] -async fn decode_video_frames_from_path( - input_path: &std::path::Path, - input_bytes: usize, - cfg: VideoFetchConfig, -) -> Result { - match video_decode_backend_override() { - Some("ffmpeg") => decode_video_with_ffmpeg(input_path, input_bytes, cfg).await, - Some("opencv") => { - #[cfg(feature = "opencv-video")] - { - let input_path = input_path.to_path_buf(); - task::spawn_blocking(move || { - decode_video_with_opencv_logged(&input_path, input_bytes, cfg) - }) - .await - .map_err(MediaConnectorError::Blocking)? - } - #[cfg(not(feature = "opencv-video"))] - { - Err(MediaConnectorError::VideoDecode( - "SMG_VIDEO_DECODE_BACKEND=opencv requires the opencv-video feature".to_string(), - )) - } - } - Some(backend) => Err(MediaConnectorError::VideoDecode(format!( - "unsupported SMG_VIDEO_DECODE_BACKEND={backend}; expected auto, opencv, or ffmpeg" - ))), - None => { - #[cfg(feature = "opencv-video")] - { - // OpenCV samples by frame index while the FFmpeg fallback uses an - // fps filter, so the fallback can select a different frame set. - let opencv_input_path = input_path.to_path_buf(); - let opencv_result = task::spawn_blocking(move || { - decode_video_with_opencv_logged(&opencv_input_path, input_bytes, cfg) - }) - .await - .map_err(MediaConnectorError::Blocking)?; - - match opencv_result { - Ok(frames) => Ok(frames), - Err(opencv_error) => { - if log_video_decode_timing_enabled() { - info!( - error = %opencv_error, - "smg_mm_timing video_decode_auto_opencv_fallback" - ); - } - - match decode_video_with_ffmpeg(input_path, input_bytes, cfg).await { - Ok(frames) => Ok(frames), - Err(ffmpeg_error) => Err(MediaConnectorError::VideoDecode(format!( - "OpenCV decode failed: {opencv_error}; ffmpeg fallback failed: {ffmpeg_error}" - ))), - } - } - } - } - - #[cfg(not(feature = "opencv-video"))] - { - decode_video_with_ffmpeg(input_path, input_bytes, cfg).await - } - } - } -} - -#[cfg(feature = "opencv-video")] -fn decode_video_with_opencv_logged( - input_path: &std::path::Path, - input_bytes: usize, - cfg: VideoFetchConfig, -) -> Result { - let started = Instant::now(); - let result = decode_video_with_opencv_file(input_path, cfg); - match &result { - Ok(_) => log_video_decode_backend_timing("opencv", started, input_bytes, cfg, None), - Err(error) => { - log_video_decode_backend_timing("opencv", started, input_bytes, cfg, Some(error)); - } - } - result -} - -#[cfg(feature = "opencv-video")] -fn decode_video_with_opencv_bytes_logged( - bytes: Bytes, - input_bytes: usize, - cfg: VideoFetchConfig, -) -> Result { - let started = Instant::now(); - let result = decode_video_with_opencv_bytes(bytes, cfg); - match &result { - Ok(_) => log_video_decode_backend_timing("opencv_buffer", started, input_bytes, cfg, None), - Err(error) => { - log_video_decode_backend_timing( - "opencv_buffer", - started, - input_bytes, - cfg, - Some(error), - ); - } - } - result -} - -fn video_decode_backend_override() -> Option<&'static str> { - VIDEO_DECODE_BACKEND - .get_or_init(|| { - let backend = std::env::var("SMG_VIDEO_DECODE_BACKEND") - .ok()? - .trim() - .to_ascii_lowercase(); - match backend.as_str() { - "" | "auto" => None, - _ => Some(backend), - } - }) - .as_deref() -} - -fn log_video_decode_timing_enabled() -> bool { - *LOG_VIDEO_DECODE_TIMING.get_or_init(|| { - std::env::var("SMG_LOG_MM_TIMING") - .map(|value| { - matches!( - value.trim().to_ascii_lowercase().as_str(), - "1" | "true" | "yes" | "on" - ) - }) - .unwrap_or(false) - }) -} - -fn log_video_decode_backend_timing( - backend: &str, - started: Instant, - input_bytes: usize, - cfg: VideoFetchConfig, - error: Option<&MediaConnectorError>, -) { - if !log_video_decode_timing_enabled() { - return; - } - let elapsed_ms = started.elapsed().as_secs_f64() * 1000.0; - match error { - Some(error) => info!( - backend, - ok = false, - input_bytes, - min_frames = cfg.min_frames, - max_frames = cfg.max_frames, - sample_fps = cfg.sample_fps, - elapsed_ms, - error = %error, - "smg_mm_timing video_decode_backend" - ), - None => info!( - backend, - ok = true, - input_bytes, - min_frames = cfg.min_frames, - max_frames = cfg.max_frames, - sample_fps = cfg.sample_fps, - elapsed_ms, - "smg_mm_timing video_decode_backend" - ), - } -} - -#[cfg(feature = "opencv-video")] -fn decode_video_with_opencv_file( - input_path: &std::path::Path, - cfg: VideoFetchConfig, -) -> Result { - let input = input_path.to_str().ok_or_else(|| { - MediaConnectorError::VideoDecode(format!( - "OpenCV video path is not valid UTF-8: {}", - input_path.display() - )) - })?; - - let active_decode = ActiveOpenCvDecode::enter(); - let decoder_threads = opencv_decoder_threads(active_decode.count()); - let capture = open_opencv_video_capture(input, decoder_threads)?; - decode_video_from_opencv_capture(capture, cfg) -} - -#[cfg(feature = "opencv-video")] -fn decode_video_with_opencv_bytes( - bytes: Bytes, - cfg: VideoFetchConfig, -) -> Result { - let active_decode = ActiveOpenCvDecode::enter(); - let decoder_threads = opencv_decoder_threads(active_decode.count()); - let capture = open_opencv_video_capture_from_buffer(bytes, decoder_threads)?; - decode_video_from_opencv_capture(capture, cfg) -} - -#[cfg(feature = "opencv-video")] -trait OpenCvCaptureOwner { - fn capture_mut(&mut self) -> &mut videoio::VideoCapture; -} - -#[cfg(feature = "opencv-video")] -impl OpenCvCaptureOwner for videoio::VideoCapture { - fn capture_mut(&mut self) -> &mut videoio::VideoCapture { - self - } -} - -#[cfg(feature = "opencv-video")] -impl OpenCvCaptureOwner for crate::opencv_buffer::BufferedCapture { - fn capture_mut(&mut self) -> &mut videoio::VideoCapture { - self.capture_mut() - } -} - -#[cfg(feature = "opencv-video")] -fn decode_video_from_opencv_capture( - mut capture: C, - cfg: VideoFetchConfig, -) -> Result -where - C: OpenCvCaptureOwner, -{ - let capture = capture.capture_mut(); - let total_frames = capture - .get(videoio::CAP_PROP_FRAME_COUNT) - .map_err(opencv_decode_error)? - .round() - .max(0.0) as usize; - if total_frames == 0 { - return Err(MediaConnectorError::VideoDecode( - "OpenCV reported zero video frames".to_string(), - )); - } - - let fps = capture - .get(videoio::CAP_PROP_FPS) - .map_err(opencv_decode_error)?; - let frame_indices = opencv_frame_indices(total_frames, fps, cfg); - if frame_indices.is_empty() { - return Err(MediaConnectorError::VideoDecode( - "OpenCV video sampling produced no frame indices".to_string(), - )); - } - - let sampled_frame_counts = counted_frame_indices(&frame_indices); - let unique_frame_count = sampled_frame_counts.len(); - let mut rgb_output = None; - let mut frames = Vec::new(); - frames.try_reserve(frame_indices.len()).map_err(|e| { - MediaConnectorError::VideoDecode(format!( - "failed to reserve {} decoded video frame records: {e}", - frame_indices.len() - )) - })?; - let mut bgr_frame = Mat::default(); - - let timeout = video_process_timeout(); - let started = Instant::now(); - // Advance to each sampled frame by SEQUENTIALLY grabbing the intervening frames - // (cheap decode-without-retrieve) and `read`ing only the sampled ones, instead of - // calling `set(CAP_PROP_POS_FRAMES)` per frame. OpenCV's POS_FRAMES set flushes/ - // re-seeks the decoder on every call (~10 ms/frame even for adjacent frames); - // sequential grab is ~1-2 ms/frame. This matches vLLM's OpenCV video backend and is - // verified bit-exact vs the old per-frame seek on both dense and sparse (non-keyframe) - // sampling, so accuracy is unchanged. `sampled_frame_counts` is monotonic. - // Index of the most recently decoded frame (-1 = nothing read yet). - let mut decoded_pos: i64 = -1; - for (idx, repeat_count) in sampled_frame_counts { - if started.elapsed() >= timeout { - return Err(MediaConnectorError::VideoDecode(format!( - "OpenCV timed out after {:.3} seconds", - timeout.as_secs_f64() - ))); - } - - // Skip-decode the frames between the current position and `idx` so the - // following `read` lands on `idx` without a decoder flush/seek. - while decoded_pos + 1 < idx as i64 { - if started.elapsed() >= timeout { - return Err(MediaConnectorError::VideoDecode(format!( - "OpenCV timed out after {:.3} seconds", - timeout.as_secs_f64() - ))); - } - if !capture.grab().map_err(opencv_decode_error)? { - return Err(MediaConnectorError::VideoDecode(format!( - "OpenCV could not grab intervening frame to reach sampled frame {idx}" - ))); - } - decoded_pos += 1; - } - - let read_successful = capture.read(&mut bgr_frame).map_err(opencv_decode_error)?; - decoded_pos = idx as i64; - if !read_successful || bgr_frame.empty() { - continue; - } - - let decoded_width = u32::try_from(bgr_frame.cols()).map_err(|_| { - MediaConnectorError::VideoDecode(format!( - "OpenCV produced invalid RGB frame width: {}", - bgr_frame.cols() - )) - })?; - let decoded_height = u32::try_from(bgr_frame.rows()).map_err(|_| { - MediaConnectorError::VideoDecode(format!( - "OpenCV produced invalid RGB frame height: {}", - bgr_frame.rows() - )) - })?; - if rgb_output.is_none() { - let frame_size = rawvideo_frame_size(decoded_width, decoded_height)?; - let decoded_bytes = frame_size.checked_mul(unique_frame_count).ok_or_else(|| { - MediaConnectorError::VideoDecode( - "decoded video byte size overflow while reserving RGB frames".to_string(), - ) - })?; - ensure_decoded_byte_limit(decoded_bytes)?; - rgb_output = Some( - crate::opencv_buffer::RgbOutputBuffer::with_capacity(decoded_bytes) - .map_err(MediaConnectorError::VideoDecode)?, - ); - } - let output = rgb_output.as_mut().ok_or_else(|| { - MediaConnectorError::VideoDecode("missing OpenCV RGB output buffer".to_string()) - })?; - let frame_size = rawvideo_frame_size(decoded_width, decoded_height)?; - let new_len = output.len().checked_add(frame_size).ok_or_else(|| { - MediaConnectorError::VideoDecode( - "decoded video byte size overflow while appending RGB frame".to_string(), - ) - })?; - ensure_decoded_byte_limit(new_len)?; - let (offset, len) = output - .push_bgr(&bgr_frame, decoded_width, decoded_height) - .map_err(MediaConnectorError::VideoDecode)?; - let frame = DecodedRgbFrame { - width: decoded_width, - height: decoded_height, - offset, - len, - }; - for _ in 0..repeat_count { - frames.push(frame.clone()); - } - } - - if frames.is_empty() { - return Err(MediaConnectorError::VideoDecode( - "OpenCV produced no readable sampled frames".to_string(), - )); - } - if frames.len() != frame_indices.len() { - return Err(MediaConnectorError::VideoDecode(format!( - "OpenCV produced {} sampled frames, expected {}", - frames.len(), - frame_indices.len() - ))); - } - - let data = rgb_output - .ok_or_else(|| { - MediaConnectorError::VideoDecode("OpenCV produced no RGB output".to_string()) - })? - .into_bytes(); - let sample_fps = effective_sample_fps( - (fps.is_finite() && fps > 0.0).then_some(total_frames as f64 / fps), - cfg, - ); - Ok(DecodedVideoFrames::Rgb { - video: DecodedRgbVideo::new(data, frames), - sample_fps, - }) -} - -#[cfg(feature = "opencv-video")] -fn open_opencv_video_capture_from_buffer( - bytes: Bytes, - decoder_threads: i32, -) -> Result { - crate::opencv_buffer::open_capture(bytes, decoder_threads).map_err(|error| { - MediaConnectorError::VideoDecode(format!("OpenCV could not open video buffer: {error}")) - }) -} - -#[cfg(feature = "opencv-video")] -fn open_opencv_video_capture( - input: &str, - decoder_threads: i32, -) -> Result { - // CAP_PROP_N_THREADS has ID 70. Referencing the numeric ID keeps builds - // compatible with pre-4.8 headers; unsupported backends reject it and use - // the parameter-free fallback below. - const CAP_PROP_N_THREADS: i32 = 70; - let params = Vector::from_slice(&[CAP_PROP_N_THREADS, decoder_threads]); - if let Ok(capture) = - videoio::VideoCapture::from_file_with_params(input, videoio::CAP_FFMPEG, ¶ms) - { - if capture.is_opened().map_err(opencv_decode_error)? { - return Ok(capture); - } - } - - for backend in [videoio::CAP_FFMPEG, videoio::CAP_ANY] { - let Ok(capture) = videoio::VideoCapture::from_file(input, backend) else { - continue; - }; - if capture.is_opened().map_err(opencv_decode_error)? { - return Ok(capture); - } - } - - Err(MediaConnectorError::VideoDecode(format!( - "OpenCV could not open video: {input}" - ))) -} - -#[cfg(feature = "opencv-video")] -struct ActiveOpenCvDecode { - count: usize, -} - -#[cfg(feature = "opencv-video")] -impl ActiveOpenCvDecode { - fn enter() -> Self { - ACTIVE_OPENCV_DECODES.fetch_add(1, Ordering::AcqRel); - // Let a burst of decode tasks become visible before dividing the CPU - // budget. The fixed window also covers blocking-pool ramp-up, where - // arrivals may briefly appear stable before the full burst. - std::thread::sleep(OPENCV_DECODE_BURST_COALESCE); - Self { - count: ACTIVE_OPENCV_DECODES.load(Ordering::Acquire), - } - } - - fn count(&self) -> usize { - self.count - } -} - -#[cfg(feature = "opencv-video")] -impl Drop for ActiveOpenCvDecode { - fn drop(&mut self) { - ACTIVE_OPENCV_DECODES.fetch_sub(1, Ordering::AcqRel); - } -} - -#[cfg(feature = "opencv-video")] -fn opencv_decoder_threads(active_decodes: usize) -> i32 { - let available = *AVAILABLE_OPENCV_CPUS.get_or_init(|| { - std::thread::available_parallelism() - .map(|parallelism| parallelism.get()) - .unwrap_or(1) - }); - adaptive_opencv_decoder_threads(available, active_decodes) -} - -#[cfg(feature = "opencv-video")] -fn adaptive_opencv_decoder_threads(available_cpus: usize, active_decodes: usize) -> i32 { - let available_cpus = available_cpus.max(1); - let active_decodes = active_decodes.max(1); - - // Once eight or more independent decoders fill the CPU quota, codec-level - // threading only adds scheduler contention. - if active_decodes >= OPENCV_LOW_CONCURRENCY_LIMIT && active_decodes >= available_cpus { - return 1; - } - - let (decoder_budget, max_threads) = if active_decodes <= OPENCV_LOW_CONCURRENCY_LIMIT { - let max_threads = if active_decodes <= 2 { 16 } else { 8 }; - ( - available_cpus.saturating_mul(OPENCV_LOW_CONCURRENCY_CPU_MULTIPLIER), - max_threads, - ) - } else { - // Independent decoders supply request-level parallelism at high - // concurrency. Reserve roughly one seventh of the CPU quota for frame - // copies, request handling, and other non-decoder work. - ( - available_cpus - .saturating_mul(OPENCV_HIGH_CONCURRENCY_CPU_BUDGET_NUMERATOR) - .div_ceil(OPENCV_HIGH_CONCURRENCY_CPU_BUDGET_DENOMINATOR), - MAX_OPENCV_DECODER_THREADS, - ) - }; - - (decoder_budget.max(1) / active_decodes).clamp(1, max_threads) as i32 -} - -#[cfg(feature = "opencv-video")] -fn opencv_frame_indices(total_frames: usize, fps: f64, cfg: VideoFetchConfig) -> Vec { - let mut target_frames = if fps.is_finite() && fps > 0.0 { - let duration = total_frames as f64 / fps; - (duration * cfg.sample_fps as f64).round() as usize - } else { - cfg.max_frames - }; - target_frames = target_frames.clamp(cfg.min_frames, cfg.max_frames); - target_frames = target_frames.max(1); - if target_frames == 1 { - return vec![0]; - } - - let last = (total_frames - 1) as f64; - let denom = (target_frames - 1) as f64; - (0..target_frames) - .map(|idx| ((idx as f64 * last) / denom).floor() as usize) - .collect() -} - -#[cfg(feature = "opencv-video")] -fn counted_frame_indices(frame_indices: &[usize]) -> Vec<(usize, usize)> { - let mut counts = Vec::new(); - for &idx in frame_indices { - if let Some((last_idx, count)) = counts.last_mut() { - if *last_idx == idx { - *count += 1; - continue; - } - } - counts.push((idx, 1)); - } - counts -} - -#[cfg(feature = "opencv-video")] -fn opencv_decode_error(err: opencv::Error) -> MediaConnectorError { - MediaConnectorError::VideoDecode(format!("OpenCV video decode failed: {err}")) -} - -async fn decode_video_with_ffmpeg( - input_path: &std::path::Path, - input_bytes: usize, - cfg: VideoFetchConfig, -) -> Result { - if let Ok(metadata) = probe_video_metadata(input_path).await { - let sample_fps = effective_sample_fps(metadata.duration_seconds, cfg); - let started = Instant::now(); - match decode_video_with_ffmpeg_ppm(input_path, cfg, metadata).await { - Ok(rgb_video) => { - log_video_decode_backend_timing("ffmpeg_ppm_file", started, input_bytes, cfg, None); - return Ok(DecodedVideoFrames::Rgb { - video: rgb_video, - sample_fps, - }); - } - Err(error) => { - log_video_decode_backend_timing( - "ffmpeg_ppm_file", - started, - input_bytes, - cfg, - Some(&error), - ); - } - } - - let started = Instant::now(); - match decode_video_with_ffmpeg_raw(input_path, cfg, metadata).await { - Ok(rgb_video) => { - log_video_decode_backend_timing("ffmpeg_raw_file", started, input_bytes, cfg, None); - return Ok(DecodedVideoFrames::Rgb { - video: rgb_video, - sample_fps, - }); - } - Err(error) => { - log_video_decode_backend_timing( - "ffmpeg_raw_file", - started, - input_bytes, - cfg, - Some(&error), - ); - } - } - } - - let started = Instant::now(); - match decode_video_with_ffmpeg_png(input_path, cfg).await { - Ok((frames, sample_fps)) => { - log_video_decode_backend_timing("ffmpeg_png_file", started, input_bytes, cfg, None); - Ok(DecodedVideoFrames::Images { frames, sample_fps }) - } - Err(error) => { - log_video_decode_backend_timing( - "ffmpeg_png_file", - started, - input_bytes, - cfg, - Some(&error), - ); - Err(error) - } - } -} - -fn write_temp_video_file(bytes: &[u8]) -> Result { - let started = Instant::now(); - let mut input_file = tempfile::Builder::new() - .prefix("smg-video-") - .suffix(video_temp_suffix(bytes)) - .tempfile()?; - input_file.write_all(bytes)?; - input_file.flush()?; - if log_video_decode_timing_enabled() { - info!( - nbytes = bytes.len(), - elapsed_ms = started.elapsed().as_secs_f64() * 1000.0, - suffix = video_temp_suffix(bytes), - "smg_mm_timing video_tempfile_write" - ); - } - Ok(input_file) -} - -fn video_temp_suffix(bytes: &[u8]) -> &'static str { - if bytes.len() >= 12 && bytes.get(4..8) == Some(b"ftyp") { - return ".mp4"; - } - if bytes.starts_with(&[0x1a, 0x45, 0xdf, 0xa3]) { - return ".webm"; - } - if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && bytes.get(8..12) == Some(b"AVI ") { - return ".avi"; - } - if bytes.starts_with(b"OggS") { - return ".ogv"; - } - if bytes.starts_with(&[0x00, 0x00, 0x01, 0xba]) { - return ".mpg"; - } - ".video" -} - -fn video_process_timeout() -> Duration { - *VIDEO_PROCESS_TIMEOUT.get_or_init(|| { - std::env::var("SMG_VIDEO_PROCESS_TIMEOUT_SECS") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|seconds| seconds.is_finite() && *seconds > 0.0) - .map(Duration::from_secs_f64) - .unwrap_or(DEFAULT_VIDEO_PROCESS_TIMEOUT) - }) -} - -fn video_max_decoded_bytes() -> usize { - env_byte_limit( - &VIDEO_MAX_DECODED_BYTES, - "SMG_VIDEO_MAX_DECODED_BYTES", - DEFAULT_VIDEO_MAX_DECODED_BYTES, - ) -} - -fn ensure_decoded_byte_limit(bytes: usize) -> Result<(), MediaConnectorError> { - let limit = video_max_decoded_bytes(); - if bytes > limit { - return Err(MediaConnectorError::VideoDecode(format!( - "decoded video RGB payload would be {bytes} bytes, exceeding SMG_VIDEO_MAX_DECODED_BYTES={limit}" - ))); - } - Ok(()) -} - -fn checked_decoded_rgb_bytes( - frame_count: usize, - frame_size: usize, -) -> Result { - let bytes = frame_count.checked_mul(frame_size).ok_or_else(|| { - MediaConnectorError::VideoDecode(format!( - "decoded video byte size overflow for {frame_count} frames of {frame_size} bytes" - )) - })?; - ensure_decoded_byte_limit(bytes)?; - Ok(bytes) -} - -async fn run_video_command_output( - mut command: Command, - program: &'static str, -) -> Result { - command - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .kill_on_drop(true); - let child = command.spawn().map_err(|e| { - if e.kind() == std::io::ErrorKind::NotFound { - MediaConnectorError::VideoDecode(format!( - "{program} executable not found; install {program} to decode video_url inputs" - )) - } else { - MediaConnectorError::Io(e) - } - })?; - - let timeout = video_process_timeout(); - match time::timeout(timeout, child.wait_with_output()).await { - Ok(Ok(output)) => Ok(output), - Ok(Err(error)) => Err(MediaConnectorError::Io(error)), - Err(_) => Err(MediaConnectorError::VideoDecode(format!( - "{program} timed out after {:.3} seconds", - timeout.as_secs_f64() - ))), - } -} - -async fn decode_video_with_ffmpeg_ppm( - input_path: &std::path::Path, - cfg: VideoFetchConfig, - metadata: VideoMetadata, -) -> Result { - let fps_filter = fps_filter_for_metadata(metadata, cfg); - let max_frames = cfg.max_frames.to_string(); - let frame_size = rawvideo_frame_size(metadata.width, metadata.height)?; - let target_frames = expected_sampled_frame_count(metadata, cfg); - let decoded_bytes = checked_decoded_rgb_bytes(target_frames, frame_size)?; - let output_limit = decoded_bytes - .checked_add(target_frames.saturating_mul(64)) - .unwrap_or_else(video_max_decoded_bytes) - .min(video_max_decoded_bytes()) - .to_string(); - let mut command = Command::new("ffmpeg"); - command - .args(["-hide_banner", "-loglevel", "error", "-nostdin", "-i"]) - .arg(input_path) - .args([ - "-vf", - &fps_filter, - "-frames:v", - &max_frames, - "-fs", - &output_limit, - "-f", - "image2pipe", - "-vcodec", - "ppm", - "-pix_fmt", - "rgb24", - "pipe:1", - ]); - let output = run_video_command_output(command, "ffmpeg").await?; - - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - return Err(MediaConnectorError::VideoDecode(format!( - "ffmpeg failed: {stderr}" - ))); - } - - parse_ppm_rgb_video(Bytes::from(output.stdout)) -} - -async fn decode_video_with_ffmpeg_raw( - input_path: &std::path::Path, - cfg: VideoFetchConfig, - metadata: VideoMetadata, -) -> Result { - let fps_filter = fps_filter_for_metadata(metadata, cfg); - let max_frames = cfg.max_frames.to_string(); - let frame_size = rawvideo_frame_size(metadata.width, metadata.height)?; - let target_frames = expected_sampled_frame_count(metadata, cfg); - let decoded_bytes = checked_decoded_rgb_bytes(target_frames, frame_size)?; - let output_limit = decoded_bytes.to_string(); - let mut command = Command::new("ffmpeg"); - // Rawvideo has no per-frame header, so we interpret stdout using ffprobe's - // coded stream dimensions. Disable FFmpeg autorotation here; otherwise a - // display-matrix rotation can swap output width/height and corrupt framing. - command - .args([ - "-hide_banner", - "-loglevel", - "error", - "-nostdin", - "-noautorotate", - "-i", - ]) - .arg(input_path) - .args([ - "-vf", - &fps_filter, - "-frames:v", - &max_frames, - "-fs", - &output_limit, - "-f", - "rawvideo", - "-pix_fmt", - "rgb24", - "pipe:1", - ]); - let output = run_video_command_output(command, "ffmpeg").await?; - - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - return Err(MediaConnectorError::VideoDecode(format!( - "ffmpeg failed: {stderr}" - ))); - } - - let frame_count = output.stdout.len() / frame_size; - checked_decoded_rgb_bytes(frame_count, frame_size)?; - let mut frames = Vec::new(); - frames.try_reserve(frame_count).map_err(|e| { - MediaConnectorError::VideoDecode(format!( - "failed to reserve {frame_count} decoded video frame records: {e}" - )) - })?; - for idx in 0..frame_count { - frames.push(DecodedRgbFrame { - width: metadata.width, - height: metadata.height, - offset: idx * frame_size, - len: frame_size, - }); - } - let remainder = output.stdout.len() % frame_size; - if remainder != 0 { - return Err(MediaConnectorError::VideoDecode(format!( - "ffmpeg rawvideo output has trailing partial frame: {remainder} bytes" - ))); - } - if frames.is_empty() { - return Err(MediaConnectorError::VideoDecode( - "ffmpeg produced no frames".to_string(), - )); - } - Ok(DecodedRgbVideo::new(Bytes::from(output.stdout), frames)) -} - -async fn decode_video_with_ffmpeg_png( - input_path: &std::path::Path, - cfg: VideoFetchConfig, -) -> Result<(Vec, f32), MediaConnectorError> { - let (fps_filter, sample_fps) = sampling_filter_for_video(input_path, cfg).await; - let max_frames = cfg.max_frames.to_string(); - let output_limit = video_max_decoded_bytes().to_string(); - let mut command = Command::new("ffmpeg"); - command - .args(["-hide_banner", "-loglevel", "error", "-nostdin", "-i"]) - .arg(input_path) - .args([ - "-vf", - &fps_filter, - "-frames:v", - &max_frames, - "-fs", - &output_limit, - "-f", - "image2pipe", - "-vcodec", - "png", - "pipe:1", - ]); - let output = run_video_command_output(command, "ffmpeg").await?; - - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - return Err(MediaConnectorError::VideoDecode(format!( - "ffmpeg failed: {stderr}" - ))); - } - - let pngs = split_png_stream(&output.stdout)?; - let mut frames = Vec::with_capacity(pngs.len()); - let mut decoded_bytes = 0usize; - for png in pngs { - let image = image::load_from_memory(png)?; - let frame_size = rawvideo_frame_size(image.width(), image.height())?; - decoded_bytes = decoded_bytes.checked_add(frame_size).ok_or_else(|| { - MediaConnectorError::VideoDecode("PNG decoded byte size overflow".to_string()) - })?; - ensure_decoded_byte_limit(decoded_bytes)?; - frames.push(image); - } - if frames.is_empty() { - return Err(MediaConnectorError::VideoDecode( - "ffmpeg produced no frames".to_string(), - )); - } - Ok((frames, sample_fps)) -} - -#[derive(Debug, Clone, Copy)] -struct VideoMetadata { - width: u32, - height: u32, - duration_seconds: Option, -} - -#[derive(Debug, Clone, Copy)] -struct ProbedVideoInfo { - width: Option, - height: Option, - duration_seconds: Option, -} - -async fn probe_video_metadata( - input_path: &std::path::Path, -) -> Result { - let info = probe_video_info(input_path).await?; - let width = info.width.ok_or_else(|| { - MediaConnectorError::VideoDecode("ffprobe did not return video width".to_string()) - })?; - let height = info.height.ok_or_else(|| { - MediaConnectorError::VideoDecode("ffprobe did not return video height".to_string()) - })?; - Ok(VideoMetadata { - width, - height, - duration_seconds: info.duration_seconds, - }) -} - -async fn probe_video_info( - input_path: &std::path::Path, -) -> Result { - let mut command = Command::new("ffprobe"); - command - .args([ - "-v", - "error", - "-select_streams", - "v:0", - "-show_entries", - "stream=width,height,duration,duration_ts,time_base:format=duration", - "-of", - "json", - ]) - .arg(input_path); - let output = run_video_command_output(command, "ffprobe").await?; - - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - return Err(MediaConnectorError::VideoDecode(format!( - "ffprobe failed: {stderr}" - ))); - } - - parse_ffprobe_video_info(&output.stdout) -} - -fn parse_ffprobe_video_info(stdout: &[u8]) -> Result { - let probe: serde_json::Value = serde_json::from_slice(stdout).map_err(|error| { - MediaConnectorError::VideoDecode(format!("failed to parse ffprobe output: {error}")) - })?; - let video_stream = probe - .get("streams") - .and_then(serde_json::Value::as_array) - .and_then(|streams| streams.first()); - - let width = video_stream - .and_then(|stream| stream.get("width")) - .and_then(json_u32); - let height = video_stream - .and_then(|stream| stream.get("height")) - .and_then(json_u32); - let stream_duration = video_stream - .and_then(|stream| stream.get("duration")) - .and_then(json_positive_f64); - let stream_time_base_duration = video_stream.and_then(|stream| { - let duration_ts = stream.get("duration_ts").and_then(json_positive_f64)?; - let time_base = stream - .get("time_base") - .and_then(serde_json::Value::as_str) - .and_then(parse_time_base)?; - let duration = duration_ts * time_base; - (duration.is_finite() && duration > 0.0).then_some(duration) - }); - let format_duration = probe - .get("format") - .and_then(|format| format.get("duration")) - .and_then(json_positive_f64); - - Ok(ProbedVideoInfo { - width, - height, - duration_seconds: stream_duration - .or(stream_time_base_duration) - .or(format_duration), - }) -} - -fn json_u32(value: &serde_json::Value) -> Option { - value - .as_u64() - .and_then(|value| u32::try_from(value).ok()) - .or_else(|| value.as_str()?.parse::().ok()) -} - -fn json_positive_f64(value: &serde_json::Value) -> Option { - value - .as_f64() - .or_else(|| value.as_str()?.parse::().ok()) - .filter(|value| value.is_finite() && *value > 0.0) -} - -fn parse_time_base(value: &str) -> Option { - let (numerator, denominator) = value.split_once('/')?; - let numerator = numerator.parse::().ok()?; - let denominator = denominator.parse::().ok()?; - let time_base = numerator / denominator; - (time_base.is_finite() && time_base > 0.0).then_some(time_base) -} - -fn fps_filter_for_metadata(metadata: VideoMetadata, cfg: VideoFetchConfig) -> String { - if let Some(duration) = metadata.duration_seconds { - if let Some(filter) = fps_filter_for_duration(duration, cfg) { - return filter; - } - } - - format!("fps={}", cfg.sample_fps) -} - -fn expected_sampled_frame_count(metadata: VideoMetadata, cfg: VideoFetchConfig) -> usize { - if let Some(duration) = metadata.duration_seconds { - if duration.is_finite() && duration > 0.0 { - return (duration * cfg.sample_fps as f64) - .round() - .clamp(cfg.min_frames as f64, cfg.max_frames as f64) as usize; - } - } - cfg.max_frames -} - -fn effective_sample_fps(duration_seconds: Option, cfg: VideoFetchConfig) -> f32 { - duration_seconds - .filter(|duration| duration.is_finite() && *duration > 0.0) - .map(|duration| { - let target_frames = (duration * cfg.sample_fps as f64) - .round() - .clamp(cfg.min_frames as f64, cfg.max_frames as f64); - (target_frames / duration) as f32 - }) - .filter(|fps| fps.is_finite() && *fps > 0.0) - .unwrap_or(cfg.sample_fps) -} - -fn fps_filter_for_duration(duration: f64, cfg: VideoFetchConfig) -> Option { - if !duration.is_finite() || duration <= 0.0 { - return None; - } - let target_frames = (duration * cfg.sample_fps as f64) - .round() - .clamp(cfg.min_frames as f64, cfg.max_frames as f64); - let fps = (target_frames / duration).max(f64::EPSILON); - Some(format!("fps={fps:.6}")) -} - -async fn sampling_filter_for_video( - input_path: &std::path::Path, - cfg: VideoFetchConfig, -) -> (String, f32) { - if let Ok(duration) = probe_video_duration_seconds(input_path).await { - if let Some(filter) = fps_filter_for_duration(duration, cfg) { - return (filter, effective_sample_fps(Some(duration), cfg)); - } - } - - (format!("fps={}", cfg.sample_fps), cfg.sample_fps) -} - -async fn probe_video_duration_seconds( - input_path: &std::path::Path, -) -> Result { - match probe_video_info(input_path).await { - Ok(ProbedVideoInfo { - duration_seconds: Some(duration), - .. - }) => Ok(duration), - Ok(_) | Err(_) => probe_video_duration_seconds_with_ffmpeg(input_path).await, - } -} - -async fn probe_video_duration_seconds_with_ffmpeg( - input_path: &std::path::Path, -) -> Result { - let mut command = Command::new("ffmpeg"); - command - .args(["-hide_banner", "-nostdin", "-i"]) - .arg(input_path); - let output = run_video_command_output(command, "ffmpeg").await?; - - let stderr = String::from_utf8_lossy(&output.stderr); - parse_ffmpeg_duration_seconds(&stderr).ok_or_else(|| { - MediaConnectorError::VideoDecode("failed to parse ffmpeg duration".to_string()) - }) -} - -fn parse_ffmpeg_duration_seconds(stderr: &str) -> Option { - let marker = "Duration:"; - let start = stderr.find(marker)? + marker.len(); - let duration = stderr[start..].trim_start().split(',').next()?.trim(); - let mut parts = duration.split(':'); - let hours = parts.next()?.parse::().ok()?; - let minutes = parts.next()?.parse::().ok()?; - let seconds = parts.next()?.parse::().ok()?; - Some(hours * 3600.0 + minutes * 60.0 + seconds) -} - -fn split_png_stream(bytes: &[u8]) -> Result, MediaConnectorError> { - const PNG_SIG: &[u8; 8] = b"\x89PNG\r\n\x1a\n"; - const IEND: &[u8; 4] = b"IEND"; - - let mut frames = Vec::new(); - let mut pos = 0; - while pos < bytes.len() { - let Some(rel_start) = bytes[pos..] - .windows(PNG_SIG.len()) - .position(|w| w == PNG_SIG) - else { - break; - }; - let start = pos + rel_start; - let mut cursor = start + PNG_SIG.len(); - - loop { - let remaining = bytes.len() - cursor; - if remaining < 12 { - return Err(MediaConnectorError::VideoDecode( - "truncated PNG frame in ffmpeg output".to_string(), - )); - } - let mut len_bytes = [0_u8; 4]; - len_bytes.copy_from_slice(&bytes[cursor..cursor + 4]); - let len = u32::from_be_bytes(len_bytes) as usize; - let chunk_type = &bytes[cursor + 4..cursor + 8]; - if remaining - 12 < len { - return Err(MediaConnectorError::VideoDecode( - "truncated PNG chunk in ffmpeg output".to_string(), - )); - } - cursor += 12 + len; - if chunk_type == IEND { - frames.push(&bytes[start..cursor]); - pos = cursor; - break; - } - } - } - - Ok(frames) -} - -#[cfg(test)] -fn parse_ppm_stream(bytes: &[u8]) -> Result, MediaConnectorError> { - let layouts = parse_ppm_frame_layout(bytes)?; - let mut frames = Vec::with_capacity(layouts.len()); - for layout in layouts { - let end = layout.offset.checked_add(layout.len).ok_or_else(|| { - MediaConnectorError::VideoDecode("PPM frame size overflow".to_string()) - })?; - let image = image::RgbImage::from_raw( - layout.width, - layout.height, - bytes[layout.offset..end].to_vec(), - ) - .ok_or_else(|| { - MediaConnectorError::VideoDecode(format!( - "failed to build RGB frame from {} bytes for {}x{} video", - layout.len, layout.width, layout.height - )) - })?; - frames.push(image::DynamicImage::ImageRgb8(image)); - } - Ok(frames) -} - -fn parse_ppm_rgb_video(bytes: Bytes) -> Result { - let layouts = parse_ppm_frame_layout(&bytes)?; - let decoded_bytes = layouts.iter().try_fold(0usize, |total, frame| { - total.checked_add(frame.len).ok_or_else(|| { - MediaConnectorError::VideoDecode("PPM decoded byte size overflow".to_string()) - }) - })?; - ensure_decoded_byte_limit(decoded_bytes)?; - Ok(DecodedRgbVideo::new(bytes, layouts)) -} - -fn parse_ppm_frame_layout(bytes: &[u8]) -> Result, MediaConnectorError> { - let mut frames = Vec::new(); - let mut pos = 0; - - while pos < bytes.len() { - skip_ppm_whitespace_and_comments(bytes, &mut pos); - if pos >= bytes.len() { - break; - } - - let magic = read_ppm_token(bytes, &mut pos)?.ok_or_else(|| { - MediaConnectorError::VideoDecode("truncated PPM frame header".to_string()) - })?; - if magic != b"P6" { - return Err(MediaConnectorError::VideoDecode(format!( - "unsupported PPM magic: {}", - String::from_utf8_lossy(magic) - ))); - } - let width = parse_ppm_u32(bytes, &mut pos, "width")?; - let height = parse_ppm_u32(bytes, &mut pos, "height")?; - let max_value = parse_ppm_u32(bytes, &mut pos, "max value")?; - if width == 0 || height == 0 { - return Err(MediaConnectorError::VideoDecode( - "PPM frame dimensions must be non-zero".to_string(), - )); - } - if max_value != 255 { - return Err(MediaConnectorError::VideoDecode(format!( - "unsupported PPM max value: {max_value}" - ))); - } - if pos >= bytes.len() || !bytes[pos].is_ascii_whitespace() { - return Err(MediaConnectorError::VideoDecode( - "PPM header is not followed by pixel data".to_string(), - )); - } - pos += 1; - - let frame_size = (width as usize) - .checked_mul(height as usize) - .and_then(|pixels| pixels.checked_mul(3)) - .ok_or_else(|| { - MediaConnectorError::VideoDecode(format!( - "PPM frame dimensions are too large: {width}x{height}" - )) - })?; - let end = pos.checked_add(frame_size).ok_or_else(|| { - MediaConnectorError::VideoDecode("PPM frame size overflow".to_string()) - })?; - if end > bytes.len() { - return Err(MediaConnectorError::VideoDecode( - "truncated PPM frame pixel data".to_string(), - )); - } - frames.push(DecodedRgbFrame { - width, - height, - offset: pos, - len: frame_size, - }); - pos = end; - } - - if frames.is_empty() { - return Err(MediaConnectorError::VideoDecode( - "ffmpeg produced no frames".to_string(), - )); - } - - Ok(frames) -} - -fn rawvideo_frame_size(width: u32, height: u32) -> Result { - let frame_size = (width as usize) - .checked_mul(height as usize) - .and_then(|pixels| pixels.checked_mul(3)) - .ok_or_else(|| { - MediaConnectorError::VideoDecode(format!( - "video frame dimensions are too large: {width}x{height}" - )) - })?; - if frame_size == 0 { - return Err(MediaConnectorError::VideoDecode( - "video frame dimensions must be non-zero".to_string(), - )); - } - Ok(frame_size) -} - -fn parse_ppm_u32(bytes: &[u8], pos: &mut usize, field: &str) -> Result { - let token = read_ppm_token(bytes, pos)? - .ok_or_else(|| MediaConnectorError::VideoDecode(format!("truncated PPM {field} header")))?; - std::str::from_utf8(token) - .ok() - .and_then(|value| value.parse::().ok()) - .ok_or_else(|| { - MediaConnectorError::VideoDecode(format!( - "invalid PPM {field}: {}", - String::from_utf8_lossy(token) - )) - }) -} - -fn read_ppm_token<'a>( - bytes: &'a [u8], - pos: &mut usize, -) -> Result, MediaConnectorError> { - skip_ppm_whitespace_and_comments(bytes, pos); - if *pos >= bytes.len() { - return Ok(None); - } - - let start = *pos; - while *pos < bytes.len() && !bytes[*pos].is_ascii_whitespace() { - if bytes[*pos] == b'#' { - return Err(MediaConnectorError::VideoDecode( - "unexpected PPM comment inside token".to_string(), - )); - } - *pos += 1; - } - Ok(Some(&bytes[start..*pos])) -} - -fn skip_ppm_whitespace_and_comments(bytes: &[u8], pos: &mut usize) { - loop { - while *pos < bytes.len() && bytes[*pos].is_ascii_whitespace() { - *pos += 1; - } - if *pos < bytes.len() && bytes[*pos] == b'#' { - while *pos < bytes.len() && bytes[*pos] != b'\n' { - *pos += 1; - } - continue; - } - break; - } -} - -#[cfg(test)] -mod tests { - use std::io::Write as _; - - use bytes::Bytes; - use futures::stream; - - use super::{ - checked_payload_length, collect_http_body_with_limit, decode_base64_with_limit, - effective_sample_fps, ensure_input_byte_limit, expected_sampled_frame_count, - fps_filter_for_metadata, parse_ffmpeg_duration_seconds, parse_ffprobe_video_info, - parse_ppm_stream, read_file_with_limit, split_png_stream, video_temp_suffix, - MediaConnectorError, VideoFetchConfig, VideoMetadata, - }; - - const TINY_PNG: &[u8] = &[ - 137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 4, - 0, 0, 0, 181, 28, 12, 2, 0, 0, 0, 11, 73, 68, 65, 84, 120, 218, 99, 96, 96, 0, 0, 0, 3, 0, - 1, 43, 9, 141, 84, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, - ]; - - #[test] - fn splits_concatenated_png_stream() { - let mut stream = Vec::new(); - stream.extend_from_slice(TINY_PNG); - stream.extend_from_slice(TINY_PNG); - - let frames = match split_png_stream(&stream) { - Ok(frames) => frames, - Err(err) => panic!("split png stream failed: {err}"), - }; - assert_eq!(frames.len(), 2); - assert_eq!(frames[0], TINY_PNG); - assert_eq!(frames[1], TINY_PNG); - } - - #[test] - fn parses_ffmpeg_duration() { - let stderr = "Input #0, mov,mp4,m4a,3gp,3g2,mj2, from 'video.mp4':\n Duration: 00:01:23.45, start: 0.000000, bitrate: 123 kb/s"; - assert_eq!(parse_ffmpeg_duration_seconds(stderr), Some(83.45)); - } - - #[test] - fn ffprobe_metadata_prefers_short_video_stream_over_long_container() { - let output = br#"{ - "streams": [{ - "width": 320, - "height": 240, - "duration": "1.000000", - "duration_ts": 30, - "time_base": "1/30" - }], - "format": {"duration": "120.000000"} - }"#; - let info = parse_ffprobe_video_info(output).expect("valid ffprobe output"); - assert_eq!(info.duration_seconds, Some(1.0)); - - let cfg = VideoFetchConfig { - min_frames: 4, - max_frames: 8, - sample_fps: 2.0, - }; - let metadata = VideoMetadata { - width: info.width.expect("video width"), - height: info.height.expect("video height"), - duration_seconds: info.duration_seconds, - }; - assert_eq!(expected_sampled_frame_count(metadata, cfg), 4); - assert_eq!(fps_filter_for_metadata(metadata, cfg), "fps=4.000000"); - } - - #[test] - fn ffprobe_metadata_uses_stream_time_base_before_container_duration() { - let output = br#"{ - "streams": [{ - "width": "640", - "height": "360", - "duration": "N/A", - "duration_ts": 45, - "time_base": "1/30" - }], - "format": {"duration": "90.000000"} - }"#; - let info = parse_ffprobe_video_info(output).expect("valid ffprobe output"); - assert_eq!(info.width, Some(640)); - assert_eq!(info.height, Some(360)); - assert_eq!(info.duration_seconds, Some(1.5)); - } - - #[test] - fn detects_video_temp_suffix_from_container_header() { - let mut mp4 = vec![0; 12]; - mp4[4..8].copy_from_slice(b"ftyp"); - assert_eq!(video_temp_suffix(&mp4), ".mp4"); - assert_eq!(video_temp_suffix(&[0x1a, 0x45, 0xdf, 0xa3]), ".webm"); - assert_eq!(video_temp_suffix(b"RIFF....AVI "), ".avi"); - assert_eq!(video_temp_suffix(b"OggS"), ".ogv"); - assert_eq!(video_temp_suffix(&[0x00, 0x00, 0x01, 0xba]), ".mpg"); - assert_eq!(video_temp_suffix(b"unknown"), ".video"); - } - - #[test] - fn parses_concatenated_ppm_stream() { - let stream = b"P6\n2 1\n255\n\x01\x02\x03\x04\x05\x06P6\n# comment\n1 2\n255\n\x07\x08\x09\x0a\x0b\x0c"; - - let frames = match parse_ppm_stream(stream) { - Ok(frames) => frames, - Err(err) => panic!("parse ppm stream failed: {err}"), - }; - assert_eq!(frames.len(), 2); - assert_eq!(frames[0].width(), 2); - assert_eq!(frames[0].height(), 1); - assert_eq!(frames[1].width(), 1); - assert_eq!(frames[1].height(), 2); - } - - #[test] - fn effective_sample_fps_tracks_min_and_max_frame_clamps() { - let cfg = VideoFetchConfig { - min_frames: 4, - max_frames: 8, - sample_fps: 2.0, - }; - - assert_eq!(effective_sample_fps(Some(1.0), cfg), 4.0); - assert!((effective_sample_fps(Some(10.0), cfg) - 0.8).abs() < 1e-6); - assert_eq!(effective_sample_fps(Some(3.0), cfg), 2.0); - assert_eq!(effective_sample_fps(None, cfg), 2.0); - } - - #[test] - fn rejects_truncated_ppm_stream() { - assert!(parse_ppm_stream(b"P6\n2 1\n255\n\x01\x02").is_err()); - } - - #[test] - fn rejects_invalid_ppm_header() { - assert!(parse_ppm_stream(b"P3\n1 1\n255\n\x01\x02\x03").is_err()); - assert!(parse_ppm_stream(b"P6\n1 1\n65535\n\x01\x02\x03").is_err()); - } - - #[test] - fn rejects_zero_dimension_ppm_stream() { - assert!(parse_ppm_stream(b"P6\n0 1\n255\n").is_err()); - assert!(parse_ppm_stream(b"P6\n1 0\n255\n").is_err()); - } - - #[test] - fn rejects_overflowing_ppm_frame_size() { - assert!(parse_ppm_stream(b"P6\n4294967295 4294967295\n255\n").is_err()); - } - - #[test] - fn enforces_media_payload_limit_for_each_label() { - for media in ["image", "video", "audio"] { - assert!(ensure_input_byte_limit(4, 4, media).is_ok()); - assert!(matches!( - ensure_input_byte_limit(5, 4, media), - Err(MediaConnectorError::PayloadTooLarge { media: actual, limit: 4 }) - if actual == media - )); - } - - assert!(checked_payload_length(usize::MAX, 1, usize::MAX, "audio").is_err()); - } - - #[test] - fn enforces_base64_payload_limit_before_and_after_decode() { - for media in ["image", "video", "audio"] { - assert_eq!(decode_base64_with_limit("AAAA", 3, media).unwrap().len(), 3); - assert!(matches!( - decode_base64_with_limit("AAAA", 2, media), - Err(MediaConnectorError::PayloadTooLarge { media: actual, limit: 2 }) - if actual == media - )); - } - - assert!(matches!( - decode_base64_with_limit("AAAAAAAAAAAAAAAA", 8, "audio"), - Err(MediaConnectorError::PayloadTooLarge { - media: "audio", - limit: 8 - }) - )); - assert!(matches!( - decode_base64_with_limit("AAAAAAAAAAAA", 8, "audio"), - Err(MediaConnectorError::PayloadTooLarge { - media: "audio", - limit: 8 - }) - )); - } - - #[tokio::test] - async fn reads_files_at_limit_and_rejects_limit_plus_one() -> Result<(), MediaConnectorError> { - let mut file = tempfile::NamedTempFile::new()?; - file.write_all(b"12345")?; - file.flush()?; - - for media in ["image", "video", "audio"] { - let bytes = read_file_with_limit(file.path(), 5, media).await?; - assert_eq!(bytes, Bytes::from_static(b"12345")); - - assert!(matches!( - read_file_with_limit(file.path(), 4, media).await, - Err(MediaConnectorError::PayloadTooLarge { media: actual, limit: 4 }) - if actual == media - )); - } - Ok(()) - } - - #[tokio::test] - async fn collects_http_body_with_content_and_streaming_limits( - ) -> Result<(), MediaConnectorError> { - let known_oversized = reqwest::Response::from(http::Response::new(reqwest::Body::from( - Bytes::from_static(b"12345"), - ))); - assert!(matches!( - collect_http_body_with_limit(known_oversized, 4, "image").await, - Err(MediaConnectorError::PayloadTooLarge { - media: "image", - limit: 4 - }) - )); - - let oversized_stream = stream::iter([ - Ok::<_, std::io::Error>(Bytes::from_static(b"123")), - Ok(Bytes::from_static(b"45")), - ]); - let unknown_oversized = reqwest::Response::from(http::Response::new( - reqwest::Body::wrap_stream(oversized_stream), - )); - assert!(matches!( - collect_http_body_with_limit(unknown_oversized, 4, "video").await, - Err(MediaConnectorError::PayloadTooLarge { - media: "video", - limit: 4 - }) - )); - - let within_limit_stream = stream::iter([ - Ok::<_, std::io::Error>(Bytes::from_static(b"12")), - Ok(Bytes::from_static(b"34")), - ]); - let within_limit = reqwest::Response::from(http::Response::new( - reqwest::Body::wrap_stream(within_limit_stream), - )); - let body = collect_http_body_with_limit(within_limit, 4, "audio").await?; - assert_eq!(body, Bytes::from_static(b"1234")); - Ok(()) - } - - #[cfg(feature = "opencv-video")] - #[test] - fn opencv_sampling_preserves_min_frames_for_short_clips() { - let cfg = VideoFetchConfig { - min_frames: 4, - max_frames: 8, - sample_fps: 2.0, - }; - let indices = super::opencv_frame_indices(1, 30.0, cfg); - assert_eq!(indices, vec![0, 0, 0, 0]); - assert_eq!(super::counted_frame_indices(&indices), vec![(0, 4)]); - } - - #[cfg(feature = "opencv-video")] - #[test] - fn opencv_decoder_threads_share_cpu_budget_across_active_decodes() { - assert_eq!(super::adaptive_opencv_decoder_threads(224, 1), 16); - assert_eq!(super::adaptive_opencv_decoder_threads(2, 1), 4); - assert_eq!(super::adaptive_opencv_decoder_threads(4, 2), 4); - assert_eq!(super::adaptive_opencv_decoder_threads(8, 4), 4); - assert_eq!(super::adaptive_opencv_decoder_threads(8, 8), 1); - assert_eq!(super::adaptive_opencv_decoder_threads(8, 9), 1); - assert_eq!(super::adaptive_opencv_decoder_threads(16, 8), 4); - assert_eq!(super::adaptive_opencv_decoder_threads(16, 16), 1); - assert_eq!(super::adaptive_opencv_decoder_threads(224, 8), 8); - assert_eq!(super::adaptive_opencv_decoder_threads(224, 32), 6); - assert_eq!(super::adaptive_opencv_decoder_threads(8, 32), 1); - assert_eq!(super::adaptive_opencv_decoder_threads(1, 0), 2); - } -} diff --git a/crates/multimodal/src/opencv_buffer.rs b/crates/multimodal/src/opencv_buffer.rs deleted file mode 100644 index 5f97c5a98..000000000 --- a/crates/multimodal/src/opencv_buffer.rs +++ /dev/null @@ -1,178 +0,0 @@ -//! Safe wrapper for OpenCV's buffered video capture constructor. -#![allow(unsafe_code)] - -use std::ffi::{c_char, c_void, CStr}; - -use bytes::Bytes; -use opencv::{ - core::{Mat, CV_8UC3}, - imgproc, - prelude::*, - traits::OpenCVFromExtern, - videoio, -}; - -unsafe extern "C" { - fn smg_opencv_capture_from_buffer( - data: *const u8, - size: usize, - decoder_threads: i32, - error: *mut c_char, - error_capacity: usize, - ) -> *mut c_void; -} - -pub(crate) struct BufferedCapture { - capture: videoio::VideoCapture, - _bytes: Bytes, -} - -pub(crate) struct RgbOutputBuffer { - data: Vec, -} - -impl RgbOutputBuffer { - pub(crate) fn with_capacity(capacity: usize) -> Result { - let mut data = Vec::new(); - data.try_reserve_exact(capacity).map_err(|error| { - format!("failed to reserve {capacity} decoded video bytes: {error}") - })?; - Ok(Self { data }) - } - - pub(crate) fn len(&self) -> usize { - self.data.len() - } - - pub(crate) fn push_bgr( - &mut self, - bgr_frame: &Mat, - width: u32, - height: u32, - ) -> Result<(usize, usize), String> { - let width_i32 = i32::try_from(width) - .map_err(|_| format!("OpenCV RGB output width does not fit i32: {width}"))?; - let height_i32 = i32::try_from(height) - .map_err(|_| format!("OpenCV RGB output height does not fit i32: {height}"))?; - let frame_size = usize::try_from(width) - .ok() - .and_then(|width| { - usize::try_from(height) - .ok() - .and_then(|height| width.checked_mul(height)) - }) - .and_then(|pixels| pixels.checked_mul(3)) - .ok_or_else(|| "OpenCV RGB frame byte size overflow".to_string())?; - self.data.try_reserve(frame_size).map_err(|error| { - format!("failed to reserve {frame_size} decoded video bytes: {error}") - })?; - let offset = self.data.len(); - let output_data = self.data.spare_capacity_mut().as_mut_ptr().cast::(); - // SAFETY: `try_reserve` made at least `frame_size` writable bytes - // available, and the Vec cannot move while `output` borrows that region. - let mut output = unsafe { - Mat::new_rows_cols_with_data_unsafe_def( - height_i32, - width_i32, - CV_8UC3, - output_data.cast(), - ) - } - .map_err(|error| error.to_string())?; - imgproc::cvt_color_def(bgr_frame, &mut output, imgproc::COLOR_BGR2RGB) - .map_err(|error| error.to_string())?; - if output.data() != output_data - || output.rows() != height_i32 - || output.cols() != width_i32 - || output.typ() != CV_8UC3 - { - return Err("OpenCV replaced the caller-provided RGB output buffer".to_string()); - } - drop(output); - // SAFETY: cvtColor successfully initialized exactly `frame_size` bytes in - // the caller-provided output region, whose identity was checked above. - unsafe { self.data.set_len(offset + frame_size) }; - Ok((offset, frame_size)) - } - - pub(crate) fn into_bytes(self) -> Bytes { - Bytes::from(self.data) - } -} - -impl BufferedCapture { - pub(crate) fn capture_mut(&mut self) -> &mut videoio::VideoCapture { - &mut self.capture - } -} - -pub(crate) fn open_capture(bytes: Bytes, decoder_threads: i32) -> Result { - let mut error = [0 as c_char; 512]; - // SAFETY: `BufferedCapture` owns `bytes` for at least as long as the capture. - let capture = unsafe { - smg_opencv_capture_from_buffer( - bytes.as_ptr(), - bytes.len(), - decoder_threads, - error.as_mut_ptr(), - error.len(), - ) - }; - if capture.is_null() { - // SAFETY: the bridge always writes a NUL-terminated message on failure. - return Err(unsafe { CStr::from_ptr(error.as_ptr()) } - .to_string_lossy() - .into_owned()); - } - - Ok(BufferedCapture { - // SAFETY: the bridge returns a heap-allocated cv::VideoCapture compatible - // with the opencv crate's generated ownership wrapper. - capture: unsafe { videoio::VideoCapture::opencv_from_extern(capture) }, - _bytes: bytes, - }) -} - -#[cfg(test)] -mod tests { - use opencv::core::Vec3b; - - use super::*; - - #[test] - fn rgb_output_buffer_writes_frames_directly_in_rgb_order() { - let pixels = [Vec3b::from([1, 2, 3]), Vec3b::from([4, 5, 6])]; - let bgr = Mat::new_rows_cols_with_data(1, 2, &pixels) - .unwrap() - .try_clone() - .unwrap(); - let mut output = RgbOutputBuffer::with_capacity(12).unwrap(); - - assert_eq!(output.push_bgr(&bgr, 2, 1).unwrap(), (0, 6)); - assert_eq!(output.push_bgr(&bgr, 2, 1).unwrap(), (6, 6)); - - let bytes = output.into_bytes(); - assert_eq!(&bytes[..], &[3, 2, 1, 6, 5, 4, 3, 2, 1, 6, 5, 4]); - } - - #[test] - fn rgb_output_buffer_supports_dimension_changes() { - let wide_pixels = [Vec3b::from([1, 2, 3]), Vec3b::from([4, 5, 6])]; - let wide = Mat::new_rows_cols_with_data(1, 2, &wide_pixels) - .unwrap() - .try_clone() - .unwrap(); - let square_pixels = [Vec3b::from([7, 8, 9])]; - let square = Mat::new_rows_cols_with_data(1, 1, &square_pixels) - .unwrap() - .try_clone() - .unwrap(); - let mut output = RgbOutputBuffer::with_capacity(1).unwrap(); - - assert_eq!(output.push_bgr(&wide, 2, 1).unwrap(), (0, 6)); - assert_eq!(output.push_bgr(&square, 1, 1).unwrap(), (6, 3)); - - let bytes = output.into_bytes(); - assert_eq!(&bytes[..], &[3, 2, 1, 6, 5, 4, 9, 8, 7]); - } -} diff --git a/crates/multimodal/src/opencv_buffer_capture.cpp b/crates/multimodal/src/opencv_buffer_capture.cpp deleted file mode 100644 index 4080ff35c..000000000 --- a/crates/multimodal/src/opencv_buffer_capture.cpp +++ /dev/null @@ -1,122 +0,0 @@ -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -namespace { - -void set_error(char* output, size_t capacity, const char* message) { - if (output == nullptr || capacity == 0) { - return; - } - std::snprintf(output, capacity, "%s", message); -} - -} // namespace - -#if CV_VERSION_MAJOR > 4 || \ - (CV_VERSION_MAJOR == 4 && CV_VERSION_MINOR >= 11) - -namespace { - -class MemoryStreamReader final : public cv::IStreamReader { - public: - MemoryStreamReader(const uint8_t* data, size_t size) - : data_(data), size_(size) {} - - long long read(char* buffer, long long size) override { - if (size <= 0 || position_ >= size_) { - return 0; - } - const size_t count = - std::min(static_cast(size), size_ - position_); - std::memcpy(buffer, data_ + position_, count); - position_ += count; - return static_cast(count); - } - - long long seek(long long offset, int origin) override { - if (size_ > static_cast(std::numeric_limits::max())) { - return -1; - } - - const long long end = static_cast(size_); - long long base = 0; - if (origin == SEEK_CUR) { - base = static_cast(position_); - } else if (origin == SEEK_END) { - base = end; - } else if (origin != SEEK_SET) { - return -1; - } - - if (offset > end - base || offset < -base) { - return -1; - } - const long long next = base + offset; - position_ = static_cast(next); - return next; - } - - private: - const uint8_t* data_; - size_t size_; - size_t position_ = 0; -}; - -} // namespace - -extern "C" void* smg_opencv_capture_from_buffer(const uint8_t* data, - size_t size, - int decoder_threads, - char* error, - size_t error_capacity) { - try { - if (data == nullptr || size == 0) { - set_error(error, error_capacity, "video buffer is empty"); - return nullptr; - } - - for (const auto backend : - cv::videoio_registry::getStreamBufferedBackends()) { - if (!cv::videoio_registry::hasBackend(backend)) { - continue; - } - cv::Ptr reader = - cv::makePtr(data, size); - auto* capture = new cv::VideoCapture( - reader, static_cast(backend), - std::vector{cv::CAP_PROP_N_THREADS, decoder_threads}); - if (capture->isOpened()) { - return capture; - } - delete capture; - } - set_error(error, error_capacity, - "OpenCV has no usable buffered video backend"); - } catch (const std::exception& exception) { - set_error(error, error_capacity, exception.what()); - } catch (...) { - set_error(error, error_capacity, "unknown OpenCV buffered capture error"); - } - return nullptr; -} - -#else - -extern "C" void* smg_opencv_capture_from_buffer(const uint8_t*, size_t, int, - char* error, - size_t error_capacity) { - set_error(error, error_capacity, - "buffered video capture requires OpenCV 4.11 or newer"); - return nullptr; -} - -#endif diff --git a/crates/multimodal/src/registry/kimi_k25.rs b/crates/multimodal/src/registry/kimi_k25.rs deleted file mode 100644 index 2abbff858..000000000 --- a/crates/multimodal/src/registry/kimi_k25.rs +++ /dev/null @@ -1,214 +0,0 @@ -use std::collections::HashMap; - -use serde_json::{json, Value}; - -use crate::{ - encoder_inputs::PreprocessedEncoderInputs, - registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, - types::{FieldLayout, Modality, PromptReplacement, TokenId}, -}; - -pub(super) struct KimiK25VisionSpec; - -impl KimiK25VisionSpec { - /// The repeated pad token (`<|media_pad|>`) — `media_placeholder_token_id` in config. - fn pad_token_id(metadata: &ModelMetadata) -> RegistryResult { - metadata - .config_u32(&["media_placeholder_token_id"]) - .map(|v| v as TokenId) - .ok_or_else(|| ModelRegistryError::MissingConfigField { - field: "media_placeholder_token_id".to_string(), - }) - } -} - -impl ModelProcessorSpec for KimiK25VisionSpec { - fn name(&self) -> &'static str { - "kimi_k25" - } - - fn matches(&self, metadata: &ModelMetadata) -> bool { - let id = metadata.model_id.to_ascii_lowercase(); - id.contains("kimi") && id.contains("k2") - || metadata - .config_model_type() - .is_some_and(|mt| mt == "kimi_k25") - } - - fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok("<|media_pad|>".to_string()) - } - - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { - Self::pad_token_id(metadata) - } - - fn modality_limits( - &self, - _metadata: &ModelMetadata, - ) -> RegistryResult> { - Ok(HashMap::from([(Modality::Image, 10)])) - } - - fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(json!({})) - } - - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - let pad_token_id = Self::pad_token_id(metadata)?; - let placeholder_token = self.placeholder_token(metadata)?; - Ok(preprocessed - .feature_token_counts - .iter() - .map(|&num_tokens| { - PromptReplacement::repeated( - Modality::Image, - &placeholder_token, - pad_token_id, - num_tokens, - ) - }) - .collect()) - } - - fn field_layouts(&self) -> HashMap { - // Kimi-K2.5 uses NaViT-style patchification: - // encoder_input is [total_patches, patch_features], split by patches_per_image. - // grid_thws is [num_images, 3] with (temporal, height, width) grid dimensions. - HashMap::from([ - ( - "pixel_values".to_string(), - FieldLayout::flat("patches_per_image"), - ), - ("grid_thws".to_string(), FieldLayout::Batched), - ("patches_per_image".to_string(), FieldLayout::Batched), - ]) - } - - fn keep_on_cpu_keys(&self) -> Vec { - vec!["grid_thws".to_string()] - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use crate::{ - registry::{test_helpers::*, ModelMetadata, ModelRegistry}, - types::ImageSize, - }; - - #[test] - fn kimi_k25_matches_model_id() { - let tokenizer = TestTokenizer::new(&[("<|media_pad|>", 163605)]); - let config = json!({ - "model_type": "kimi_k25", - "media_placeholder_token_id": 163605 - }); - let metadata = ModelMetadata { - model_id: "moonshotai/Kimi-K2.5", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("kimi_k25 spec"); - assert_eq!(spec.name(), "kimi_k25"); - } - - #[test] - fn kimi_k25_prompt_replacements() { - let tokenizer = TestTokenizer::new(&[("<|media_pad|>", 163605)]); - let config = json!({ - "model_type": "kimi_k25", - "media_placeholder_token_id": 163605 - }); - let metadata = ModelMetadata { - model_id: "moonshotai/Kimi-K2.5", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("kimi_k25 spec"); - - let replacements = spec - .prompt_replacements( - &metadata, - &test_preprocessed_with_tokens(&[ImageSize::new(448, 448)], &[256]), - ) - .unwrap(); - - // 256 pad tokens (no start/end wrapper — SGLang handles that in the chat template) - assert_eq!(replacements[0].tokens.len(), 256); - assert!(replacements[0].tokens.iter().all(|&t| t == 163605)); - } - - #[test] - fn kimi_k25_prompt_replacements_multiple_images() { - let tokenizer = TestTokenizer::new(&[("<|media_pad|>", 163605)]); - let config = json!({ - "model_type": "kimi_k25", - "media_placeholder_token_id": 163605 - }); - let metadata = ModelMetadata { - model_id: "moonshotai/Kimi-K2.5", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("kimi_k25 spec"); - - let replacements = spec - .prompt_replacements( - &metadata, - &test_preprocessed_with_tokens( - &[ImageSize::new(448, 448), ImageSize::new(224, 224)], - &[256, 64], - ), - ) - .unwrap(); - - assert_eq!(replacements.len(), 2); - assert_eq!(replacements[0].tokens.len(), 256); - assert_eq!(replacements[1].tokens.len(), 64); - assert!(replacements[1].tokens.iter().all(|&t| t == 163605)); - } - - #[test] - fn kimi_k25_matches_kimi_k2_variant() { - let tokenizer = TestTokenizer::new(&[("<|media_pad|>", 163605)]); - let config = json!({ - "model_type": "kimi_k25", - "media_placeholder_token_id": 163605 - }); - let metadata = ModelMetadata { - model_id: "moonshotai/Kimi-K2-VL", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata); - assert!(spec.is_some(), "Should match Kimi-K2 variants"); - } - - #[test] - fn kimi_k25_does_not_match_kimi_k1() { - let tokenizer = TestTokenizer::new(&[("<|media_pad|>", 163605)]); - let config = json!({ - "model_type": "kimi_k1", - "media_placeholder_token_id": 163605 - }); - let metadata = ModelMetadata { - model_id: "moonshotai/Kimi-K1-VL", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata); - assert!(spec.is_none(), "Should not match Kimi-K1"); - } -} diff --git a/crates/multimodal/src/registry/llama4.rs b/crates/multimodal/src/registry/llama4.rs deleted file mode 100644 index 3e04058a1..000000000 --- a/crates/multimodal/src/registry/llama4.rs +++ /dev/null @@ -1,296 +0,0 @@ -use std::collections::HashMap; - -use serde_json::{json, Value}; - -use crate::{ - encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}, - registry::{ModelMetadata, ModelProcessorSpec, RegistryResult}, - types::{FieldLayout, Modality, PromptReplacement, TokenId}, -}; - -pub(super) struct Llama4Spec; - -impl Llama4Spec { - fn patch_size(metadata: &ModelMetadata) -> u32 { - metadata - .config_u32(&["vision_config", "patch_size"]) - .unwrap_or(14) - } - - fn tile_size(metadata: &ModelMetadata) -> u32 { - metadata - .config_u32(&["vision_config", "image_size"]) - .filter(|v| *v > 0) - .unwrap_or(336) - } - - fn pixel_shuffle_ratio(metadata: &ModelMetadata) -> f64 { - metadata - .config - .get("vision_config") - .and_then(|v| v.get("pixel_shuffle_ratio")) - .and_then(|v| v.as_f64()) - .unwrap_or(0.5) - } - - fn tokens_per_tile(metadata: &ModelMetadata) -> usize { - let tile = Self::tile_size(metadata) as usize; - let patch = Self::patch_size(metadata) as usize; - if patch == 0 { - return 0; - } - let patches = (tile / patch).pow(2); - // Pixel shuffle reduces spatial dims by ratio, so token count by ratio^2 - let ratio = Self::pixel_shuffle_ratio(metadata); - let downsample = (1.0 / (ratio * ratio)).round().max(1.0) as usize; - patches / downsample - } - - /// Extract per-image `(h_tiles, w_tiles)` from the preprocessor's - /// `aspect_ratios` tensor. Falls back to deriving tile counts from - /// the original image sizes when aspect_ratios are unavailable. - fn extract_aspect_ratios( - preprocessed: &PreprocessedEncoderInputs, - tile_size: usize, - ) -> Vec<(usize, usize)> { - if let Some(ModelSpecificValue::IntTensor { data, shape }) = - preprocessed.model_specific.get("aspect_ratios") - { - if shape.len() == 2 && shape[1] == 2 && data.len() == shape[0] * 2 { - return data - .chunks_exact(2) - .map(|chunk| (chunk[0] as usize, chunk[1] as usize)) - .collect(); - } - } - // Fallback: derive from original image sizes (height, width). - preprocessed - .item_sizes - .iter() - .map(|&(h, w)| { - let h_tiles = (h as usize).div_ceil(tile_size); - let w_tiles = (w as usize).div_ceil(tile_size); - (h_tiles, w_tiles) - }) - .collect() - } -} - -impl ModelProcessorSpec for Llama4Spec { - fn name(&self) -> &'static str { - "llama4" - } - - fn matches(&self, metadata: &ModelMetadata) -> bool { - let id = metadata.model_id.to_ascii_lowercase(); - id.contains("llama-4") - || id.contains("llama4") - || metadata - .config_model_type() - .is_some_and(|mt| mt == "llama4") - } - - fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok("<|image|>".to_string()) - } - - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { - if let Some(value) = metadata.config_u32(&["image_token_index"]) { - return Ok(value as TokenId); - } - metadata.token_id("<|image|>") - } - - fn modality_limits( - &self, - _metadata: &ModelMetadata, - ) -> RegistryResult> { - Ok(HashMap::from([(Modality::Image, 8)])) - } - - fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(json!({})) - } - - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - let patch_token_id = self.placeholder_token_id(metadata)?; - let placeholder = self.placeholder_token(metadata)?; - let tokens_per_tile = Self::tokens_per_tile(metadata); - let tile_size = Self::tile_size(metadata) as usize; - - // Structural token IDs matching HF _prompt_split_image format. - let image_start_id = metadata.token_id("<|image_start|>")?; - let image_end_id = metadata.token_id("<|image_end|>")?; - let image_id = metadata.token_id("<|image|>")?; - let tile_x_sep_id = metadata.token_id("<|tile_x_separator|>")?; - let tile_y_sep_id = metadata.token_id("<|tile_y_separator|>")?; - - // Extract aspect_ratios from preprocessor output (computed by - // get_best_fit, respecting max_patches cap). This is the Llama 4 - // analog of vLLM's out_mm_kwargs["image"][i]["aspect_ratios"].data. - let aspect_ratios = Self::extract_aspect_ratios(preprocessed, tile_size); - - Ok(aspect_ratios - .iter() - .map(|&(h_tiles, w_tiles)| { - let num_tiles = h_tiles * w_tiles; - - let mut tokens = Vec::new(); - - // <|image_start|> - tokens.push(image_start_id); - - // Grid tiles with separators (only for multi-tile images) - if num_tiles > 1 { - for _row in 0..h_tiles { - for col in 0..w_tiles { - tokens.extend(std::iter::repeat_n(patch_token_id, tokens_per_tile)); - if col < w_tiles - 1 { - tokens.push(tile_x_sep_id); - } - } - tokens.push(tile_y_sep_id); - } - } - - // Global/cover tile: <|image|> + <|patch|> * tokens_per_tile - tokens.push(image_id); - tokens.extend(std::iter::repeat_n(patch_token_id, tokens_per_tile)); - - // <|image_end|> - tokens.push(image_end_id); - - PromptReplacement::sequence(Modality::Image, &placeholder, tokens) - }) - .collect()) - } - - fn field_layouts(&self) -> HashMap { - // encoder_input is [total_tiles, C, H, W] — variable tiles per image. - // aspect_ratios and patches_per_image are [num_images, ...]. - HashMap::from([ - ( - "pixel_values".to_string(), - FieldLayout::flat("patches_per_image"), - ), - ("aspect_ratios".to_string(), FieldLayout::Batched), - ("patches_per_image".to_string(), FieldLayout::Batched), - ]) - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use crate::{ - registry::{test_helpers::*, ModelMetadata, ModelRegistry}, - types::ImageSize, - }; - - #[test] - fn llama4_single_tile_token_count() { - let tokenizer = TestTokenizer::new(&[ - ("<|image|>", 200090), - ("<|image_start|>", 200088), - ("<|image_end|>", 200089), - ("<|patch|>", 200092), - ("<|tile_x_separator|>", 200093), - ("<|tile_y_separator|>", 200094), - ]); - let config = json!({ - "model_type": "llama4", - "image_token_index": 200092, - "vision_config": {"image_size": 336, "patch_size": 14} - }); - let metadata = ModelMetadata { - model_id: "/models/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("llama4 spec"); - assert_eq!(spec.name(), "llama4"); - - // Single tile (336x336): <|image_start|> <|image|> <|patch|>*144 <|image_end|> - // = 1 + 1 + 144 + 1 = 147 tokens - let pp = test_preprocessed_with_aspects(&[ImageSize::new(336, 336)], &[(1, 1)]); - let replacements = spec.prompt_replacements(&metadata, &pp).unwrap(); - assert_eq!(replacements[0].tokens.len(), 147); - assert_eq!(replacements[0].tokens[0], 200088); // <|image_start|> - assert_eq!(replacements[0].tokens[1], 200090); // <|image|> - assert_eq!(replacements[0].tokens[2], 200092); // <|patch|> (first) - assert_eq!(replacements[0].tokens[145], 200092); // <|patch|> (last) - assert_eq!(replacements[0].tokens[146], 200089); // <|image_end|> - } - - #[test] - fn llama4_multi_tile_adds_global() { - let tokenizer = TestTokenizer::new(&[ - ("<|image|>", 200090), - ("<|image_start|>", 200088), - ("<|image_end|>", 200089), - ("<|patch|>", 200092), - ("<|tile_x_separator|>", 200093), - ("<|tile_y_separator|>", 200094), - ]); - let config = json!({ - "model_type": "llama4", - "image_token_index": 200092, - "vision_config": {"image_size": 336, "patch_size": 14} - }); - let metadata = ModelMetadata { - model_id: "Llama-4-Scout-Vision", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("llama4 spec"); - - // 672x672 = 2x2 tiles + 1 global: - // <|image_start|> = 1 - // row0: <|patch|>*144 <|tile_x_sep|> <|patch|>*144 <|tile_y_sep|> = 290 - // row1: <|patch|>*144 <|tile_x_sep|> <|patch|>*144 <|tile_y_sep|> = 290 - // <|image|> <|patch|>*144 = 145 - // <|image_end|> = 1 - // Total = 1 + 290 + 290 + 145 + 1 = 727 - let pp = test_preprocessed_with_aspects(&[ImageSize::new(672, 672)], &[(2, 2)]); - let replacements = spec.prompt_replacements(&metadata, &pp).unwrap(); - assert_eq!(replacements[0].tokens.len(), 727); - // Verify structure: starts with image_start, ends with image_end - assert_eq!(replacements[0].tokens[0], 200088); // <|image_start|> - assert_eq!(*replacements[0].tokens.last().unwrap(), 200089); // <|image_end|> - // The token before the last patch block is <|image|> (global tile marker) - // Position: 1 + 290 + 290 = 581 - assert_eq!(replacements[0].tokens[581], 200090); // <|image|> - } - - #[test] - fn llama4_matches_alias_via_model_type() { - let tokenizer = TestTokenizer::new(&[ - ("<|image|>", 200090), - ("<|image_start|>", 200088), - ("<|image_end|>", 200089), - ("<|patch|>", 200092), - ("<|tile_x_separator|>", 200093), - ("<|tile_y_separator|>", 200094), - ]); - let config = json!({ - "model_type": "llama4", - "image_token_index": 200092, - "vision_config": {"image_size": 336, "patch_size": 14} - }); - let metadata = ModelMetadata { - model_id: "custom-model", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("llama4 alias"); - assert_eq!(spec.name(), "llama4"); - } -} diff --git a/crates/multimodal/src/registry/llava.rs b/crates/multimodal/src/registry/llava.rs deleted file mode 100644 index 117091c00..000000000 --- a/crates/multimodal/src/registry/llava.rs +++ /dev/null @@ -1,220 +0,0 @@ -use std::collections::HashMap; - -use serde_json::{json, Value}; - -use crate::{ - encoder_inputs::PreprocessedEncoderInputs, - registry::{ModelMetadata, ModelProcessorSpec, RegistryResult}, - types::{FieldLayout, Modality, PromptReplacement, TokenId}, -}; - -pub(super) struct LlavaSpec; -pub(super) struct LlavaNextSpec; - -impl ModelProcessorSpec for LlavaSpec { - fn name(&self) -> &'static str { - "llava" - } - - fn matches(&self, metadata: &ModelMetadata) -> bool { - // Match plain "llava" but not "llava_next" (handled by LlavaNextSpec). - let model_type = metadata.config_model_type(); - if model_type.is_some_and(|mt| mt == "llava_next") { - return false; - } - let model_id_lower = metadata.model_id.to_ascii_lowercase(); - if model_id_lower.contains("llava-next") || model_id_lower.contains("llava_next") { - return false; - } - model_id_lower.contains("llava") || model_type.is_some_and(|mt| mt == "llava") - } - - fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok("".to_string()) - } - - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { - if let Some(value) = metadata.config_u32(&["image_token_index"]) { - return Ok(value as TokenId); - } - metadata.token_id("") - } - - fn modality_limits( - &self, - _metadata: &ModelMetadata, - ) -> RegistryResult> { - Ok(HashMap::from([(Modality::Image, 4)])) - } - - fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(json!({})) - } - - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - let token_id = self.placeholder_token_id(metadata)?; - let token = self.placeholder_token(metadata)?; - Ok(preprocessed - .feature_token_counts - .iter() - .map(|&count| PromptReplacement::repeated(Modality::Image, &token, token_id, count)) - .collect()) - } -} - -impl ModelProcessorSpec for LlavaNextSpec { - fn name(&self) -> &'static str { - "llava_next" - } - - fn matches(&self, metadata: &ModelMetadata) -> bool { - metadata - .config_model_type() - .is_some_and(|mt| mt == "llava_next") - } - - fn placeholder_token(&self, metadata: &ModelMetadata) -> RegistryResult { - LlavaSpec.placeholder_token(metadata) - } - - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { - LlavaSpec.placeholder_token_id(metadata) - } - - fn modality_limits( - &self, - metadata: &ModelMetadata, - ) -> RegistryResult> { - LlavaSpec.modality_limits(metadata) - } - - fn processor_kwargs(&self, metadata: &ModelMetadata) -> RegistryResult { - LlavaSpec.processor_kwargs(metadata) - } - - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - // LLaVA-Next token counts differ from plain LLaVA because of - // anyres multi-crop + spatial_unpad. The correct per-image counts - // are already computed by LlavaNextProcessor::calculate_num_tokens - // and stored in preprocessed.feature_token_counts. - let token_id = LlavaSpec.placeholder_token_id(metadata)?; - let token = LlavaSpec.placeholder_token(metadata)?; - Ok(preprocessed - .feature_token_counts - .iter() - .map(|&count| PromptReplacement::repeated(Modality::Image, &token, token_id, count)) - .collect()) - } - - fn field_layouts(&self) -> HashMap { - // encoder_input is [num_images, max_patches, C, H, W] (5D, batched). - // image_sizes is [num_images, 2] (batched), matching HF/vLLM kwargs. - HashMap::from([ - ("pixel_values".to_string(), FieldLayout::Batched), - ("image_sizes".to_string(), FieldLayout::Batched), - ]) - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use crate::{ - registry::{test_helpers::*, ModelMetadata, ModelRegistry}, - types::ImageSize, - }; - - #[test] - fn llava_prompt_replacement_uses_preprocessed_tokens() { - let tokenizer = TestTokenizer::new(&[("", 32000)]); - let config = json!({ - "model_type": "llava", - "image_token_index": 32000, - "vision_config": {"patch_size": 14} - }); - let metadata = ModelMetadata { - model_id: "llava-v1.5", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("llava spec"); - // Token count comes from preprocessed.feature_token_counts (set by - // LlavaProcessor::calculate_num_tokens), not from image dimensions. - let preprocessed = test_preprocessed_with_tokens(&[ImageSize::new(336, 336)], &[576]); - let replacements = spec.prompt_replacements(&metadata, &preprocessed).unwrap(); - assert_eq!(replacements[0].tokens.len(), 576); - } - - #[test] - fn llava_prompt_replacement_uses_per_image_token_counts() { - let tokenizer = TestTokenizer::new(&[("", 32000)]); - let config = json!({ - "model_type": "llava", - "image_token_index": 32000, - "vision_config": {"patch_size": 14} - }); - let metadata = ModelMetadata { - model_id: "llava-v1.5", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("llava spec"); - let preprocessed = test_preprocessed_with_tokens( - &[ImageSize::new(336, 336), ImageSize::new(448, 448)], - &[576, 1024], - ); - - let replacements = spec.prompt_replacements(&metadata, &preprocessed).unwrap(); - - assert_eq!(replacements.len(), 2); - assert_eq!(replacements[0].tokens.len(), 576); - assert_eq!(replacements[1].tokens.len(), 1024); - } - - #[test] - fn llava_matches_alias_via_model_type() { - let tokenizer = TestTokenizer::new(&[("", 32000)]); - let config = json!({ - "model_type": "llava", - "image_token_index": 32000, - "vision_config": {"patch_size": 14} - }); - let metadata = ModelMetadata { - model_id: "custom-model", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("llava alias"); - assert_eq!(spec.name(), "llava"); - } - - #[test] - fn llava_spec_has_no_audio_processor() { - use crate::vision::PreProcessorConfig; - - let tokenizer = TestTokenizer::new(&[("", 32000)]); - let config = json!({"model_type": "llava", "image_token_index": 32000}); - let metadata = ModelMetadata { - model_id: "llava-v1.5", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("llava spec"); - assert!(spec - .audio_processor(&config, &PreProcessorConfig::default()) - .is_none()); - } -} diff --git a/crates/multimodal/src/registry/mod.rs b/crates/multimodal/src/registry/mod.rs deleted file mode 100644 index 6fb5a1902..000000000 --- a/crates/multimodal/src/registry/mod.rs +++ /dev/null @@ -1,202 +0,0 @@ -mod kimi_k25; -mod llama4; -mod llava; -mod phi3_v; -mod qwen3_asr; -mod qwen3_omni; -mod qwen3_vl; -mod qwen_vl; -mod traits; - -use kimi_k25::KimiK25VisionSpec; -use llama4::Llama4Spec; -use llava::{LlavaNextSpec, LlavaSpec}; -use once_cell::sync::Lazy; -use phi3_v::Phi3VisionSpec; -use qwen3_asr::Qwen3AsrSpec; -use qwen3_omni::Qwen3OmniSpec; -use qwen3_vl::Qwen3VLVisionSpec; -use qwen_vl::QwenVLVisionSpec; -// Re-export public API from traits. -pub use traits::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}; - -pub struct ModelRegistry { - specs: Vec, -} - -impl ModelRegistry { - pub fn new() -> Self { - Self { - specs: vec![ - LazySpec::new(|| Box::new(KimiK25VisionSpec)), - LazySpec::new(|| Box::new(Llama4Spec)), - // LlavaNext must be registered before Llava so "llava_next" model_type matches first. - LazySpec::new(|| Box::new(LlavaNextSpec)), - LazySpec::new(|| Box::new(LlavaSpec)), - LazySpec::new(|| Box::new(Qwen3AsrSpec)), - LazySpec::new(|| Box::new(Qwen3OmniSpec)), - // Qwen3-VL must be registered before QwenVL so "qwen3" matches first. - LazySpec::new(|| Box::new(Qwen3VLVisionSpec)), - LazySpec::new(|| Box::new(QwenVLVisionSpec)), - LazySpec::new(|| Box::new(Phi3VisionSpec)), - ], - } - } - - pub fn lookup<'a>(&'a self, metadata: &ModelMetadata) -> Option<&'a dyn ModelProcessorSpec> { - for spec in &self.specs { - let spec_ref = spec.get(); - if spec_ref.matches(metadata) { - return Some(spec_ref); - } - } - None - } -} - -impl Default for ModelRegistry { - fn default() -> Self { - Self::new() - } -} - -struct LazySpec { - inner: Lazy>, -} - -impl LazySpec { - fn new(factory: fn() -> Box) -> Self { - Self { - inner: Lazy::new(factory), - } - } - - fn get(&self) -> &dyn ModelProcessorSpec { - self.inner.as_ref() - } -} - -#[cfg(test)] -pub(super) mod test_helpers { - use std::collections::HashMap; - - use llm_tokenizer::{Decoder, Encoder, Encoding, SpecialTokens, TokenizerTrait}; - use once_cell::sync::Lazy; - - use crate::{ - encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}, - types::ImageSize, - }; - - pub struct TestTokenizer { - vocab: HashMap, - } - - impl TestTokenizer { - pub fn new(pairs: &[(&str, u32)]) -> Self { - let vocab = pairs - .iter() - .map(|(token, id)| ((*token).to_string(), *id)) - .collect(); - Self { vocab } - } - } - - impl Encoder for TestTokenizer { - fn encode(&self, _input: &str, _add_special_tokens: bool) -> anyhow::Result { - Ok(Encoding::Plain(Vec::new())) - } - - fn encode_batch( - &self, - inputs: &[&str], - add_special_tokens: bool, - ) -> anyhow::Result> { - inputs - .iter() - .map(|_| self.encode("", add_special_tokens)) - .collect() - } - } - - impl Decoder for TestTokenizer { - fn decode(&self, _token_ids: &[u32], _skip_special_tokens: bool) -> anyhow::Result { - Ok(String::new()) - } - } - - impl TokenizerTrait for TestTokenizer { - fn vocab_size(&self) -> usize { - self.vocab.len() - } - - fn get_special_tokens(&self) -> &SpecialTokens { - static TOKENS: Lazy = Lazy::new(|| SpecialTokens { - bos_token: None, - eos_token: None, - unk_token: None, - sep_token: None, - pad_token: None, - cls_token: None, - mask_token: None, - additional_special_tokens: vec![], - }); - &TOKENS - } - - fn token_to_id(&self, token: &str) -> Option { - self.vocab.get(token).copied() - } - - fn id_to_token(&self, id: u32) -> Option { - self.vocab - .iter() - .find(|(_, &v)| v == id) - .map(|(k, _)| k.clone()) - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } - } - - pub fn test_preprocessed_with_tokens( - item_sizes: &[ImageSize], - feature_token_counts: &[usize], - ) -> PreprocessedEncoderInputs { - let sizes: Vec<(u32, u32)> = item_sizes.iter().map(|s| (s.height, s.width)).collect(); - PreprocessedEncoderInputs { - encoder_input: ndarray::ArrayD::zeros(vec![1, 3, 336, 336]), - feature_token_counts: feature_token_counts.to_vec(), - item_sizes: sizes, - model_specific: HashMap::new(), - } - } - - /// Build `PreprocessedEncoderInputs` with explicit aspect_ratios (for Llama4 tests). - pub fn test_preprocessed_with_aspects( - item_sizes: &[ImageSize], - aspect_ratios: &[(i64, i64)], - ) -> PreprocessedEncoderInputs { - let sizes: Vec<(u32, u32)> = item_sizes.iter().map(|s| (s.height, s.width)).collect(); - let flat: Vec = aspect_ratios - .iter() - .flat_map(|&(h, w)| vec![h, w]) - .collect(); - let batch = aspect_ratios.len(); - let mut model_specific = HashMap::new(); - model_specific.insert( - "aspect_ratios".to_string(), - ModelSpecificValue::IntTensor { - data: flat, - shape: vec![batch, 2], - }, - ); - PreprocessedEncoderInputs { - encoder_input: ndarray::ArrayD::zeros(vec![1, 3, 336, 336]), - feature_token_counts: vec![0; sizes.len()], - item_sizes: sizes, - model_specific, - } - } -} diff --git a/crates/multimodal/src/registry/phi3_v.rs b/crates/multimodal/src/registry/phi3_v.rs deleted file mode 100644 index b5229d18d..000000000 --- a/crates/multimodal/src/registry/phi3_v.rs +++ /dev/null @@ -1,116 +0,0 @@ -use std::collections::HashMap; - -use serde_json::{json, Value}; - -use crate::{ - encoder_inputs::PreprocessedEncoderInputs, - registry::{ModelMetadata, ModelProcessorSpec, RegistryResult}, - types::{FieldLayout, Modality, PromptReplacement, TokenId}, -}; - -pub(super) struct Phi3VisionSpec; - -impl ModelProcessorSpec for Phi3VisionSpec { - fn name(&self) -> &'static str { - "phi3_v" - } - - fn matches(&self, metadata: &ModelMetadata) -> bool { - let id = metadata.model_id.to_ascii_lowercase(); - id.contains("phi") && id.contains("vision") - || metadata - .config_model_type() - .is_some_and(|mt| mt == "phi3_v") - } - - fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok("<|image|>".to_owned()) - } - - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { - metadata.token_id("<|image|>") - } - - fn modality_limits( - &self, - _metadata: &ModelMetadata, - ) -> RegistryResult> { - Ok(HashMap::from([(Modality::Image, 4)])) - } - - fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(json!({})) - } - - fn field_layouts(&self) -> HashMap { - HashMap::from([ - ("pixel_values".to_string(), FieldLayout::Batched), - ("image_sizes".to_string(), FieldLayout::Batched), - ]) - } - - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - let token_id = self.placeholder_token_id(metadata)?; - let token = self.placeholder_token(metadata)?; - Ok(preprocessed - .feature_token_counts - .iter() - .map(|&count| PromptReplacement::repeated(Modality::Image, &token, token_id, count)) - .collect()) - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use crate::{ - registry::{test_helpers::*, ModelMetadata, ModelRegistry}, - types::ImageSize, - }; - - #[test] - fn phi3_uses_num_img_tokens() { - let tokenizer = TestTokenizer::new(&[("<|image|>", 555)]); - let config = json!({ - "model_type": "phi3_v", - "img_processor": {"num_img_tokens": 144} - }); - let metadata = ModelMetadata { - model_id: "Phi-3-vision", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("phi3 spec"); - let replacements = spec - .prompt_replacements( - &metadata, - &test_preprocessed_with_tokens(&[ImageSize::new(336, 336)], &[144]), - ) - .unwrap(); - assert_eq!(replacements[0].tokens.len(), 144); - assert_eq!(replacements[0].tokens[0], 555); - } - - #[test] - fn phi3_matches_alias_via_model_type() { - let tokenizer = TestTokenizer::new(&[("<|image|>", 555)]); - let config = json!({ - "model_type": "phi3_v", - "img_processor": {"num_img_tokens": 144} - }); - let metadata = ModelMetadata { - model_id: "custom-model", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("phi3 alias"); - assert_eq!(spec.name(), "phi3_v"); - } -} diff --git a/crates/multimodal/src/registry/qwen3_asr.rs b/crates/multimodal/src/registry/qwen3_asr.rs deleted file mode 100644 index bdcf3c2f4..000000000 --- a/crates/multimodal/src/registry/qwen3_asr.rs +++ /dev/null @@ -1,259 +0,0 @@ -use std::collections::HashMap; - -use serde_json::{json, Value}; - -use crate::{ - audio::{AudioPreProcessor, Qwen3AudioProcessor}, - encoder_inputs::PreprocessedEncoderInputs, - registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, - types::{EncoderFieldLayouts, FieldLayout, Modality, PromptReplacement, TokenId}, - vision::PreProcessorConfig, -}; - -const AUDIO_PAD_TOKEN: &str = "<|audio_pad|>"; - -pub(super) struct Qwen3AsrSpec; - -impl Qwen3AsrSpec { - fn audio_token_id(metadata: &ModelMetadata) -> RegistryResult { - metadata - .config_u32(&["thinker_config", "audio_token_id"]) - .or_else(|| metadata.config_u32(&["audio_token_id"])) - .map(|value| value as TokenId) - .map_or_else(|| metadata.token_id(AUDIO_PAD_TOKEN), Ok) - } -} - -impl ModelProcessorSpec for Qwen3AsrSpec { - fn name(&self) -> &'static str { - "qwen3_asr" - } - - fn matches(&self, metadata: &ModelMetadata) -> bool { - let model_id = metadata.model_id.to_ascii_lowercase(); - model_id.contains("qwen3-asr") - || model_id.contains("qwen3_asr") - || metadata - .config_model_type() - .is_some_and(|model_type| model_type == "qwen3_asr") - } - - fn placeholder_token(&self, metadata: &ModelMetadata) -> RegistryResult { - self.placeholder_token_for(metadata, Modality::Audio) - } - - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { - self.placeholder_token_id_for(metadata, Modality::Audio) - } - - fn placeholder_token_for( - &self, - metadata: &ModelMetadata, - modality: Modality, - ) -> RegistryResult { - match modality { - Modality::Audio => { - let token_id = Self::audio_token_id(metadata)?; - match metadata.tokenizer.id_to_token(token_id as u32) { - Some(token) => Ok(token), - None => { - metadata.token_id(AUDIO_PAD_TOKEN)?; - Ok(AUDIO_PAD_TOKEN.to_string()) - } - } - } - Modality::Image | Modality::Video | Modality::ImageEmbeds => { - Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }) - } - } - } - - fn placeholder_token_id_for( - &self, - metadata: &ModelMetadata, - modality: Modality, - ) -> RegistryResult { - match modality { - Modality::Audio => Self::audio_token_id(metadata), - Modality::Image | Modality::Video | Modality::ImageEmbeds => { - Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }) - } - } - } - - fn modality_limits( - &self, - _metadata: &ModelMetadata, - ) -> RegistryResult> { - Ok(HashMap::from([(Modality::Audio, 10)])) - } - - fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(json!({})) - } - - fn audio_processor( - &self, - model_config: &Value, - preprocessor_config: &PreProcessorConfig, - ) -> Option> { - Some(Box::new(Qwen3AudioProcessor::from_configs( - model_config, - preprocessor_config, - ))) - } - - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - self.prompt_replacements_for(metadata, preprocessed, Modality::Audio) - } - - fn prompt_replacements_for( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - modality: Modality, - ) -> RegistryResult> { - match modality { - Modality::Audio => { - let token_id = Self::audio_token_id(metadata)?; - let token = self.placeholder_token_for(metadata, Modality::Audio)?; - Ok(preprocessed - .feature_token_counts - .iter() - .map(|&count| { - PromptReplacement::repeated(Modality::Audio, &token, token_id, count) - }) - .collect()) - } - Modality::Image | Modality::Video | Modality::ImageEmbeds => { - Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }) - } - } - } - - fn encoder_field_layouts_for(&self, modality: Modality) -> EncoderFieldLayouts { - match modality { - Modality::Audio => EncoderFieldLayouts::new( - FieldLayout::Batched, - HashMap::from([ - ("feature_attention_mask".to_string(), FieldLayout::Batched), - ("audio_feature_lengths".to_string(), FieldLayout::Batched), - ]), - ), - Modality::Image | Modality::Video | Modality::ImageEmbeds => { - EncoderFieldLayouts::default() - } - } - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - use crate::{ - registry::{test_helpers::*, ModelRegistry}, - types::ImageSize, - }; - - #[test] - fn asr_matches_and_expands_nested_audio_token() { - let tokenizer = TestTokenizer::new(&[(AUDIO_PAD_TOKEN, 151676)]); - let config = json!({ - "model_type": "qwen3_asr", - "thinker_config": {"audio_token_id": 151676} - }); - let metadata = ModelMetadata { - model_id: "Qwen/Qwen3-ASR-1.7B", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).unwrap(); - assert_eq!(spec.name(), "qwen3_asr"); - assert_eq!( - spec.placeholder_token(&metadata).unwrap(), - spec.placeholder_token_for(&metadata, Modality::Audio) - .unwrap() - ); - assert_eq!( - spec.placeholder_token_id(&metadata).unwrap(), - spec.placeholder_token_id_for(&metadata, Modality::Audio) - .unwrap() - ); - - let replacements = spec - .prompt_replacements_for( - &metadata, - &test_preprocessed_with_tokens(&[ImageSize::new(128, 100)], &[13]), - Modality::Audio, - ) - .unwrap(); - assert_eq!(replacements[0].tokens, vec![151676; 13]); - assert_eq!( - spec.encoder_field_layouts_for(Modality::Audio) - .encoder_input, - FieldLayout::Batched - ); - assert_eq!( - spec.modality_limits(&metadata).unwrap(), - HashMap::from([(Modality::Audio, 10)]) - ); - } - - #[test] - fn asr_spec_builds_qwen_audio_processor() { - use std::sync::Arc; - - use bytes::Bytes; - - use crate::{ - audio::DecodedAudio, - types::{AudioClip, AudioSource}, - }; - - let tokenizer = TestTokenizer::new(&[(AUDIO_PAD_TOKEN, 151676)]); - let config = json!({"model_type": "qwen3_asr"}); - let metadata = ModelMetadata { - model_id: "Qwen/Qwen3-ASR-1.7B", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).unwrap(); - - let preprocessor_config = PreProcessorConfig::from_json( - r#"{"feature_size": 16, "sampling_rate": 16000, "n_fft": 400, "hop_length": 160}"#, - ) - .unwrap(); - let processor = spec - .audio_processor(&config, &preprocessor_config) - .expect("qwen3_asr spec must provide an audio processor"); - - let clip = Arc::new(AudioClip::new( - Bytes::from_static(b"audio"), - DecodedAudio { - samples: vec![0.0; 800], - sample_rate: 16_000, - }, - AudioSource::InlineBytes, - "audio-hash".to_string(), - )); - let result = processor.preprocess(&[clip]).unwrap(); - assert_eq!(result.encoder_input.shape(), &[1, 16, 5]); - } -} diff --git a/crates/multimodal/src/registry/qwen3_omni.rs b/crates/multimodal/src/registry/qwen3_omni.rs deleted file mode 100644 index 1ec927c43..000000000 --- a/crates/multimodal/src/registry/qwen3_omni.rs +++ /dev/null @@ -1,365 +0,0 @@ -use std::collections::HashMap; - -use serde_json::{json, Value}; - -use crate::{ - audio::{AudioPreProcessor, Qwen3AudioProcessor}, - encoder_inputs::PreprocessedEncoderInputs, - registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, - types::{EncoderFieldLayouts, FieldLayout, Modality, PromptReplacement, TokenId}, - vision::PreProcessorConfig, -}; - -const IMAGE_PAD_TOKEN: &str = "<|image_pad|>"; -const VIDEO_PAD_TOKEN: &str = "<|video_pad|>"; -const AUDIO_PAD_TOKEN: &str = "<|audio_pad|>"; - -pub(super) struct Qwen3OmniSpec; - -impl Qwen3OmniSpec { - fn token_id(metadata: &ModelMetadata, field: &str) -> RegistryResult { - metadata - .config_u32(&["thinker_config", field]) - .or_else(|| metadata.config_u32(&[field])) - .map(|value| value as TokenId) - .ok_or_else(|| ModelRegistryError::MissingConfigField { - field: format!("thinker_config.{field}"), - }) - } - - fn token(metadata: &ModelMetadata, field: &str, fallback: &str) -> RegistryResult { - let token_id = Self::token_id(metadata, field)?; - if let Some(token) = metadata.tokenizer.id_to_token(token_id as u32) { - return Ok(token); - } - metadata.token_id(fallback)?; - Ok(fallback.to_string()) - } - - fn replacements( - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - modality: Modality, - field: &str, - fallback: &str, - ) -> RegistryResult> { - let token_id = Self::token_id(metadata, field)?; - let token = Self::token(metadata, field, fallback)?; - Ok(preprocessed - .feature_token_counts - .iter() - .map(|&count| PromptReplacement::repeated(modality, &token, token_id, count)) - .collect()) - } -} - -impl ModelProcessorSpec for Qwen3OmniSpec { - fn name(&self) -> &'static str { - "qwen3_omni" - } - - fn matches(&self, metadata: &ModelMetadata) -> bool { - let model_id = metadata.model_id.to_ascii_lowercase(); - model_id.contains("qwen3-omni") - || model_id.contains("qwen3_omni") - || metadata.config_model_type().is_some_and(|model_type| { - model_type == "qwen3_omni_moe" || model_type == "qwen3_omni_moe_thinker" - }) - } - - fn placeholder_token(&self, metadata: &ModelMetadata) -> RegistryResult { - self.placeholder_token_for(metadata, Modality::Image) - } - - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { - self.placeholder_token_id_for(metadata, Modality::Image) - } - - fn placeholder_token_for( - &self, - metadata: &ModelMetadata, - modality: Modality, - ) -> RegistryResult { - match modality { - Modality::Image => Self::token(metadata, "image_token_id", IMAGE_PAD_TOKEN), - Modality::Video => Self::token(metadata, "video_token_id", VIDEO_PAD_TOKEN), - Modality::Audio => Self::token(metadata, "audio_token_id", AUDIO_PAD_TOKEN), - Modality::ImageEmbeds => Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }), - } - } - - fn placeholder_token_id_for( - &self, - metadata: &ModelMetadata, - modality: Modality, - ) -> RegistryResult { - match modality { - Modality::Image => Self::token_id(metadata, "image_token_id"), - Modality::Video => Self::token_id(metadata, "video_token_id"), - Modality::Audio => Self::token_id(metadata, "audio_token_id"), - Modality::ImageEmbeds => Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }), - } - } - - fn modality_limits( - &self, - _metadata: &ModelMetadata, - ) -> RegistryResult> { - Ok(HashMap::from([ - (Modality::Image, 10), - (Modality::Video, 1), - (Modality::Audio, 10), - ])) - } - - fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(json!({"use_audio_in_video": false})) - } - - fn audio_processor( - &self, - model_config: &Value, - preprocessor_config: &PreProcessorConfig, - ) -> Option> { - Some(Box::new(Qwen3AudioProcessor::from_configs( - model_config, - preprocessor_config, - ))) - } - - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - self.prompt_replacements_for(metadata, preprocessed, Modality::Image) - } - - fn prompt_replacements_for( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - modality: Modality, - ) -> RegistryResult> { - match modality { - Modality::Image => Self::replacements( - metadata, - preprocessed, - modality, - "image_token_id", - IMAGE_PAD_TOKEN, - ), - Modality::Video => Self::replacements( - metadata, - preprocessed, - modality, - "video_token_id", - VIDEO_PAD_TOKEN, - ), - Modality::Audio => Self::replacements( - metadata, - preprocessed, - modality, - "audio_token_id", - AUDIO_PAD_TOKEN, - ), - Modality::ImageEmbeds => Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }), - } - } - - fn encoder_field_layouts_for(&self, modality: Modality) -> EncoderFieldLayouts { - match modality { - Modality::Image => EncoderFieldLayouts::new( - FieldLayout::flat("patches_per_image"), - HashMap::from([ - ("image_grid_thw".to_string(), FieldLayout::Batched), - ("patches_per_image".to_string(), FieldLayout::Batched), - ]), - ), - Modality::Video => EncoderFieldLayouts::new( - FieldLayout::flat("patches_per_video"), - HashMap::from([ - ("video_grid_thw".to_string(), FieldLayout::Batched), - ("patches_per_video".to_string(), FieldLayout::Batched), - ("video_second_per_grid".to_string(), FieldLayout::Batched), - ]), - ), - Modality::Audio => EncoderFieldLayouts::new( - FieldLayout::Batched, - HashMap::from([ - ("feature_attention_mask".to_string(), FieldLayout::Batched), - ("audio_feature_lengths".to_string(), FieldLayout::Batched), - ]), - ), - Modality::ImageEmbeds => EncoderFieldLayouts::default(), - } - } - - fn keep_on_cpu_keys(&self) -> Vec { - vec!["image_grid_thw".to_string(), "video_grid_thw".to_string()] - } - - fn keep_on_cpu_keys_for(&self, modality: Modality) -> Vec { - match modality { - Modality::Image => vec!["image_grid_thw".to_string()], - Modality::Video => vec!["video_grid_thw".to_string()], - Modality::Audio | Modality::ImageEmbeds => vec![], - } - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - use crate::{ - registry::{test_helpers::*, ModelRegistry}, - types::ImageSize, - }; - - fn omni_tokenizer() -> TestTokenizer { - TestTokenizer::new(&[ - (AUDIO_PAD_TOKEN, 151675), - (IMAGE_PAD_TOKEN, 151655), - (VIDEO_PAD_TOKEN, 151656), - ]) - } - - #[test] - fn omni_accepts_mixed_modalities_and_uses_nested_tokens() { - let tokenizer = omni_tokenizer(); - let config = json!({ - "model_type": "qwen3_omni_moe", - "thinker_config": { - "audio_token_id": 151675, - "image_token_id": 151655, - "video_token_id": 151656 - } - }); - let metadata = ModelMetadata { - model_id: "Qwen/Qwen3-Omni-30B-A3B-Thinking", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).unwrap(); - assert_eq!(spec.name(), "qwen3_omni"); - assert_eq!( - spec.placeholder_token(&metadata).unwrap(), - spec.placeholder_token_for(&metadata, Modality::Image) - .unwrap() - ); - assert_eq!( - spec.placeholder_token_id(&metadata).unwrap(), - spec.placeholder_token_id_for(&metadata, Modality::Image) - .unwrap() - ); - spec.validate_media_request( - &metadata, - &[ - (Modality::Image, 2), - (Modality::Video, 1), - (Modality::Audio, 2), - ], - ) - .unwrap(); - - for (modality, expected) in [ - (Modality::Image, 151655), - (Modality::Video, 151656), - (Modality::Audio, 151675), - ] { - let replacements = spec - .prompt_replacements_for( - &metadata, - &test_preprocessed_with_tokens(&[ImageSize::new(32, 32)], &[3]), - modality, - ) - .unwrap(); - assert_eq!(replacements[0].tokens, vec![expected; 3]); - } - } - - #[test] - fn omni_layouts_are_modality_specific() { - let image = Qwen3OmniSpec.encoder_field_layouts_for(Modality::Image); - assert_eq!(image.encoder_input, FieldLayout::flat("patches_per_image")); - assert!(image.model_specific.contains_key("image_grid_thw")); - - let video = Qwen3OmniSpec.encoder_field_layouts_for(Modality::Video); - assert_eq!(video.encoder_input, FieldLayout::flat("patches_per_video")); - assert!(video.model_specific.contains_key("video_grid_thw")); - - let audio = Qwen3OmniSpec.encoder_field_layouts_for(Modality::Audio); - assert_eq!(audio.encoder_input, FieldLayout::Batched); - assert!(audio.model_specific.contains_key("feature_attention_mask")); - } - - #[test] - fn omni_keep_on_cpu_keys_are_modality_specific() { - assert_eq!( - Qwen3OmniSpec.keep_on_cpu_keys_for(Modality::Image), - vec!["image_grid_thw"] - ); - assert_eq!( - Qwen3OmniSpec.keep_on_cpu_keys_for(Modality::Video), - vec!["video_grid_thw"] - ); - assert!(Qwen3OmniSpec - .keep_on_cpu_keys_for(Modality::Audio) - .is_empty()); - } - - #[test] - fn omni_spec_builds_qwen_audio_processor() { - use std::sync::Arc; - - use bytes::Bytes; - - use crate::{ - audio::DecodedAudio, - types::{AudioClip, AudioSource}, - vision::PreProcessorConfig, - }; - - let tokenizer = omni_tokenizer(); - let config = json!({"model_type": "qwen3_omni_moe"}); - let metadata = ModelMetadata { - model_id: "Qwen/Qwen3-Omni-30B-A3B-Thinking", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).unwrap(); - - let preprocessor_config = PreProcessorConfig::from_json( - r#"{"feature_size": 16, "sampling_rate": 16000, "n_fft": 400, "hop_length": 160}"#, - ) - .unwrap(); - let processor = spec - .audio_processor(&config, &preprocessor_config) - .expect("qwen3_omni spec must provide an audio processor"); - - let clip = Arc::new(AudioClip::new( - Bytes::from_static(b"audio"), - DecodedAudio { - samples: vec![0.0; 800], - sample_rate: 16_000, - }, - AudioSource::InlineBytes, - "audio-hash".to_string(), - )); - let result = processor.preprocess(&[clip]).unwrap(); - assert_eq!(result.encoder_input.shape(), &[1, 16, 5]); - } -} diff --git a/crates/multimodal/src/registry/qwen3_vl.rs b/crates/multimodal/src/registry/qwen3_vl.rs deleted file mode 100644 index 3e2ff0342..000000000 --- a/crates/multimodal/src/registry/qwen3_vl.rs +++ /dev/null @@ -1,517 +0,0 @@ -use std::collections::HashMap; - -use llm_tokenizer::Encoding; -use serde_json::{json, Value}; - -use crate::{ - encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}, - registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, - types::{FieldLayout, Modality, PromptReplacement, TokenId}, -}; - -pub(super) struct Qwen3VLVisionSpec; - -impl Qwen3VLVisionSpec { - fn image_pad_token_id(metadata: &ModelMetadata) -> RegistryResult { - metadata - .config_u32(&["image_token_id"]) - .map(|v| v as TokenId) - .ok_or_else(|| ModelRegistryError::MissingConfigField { - field: "image_token_id".to_string(), - }) - } - - fn video_pad_token_id(metadata: &ModelMetadata) -> RegistryResult { - metadata - .config_u32(&["video_token_id"]) - .map(|v| v as TokenId) - .ok_or_else(|| ModelRegistryError::MissingConfigField { - field: "video_token_id".to_string(), - }) - } - - fn vision_start_token_id(metadata: &ModelMetadata) -> Option { - metadata - .config_u32(&["vision_start_token_id"]) - .map(|v| v as TokenId) - } - - fn vision_end_token_id(metadata: &ModelMetadata) -> Option { - metadata - .config_u32(&["vision_end_token_id"]) - .map(|v| v as TokenId) - } - - fn token_for_id( - metadata: &ModelMetadata, - token_id: TokenId, - field: &str, - ) -> RegistryResult { - metadata - .tokenizer - .id_to_token(token_id as u32) - .ok_or_else(|| ModelRegistryError::TokenNotFound { - token: format!("{field}:{token_id}"), - }) - } - - fn video_grid_t(preprocessed: &PreprocessedEncoderInputs) -> Option { - match preprocessed.model_specific.get("video_grid_thw") { - Some(ModelSpecificValue::IntTensor { data, shape }) - if shape == &[1, 3] && !data.is_empty() => - { - usize::try_from(data[0]).ok() - } - _ => None, - } - } - - fn encode_plain_text(metadata: &ModelMetadata, text: &str) -> Vec { - metadata - .tokenizer - .encode(text, false) - .ok() - .map(|encoding| match encoding { - Encoding::Hf(inner) => inner - .get_ids() - .iter() - .map(|&id| id as TokenId) - .collect::>(), - Encoding::Plain(ids) | Encoding::Tiktoken(ids) => { - ids.into_iter().map(|id| id as TokenId).collect() - } - }) - .unwrap_or_default() - } - - /// Build the per-frame video placeholder body for the Qwen3-VL family. - /// - /// Qwen3-VL lays out video as one `<|vision_start|> .. <|vision_end|>` block - /// per temporal frame with a `` timestamp between frames. The chat - /// template already supplies the outer `<|vision_start|>`/`<|vision_end|>`, so - /// this emits only the inner per-frame structure (hence the `grid_idx > 0` - /// guards that reuse the template's opener/closer for the first/last frame). - /// Returns `None` when the layout can't apply (single-frame or ragged token - /// counts), leaving the caller to fall back to a flat pad block. - fn per_frame_video_tokens( - metadata: &ModelMetadata, - pad_token_id: TokenId, - num_tokens: usize, - grid_t: usize, - ) -> Option> { - if grid_t <= 1 || num_tokens == 0 || !num_tokens.is_multiple_of(grid_t) { - return None; - } - let vision_start = Self::vision_start_token_id(metadata)?; - let vision_end = Self::vision_end_token_id(metadata)?; - let tokens_per_grid = num_tokens / grid_t; - let mut tokens = Vec::with_capacity(num_tokens + (grid_t.saturating_sub(1)) * 8); - let temporal_patch_size = metadata - .config_u32(&["vision_config", "temporal_patch_size"]) - .unwrap_or(2) as f64; - // SMG currently samples Qwen videos at the HF default 2 fps. Match HF's - // prompt timestamp convention: timestamp each temporal patch by the - // average frame time and format it with one decimal place. - let sample_fps = 2.0_f64; - - for grid_idx in 0..grid_t { - let seconds = (grid_idx as f64 * temporal_patch_size - + (temporal_patch_size - 1.0) / 2.0) - / sample_fps; - if grid_idx > 0 { - tokens.push(vision_end); - } - tokens.extend(Self::encode_plain_text( - metadata, - &format!("<{seconds:.1} seconds>"), - )); - if grid_idx > 0 { - tokens.push(vision_start); - } - tokens.extend(std::iter::repeat_n(pad_token_id, tokens_per_grid)); - } - - Some(tokens) - } -} - -impl ModelProcessorSpec for Qwen3VLVisionSpec { - fn name(&self) -> &'static str { - "qwen3_vl" - } - - fn matches(&self, metadata: &ModelMetadata) -> bool { - let id = metadata.model_id.to_ascii_lowercase(); - let model_type = metadata.config_model_type(); - let is_qwen3_vl = id.contains("qwen3") && id.contains("vl") - || model_type.is_some_and(|mt| mt == "qwen3_vl"); - let is_qwen3_5 = id.contains("qwen3.5") - || id.contains("qwen3.6") - || model_type.is_some_and(|mt| mt == "qwen3_5" || mt == "qwen3_5_moe"); - is_qwen3_vl || is_qwen3_5 - } - - fn placeholder_token(&self, metadata: &ModelMetadata) -> RegistryResult { - Self::token_for_id( - metadata, - Self::image_pad_token_id(metadata)?, - "image_token_id", - ) - } - - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { - Self::image_pad_token_id(metadata) - } - - fn placeholder_token_for( - &self, - metadata: &ModelMetadata, - modality: Modality, - ) -> RegistryResult { - match modality { - Modality::Image => self.placeholder_token(metadata), - Modality::Video => Self::token_for_id( - metadata, - Self::video_pad_token_id(metadata)?, - "video_token_id", - ), - _ => Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }), - } - } - - fn placeholder_token_id_for( - &self, - metadata: &ModelMetadata, - modality: Modality, - ) -> RegistryResult { - match modality { - Modality::Image => Self::image_pad_token_id(metadata), - Modality::Video => Self::video_pad_token_id(metadata), - _ => Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }), - } - } - - fn modality_limits( - &self, - metadata: &ModelMetadata, - ) -> RegistryResult> { - let mut limits = HashMap::from([(Modality::Image, 10)]); - if metadata.config_u32(&["video_token_id"]).is_some() { - limits.insert(Modality::Video, 1); - } - Ok(limits) - } - - fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(json!({})) - } - - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - let pad_token_id = Self::image_pad_token_id(metadata)?; - let placeholder_token = self.placeholder_token(metadata)?; - // The chat template already wraps each image with <|vision_start|> ... <|vision_end|>, - // so we only expand the single <|image_pad|> placeholder to N pad tokens. - Ok(preprocessed - .feature_token_counts - .iter() - .map(|&num_tokens| { - let tokens = vec![pad_token_id; num_tokens]; - PromptReplacement::sequence(Modality::Image, &placeholder_token, tokens) - }) - .collect()) - } - - fn prompt_replacements_for( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - modality: Modality, - ) -> RegistryResult> { - match modality { - Modality::Image => self.prompt_replacements(metadata, preprocessed), - Modality::Video => { - let pad_token_id = Self::video_pad_token_id(metadata)?; - let placeholder_token = self.placeholder_token_for(metadata, Modality::Video)?; - let video_grid_t = Self::video_grid_t(preprocessed); - Ok(preprocessed - .feature_token_counts - .iter() - .map(|&num_tokens| { - // Every Qwen3-VL model routed to this spec (base VL and the - // 3.5/3.6 family) needs the per-frame video layout: vLLM's - // mrope pass scans for one <|vision_start|> per temporal - // frame, so a single flat block crashes any multi-frame - // video. Fall back to a flat block only when the per-frame - // layout can't be built (single-frame or unknown grid_t). - let tokens = video_grid_t - .and_then(|grid_t| { - Self::per_frame_video_tokens( - metadata, - pad_token_id, - num_tokens, - grid_t, - ) - }) - .unwrap_or_else(|| vec![pad_token_id; num_tokens]); - // The chat template wraps the placeholder as - // <|vision_start|><|video_pad|><|vision_end|>; the leading - // <|vision_start|> belongs to the placeholder range so vLLM's - // per-frame video mrope finds one marker per frame starting at - // the range offset (it scans even for a single frame). - PromptReplacement::sequence(Modality::Video, &placeholder_token, tokens) - .with_structural_prefix(1) - }) - .collect()) - } - _ => Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }), - } - } - - fn field_layouts(&self) -> HashMap { - // encoder_input is patchified: [total_patches, patch_features]. - // patches_per_image tells how many patches belong to each image. - // image_grid_thw is [num_images, 3]. - HashMap::from([ - ( - "pixel_values".to_string(), - FieldLayout::flat("patches_per_image"), - ), - ("image_grid_thw".to_string(), FieldLayout::Batched), - ("patches_per_image".to_string(), FieldLayout::Batched), - ("video_grid_thw".to_string(), FieldLayout::Batched), - ("patches_per_video".to_string(), FieldLayout::Batched), - ("video_second_per_grid".to_string(), FieldLayout::Batched), - ]) - } - - fn keep_on_cpu_keys(&self) -> Vec { - vec!["image_grid_thw".to_string(), "video_grid_thw".to_string()] - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use crate::{ - encoder_inputs::ModelSpecificValue, - registry::{test_helpers::*, ModelMetadata, ModelRegistry}, - types::ImageSize, - }; - - #[test] - fn qwen3_vl_pad_only_replacement() { - let tokenizer = TestTokenizer::new(&[("", 999), ("<|image_pad|>", 151655)]); - let config = json!({ - "model_type": "qwen3_vl", - "vision_start_token_id": 151652, - "image_token_id": 151655, - "vision_end_token_id": 151653, - "vision_config": {"patch_size": 16} - }); - let metadata = ModelMetadata { - model_id: "Qwen3-VL-7B", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("qwen3 spec"); - assert_eq!(spec.name(), "qwen3_vl"); - // 448/16 = 28 grid, merge_size=2 => (28*28)/4 = 196 tokens - let replacements = spec - .prompt_replacements( - &metadata, - &test_preprocessed_with_tokens(&[ImageSize::new(448, 448)], &[196]), - ) - .unwrap(); - // Only pad tokens — vision_start/vision_end are already in the chat template - assert_eq!(replacements[0].tokens.len(), 196); - assert_eq!(replacements[0].tokens[0], 151655); // pad (image_token_id) - assert_eq!(*replacements[0].tokens.last().unwrap(), 151655); // pad - } - - #[test] - fn qwen3_vl_video_pad_replacement() { - let tokenizer = TestTokenizer::new(&[("<|video_pad|>", 151656)]); - let config = json!({ - "model_type": "qwen3_5", - "image_token_id": 151655, - "video_token_id": 151656, - }); - let metadata = ModelMetadata { - model_id: "Qwen3.5-VL-7B", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("qwen3.5 spec"); - let replacements = spec - .prompt_replacements_for( - &metadata, - &test_preprocessed_with_tokens(&[ImageSize::new(448, 448)], &[128]), - crate::types::Modality::Video, - ) - .unwrap(); - - assert_eq!(replacements[0].modality, crate::types::Modality::Video); - assert_eq!(replacements[0].tokens.len(), 128); - assert_eq!(replacements[0].tokens[0], 151656); - } - - #[test] - fn qwen3_5_video_replacement_splits_temporal_grid() { - let tokenizer = TestTokenizer::new(&[ - ("<|video_pad|>", 151656), - ("<|vision_start|>", 151652), - ("<|vision_end|>", 151653), - ]); - let config = json!({ - "model_type": "qwen3_5", - "image_token_id": 151655, - "video_token_id": 151656, - "vision_start_token_id": 151652, - "vision_end_token_id": 151653, - }); - let metadata = ModelMetadata { - model_id: "Qwen/Qwen3.5-4B", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("qwen3.5 spec"); - let preprocessed = test_preprocessed_with_tokens(&[ImageSize::new(320, 256)], &[160]) - .with_extra( - "video_grid_thw", - ModelSpecificValue::int_2d(vec![2, 16, 20], 1, 3), - ); - let replacements = spec - .prompt_replacements_for(&metadata, &preprocessed, crate::types::Modality::Video) - .unwrap(); - - let tokens = &replacements[0].tokens; - assert_eq!(tokens.len(), 162); - assert!(tokens[..80].iter().all(|&token| token == 151656)); - assert_eq!(tokens[80], 151653); - assert_eq!(tokens[81], 151652); - assert!(tokens[82..].iter().all(|&token| token == 151656)); - } - - #[test] - fn qwen3_vl_video_splits_temporal_grid() { - // Base Qwen3-VL (not the 3.5/3.6 family) must ALSO emit one vision block - // per temporal frame. vLLM's mrope pass scans for a <|vision_start|> per - // frame, so a flat single block crashes any multi-frame video. Regression - // guard for the is_qwen3_5-only gate that previously left base VL flat. - let tokenizer = TestTokenizer::new(&[ - ("<|video_pad|>", 151656), - ("<|vision_start|>", 151652), - ("<|vision_end|>", 151653), - ]); - let config = json!({ - "model_type": "qwen3_vl", - "image_token_id": 151655, - "video_token_id": 151656, - "vision_start_token_id": 151652, - "vision_end_token_id": 151653, - }); - let metadata = ModelMetadata { - model_id: "Qwen/Qwen3-VL-8B-Instruct", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("qwen3_vl spec"); - assert_eq!(spec.name(), "qwen3_vl"); - let preprocessed = test_preprocessed_with_tokens(&[ImageSize::new(320, 256)], &[160]) - .with_extra( - "video_grid_thw", - ModelSpecificValue::int_2d(vec![2, 16, 20], 1, 3), - ); - let replacements = spec - .prompt_replacements_for(&metadata, &preprocessed, crate::types::Modality::Video) - .unwrap(); - - // Two temporal frames -> a vision_end/vision_start seam splits the 160 - // pads into two 80-token halves (mirrors the 3.5 case above). - let tokens = &replacements[0].tokens; - assert_eq!(tokens.len(), 162); - assert!(tokens[..80].iter().all(|&token| token == 151656)); - assert_eq!(tokens[80], 151653); - assert_eq!(tokens[81], 151652); - assert!(tokens[82..].iter().all(|&token| token == 151656)); - } - - #[test] - fn qwen2_vl_does_not_match_qwen3() { - let tokenizer = TestTokenizer::new(&[("", 999)]); - let config = json!({ - "model_type": "qwen3_vl", - "vision_start_token_id": 151652, - "image_token_id": 151655, - "vision_end_token_id": 151653, - "vision_config": {"patch_size": 16} - }); - let metadata = ModelMetadata { - model_id: "Qwen3-VL-7B", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("should match qwen3"); - // Must match qwen3_vl spec, not qwen_vl - assert_eq!(spec.name(), "qwen3_vl"); - } - - #[test] - fn qwen3_vl_matches_alias_via_model_type() { - let tokenizer = TestTokenizer::new(&[("<|image_pad|>", 151655)]); - let config = json!({ - "model_type": "qwen3_vl", - "vision_start_token_id": 151652, - "image_token_id": 151655, - "vision_end_token_id": 151653 - }); - let metadata = ModelMetadata { - model_id: "custom-model", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry - .lookup(&metadata) - .expect("should match qwen3 alias"); - assert_eq!(spec.name(), "qwen3_vl"); - } - - #[test] - fn qwen3_5_matches_alias_via_model_type() { - let tokenizer = TestTokenizer::new(&[("<|image_pad|>", 151655)]); - let config = json!({ - "model_type": "qwen3_5_moe", - "image_token_id": 151655, - }); - let metadata = ModelMetadata { - model_id: "custom-model", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry - .lookup(&metadata) - .expect("should match qwen3.5 alias"); - assert_eq!(spec.name(), "qwen3_vl"); - } -} diff --git a/crates/multimodal/src/registry/qwen_vl.rs b/crates/multimodal/src/registry/qwen_vl.rs deleted file mode 100644 index e991a6a94..000000000 --- a/crates/multimodal/src/registry/qwen_vl.rs +++ /dev/null @@ -1,153 +0,0 @@ -use std::collections::HashMap; - -use serde_json::{json, Value}; - -use crate::{ - encoder_inputs::PreprocessedEncoderInputs, - registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, - types::{FieldLayout, Modality, PromptReplacement, TokenId}, -}; - -pub(super) struct QwenVLVisionSpec; - -impl QwenVLVisionSpec { - fn pad_token_id(metadata: &ModelMetadata) -> RegistryResult { - metadata - .config_u32(&["vision_token_id"]) - .map(|v| v as TokenId) - .ok_or_else(|| ModelRegistryError::MissingConfigField { - field: "vision_token_id".to_string(), - }) - } -} - -impl ModelProcessorSpec for QwenVLVisionSpec { - fn name(&self) -> &'static str { - "qwen_vl" - } - - fn matches(&self, metadata: &ModelMetadata) -> bool { - let id = metadata.model_id.to_ascii_lowercase(); - id.contains("qwen") && id.contains("vl") - || metadata - .config_model_type() - .is_some_and(|mt| mt == "qwen2_vl") - } - - fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok("".to_string()) - } - - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { - // Must match pad_token_id (vision_token_id) — this is the repeated token - // in the expanded sequence. image_token_id is a distinct token in Qwen2-VL. - Self::pad_token_id(metadata) - } - - fn modality_limits( - &self, - _metadata: &ModelMetadata, - ) -> RegistryResult> { - Ok(HashMap::from([(Modality::Image, 10)])) - } - - fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(json!({})) - } - - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - let pad_token_id = Self::pad_token_id(metadata)?; - let placeholder_token = self.placeholder_token(metadata)?; - // The chat template already wraps each image with <|vision_start|> ... <|vision_end|>, - // so we only expand the single placeholder to N pad tokens. - Ok(preprocessed - .feature_token_counts - .iter() - .map(|&num_tokens| { - let tokens = vec![pad_token_id; num_tokens]; - PromptReplacement::sequence(Modality::Image, &placeholder_token, tokens) - }) - .collect()) - } - - fn field_layouts(&self) -> HashMap { - // encoder_input is patchified: [total_patches, patch_features]. - // patches_per_image tells how many patches belong to each image. - // image_grid_thw is [num_images, 3]. - HashMap::from([ - ( - "pixel_values".to_string(), - FieldLayout::flat("patches_per_image"), - ), - ("image_grid_thw".to_string(), FieldLayout::Batched), - ("patches_per_image".to_string(), FieldLayout::Batched), - ]) - } - - fn keep_on_cpu_keys(&self) -> Vec { - vec!["image_grid_thw".to_string()] - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use crate::{ - registry::{test_helpers::*, ModelMetadata, ModelRegistry}, - types::ImageSize, - }; - - #[test] - fn qwen_vision_uses_config_token_ids() { - let tokenizer = TestTokenizer::new(&[("", 999)]); - let config = json!({ - "model_type": "qwen2_vl", - "vision_start_token_id": 151652, - "vision_token_id": 151654, - "image_token_id": 151655, - "vision_config": {"patch_size": 14} - }); - let metadata = ModelMetadata { - model_id: "Qwen2-VL-7B", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("qwen spec"); - // 448/14 = 32 grid, merge_size=2 => (32*32)/4 = 256 tokens - let replacements = spec - .prompt_replacements( - &metadata, - &test_preprocessed_with_tokens(&[ImageSize::new(448, 448)], &[256]), - ) - .unwrap(); - // Only pad tokens — vision_start/vision_end are already in the chat template - assert_eq!(replacements[0].tokens.len(), 256); - assert_eq!(replacements[0].tokens[0], 151654); // pad (vision_token_id) - assert_eq!(*replacements[0].tokens.last().unwrap(), 151654); - } - - #[test] - fn qwen_vl_matches_alias_via_model_type() { - let tokenizer = TestTokenizer::new(&[("", 999)]); - let config = json!({ - "model_type": "qwen2_vl", - "vision_start_token_id": 151652, - "vision_token_id": 151654, - "image_token_id": 151655 - }); - let metadata = ModelMetadata { - model_id: "custom-model", - tokenizer: &tokenizer, - config: &config, - }; - let registry = ModelRegistry::new(); - let spec = registry.lookup(&metadata).expect("should match qwen alias"); - assert_eq!(spec.name(), "qwen_vl"); - } -} diff --git a/crates/multimodal/src/registry/traits.rs b/crates/multimodal/src/registry/traits.rs deleted file mode 100644 index 96c9059e2..000000000 --- a/crates/multimodal/src/registry/traits.rs +++ /dev/null @@ -1,339 +0,0 @@ -use std::collections::HashMap; - -use llm_tokenizer::TokenizerTrait; -use serde_json::Value; -use thiserror::Error; - -use crate::{ - audio::AudioPreProcessor, - encoder_inputs::PreprocessedEncoderInputs, - types::{EncoderFieldLayouts, FieldLayout, Modality, PromptReplacement, TokenId}, - vision::PreProcessorConfig, -}; - -#[derive(Debug, Error, PartialEq, Eq)] -pub enum ModelRegistryError { - #[error("unsupported model: {0}")] - UnsupportedModel(String), - #[error("token '{token}' not found in tokenizer vocabulary")] - TokenNotFound { token: String }, - #[error("missing config field '{field}'")] - MissingConfigField { field: String }, - #[error("modality {modality} is not supported by model spec {spec}")] - UnsupportedModality { - spec: &'static str, - modality: Modality, - }, - #[error("model spec {spec} supports at most {limit} {modality} inputs; got {requested}")] - ModalityLimitExceeded { - spec: &'static str, - modality: Modality, - limit: usize, - requested: usize, - }, - #[error("modality {modality} appears more than once in the request for model spec {spec}")] - DuplicateModality { - spec: &'static str, - modality: Modality, - }, -} - -pub type RegistryResult = Result; - -/// Metadata about the current model used to derive tokenizer/config dependent fields. -pub struct ModelMetadata<'a> { - pub model_id: &'a str, - pub tokenizer: &'a dyn TokenizerTrait, - pub config: &'a Value, -} - -impl<'a> ModelMetadata<'a> { - pub fn token_id(&self, token: &str) -> RegistryResult { - self.tokenizer - .token_to_id(token) - .map(|id| id as TokenId) - .ok_or_else(|| ModelRegistryError::TokenNotFound { - token: token.to_string(), - }) - } - - pub fn config_u32(&self, path: &[&str]) -> Option { - Self::find_value(self.config, path).and_then(|value| value.as_u64().map(|v| v as u32)) - } - - pub fn config_model_type(&self) -> Option<&str> { - Self::find_value(self.config, &["model_type"]).and_then(Value::as_str) - } - - fn find_value<'v>(value: &'v Value, path: &[&str]) -> Option<&'v Value> { - let mut current = value; - for key in path { - current = current.get(*key)?; - } - Some(current) - } -} - -pub trait ModelProcessorSpec: Send + Sync { - fn name(&self) -> &'static str; - fn matches(&self, metadata: &ModelMetadata) -> bool; - fn placeholder_token(&self, metadata: &ModelMetadata) -> RegistryResult; - fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult; - fn placeholder_token_for( - &self, - metadata: &ModelMetadata, - modality: Modality, - ) -> RegistryResult { - match modality { - Modality::Image => self.placeholder_token(metadata), - _ => Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }), - } - } - fn placeholder_token_id_for( - &self, - metadata: &ModelMetadata, - modality: Modality, - ) -> RegistryResult { - match modality { - Modality::Image => self.placeholder_token_id(metadata), - _ => Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }), - } - } - fn modality_limits(&self, metadata: &ModelMetadata) - -> RegistryResult>; - - /// Validate the active modalities and item counts in one media request. - /// - /// Any subset of the modalities declared by [`Self::modality_limits`] is - /// accepted. Each modality may appear once in `requested`; zero-count - /// entries are ignored. - fn validate_media_request( - &self, - metadata: &ModelMetadata, - requested: &[(Modality, usize)], - ) -> RegistryResult<()> { - let limits = self.modality_limits(metadata)?; - let mut active = Vec::with_capacity(requested.len()); - - for &(modality, count) in requested { - if count == 0 { - continue; - } - if active.contains(&modality) { - return Err(ModelRegistryError::DuplicateModality { - spec: self.name(), - modality, - }); - } - active.push(modality); - - let Some(&limit) = limits.get(&modality) else { - return Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }); - }; - if count > limit { - return Err(ModelRegistryError::ModalityLimitExceeded { - spec: self.name(), - modality, - limit, - requested: count, - }); - } - } - - Ok(()) - } - - fn processor_kwargs(&self, metadata: &ModelMetadata) -> RegistryResult; - - /// Build the audio preprocessor for this model, if it supports audio. - /// - /// This is the single source of truth for audio-processor selection: the - /// same spec that owns a model's prompt/placeholder logic also owns its - /// audio preprocessor factory, so there is no separate string-keyed - /// registry to keep in sync. Audio-less specs use the default (`None`). - /// - /// The processor is built from the current model config because its feature - /// shapes and quantization parameters can be checkpoint-specific. - fn audio_processor( - &self, - _model_config: &Value, - _preprocessor_config: &PreProcessorConfig, - ) -> Option> { - None - } - - /// Compute per-media prompt replacement token sequences. - /// - /// Receives the full preprocessed output so each model can extract whatever - /// metadata it needs (e.g. aspect_ratios for tile-based models). This - /// mirrors vLLM's `_get_prompt_updates(out_mm_kwargs)` pattern. - fn prompt_replacements( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult>; - fn prompt_replacements_for( - &self, - metadata: &ModelMetadata, - preprocessed: &PreprocessedEncoderInputs, - modality: Modality, - ) -> RegistryResult> { - match modality { - Modality::Image => self.prompt_replacements(metadata, preprocessed), - _ => Err(ModelRegistryError::UnsupportedModality { - spec: self.name(), - modality, - }), - } - } - - /// Declare how each tensor's first dimension maps to media items. - /// - /// Keys not listed are treated as shared (replicated across all media items). - /// The `"pixel_values"` key mirrors HF/vLLM vision kwargs and should be - /// included when the primary encoder input differs from batched layout. - fn field_layouts(&self) -> HashMap { - // Default: encoder_input is batched (most models). - HashMap::from([("pixel_values".to_string(), FieldLayout::Batched)]) - } - - /// Declare the neutral primary/side-tensor layout contract for one modality. - /// - /// The default converts the legacy HF/vLLM-shaped field map so existing - /// vision specs remain source-compatible. New multimodal specs should - /// override this method and keep backend-specific field names at adapters. - fn encoder_field_layouts_for(&self, _modality: Modality) -> EncoderFieldLayouts { - EncoderFieldLayouts::from_legacy_fields(self.field_layouts()) - } - - /// Tensor keys that should remain on CPU (not transferred to GPU). - /// - /// In vLLM, certain model-specific tensors are marked `keep_on_cpu=True` - /// in their `MultiModalFieldConfig`. This method mirrors that per-model - /// knowledge so the router can send the hint via gRPC, avoiding the need - /// for the backend to instantiate a Python processor just to query it. - fn keep_on_cpu_keys(&self) -> Vec { - vec![] - } - - /// Tensor keys that should remain on CPU for one modality. - /// - /// The default preserves the legacy model-wide declaration. - fn keep_on_cpu_keys_for(&self, _modality: Modality) -> Vec { - self.keep_on_cpu_keys() - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - use crate::registry::test_helpers::TestTokenizer; - - struct TestSpec; - - impl ModelProcessorSpec for TestSpec { - fn name(&self) -> &'static str { - "test" - } - - fn matches(&self, _metadata: &ModelMetadata) -> bool { - true - } - - fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok("".to_string()) - } - - fn placeholder_token_id(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(1) - } - - fn modality_limits( - &self, - _metadata: &ModelMetadata, - ) -> RegistryResult> { - Ok(HashMap::from([(Modality::Image, 2), (Modality::Audio, 1)])) - } - - fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { - Ok(json!({})) - } - - fn prompt_replacements( - &self, - _metadata: &ModelMetadata, - _preprocessed: &PreprocessedEncoderInputs, - ) -> RegistryResult> { - Ok(vec![]) - } - } - - fn validate( - spec: &dyn ModelProcessorSpec, - requested: &[(Modality, usize)], - ) -> RegistryResult<()> { - let tokenizer = TestTokenizer::new(&[]); - let config = json!({}); - let metadata = ModelMetadata { - model_id: "test-model", - tokenizer: &tokenizer, - config: &config, - }; - spec.validate_media_request(&metadata, requested) - } - - #[test] - fn validation_accepts_any_declared_modality_subset() { - assert_eq!(validate(&TestSpec, &[(Modality::Image, 2)]), Ok(())); - assert_eq!( - validate(&TestSpec, &[(Modality::Image, 1), (Modality::Audio, 1)]), - Ok(()) - ); - } - - #[test] - fn validation_rejects_undeclared_modality() { - assert_eq!( - validate(&TestSpec, &[(Modality::Video, 1)]), - Err(ModelRegistryError::UnsupportedModality { - spec: "test", - modality: Modality::Video, - }) - ); - } - - #[test] - fn validation_rejects_count_above_limit() { - assert_eq!( - validate(&TestSpec, &[(Modality::Image, 3)]), - Err(ModelRegistryError::ModalityLimitExceeded { - spec: "test", - modality: Modality::Image, - limit: 2, - requested: 3, - }) - ); - } - - #[test] - fn validation_rejects_duplicate_modality_counts() { - assert_eq!( - validate(&TestSpec, &[(Modality::Image, 1), (Modality::Image, 1)]), - Err(ModelRegistryError::DuplicateModality { - spec: "test", - modality: Modality::Image, - }) - ); - } -} diff --git a/crates/multimodal/src/tracker.rs b/crates/multimodal/src/tracker.rs deleted file mode 100644 index 10fd4386f..000000000 --- a/crates/multimodal/src/tracker.rs +++ /dev/null @@ -1,164 +0,0 @@ -use std::{collections::HashMap, sync::Arc}; - -use tokio::task::JoinHandle; - -use super::{ - error::{MultiModalError, MultiModalResult}, - media::{ImageFetchConfig, MediaConnector, MediaSource, VideoFetchConfig}, - types::{ - ImageDetail, MediaContentPart, Modality, MultiModalData, MultiModalUUIDs, TrackedMedia, - }, -}; - -type PendingTask = JoinHandle>; - -#[derive(Debug)] -pub struct TrackerOutput { - pub data: MultiModalData, - pub uuids: MultiModalUUIDs, -} - -pub struct AsyncMultiModalTracker { - media_connector: Arc, - pending: HashMap>, - uuids: MultiModalUUIDs, -} - -impl AsyncMultiModalTracker { - pub fn new(media_connector: Arc) -> Self { - Self { - media_connector, - pending: HashMap::new(), - uuids: HashMap::new(), - } - } - - pub fn push_part(&mut self, part: MediaContentPart) -> MultiModalResult<()> { - match part { - MediaContentPart::Text { .. } => {} - MediaContentPart::ImageUrl { url, detail, uuid } => { - let source = match url::Url::parse(&url) { - Ok(parsed) if parsed.scheme() == "data" => MediaSource::DataUrl(url), - _ => MediaSource::Url(url), - }; - self.enqueue_image(source, detail.unwrap_or_default(), uuid); - } - MediaContentPart::ImageData { - data, - mime_type: _, - uuid, - detail, - } => { - self.enqueue_image( - MediaSource::InlineBytes(data), - detail.unwrap_or_default(), - uuid, - ); - } - MediaContentPart::ImageEmbeds { .. } => { - return Err(MultiModalError::UnsupportedContent("image_embeds")); - } - MediaContentPart::AudioUrl { url, uuid } => { - let source = match url::Url::parse(&url) { - Ok(parsed) if parsed.scheme() == "data" => MediaSource::DataUrl(url), - _ => MediaSource::Url(url), - }; - self.enqueue_audio(source, uuid); - } - MediaContentPart::AudioData { - data, - mime_type: _, - uuid, - } => { - self.enqueue_audio(MediaSource::InlineBytes(data), uuid); - } - MediaContentPart::VideoUrl { url, uuid } => { - let source = match url::Url::parse(&url) { - Ok(parsed) if parsed.scheme() == "data" => MediaSource::DataUrl(url), - _ => MediaSource::Url(url), - }; - self.enqueue_video(source, uuid); - } - MediaContentPart::VideoData { - data, - mime_type: _, - uuid, - } => { - self.enqueue_video(MediaSource::InlineBytes(data), uuid); - } - } - Ok(()) - } - - pub async fn finalize(mut self) -> MultiModalResult { - let mut data = MultiModalData::new(); - for (modality, tasks) in self.pending.drain() { - let mut items = Vec::with_capacity(tasks.len()); - for task in tasks { - let media = task.await??; - items.push(media); - } - data.insert(modality, items); - } - - Ok(TrackerOutput { - data, - uuids: self.uuids, - }) - } - - fn enqueue_image(&mut self, source: MediaSource, detail: ImageDetail, uuid: Option) { - let modality = Modality::Image; - self.uuids.entry(modality).or_default().push(uuid); - - let connector = Arc::clone(&self.media_connector); - #[expect( - clippy::disallowed_methods, - reason = "spawn handle is stored in self.pending and awaited in finalize(); fire-and-forget is intentional for concurrent media fetching" - )] - let handle = tokio::spawn(async move { - let frame = connector - .fetch_image(source, ImageFetchConfig { detail }) - .await?; - Ok(TrackedMedia::Image(frame)) - }); - - self.pending.entry(modality).or_default().push(handle); - } - - fn enqueue_video(&mut self, source: MediaSource, uuid: Option) { - let modality = Modality::Video; - self.uuids.entry(modality).or_default().push(uuid); - - let connector = Arc::clone(&self.media_connector); - #[expect( - clippy::disallowed_methods, - reason = "spawn handle is stored in self.pending and awaited in finalize(); fire-and-forget is intentional for concurrent media fetching" - )] - let handle = tokio::spawn(async move { - let clip = connector - .fetch_video(source, VideoFetchConfig::default()) - .await?; - Ok(TrackedMedia::Video(clip)) - }); - - self.pending.entry(modality).or_default().push(handle); - } - - fn enqueue_audio(&mut self, source: MediaSource, uuid: Option) { - let modality = Modality::Audio; - self.uuids.entry(modality).or_default().push(uuid); - - let connector = Arc::clone(&self.media_connector); - #[expect( - clippy::disallowed_methods, - reason = "spawn handle is stored in self.pending and awaited in finalize(); fire-and-forget is intentional for concurrent media fetching" - )] - let handle = tokio::spawn(async move { - let clip = connector.fetch_audio(source).await?; - Ok(TrackedMedia::Audio(clip)) - }); - - self.pending.entry(modality).or_default().push(handle); - } -} diff --git a/crates/multimodal/src/types.rs b/crates/multimodal/src/types.rs deleted file mode 100644 index fe49170e8..000000000 --- a/crates/multimodal/src/types.rs +++ /dev/null @@ -1,559 +0,0 @@ -use std::{collections::HashMap, fmt, path::PathBuf, sync::Arc}; - -use image::{DynamicImage, RgbImage}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -use crate::audio::DecodedAudio; - -/// Supported multimodal modalities. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum Modality { - Image, - ImageEmbeds, - Audio, - Video, -} - -impl fmt::Display for Modality { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Modality::Image => write!(f, "image"), - Modality::ImageEmbeds => write!(f, "image_embeds"), - Modality::Audio => write!(f, "audio"), - Modality::Video => write!(f, "video"), - } - } -} - -/// Detail level passed by OpenAI style APIs. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] -#[serde(rename_all = "snake_case")] -pub enum ImageDetail { - #[default] - Auto, - Low, - High, -} - -/// A normalized content part understood by the tracker. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum MediaContentPart { - Text { - text: String, - }, - ImageUrl { - url: String, - #[serde(skip_serializing_if = "Option::is_none")] - detail: Option, - #[serde(skip_serializing_if = "Option::is_none")] - uuid: Option, - }, - ImageData { - data: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - mime_type: Option, - #[serde(skip_serializing_if = "Option::is_none")] - uuid: Option, - #[serde(skip_serializing_if = "Option::is_none")] - detail: Option, - }, - ImageEmbeds { - payload: Value, - #[serde(skip_serializing_if = "Option::is_none")] - uuid: Option, - }, - AudioUrl { - url: String, - #[serde(skip_serializing_if = "Option::is_none")] - uuid: Option, - }, - AudioData { - data: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - mime_type: Option, - #[serde(skip_serializing_if = "Option::is_none")] - uuid: Option, - }, - VideoUrl { - url: String, - #[serde(skip_serializing_if = "Option::is_none")] - uuid: Option, - }, - VideoData { - data: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - mime_type: Option, - #[serde(skip_serializing_if = "Option::is_none")] - uuid: Option, - }, -} - -/// Image source metadata (useful for hashing & tracing). -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum ImageSource { - Url { url: String }, - DataUrl, - InlineBytes, - File { path: PathBuf }, -} - -/// Audio source metadata (useful for hashing & tracing). -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum AudioSource { - Url { url: String }, - DataUrl, - InlineBytes, - File { path: PathBuf }, -} - -/// Video source metadata (useful for hashing & tracing). -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum VideoSource { - Url { url: String }, - DataUrl, - InlineBytes, - File { path: PathBuf }, -} - -/// Concrete image payload captured by the media connector. -#[derive(Debug, Clone)] -pub struct ImageFrame { - pub image: DynamicImage, - pub raw_bytes: bytes::Bytes, - pub detail: ImageDetail, - pub source: ImageSource, - /// Blake3 hex-digest of raw_bytes, computed at decode time. - pub hash: String, -} - -/// Decoded audio payload captured by the media connector. -#[derive(Debug, Clone)] -pub struct AudioClip { - pub raw_bytes: bytes::Bytes, - pub decoded: DecodedAudio, - pub source: AudioSource, - /// Blake3 hex-digest of raw_bytes, computed at decode time. - pub hash: String, -} - -/// Decoded video payload captured by the media connector. -#[derive(Debug, Clone)] -pub struct VideoClip { - pub frames: Vec, - pub rgb_video: Option, - /// Effective frame rate after connector-side sampling and frame-count clamps. - pub sample_fps: f32, - pub raw_bytes: bytes::Bytes, - pub source: VideoSource, - /// Blake3 hex-digest of raw_bytes, computed at decode time. - pub hash: String, -} - -/// Borrowed RGB frame data for video preprocessors. -#[derive(Debug, Clone, Copy)] -pub struct RgbFrameRef<'a> { - pub width: u32, - pub height: u32, - pub data: &'a [u8], -} - -/// One decoded RGB frame inside a shared decoded-video byte buffer. -#[derive(Debug, Clone)] -pub struct DecodedRgbFrame { - pub width: u32, - pub height: u32, - pub offset: usize, - pub len: usize, -} - -/// Decoded RGB video frames backed by one shared byte buffer. -#[derive(Debug, Clone)] -pub struct DecodedRgbVideo { - pub data: bytes::Bytes, - pub frames: Vec, -} - -impl DecodedRgbVideo { - pub fn new(data: bytes::Bytes, frames: Vec) -> Self { - Self { data, frames } - } - - pub fn frame_refs(&self) -> Result>, String> { - self.frames - .iter() - .map(|frame| { - let end = frame - .offset - .checked_add(frame.len) - .ok_or_else(|| "decoded RGB frame offset overflow".to_string())?; - let data = self - .data - .get(frame.offset..end) - .ok_or_else(|| "decoded RGB frame range is out of bounds".to_string())?; - Ok(RgbFrameRef { - width: frame.width, - height: frame.height, - data, - }) - }) - .collect() - } - - pub fn to_dynamic_images(&self) -> Result, String> { - let mut images = Vec::with_capacity(self.frames.len()); - for frame in &self.frames { - let end = frame - .offset - .checked_add(frame.len) - .ok_or_else(|| "decoded RGB frame offset overflow".to_string())?; - let data = self - .data - .get(frame.offset..end) - .ok_or_else(|| "decoded RGB frame range is out of bounds".to_string())?; - let image = - RgbImage::from_raw(frame.width, frame.height, data.to_vec()).ok_or_else(|| { - format!( - "failed to build RGB frame from {} bytes for {}x{} video", - frame.len, frame.width, frame.height - ) - })?; - images.push(DynamicImage::ImageRgb8(image)); - } - Ok(images) - } -} - -impl VideoClip { - pub fn new( - frames: Vec, - raw_bytes: bytes::Bytes, - source: VideoSource, - hash: String, - ) -> Self { - Self::new_with_sample_fps(frames, raw_bytes, source, hash, 2.0) - } - - pub fn new_with_sample_fps( - frames: Vec, - raw_bytes: bytes::Bytes, - source: VideoSource, - hash: String, - sample_fps: f32, - ) -> Self { - Self { - frames, - rgb_video: None, - sample_fps, - raw_bytes, - source, - hash, - } - } - - pub fn new_rgb( - rgb_video: DecodedRgbVideo, - raw_bytes: bytes::Bytes, - source: VideoSource, - hash: String, - ) -> Self { - Self::new_rgb_with_sample_fps(rgb_video, raw_bytes, source, hash, 2.0) - } - - pub fn new_rgb_with_sample_fps( - rgb_video: DecodedRgbVideo, - raw_bytes: bytes::Bytes, - source: VideoSource, - hash: String, - sample_fps: f32, - ) -> Self { - Self { - frames: Vec::new(), - rgb_video: Some(rgb_video), - sample_fps, - raw_bytes, - source, - hash, - } - } - - pub fn frames(&self) -> &[DynamicImage] { - &self.frames - } - - pub fn rgb_video(&self) -> Option<&DecodedRgbVideo> { - self.rgb_video.as_ref() - } - - pub fn sample_fps(&self) -> f32 { - self.sample_fps - } - - pub fn materialized_frames(&self) -> Result, String> { - if !self.frames.is_empty() { - return Ok(self.frames.clone()); - } - self.rgb_video - .as_ref() - .ok_or_else(|| "video clip has no decoded frames".to_string())? - .to_dynamic_images() - } - - pub fn raw_bytes(&self) -> &[u8] { - &self.raw_bytes - } - - pub fn source(&self) -> &VideoSource { - &self.source - } -} - -impl AudioClip { - pub fn new( - raw_bytes: bytes::Bytes, - decoded: DecodedAudio, - source: AudioSource, - hash: String, - ) -> Self { - Self { - raw_bytes, - decoded, - source, - hash, - } - } - - pub fn raw_bytes(&self) -> &[u8] { - &self.raw_bytes - } - - pub fn decoded(&self) -> &DecodedAudio { - &self.decoded - } - - pub fn source(&self) -> &AudioSource { - &self.source - } -} - -impl ImageFrame { - pub fn new( - image: DynamicImage, - raw_bytes: bytes::Bytes, - detail: ImageDetail, - source: ImageSource, - hash: String, - ) -> Self { - Self { - image, - raw_bytes, - detail, - source, - hash, - } - } - - pub fn data(&self) -> &DynamicImage { - &self.image - } - - pub fn raw_bytes(&self) -> &[u8] { - &self.raw_bytes - } - - pub fn source(&self) -> &ImageSource { - &self.source - } - - pub fn size(&self) -> ImageSize { - ImageSize::new(self.image.width(), self.image.height()) - } -} - -/// Container for all supported multimodal media objects. -#[derive(Debug, Clone)] -pub enum TrackedMedia { - Image(Arc), - Audio(Arc), - Video(Arc), - /// Placeholder variants for future modalities. - Embeddings, -} - -pub type MultiModalData = HashMap>; -pub type MultiModalUUIDs = HashMap>>; - -pub type TokenId = i32; - -/// Declares how a multimodal tensor's first dimension maps to media items. -/// -/// Used by [`crate::registry::ModelProcessorSpec::encoder_field_layouts_for`] to tell the backend -/// how to split tensors for per-item scheduling (vLLM `MultiModalFieldConfig`). -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum FieldLayout { - /// First dimension equals number of media items (one slice per item). - Batched, - /// Variable-length slices per item. The sizes are stored in the tensor - /// named by `sizes_key` (e.g. `"patches_per_image"` or `"patches_per_video"`). - Flat { sizes_key: String }, -} - -impl FieldLayout { - /// Convenience constructor for `Flat`. - pub fn flat(sizes_key: impl Into) -> Self { - Self::Flat { - sizes_key: sizes_key.into(), - } - } -} - -/// Layout contract for one modality's encoder inputs. -/// -/// The primary encoder input is transported independently from named, -/// model-specific side tensors. Keeping its layout typed avoids leaking a -/// vision-specific field name into audio and other modality processors. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct EncoderFieldLayouts { - pub encoder_input: FieldLayout, - pub model_specific: HashMap, -} - -impl EncoderFieldLayouts { - pub fn new(encoder_input: FieldLayout, model_specific: HashMap) -> Self { - Self { - encoder_input, - model_specific, - } - } - - /// Convert the legacy HF/vLLM-shaped field map into the neutral contract. - /// - /// Existing vision specs use `pixel_values` for the primary encoder input. - /// New specs should construct [`Self`] directly instead. - pub fn from_legacy_fields(mut fields: HashMap) -> Self { - let encoder_input = fields - .remove("pixel_values") - .unwrap_or(FieldLayout::Batched); - Self::new(encoder_input, fields) - } -} - -impl Default for EncoderFieldLayouts { - fn default() -> Self { - Self::new(FieldLayout::Batched, HashMap::new()) - } -} - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -pub struct ImageSize { - pub width: u32, - pub height: u32, -} - -impl ImageSize { - pub fn new(width: u32, height: u32) -> Self { - Self { width, height } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize, Default)] -pub struct PlaceholderRange { - pub offset: usize, - pub length: usize, -} - -#[derive(Debug, Clone)] -pub struct PromptReplacement { - pub modality: Modality, - pub placeholder_token: String, - pub tokens: Vec, - /// Number of structural tokens the chat template emits *immediately before* - /// this placeholder (e.g. Qwen's leading `<|vision_start|>`) that belong to - /// the placeholder's range. `expand_tokens` folds them into the reported - /// [`PlaceholderRange`] without re-emitting them, so backends that scan the - /// range for structural markers see the leading marker. vLLM's video mrope - /// walks each frame from `<|vision_start|>` starting at the range offset, so - /// the offset must sit on (or before) the first marker. 0 for the common - /// case where the range is exactly the replacement. - pub structural_prefix: usize, -} - -impl PromptReplacement { - pub fn repeated( - modality: Modality, - placeholder_token: &str, - token_id: TokenId, - count: usize, - ) -> Self { - Self { - modality, - placeholder_token: placeholder_token.to_string(), - tokens: vec![token_id; count], - structural_prefix: 0, - } - } - - pub fn sequence(modality: Modality, placeholder_token: &str, sequence: Vec) -> Self { - Self { - modality, - placeholder_token: placeholder_token.to_string(), - tokens: sequence, - structural_prefix: 0, - } - } - - /// Declare that `n` template-emitted structural tokens precede this - /// placeholder and should be included in its reported range. See - /// [`Self::structural_prefix`]. - #[must_use] - pub fn with_structural_prefix(mut self, n: usize) -> Self { - self.structural_prefix = n; - self - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn placeholder_range_serializes() { - let range = PlaceholderRange { - offset: 10, - length: 4, - }; - let json = serde_json::to_string(&range).unwrap(); - assert!(json.contains("offset")); - } - - #[test] - fn prompt_replacement_builders() { - let rep = PromptReplacement::repeated(Modality::Image, "", 100, 3); - assert_eq!(rep.tokens, vec![100, 100, 100]); - } - - #[test] - fn legacy_encoder_fields_are_split_into_typed_layouts() { - let layouts = EncoderFieldLayouts::from_legacy_fields(HashMap::from([ - ( - "pixel_values".to_string(), - FieldLayout::flat("patches_per_image"), - ), - ("image_grid_thw".to_string(), FieldLayout::Batched), - ])); - - assert_eq!( - layouts.encoder_input, - FieldLayout::flat("patches_per_image") - ); - assert_eq!( - layouts.model_specific, - HashMap::from([("image_grid_thw".to_string(), FieldLayout::Batched)]) - ); - } -} diff --git a/crates/multimodal/src/vision/execution.rs b/crates/multimodal/src/vision/execution.rs deleted file mode 100644 index fa3376c62..000000000 --- a/crates/multimodal/src/vision/execution.rs +++ /dev/null @@ -1,48 +0,0 @@ -//! Shared execution primitives for CPU-bound vision preprocessing. -//! -//! This module defines task granularity without owning a process-wide thread -//! pool. Rayon remains responsible for worker lifecycle and work stealing. - -const PARALLEL_MIN_BYTES: usize = 1 << 19; -const MAX_TASKS_PER_OPERATION: usize = 8; - -pub(crate) fn scope<'scope, OP, R>(operation: OP) -> R -where - OP: FnOnce(&rayon::Scope<'scope>) -> R + Send, - R: Send, -{ - rayon::scope(operation) -} - -pub(crate) fn task_count( - output_bytes: usize, - work_items: usize, - min_items_per_task: usize, -) -> usize { - debug_assert!(min_items_per_task > 0); - if output_bytes < PARALLEL_MIN_BYTES || work_items < 2 * min_items_per_task { - return 1; - } - - (work_items / min_items_per_task) - .min(rayon::current_num_threads()) - .clamp(1, MAX_TASKS_PER_OPERATION) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn small_operations_stay_serial() { - assert_eq!(task_count(PARALLEL_MIN_BYTES - 1, 1_024, 1), 1); - assert_eq!(task_count(PARALLEL_MIN_BYTES, 63, 32), 1); - } - - #[test] - fn task_count_respects_operation_and_executor_limits() { - let tasks = task_count(PARALLEL_MIN_BYTES, usize::MAX, 1); - assert!(tasks <= MAX_TASKS_PER_OPERATION); - assert!(tasks <= rayon::current_num_threads()); - } -} diff --git a/crates/multimodal/src/vision/mod.rs b/crates/multimodal/src/vision/mod.rs deleted file mode 100644 index 56cd062fd..000000000 --- a/crates/multimodal/src/vision/mod.rs +++ /dev/null @@ -1,53 +0,0 @@ -//! Pure Rust vision processing module for multimodal models. -//! -//! This module provides vision preprocessing pipelines that match HuggingFace -//! processor outputs without requiring Python dependencies. -//! -//! # Architecture -//! -//! The vision module is structured as follows: -//! -//! - `transforms`: Core image transformations (resize, normalize, crop, etc.) -//! - `preprocessor_config`: HuggingFace config parsing -//! - `processor`: Vision processor trait and registry -//! - `processors`: Model-specific implementations (LLaVA, Qwen-VL, etc.) -//! -//! Modality-neutral encoder outputs live in [`crate::encoder_inputs`], while -//! shared errors live in [`crate::error`]. -//! -//! # Usage -//! -//! ```rust,ignore -//! use smg::multimodal::vision::{ -//! PreProcessorConfig, -//! processors::LlavaProcessor, -//! VisionPreProcessor, -//! }; -//! -//! // Load config from HuggingFace -//! let config = PreProcessorConfig::from_json(config_json)?; -//! -//! // Create processor and preprocess images -//! let processor = LlavaProcessor::new(); -//! let result = processor.preprocess(&images, &config)?; -//! ``` - -pub(crate) mod execution; -pub mod preprocessor_config; -pub mod processor; -pub mod processors; -pub(crate) mod scratch; -pub mod transforms; - -// Re-export commonly used types, including compatibility paths for shared -// preprocessing outputs. -pub use preprocessor_config::PreProcessorConfig; -pub use processor::{ - ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor, VisionProcessorRegistry, -}; -pub use processors::{ - Llama4VisionProcessor, LlavaNextProcessor, LlavaProcessor, Phi3VisionProcessor, - Phi4VisionProcessor, PixtralProcessor, Qwen2VLProcessor, Qwen3OmniVisionProcessor, - Qwen3VLProcessor, -}; -pub use transforms::TransformError; diff --git a/crates/multimodal/src/vision/preprocessor_config.rs b/crates/multimodal/src/vision/preprocessor_config.rs deleted file mode 100644 index 9485e6072..000000000 --- a/crates/multimodal/src/vision/preprocessor_config.rs +++ /dev/null @@ -1,605 +0,0 @@ -//! HuggingFace preprocessor_config.json parsing. -//! -//! This module parses the `preprocessor_config.json` files from HuggingFace model -//! repositories, providing the configuration needed for image preprocessing. - -use std::collections::HashMap; - -use image::imageops::FilterType; -use serde::{Deserialize, Deserializer}; - -use super::transforms; - -/// Struct to represent patch_size as dict {"height": x, "width": y} -#[derive(Debug, Clone, Deserialize, Default)] -pub struct PatchSize { - pub height: Option, - pub width: Option, -} - -/// Custom deserializer for patch_size that handles both integer and dict formats. -/// - Integer format: `"patch_size": 16` -> PatchSize { height: 16, width: 16 } -/// - Dict format: `"patch_size": {"height": 16, "width": 16}` -> PatchSize { height: 16, width: 16 } -fn deserialize_patch_size<'de, D>(deserializer: D) -> Result, D::Error> -where - D: Deserializer<'de>, -{ - use std::fmt; - - use serde::de::{self, MapAccess, Visitor}; - - struct PatchSizeVisitor; - - impl<'de> Visitor<'de> for PatchSizeVisitor { - type Value = Option; - - fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { - formatter.write_str("an integer, a dict with height/width, or null") - } - - fn visit_none(self) -> Result - where - E: de::Error, - { - Ok(None) - } - - fn visit_unit(self) -> Result - where - E: de::Error, - { - Ok(None) - } - - fn visit_i64(self, value: i64) -> Result - where - E: de::Error, - { - let v = value as u32; - Ok(Some(PatchSize { - height: Some(v), - width: Some(v), - })) - } - - fn visit_u64(self, value: u64) -> Result - where - E: de::Error, - { - let v = value as u32; - Ok(Some(PatchSize { - height: Some(v), - width: Some(v), - })) - } - - fn visit_map(self, mut map: M) -> Result - where - M: MapAccess<'de>, - { - let mut height = None; - let mut width = None; - - while let Some(key) = map.next_key::()? { - match key.as_str() { - "height" => height = Some(map.next_value::()?), - "width" => width = Some(map.next_value::()?), - _ => { - let _ = map.next_value::()?; - } - } - } - - Ok(Some(PatchSize { height, width })) - } - } - - deserializer.deserialize_any(PatchSizeVisitor) -} - -/// HuggingFace preprocessor_config.json structure. -/// -/// This struct captures the common fields across different vision model processors. -/// Model-specific fields are accessed via the flexible `extra` field. -#[derive(Debug, Clone, Deserialize, Default)] -pub struct PreProcessorConfig { - /// Processor class name (e.g., "CLIPImageProcessor", "Qwen2VLImageProcessor") - #[serde(default)] - pub image_processor_type: Option, - - /// Whether to convert to RGB - #[serde(default)] - pub do_convert_rgb: Option, - - /// Whether to normalize with mean/std - #[serde(default)] - pub do_normalize: Option, - - /// Whether to pad images - #[serde(default)] - pub do_pad: Option, - - /// Whether to rescale pixel values (typically by 1/255) - #[serde(default)] - pub do_rescale: Option, - - /// Whether to resize images - #[serde(default)] - pub do_resize: Option, - - /// Whether to center crop after resizing - #[serde(default)] - pub do_center_crop: Option, - - /// Per-channel normalization mean - #[serde(default, alias = "norm_mean")] - pub image_mean: Option>, - - /// Per-channel normalization std - #[serde(default, alias = "norm_std")] - pub image_std: Option>, - - /// Rescale factor (typically 1/255 = 0.00392156862745098) - #[serde(default)] - pub rescale_factor: Option, - - /// PIL resampling filter enum (0=Nearest, 1=Lanczos, 2=Bilinear, 3=Bicubic) - #[serde(default, alias = "resample")] - pub resampling: Option, - - /// Target size for resizing - /// Can be {"height": H, "width": W} or {"shortest_edge": S} - #[serde(default)] - pub size: Option>, - - /// Target size for center cropping - #[serde(default)] - pub crop_size: Option>, - - // ===================== - // Model-specific fields - // ===================== - /// Vision encoder patch size (typically 14 or 16) - /// Can be an integer or a dict {"height": x, "width": y} - #[serde(default, deserialize_with = "deserialize_patch_size")] - pub patch_size: Option, - - /// Qwen-VL: merge size for token reduction - #[serde(default)] - pub merge_size: Option, - - /// Qwen-VL: minimum total pixels - #[serde(default)] - pub min_pixels: Option, - - /// Qwen-VL: maximum total pixels - #[serde(default)] - pub max_pixels: Option, - - /// Qwen-VL: temporal patch size for video - #[serde(default)] - pub temporal_patch_size: Option, - - /// Phi3-Vision: number of image crops - #[serde(default)] - pub num_crops: Option, - - /// Phi4-Vision: dynamic HD max crops - #[serde(default)] - pub dynamic_hd: Option, - - /// LLaMA-Vision: maximum image tiles - #[serde(default)] - pub max_image_tiles: Option, - - /// Fixed number of image tokens (some model configs use this HF field name). - #[serde(default)] - pub num_img_tokens: Option, - - // ===================== - // Special tokens - // ===================== - /// Image start token - #[serde(default)] - pub im_start_token: Option, - - /// Image end token - #[serde(default)] - pub im_end_token: Option, - - /// Slice start token (for multi-crop) - #[serde(default)] - pub slice_start_token: Option, - - /// Slice end token - #[serde(default)] - pub slice_end_token: Option, - - /// Vision start token (alternative naming) - #[serde(default)] - pub vision_start_token: Option, - - /// Vision end token - #[serde(default)] - pub vision_end_token: Option, - - /// Catch-all for model-specific fields not explicitly defined - #[serde(flatten)] - pub extra: HashMap, -} - -impl PreProcessorConfig { - /// Parse from JSON string. - /// - /// Handles both standard HuggingFace format (top-level fields) and Kimi-K2.5's - /// nested format where values are under `media_proc_cfg`. - pub fn from_json(json: &str) -> Result { - let raw: serde_json::Value = serde_json::from_str(json)?; - Self::from_value(raw) - } - - /// Parse from JSON value. - /// - /// Handles both standard HuggingFace format (top-level fields) and Kimi-K2.5's - /// nested format where values are under `media_proc_cfg`. - pub fn from_value(value: serde_json::Value) -> Result { - let mut config: Self = serde_json::from_value(value.clone())?; - Self::apply_nested_media_cfg(&mut config, &value); - Ok(config) - } - - /// Extract values from nested `media_proc_cfg` (used by Kimi-K2.5 and - /// similar models) when top-level fields are missing. - fn apply_nested_media_cfg(config: &mut Self, raw: &serde_json::Value) { - let Some(media_cfg) = raw.get("media_proc_cfg") else { - return; - }; - if config.image_mean.is_none() { - config.image_mean = media_cfg - .get("image_mean") - .and_then(|v| serde_json::from_value(v.clone()).ok()); - } - if config.image_std.is_none() { - config.image_std = media_cfg - .get("image_std") - .and_then(|v| serde_json::from_value(v.clone()).ok()); - } - if config.patch_size.is_none() { - config.patch_size = media_cfg.get("patch_size").and_then(|v| { - v.as_u64().map(|ps| PatchSize { - height: Some(ps as u32), - width: Some(ps as u32), - }) - }); - } - if config.merge_size.is_none() { - config.merge_size = media_cfg - .get("merge_kernel_size") - .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"] { - if !config.extra.contains_key(key) { - if let Some(v) = media_cfg.get(key) { - config.extra.insert(key.to_string(), v.clone()); - } - } - } - } - - /// Get patch size as a simple usize. - /// - /// Returns the height value from PatchSize if available, falling back to provided default. - pub fn get_patch_size(&self, default: usize) -> usize { - self.patch_size - .as_ref() - .and_then(|p| p.height) - .map(|h| h as usize) - .unwrap_or(default) - } - - /// Whether this config changes Qwen-style processor structure or budgets. - pub(crate) fn has_structural_overrides(&self) -> bool { - self.patch_size.is_some() - || self.merge_size.is_some() - || self.min_pixels.is_some() - || self.max_pixels.is_some() - || self.temporal_patch_size.is_some() - || self.size.is_some() - } - - /// Whether the declared processor type is image-only rather than video-capable. - pub(crate) fn is_image_only_processor_type(&self) -> bool { - self.image_processor_type - .as_deref() - .map(str::to_ascii_lowercase) - .is_some_and(|processor| { - processor.contains("imageprocessor") && !processor.contains("video") - }) - } - - /// Get image mean as fixed array, with fallback to CLIP defaults. - pub fn get_image_mean(&self) -> [f64; 3] { - self.image_mean - .as_ref() - .and_then(|v| { - if v.len() >= 3 { - Some([v[0], v[1], v[2]]) - } else { - None - } - }) - .unwrap_or(Self::CLIP_MEAN) - } - - /// Get image std as fixed array, with fallback to CLIP defaults. - pub fn get_image_std(&self) -> [f64; 3] { - self.image_std - .as_ref() - .and_then(|v| { - if v.len() >= 3 { - Some([v[0], v[1], v[2]]) - } else { - None - } - }) - .unwrap_or(Self::CLIP_STD) - } - - /// Get target size from various config formats. - /// - /// Handles both `{"height": H, "width": W}` and `{"shortest_edge": S}` formats. - /// Returns (height, width). - pub fn get_target_size(&self) -> Option<(u32, u32)> { - self.size.as_ref().map(|s| { - // Try explicit height/width first - let h = s - .get("height") - .or_else(|| s.get("shortest_edge")) - .copied() - .unwrap_or(224); - let w = s - .get("width") - .or_else(|| s.get("shortest_edge")) - .copied() - .unwrap_or(224); - (h, w) - }) - } - - /// Get a scalar value from the `size` map, such as `shortest_edge` or - /// `longest_edge`. - pub fn get_size_value(&self, key: &str) -> Option { - self.size - .as_ref() - .and_then(|s| s.get(key)) - .map(|v| *v as usize) - } - - pub fn get_shortest_edge(&self) -> Option { - self.get_size_value("shortest_edge") - } - - pub fn get_longest_edge(&self) -> Option { - self.get_size_value("longest_edge") - } - - /// Get crop size. - /// - /// Returns (height, width). - pub fn get_crop_size(&self) -> Option<(u32, u32)> { - self.crop_size.as_ref().map(|s| { - let h = s.get("height").copied().unwrap_or(224); - let w = s.get("width").copied().unwrap_or(224); - (h, w) - }) - } - - /// Get the interpolation filter for resizing. - pub fn get_filter(&self) -> FilterType { - transforms::pil_to_filter(self.resampling) - } - - /// Check if normalization should be applied. - pub fn should_normalize(&self) -> bool { - self.do_normalize.unwrap_or(true) - } - - /// Check if rescaling should be applied. - pub fn should_rescale(&self) -> bool { - self.do_rescale.unwrap_or(false) - } - - /// Check if resizing should be applied. - pub fn should_resize(&self) -> bool { - self.do_resize.unwrap_or(true) - } - - /// Check if center cropping should be applied. - pub fn should_center_crop(&self) -> bool { - self.do_center_crop.unwrap_or(false) - } - - /// Get rescale factor with default. - pub fn get_rescale_factor(&self) -> f64 { - self.rescale_factor.unwrap_or(1.0 / 255.0) - } - - /// Get a typed extra field. - pub fn get_extra(&self, key: &str) -> Option { - self.extra - .get(key) - .and_then(|v| serde_json::from_value(v.clone()).ok()) - } - - // 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]; - - pub const IMAGENET_MEAN: [f64; 3] = [0.485, 0.456, 0.406]; - pub const IMAGENET_STD: [f64; 3] = [0.229, 0.224, 0.225]; - - pub const SIGLIP_MEAN: [f64; 3] = [0.5, 0.5, 0.5]; - pub const SIGLIP_STD: [f64; 3] = [0.5, 0.5, 0.5]; -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_parse_clip_config() { - let json = r#"{ - "do_center_crop": true, - "do_normalize": true, - "do_resize": true, - "image_mean": [0.48145466, 0.4578275, 0.40821073], - "image_std": [0.26862954, 0.26130258, 0.27577711], - "resample": 3, - "size": {"shortest_edge": 224} - }"#; - - let config = PreProcessorConfig::from_json(json).unwrap(); - - assert!(config.should_normalize()); - assert!(config.should_center_crop()); - assert!(config.should_resize()); - assert_eq!(config.resampling, Some(3)); - - let (h, w) = config.get_target_size().unwrap(); - assert_eq!(h, 224); - assert_eq!(w, 224); - - let mean = config.get_image_mean(); - assert!((mean[0] - 0.48145466).abs() < 1e-6); - } - - #[test] - fn test_parse_qwen_vl_config() { - let json = r#"{ - "do_normalize": true, - "do_rescale": true, - "do_resize": true, - "image_mean": [0.48145466, 0.4578275, 0.40821073], - "image_std": [0.26862954, 0.26130258, 0.27577711], - "min_pixels": 200704, - "max_pixels": 1003520, - "patch_size": 14, - "merge_size": 2, - "temporal_patch_size": 2, - "rescale_factor": 0.00392156862745098 - }"#; - - let config = PreProcessorConfig::from_json(json).unwrap(); - - assert_eq!(config.min_pixels, Some(200704)); - assert_eq!(config.max_pixels, Some(1003520)); - assert_eq!(config.get_patch_size(0), 14); - assert_eq!(config.merge_size, Some(2)); - assert!((config.get_rescale_factor() - 1.0 / 255.0).abs() < 1e-10); - } - - #[test] - fn test_parse_size_formats() { - // Height/width format - let json1 = r#"{"size": {"height": 336, "width": 336}}"#; - let config1 = PreProcessorConfig::from_json(json1).unwrap(); - assert_eq!(config1.get_target_size(), Some((336, 336))); - - // Shortest edge format - let json2 = r#"{"size": {"shortest_edge": 224}}"#; - let config2 = PreProcessorConfig::from_json(json2).unwrap(); - assert_eq!(config2.get_target_size(), Some((224, 224))); - } - - #[test] - fn test_defaults() { - let config = PreProcessorConfig::default(); - - // Should use CLIP defaults when not specified - let mean = config.get_image_mean(); - assert!((mean[0] - PreProcessorConfig::CLIP_MEAN[0]).abs() < 1e-6); - - // Default behaviors - assert!(config.should_normalize()); // true by default - assert!(!config.should_rescale()); // false by default - assert!(config.should_resize()); // true by default - assert!(!config.should_center_crop()); // false by default - } - - #[test] - fn test_image_only_processor_type_detection() { - for (processor_type, expected) in [ - (Some("Qwen3VLImageProcessor"), true), - (Some("qWeN3vLiMaGePrOcEsSoR"), true), - (Some("Qwen3VLVideoProcessor"), false), - (Some("Qwen3VLImageProcessorVideo"), false), - (None, false), - ] { - let config = PreProcessorConfig { - image_processor_type: processor_type.map(str::to_owned), - ..Default::default() - }; - assert_eq!(config.is_image_only_processor_type(), expected); - } - } - - #[test] - fn test_filter_conversion() { - let json = r#"{"resampling": 3}"#; - let config = PreProcessorConfig::from_json(json).unwrap(); - assert!(matches!(config.get_filter(), FilterType::CatmullRom)); - } - - #[test] - fn test_extra_fields() { - let json = r#"{ - "custom_field": 42, - "nested": {"foo": "bar"} - }"#; - - let config = PreProcessorConfig::from_json(json).unwrap(); - - let custom: Option = config.get_extra("custom_field"); - assert_eq!(custom, Some(42)); - - let nested: Option> = config.get_extra("nested"); - assert_eq!( - nested.as_ref().unwrap().get("foo"), - Some(&"bar".to_string()) - ); - } - - #[test] - fn test_parse_kimi_nested_media_proc_cfg() { - let json = r#"{ - "auto_map": { - "AutoProcessor": "kimi_k25_processor.KimiK25Processor" - }, - "media_proc_cfg": { - "in_patch_limit": 16384, - "patch_size": 14, - "image_mean": [0.5, 0.5, 0.5], - "image_std": [0.5, 0.5, 0.5], - "merge_kernel_size": 2, - "patch_limit_on_one_side": 512 - } - }"#; - - let config = PreProcessorConfig::from_json(json).unwrap(); - - // image_mean/std should be extracted from media_proc_cfg - let mean = config.get_image_mean(); - assert!((mean[0] - 0.5).abs() < 1e-6); - assert!((mean[1] - 0.5).abs() < 1e-6); - assert!((mean[2] - 0.5).abs() < 1e-6); - - let std = config.get_image_std(); - assert!((std[0] - 0.5).abs() < 1e-6); - - assert_eq!(config.get_patch_size(0), 14); - assert_eq!(config.merge_size, Some(2)); - } -} diff --git a/crates/multimodal/src/vision/processor.rs b/crates/multimodal/src/vision/processor.rs deleted file mode 100644 index de025dddb..000000000 --- a/crates/multimodal/src/vision/processor.rs +++ /dev/null @@ -1,422 +0,0 @@ -//! Vision processor trait and registry. -//! -//! Shared encoder output types live in [`crate::encoder_inputs`] and are re-exported -//! here for compatibility. - -use std::collections::HashMap; - -use image::DynamicImage; - -use super::{preprocessor_config::PreProcessorConfig, transforms::TransformError}; -pub use crate::encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}; -use crate::types::RgbFrameRef; - -/// Helper to extract a dimension from encoder_input given an ndim-dependent axis index. -/// Returns `Err` if the ndim is not 4 or 5. -fn dim_for_ndim( - ndim: usize, - axis_4d: usize, - axis_5d: usize, - shape: &[usize], -) -> Result { - match ndim { - 4 => Ok(shape[axis_4d]), - 5 => Ok(shape[axis_5d]), - _ => Err(TransformError::InvalidShape { - expected: format!("4D or 5D encoder_input tensor, got {ndim}D"), - actual: shape.to_vec(), - }), - } -} - -impl PreprocessedEncoderInputs { - /// Get the number of channels. - /// - /// For 4D tensors [B, C, H, W], returns shape[1]. - /// For 5D tensors [B, N, C, H, W] (Phi3-Vision), returns shape[2]. - /// - /// # Errors - /// Returns `TransformError::InvalidShape` if encoder_input is not 4D or 5D. - pub fn channels(&self) -> Result { - dim_for_ndim(self.encoder_input.ndim(), 1, 2, self.encoder_input.shape()) - } - - /// Get the height of processed images. - /// - /// For 4D tensors [B, C, H, W], returns shape[2]. - /// For 5D tensors [B, N, C, H, W] (Phi3-Vision), returns shape[3]. - /// - /// # Errors - /// Returns `TransformError::InvalidShape` if encoder_input is not 4D or 5D. - pub fn height(&self) -> Result { - dim_for_ndim(self.encoder_input.ndim(), 2, 3, self.encoder_input.shape()) - } - - /// Get the width of processed images. - /// - /// For 4D tensors [B, C, H, W], returns shape[3]. - /// For 5D tensors [B, N, C, H, W] (Phi3-Vision), returns shape[4]. - /// - /// # Errors - /// Returns `TransformError::InvalidShape` if encoder_input is not 4D or 5D. - pub fn width(&self) -> Result { - dim_for_ndim(self.encoder_input.ndim(), 3, 4, self.encoder_input.shape()) - } -} - -/// Trait for model-specific vision preprocessors. -/// -/// Each vision model (LLaVA, Qwen-VL, Phi3-Vision, etc.) implements this trait -/// to provide the correct preprocessing pipeline. -pub trait VisionPreProcessor: Send + Sync { - /// Default normalization mean for this model family. - fn default_mean(&self) -> [f64; 3]; - - /// Default normalization std for this model family. - fn default_std(&self) -> [f64; 3]; - - /// Preprocess a batch of images. - /// - /// # Arguments - /// * `images` - Input images to preprocess - /// * `config` - Preprocessor configuration from HuggingFace - /// - /// # Returns - /// Preprocessed encoder inputs ready for the model, or an error. - fn preprocess( - &self, - images: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result; - - /// Preprocess one decoded video clip represented as sampled frames. - /// - /// Implementations that support video should emit the same primary - /// `encoder_input` tensor shape used by the image path, plus video-specific - /// model metadata such as `video_grid_thw`. - fn preprocess_video( - &self, - _frames: &[DynamicImage], - _config: &PreProcessorConfig, - ) -> Result { - Err(TransformError::ShapeError(format!( - "{} does not support video preprocessing", - self.model_name() - ))) - } - - /// Preprocess one decoded video clip represented as borrowed RGB frame - /// buffers. Implementations can override this to avoid materializing - /// `DynamicImage` objects after media decode. - fn preprocess_video_rgb( - &self, - _frames: &[RgbFrameRef<'_>], - _config: &PreProcessorConfig, - ) -> Result { - Err(TransformError::ShapeError(format!( - "{} does not support RGB video preprocessing", - self.model_name() - ))) - } - - /// Calculate the number of vision tokens for a given image size. - /// - /// This is used to determine how many placeholder tokens to insert - /// in the text input before the image has been fully processed. - /// - /// # Arguments - /// * `width` - Image width after preprocessing - /// * `height` - Image height after preprocessing - /// * `config` - Preprocessor configuration - fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize; - - /// Get the model family name for identification. - fn model_name(&self) -> &'static str; - - /// Get the expected image size after preprocessing. - /// - /// Some models have fixed sizes, others are dynamic. - fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> { - config.get_target_size() - } -} - -/// Registry of available vision processors. -pub struct VisionProcessorRegistry { - processors: HashMap>, -} - -impl VisionProcessorRegistry { - /// Create a new empty registry. - pub fn new() -> Self { - Self { - processors: HashMap::new(), - } - } - - /// Register a processor for a model pattern. - pub fn register(&mut self, pattern: impl Into, processor: Box) { - self.processors.insert(pattern.into(), processor); - } - - /// Find a processor for the given model ID, falling back to model_type. - /// - /// Matches by substring containment (case-insensitive). - pub fn find( - &self, - model_id: &str, - model_type: Option<&str>, - ) -> Option<&dyn VisionPreProcessor> { - self.find_in_candidate(model_id) - .or_else(|| model_type.and_then(|mt| self.find_in_candidate(mt))) - } - - fn find_in_candidate(&self, candidate: &str) -> Option<&dyn VisionPreProcessor> { - let candidate = candidate.to_lowercase(); - for (pattern, processor) in &self.processors { - if candidate.contains(&pattern.to_lowercase()) { - return Some(processor.as_ref()); - } - } - None - } - - /// Get list of supported model patterns. - pub fn supported_patterns(&self) -> Vec<&str> { - self.processors.keys().map(|s| s.as_str()).collect() - } -} - -impl Default for VisionProcessorRegistry { - fn default() -> Self { - Self::new() - } -} - -impl VisionProcessorRegistry { - /// Create a registry with all built-in processors registered. - /// - /// Currently registers: - /// - `llava-next` -> LlavaNextProcessor - /// - `llava-1.5` / `llava-v1.5` -> LlavaProcessor - /// - `qwen2-vl` -> Qwen2VLProcessor - /// - `qwen2.5-vl` -> Qwen2VLProcessor (same preprocessing as Qwen2-VL) - /// - `qwen3-vl` -> Qwen3VLProcessor (patch_size=16, [0.5,0.5,0.5] normalization) - /// - `qwen3.5` / `qwen3_5` -> Qwen3VLProcessor (Qwen3.5 reuses Qwen3-VL preprocessing) - /// - `phi-3-vision` -> Phi3VisionProcessor (HD transform with 336x336 tiles) - pub fn with_defaults() -> Self { - let mut registry = Self::new(); - - // LLaVA-NeXT (v1.6+, anyres multi-crop) - registry.register( - "llava-next", - Box::new(super::processors::LlavaNextProcessor::new()), - ); - registry.register( - "llava_next", - Box::new(super::processors::LlavaNextProcessor::new()), - ); - registry.register( - "llava-v1.6", - Box::new(super::processors::LlavaNextProcessor::new()), - ); - - // Standard LLaVA (v1.5, single-patch). - // Use specific patterns so they don't accidentally match LLaVA-NeXT - // model IDs like "llava-v1.6-*". - registry.register( - "llava-1.5", - Box::new(super::processors::LlavaProcessor::new()), - ); - registry.register( - "llava-v1.5", - Box::new(super::processors::LlavaProcessor::new()), - ); - - // Register Qwen3-VL first (more specific pattern - must match before qwen2) - registry.register( - "qwen3-vl", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - registry.register( - "qwen3_vl", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - - // Qwen3-Omni uses the same patchification and normalization contract - // as Qwen3-VL for its image and video towers. - registry.register( - "qwen3-omni", - Box::new(super::processors::Qwen3OmniVisionProcessor::new()), - ); - registry.register( - "qwen3_omni", - Box::new(super::processors::Qwen3OmniVisionProcessor::new()), - ); - - // Qwen3.5 family (and Qwen3.6: same arch) reuses Qwen3-VL preprocessing. - registry.register( - "qwen3.5", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - registry.register( - "qwen3_5", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - registry.register( - "qwen3.6", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - registry.register( - "qwen3_6", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - - // Register Qwen2-VL (matches Qwen/Qwen2-VL-*, etc.) - registry.register( - "qwen2-vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); - registry.register( - "qwen2_vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); - - // Register Qwen2.5-VL (uses identical preprocessing to Qwen2-VL) - registry.register( - "qwen2.5-vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); - registry.register( - "qwen2_5-vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); - registry.register( - "qwen2_5_vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); - - // Register Phi3-Vision - registry.register( - "phi-3-vision", - Box::new(super::processors::Phi3VisionProcessor::new()), - ); - registry.register( - "phi3-vision", - Box::new(super::processors::Phi3VisionProcessor::new()), - ); - registry.register( - "phi3_v", - Box::new(super::processors::Phi3VisionProcessor::new()), - ); - - // Register LLaMA 4 Vision - registry.register( - "llama-4", - Box::new(super::processors::Llama4VisionProcessor::new()), - ); - registry.register( - "llama4", - Box::new(super::processors::Llama4VisionProcessor::new()), - ); - - // Register Kimi-K2.5 Vision - registry.register( - "kimi-k2", - Box::new(super::processors::KimiK25Processor::new()), - ); - registry.register( - "kimi_k2", - Box::new(super::processors::KimiK25Processor::new()), - ); - - registry - } -} - -#[cfg(test)] -mod tests { - use ndarray::Array4; - - use super::*; - use crate::vision::processors::LlavaProcessor; - - #[test] - fn test_preprocessed_encoder_inputs_geometry_accessors() { - let encoder_input = Array4::::zeros((2, 3, 336, 336)); - let inputs = PreprocessedEncoderInputs::new( - encoder_input, - vec![576, 576], - vec![(640, 480), (800, 600)], - ); - - assert_eq!(inputs.channels().unwrap(), 3); - assert_eq!(inputs.height().unwrap(), 336); - assert_eq!(inputs.width().unwrap(), 336); - } - - #[test] - fn test_registry_with_defaults() { - let registry = VisionProcessorRegistry::with_defaults(); - - // Should find LLaVA processor - assert!(registry.find("llava-hf/llava-1.5-7b-hf", None).is_some()); - assert!(registry.find("liuhaotian/llava-v1.5-7b", None).is_some()); - - // Should find LLaVA-NeXT processor - assert!(registry - .find("llava-hf/llava-v1.6-mistral-7b-hf", None) - .is_some()); - assert!(registry - .find("lmms-lab/llava-next-interleave-qwen-7b", None) - .is_some()); - - // Get the processor and check model name - let processor = registry.find("llava-hf/llava-1.5-7b-hf", None).unwrap(); - assert_eq!(processor.model_name(), "llava"); - } - - #[test] - fn test_registry_find() { - let mut registry = VisionProcessorRegistry::new(); - - // Create a mock processor using LlavaProcessor - registry.register("test-model", Box::new(LlavaProcessor::new())); - - assert!(registry.find("test-model-7b", None).is_some()); - assert!(registry.find("TEST-MODEL", None).is_some()); - assert!(registry.find("other-model", None).is_none()); - } - - #[test] - fn test_registry_find_falls_back_to_model_type() { - let registry = VisionProcessorRegistry::with_defaults(); - - assert!(registry.find("custom-model", None).is_none()); - - let processor = registry - .find("custom-model", Some("qwen3_vl")) - .expect("qwen3 processor by model_type"); - assert_eq!(processor.model_name(), "qwen3-vl"); - } - - #[test] - fn test_registry_find_preserves_fast_path() { - let registry = VisionProcessorRegistry::with_defaults(); - - let processor = registry - .find("Qwen3-VL-30B-A3B-Instruct", Some("qwen2_vl")) - .expect("qwen3 processor by model_id"); - assert_eq!(processor.model_name(), "qwen3-vl"); - } - - #[test] - fn test_registry_find_phi3_model_type_fallback() { - let registry = VisionProcessorRegistry::with_defaults(); - - let processor = registry - .find("custom-model", Some("phi3_v")) - .expect("phi3 processor by model_type"); - assert_eq!(processor.model_name(), "phi3-vision"); - } -} diff --git a/crates/multimodal/src/vision/processors/kimi_k25.rs b/crates/multimodal/src/vision/processors/kimi_k25.rs deleted file mode 100644 index bcdc58ad4..000000000 --- a/crates/multimodal/src/vision/processors/kimi_k25.rs +++ /dev/null @@ -1,572 +0,0 @@ -//! Kimi-K2.5 (MoonViT) image processor. -//! -//! Matches the HuggingFace `KimiK25VisionProcessor` preprocessing pipeline: -//! -//! 1. Compute scale to fit within patch limits (never upscale) -//! 2. Resize with BICUBIC interpolation -//! 3. Zero-pad to make dimensions divisible by factor (patch_size * merge_size) -//! 4. Normalize with [0.5, 0.5, 0.5] mean/std -//! 5. Extract patches as [N, C, patch_size, patch_size] -//! -//! Kimi resizes then zero-pads to make dimensions divisible by the alignment -//! factor (patch_size * merge_size). The model was trained with zero-padded -//! images, so using direct resize-to-aligned would degrade image quality. - -use image::{DynamicImage, GenericImageView}; -use ndarray::Array3; - -use crate::vision::{ - preprocessor_config::PreProcessorConfig, - processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, - scratch, - transforms::{self, TransformError}, -}; - -pub const KIMI_K25_MEAN: [f64; 3] = [0.5, 0.5, 0.5]; -pub const KIMI_K25_STD: [f64; 3] = [0.5, 0.5, 0.5]; - -pub const DEFAULT_PATCH_SIZE: usize = 14; -pub const DEFAULT_MERGE_SIZE: usize = 2; -/// Maximum total patches before merge (from preprocessor_config.json in_patch_limit) -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, -} - -impl Default for KimiK25Processor { - fn default() -> Self { - Self::new() - } -} - -impl KimiK25Processor { - pub fn new() -> Self { - Self { - patch_size: DEFAULT_PATCH_SIZE, - merge_size: DEFAULT_MERGE_SIZE, - in_patch_limit: DEFAULT_IN_PATCH_LIMIT, - patch_limit_on_one_side: DEFAULT_PATCH_LIMIT_ON_ONE_SIDE, - } - } - - 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), - } - } - - pub fn patch_size(&self) -> usize { - self.patch_size - } - - pub fn merge_size(&self) -> usize { - self.merge_size - } - - #[inline] - fn factor(&self) -> usize { - self.patch_size * self.merge_size - } - - /// Compute resize dimensions and padding, matching HF `navit_resize_image`. - /// - /// Never upscales (scale capped at 1.0). Pads with zeros to align to factor. - fn compute_resize_config(&self, width: usize, height: usize) -> ResizeConfig { - let ps = self.patch_size; - let patches_w = (width / ps).max(1) as f64; - let patches_h = (height / ps).max(1) as f64; - - let s1 = (self.in_patch_limit as f64 / (patches_w * patches_h)).sqrt(); - let s2 = (self.patch_limit_on_one_side * ps) as f64 / width as f64; - let s3 = (self.patch_limit_on_one_side * ps) as f64 / height as f64; - let scale = f64::min(1.0, f64::min(s1, f64::min(s2, s3))); - - let new_w = ((width as f64 * scale) as usize).max(1); - let new_h = ((height as f64 * scale) as usize).max(1); - let new_w = new_w.min(self.patch_limit_on_one_side * ps); - let new_h = new_h.min(self.patch_limit_on_one_side * ps); - - let factor = self.factor(); - let pad_width = (factor - new_w % factor) % factor; - let pad_height = (factor - new_h % factor) % factor; - - let token_height = (new_h + pad_height) / factor; - let token_width = (new_w + pad_width) / factor; - let num_tokens = token_height * token_width; - - ResizeConfig { - new_width: new_w, - new_height: new_h, - pad_width, - pad_height, - num_tokens, - } - } - - /// Fused resize + zero-pad + normalize into a single [C, H_padded, W_padded] tensor. - /// - /// Avoids intermediate allocations by: - /// 1. Allocating the final padded canvas directly - /// 2. Pre-filling with normalized black (bias value) - /// 3. Deinterleaving + normalizing the image region in one pass - fn resize_pad_and_normalize( - image: &DynamicImage, - cfg: &ResizeConfig, - mean: &[f64; 3], - std: &[f64; 3], - ) -> Array3 { - let canvas_h = cfg.new_height + cfg.pad_height; - let canvas_w = cfg.new_width + cfg.pad_width; - - // Resize using SIMD-accelerated BICUBIC (fast_image_resize) - let resized = transforms::resize( - image, - cfg.new_width as u32, - cfg.new_height as u32, - image::imageops::FilterType::CatmullRom, - ); - - let (img_w, img_h, raw) = transforms::rgb_bytes(&resized); - let canvas_pixels = canvas_h * canvas_w; - - // Precompute fused scale/bias: pixel/255 → normalized - // output[c][i] = raw[i*3+c] / 255.0 * (1/std[c]) + (-mean[c]/std[c]) - let scale: [f32; 3] = std::array::from_fn(|c| 1.0 / (255.0 * std[c] as f32)); - let bias: [f32; 3] = std::array::from_fn(|c| -(mean[c] as f32) / (std[c] as f32)); - - // Pooled: this per-image CHW buffer (tens of MB) is recycled by the - // caller after patch extraction, keeping its pages mapped and hot. - let mut data = scratch::take_f32(3 * canvas_pixels); - let (r_plane, rest) = data.split_at_mut(canvas_pixels); - let (g_plane, b_plane) = rest.split_at_mut(canvas_pixels); - - // Pre-fill with normalized black: (0/255 - mean) / std = bias - r_plane.fill(bias[0]); - g_plane.fill(bias[1]); - b_plane.fill(bias[2]); - - // Overwrite image region row-by-row using vectorized deinterleave - let rw = img_w.min(canvas_w); - let rh = img_h.min(canvas_h); - for y in 0..rh { - let src_row = &raw[y * img_w * 3..y * img_w * 3 + rw * 3]; - let dst_offset = y * canvas_w; - transforms::deinterleave_rgb_to_planes( - src_row, - &mut r_plane[dst_offset..dst_offset + rw], - &mut g_plane[dst_offset..dst_offset + rw], - &mut b_plane[dst_offset..dst_offset + rw], - scale, - bias, - ); - } - - #[expect( - clippy::expect_used, - reason = "data has exactly 3*canvas_h*canvas_w elements by construction" - )] - Array3::from_shape_vec((3, canvas_h, canvas_w), data) - .expect("shape matches pre-allocated buffer") - } - - /// Extract [C, patch_size, patch_size] patches from a contiguous [C, H, W] tensor. - /// - /// Uses row-based `copy_from_slice` instead of per-element indexing so the - /// compiler can auto-vectorize the inner copy. - /// Append this image's patches directly into `out` (no per-image intermediate - /// Vec): `out` is the pooled batch buffer pre-sized for the whole request. - fn extract_patches_into(tensor: &Array3, patch_size: usize, out: &mut Vec) { - let channels = tensor.shape()[0]; - let height = tensor.shape()[1]; - let width = tensor.shape()[2]; - - let grid_h = height / patch_size; - let grid_w = width / patch_size; - - // Get contiguous slice for direct row addressing - let flat = tensor.as_standard_layout(); - #[expect( - clippy::expect_used, - reason = "as_standard_layout guarantees contiguous C-order memory" - )] - let data = flat - .as_slice() - .expect("as_standard_layout guarantees contiguous memory"); - - for gh in 0..grid_h { - for gw in 0..grid_w { - let h_start = gh * patch_size; - let w_start = gw * patch_size; - for c in 0..channels { - let plane_offset = c * height * width; - for ph in 0..patch_size { - let row_start = plane_offset + (h_start + ph) * width + w_start; - out.extend_from_slice(&data[row_start..row_start + patch_size]); - } - } - } - } - } -} - -impl VisionPreProcessor for KimiK25Processor { - fn default_mean(&self) -> [f64; 3] { - KIMI_K25_MEAN - } - - fn default_std(&self) -> [f64; 3] { - KIMI_K25_STD - } - - fn preprocess( - &self, - 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) - } - - 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 { - "kimi-k2.5" - } - - fn get_processed_size(&self, _config: &PreProcessorConfig) -> Option<(u32, u32)> { - None - } -} - -#[cfg(test)] -mod tests { - use image::{Rgb, RgbImage}; - - use super::*; - use crate::vision::preprocessor_config::PatchSize; - - fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { - DynamicImage::from(RgbImage::from_pixel(width, height, color)) - } - - #[test] - fn test_defaults() { - let p = KimiK25Processor::new(); - assert_eq!(p.patch_size(), 14); - assert_eq!(p.merge_size(), 2); - assert_eq!(p.factor(), 28); - } - - #[test] - fn test_mean_std() { - let p = KimiK25Processor::new(); - assert_eq!(p.default_mean(), KIMI_K25_MEAN); - assert_eq!(p.default_std(), KIMI_K25_STD); - } - - #[test] - fn test_model_name() { - assert_eq!(KimiK25Processor::new().model_name(), "kimi-k2.5"); - } - - #[test] - fn test_resize_config_no_upscale() { - let p = KimiK25Processor::new(); - // Small image should NOT be upscaled (scale capped at 1.0) - let cfg = p.compute_resize_config(100, 100); - assert!(cfg.new_width <= 100); - assert!(cfg.new_height <= 100); - // Padded dimensions must be factor-aligned - assert_eq!((cfg.new_height + cfg.pad_height) % 28, 0); - assert_eq!((cfg.new_width + cfg.pad_width) % 28, 0); - } - - #[test] - fn test_resize_config_large_image_downscaled() { - let p = KimiK25Processor::new(); - // Large image should be downscaled - let cfg = p.compute_resize_config(4000, 3000); - // Resized dimensions should be smaller than original - assert!(cfg.new_width < 4000); - assert!(cfg.new_height < 3000); - // Per-side patch limit must be respected (HF assertion) - let padded_h = cfg.new_height + cfg.pad_height; - let padded_w = cfg.new_width + cfg.pad_width; - assert!(padded_h / 14 <= DEFAULT_PATCH_LIMIT_ON_ONE_SIDE * 2); - assert!(padded_w / 14 <= DEFAULT_PATCH_LIMIT_ON_ONE_SIDE * 2); - } - - #[test] - fn test_resize_config_matches_hf_reference() { - let p = KimiK25Processor::new(); - // 600x400 image: scale=1.0 (small enough), resize to 600x400, - // pad to (600+4=) → let's compute: - // factor=28, 400 % 28 = 400 - 14*28 = 400-392 = 8, pad_h = 28-8 = 20 - // 600 % 28 = 600 - 21*28 = 600-588 = 12, pad_w = 28-12 = 16 - let cfg = p.compute_resize_config(600, 400); - assert_eq!(cfg.new_width, 600); - assert_eq!(cfg.new_height, 400); - assert_eq!(cfg.pad_height, 20); - assert_eq!(cfg.pad_width, 16); - // Padded: 420 x 616, grid: 30 x 44, tokens: (30*44)/(2*2) = 330 - assert_eq!(cfg.num_tokens, 330); - } - - #[test] - fn test_preprocess_4d_output() { - let p = KimiK25Processor::new(); - let config = PreProcessorConfig { - do_normalize: Some(true), - image_mean: Some(KIMI_K25_MEAN.to_vec()), - image_std: Some(KIMI_K25_STD.to_vec()), - ..Default::default() - }; - - let image = create_test_image(600, 400, Rgb([128, 128, 128])); - let result = p.preprocess(&[image], &config).unwrap(); - - // 4D output: [total_patches, 3, 14, 14] - assert_eq!(result.encoder_input.ndim(), 4); - assert_eq!(result.encoder_input.shape()[1], 3); - assert_eq!(result.encoder_input.shape()[2], 14); - assert_eq!(result.encoder_input.shape()[3], 14); - - assert!(result.model_specific.contains_key("grid_thws")); - assert!(result.model_specific.contains_key("patches_per_image")); - assert!(result.feature_token_counts[0] > 0); - } - - #[test] - fn test_preprocess_multiple_images() { - let p = KimiK25Processor::new(); - let config = PreProcessorConfig::default(); - let images = vec![ - create_test_image(600, 400, Rgb([100, 100, 100])), - create_test_image(400, 600, Rgb([150, 150, 150])), - ]; - - let result = p.preprocess(&images, &config).unwrap(); - - assert_eq!(result.item_sizes.len(), 2); - assert_eq!(result.feature_token_counts.len(), 2); - assert_eq!(result.encoder_input.ndim(), 4); - assert_eq!(result.encoder_input.shape()[1], 3); - - if let Some(ModelSpecificValue::IntTensor { data, shape }) = - result.model_specific.get("grid_thws") - { - assert_eq!(shape, &[2, 3]); - assert_eq!(data.len(), 6); - } else { - panic!("Expected grid_thws to be IntTensor"); - } - - if let Some(ModelSpecificValue::IntTensor { data, .. }) = - result.model_specific.get("patches_per_image") - { - let total: i64 = data.iter().sum(); - assert_eq!(total as usize, result.encoder_input.shape()[0]); - } - } - - #[test] - fn test_calculate_num_tokens() { - let p = KimiK25Processor::new(); - let config = PreProcessorConfig::default(); - let tokens = p.calculate_num_tokens(600, 400, &config); - assert_eq!(tokens, 330); - } - - #[test] - fn test_from_preprocessor_config() { - let config = PreProcessorConfig { - patch_size: Some(PatchSize { - height: Some(14), - width: Some(14), - }), - merge_size: Some(2), - ..Default::default() - }; - let p = KimiK25Processor::from_preprocessor_config(&config); - assert_eq!(p.patch_size(), 14); - assert_eq!(p.merge_size(), 2); - } - - #[test] - fn test_zero_padding_applied() { - let p = KimiK25Processor::new(); - let config = PreProcessorConfig { - image_mean: Some(KIMI_K25_MEAN.to_vec()), - image_std: Some(KIMI_K25_STD.to_vec()), - ..Default::default() - }; - - // 100x100 white image — after normalization: (255/255 - 0.5) / 0.5 = 1.0 - // Padded region: (0/255 - 0.5) / 0.5 = -1.0 - let image = create_test_image(100, 100, Rgb([255, 255, 255])); - let result = p.preprocess(&[image], &config).unwrap(); - - let flat = result.encoder_input_flat(); - // Padded region should be normalized black (-1.0) - let has_neg_ones = flat.iter().any(|&v| (v - (-1.0)).abs() < 1e-6); - assert!( - has_neg_ones, - "Expected normalized-black padding (-1.0) in output" - ); - - // Image region should be normalized white (1.0) - let has_ones = flat.iter().any(|&v| (v - 1.0).abs() < 1e-6); - assert!( - has_ones, - "Expected normalized-white image values (1.0) in output" - ); - } - - #[test] - fn test_preprocess_tiny_image() { - // 1x1 image should not panic — padded to 28x28 - let p = KimiK25Processor::new(); - let config = PreProcessorConfig { - image_mean: Some(KIMI_K25_MEAN.to_vec()), - image_std: Some(KIMI_K25_STD.to_vec()), - ..Default::default() - }; - let image = create_test_image(1, 1, Rgb([128, 128, 128])); - let result = p.preprocess(&[image], &config).unwrap(); - assert_eq!(result.encoder_input.ndim(), 4); - assert!(result.encoder_input.shape()[0] > 0); - assert!(result.feature_token_counts[0] > 0); - } - - #[test] - fn test_preprocess_empty_batch_returns_error() { - let p = KimiK25Processor::new(); - let config = PreProcessorConfig::default(); - let result = p.preprocess(&[], &config); - assert!(result.is_err()); - } - - #[test] - fn test_from_preprocessor_config_reads_limits() { - let config = PreProcessorConfig { - patch_size: Some(PatchSize { - height: Some(14), - width: Some(14), - }), - merge_size: Some(2), - extra: [ - ("in_patch_limit".to_string(), serde_json::json!(8192)), - ( - "patch_limit_on_one_side".to_string(), - serde_json::json!(256), - ), - ] - .into_iter() - .collect(), - ..Default::default() - }; - let p = KimiK25Processor::from_preprocessor_config(&config); - assert_eq!(p.in_patch_limit, 8192); - assert_eq!(p.patch_limit_on_one_side, 256); - } -} diff --git a/crates/multimodal/src/vision/processors/llama4_vision.rs b/crates/multimodal/src/vision/processors/llama4_vision.rs deleted file mode 100644 index 033adc128..000000000 --- a/crates/multimodal/src/vision/processors/llama4_vision.rs +++ /dev/null @@ -1,727 +0,0 @@ -//! LLaMA 4 Vision image processor. -//! -//! This module implements the LLaMA 4 Vision (Llama-4-Scout, Llama-4-Maverick) image preprocessing -//! pipeline with tile-based processing similar to other dynamic resolution models. -//! -//! # Key Features -//! -//! | Feature | Value | -//! |---------|-------| -//! | Tile size | 336x336 | -//! | Default max_patches | 16 | -//! | Normalization | [0.5, 0.5, 0.5] mean/std | -//! | Interpolation | Bilinear | -//! | Global tile | Added when num_tiles > 1 | -//! -//! # Processing Pipeline -//! -//! 1. **Find supported resolutions**: Calculate valid tile configurations -//! 2. **Get best fit**: Find optimal resolution without distortion -//! 3. **Resize**: Scale to target resolution maintaining aspect ratio -//! 4. **Pad**: Add black padding (0) to reach target dimensions -//! 5. **Normalize**: Apply [0.5, 0.5, 0.5] mean/std normalization -//! 6. **Tile**: Split into (num_tiles_h * num_tiles_w, 3, 336, 336) tiles -//! 7. **Global tile**: If multiple tiles, add global view at the end -//! -//! # Token Count -//! -//! For LLaMA 4, tokens = num_tiles * (tile_size / patch_size)² -//! where patch_size is typically 14, giving 576 tokens per tile. - -use std::collections::HashSet; - -use image::{imageops::FilterType, DynamicImage, GenericImageView}; -use ndarray::{s, Array3, Array4}; - -use crate::vision::{ - preprocessor_config::PreProcessorConfig, - processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::{self, TransformError}, -}; - -/// Default normalization mean for LLaMA 4 Vision. -pub const LLAMA4_MEAN: [f64; 3] = [0.5, 0.5, 0.5]; - -/// Default normalization std for LLaMA 4 Vision. -pub const LLAMA4_STD: [f64; 3] = [0.5, 0.5, 0.5]; - -/// Default tile size for LLaMA 4 Vision. -pub const TILE_SIZE: u32 = 336; - -/// Default maximum number of patches/tiles. -pub const DEFAULT_MAX_PATCHES: usize = 16; - -/// Patch size used in vision encoder. -pub const PATCH_SIZE: usize = 14; - -/// LLaMA 4 Vision image processor. -/// -/// Implements tile-based processing with dynamic resolution selection. -#[derive(Debug, Clone)] -pub struct Llama4VisionProcessor { - /// Tile size (both height and width). - tile_size: u32, - /// Maximum number of tiles/patches. - max_patches: usize, - /// Whether to resize to max canvas (upscale aggressively). - resize_to_max_canvas: bool, - /// Normalization mean. - mean: [f64; 3], - /// Normalization std. - std: [f64; 3], -} - -impl Default for Llama4VisionProcessor { - fn default() -> Self { - Self::new() - } -} - -impl Llama4VisionProcessor { - /// Create a new LLaMA 4 Vision processor with default settings. - pub fn new() -> Self { - Self { - tile_size: TILE_SIZE, - max_patches: DEFAULT_MAX_PATCHES, - resize_to_max_canvas: false, - mean: LLAMA4_MEAN, - std: LLAMA4_STD, - } - } - - /// Create a processor with custom max_patches setting. - pub fn with_max_patches(max_patches: usize) -> Self { - Self { - tile_size: TILE_SIZE, - max_patches, - resize_to_max_canvas: false, - mean: LLAMA4_MEAN, - std: LLAMA4_STD, - } - } - - /// Create a processor from preprocessor config. - pub fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { - Self { - tile_size: config - .size - .as_ref() - .and_then(|s| s.get("height").copied()) - .unwrap_or(TILE_SIZE), - max_patches: config.max_image_tiles.unwrap_or(DEFAULT_MAX_PATCHES), - resize_to_max_canvas: false, - mean: config - .image_mean - .as_ref() - .map(|v| [v[0], v[1], v[2]]) - .unwrap_or(LLAMA4_MEAN), - std: config - .image_std - .as_ref() - .map(|v| [v[0], v[1], v[2]]) - .unwrap_or(LLAMA4_STD), - } - } - - /// Get the tile size. - pub fn tile_size(&self) -> u32 { - self.tile_size - } - - /// Get the max patches setting. - pub fn max_patches(&self) -> usize { - self.max_patches - } - - /// Get all factors of a number. - fn get_factors(n: usize) -> HashSet { - let mut factors = HashSet::new(); - for i in 1..=(n as f64).sqrt() as usize { - if n.is_multiple_of(i) { - factors.insert(i); - factors.insert(n / i); - } - } - factors - } - - /// Find all supported resolutions for the given max_patches. - /// - /// Returns list of (height, width) in pixels. - fn find_supported_resolutions(&self) -> Vec<(u32, u32)> { - let mut resolutions = Vec::new(); - let tile = self.tile_size; - - // For each possible number of chunks from max_patches down to 1 - for chunk_size in (1..=self.max_patches).rev() { - let factors = Self::get_factors(chunk_size); - for &factor in &factors { - let h_tiles = factor; - let w_tiles = chunk_size / factor; - resolutions.push((h_tiles as u32 * tile, w_tiles as u32 * tile)); - } - } - - resolutions - } - - /// Get the maximum resolution without distortion. - /// - /// Given an image size and target size, compute the largest size - /// that fits within target while maintaining aspect ratio. - fn get_max_res_without_distortion( - image_size: (u32, u32), - target_size: (u32, u32), - ) -> (u32, u32) { - let (orig_h, orig_w) = image_size; - let (target_h, target_w) = target_size; - - let scale_w = target_w as f64 / orig_w as f64; - let scale_h = target_h as f64 / orig_h as f64; - - if scale_w < scale_h { - let new_w = target_w; - let new_h = (orig_h as f64 * scale_w).floor() as u32; - (new_h.min(target_h), new_w) - } else { - let new_h = target_h; - let new_w = (orig_w as f64 * scale_h).floor() as u32; - (new_h, new_w.min(target_w)) - } - } - - /// Find the best fitting resolution from supported resolutions. - /// - /// Selects resolution that: - /// - Minimizes upscaling if possible (unless resize_to_max_canvas) - /// - Minimizes downscaling if no upscaling possible - /// - Minimizes padding area when tied - fn get_best_fit(&self, image_size: (u32, u32)) -> (u32, u32) { - let resolutions = self.find_supported_resolutions(); - let (orig_h, orig_w) = image_size; - - // Calculate scaling factors for each resolution - let scales_and_resolutions: Vec<(f64, (u32, u32))> = resolutions - .iter() - .map(|&(target_h, target_w)| { - let scale_w = target_w as f64 / orig_w as f64; - let scale_h = target_h as f64 / orig_h as f64; - // Limiting scale is the minimum (the side that constrains) - let scale = scale_w.min(scale_h); - (scale, (target_h, target_w)) - }) - .collect(); - - // Separate upscaling and downscaling options - let upscaling: Vec<_> = scales_and_resolutions - .iter() - .filter(|(s, _)| *s >= 1.0) - .copied() - .collect(); - - let selected_scale = if upscaling.is_empty() { - // No upscaling possible, pick largest downscaling (minimum reduction) - scales_and_resolutions - .iter() - .filter(|(s, _)| *s < 1.0) - .map(|(s, _)| *s) - .fold(f64::NEG_INFINITY, f64::max) - } else if self.resize_to_max_canvas { - // Pick largest upscaling - upscaling - .iter() - .map(|(s, _)| *s) - .fold(f64::NEG_INFINITY, f64::max) - } else { - // Pick smallest upscaling (minimum distortion) - upscaling - .iter() - .map(|(s, _)| *s) - .fold(f64::INFINITY, f64::min) - }; - - // Get all resolutions with the selected scale - let candidates: Vec<_> = scales_and_resolutions - .iter() - .filter(|(s, _)| (*s - selected_scale).abs() < 1e-9) - .map(|(_, res)| *res) - .collect(); - - // If multiple candidates, pick the one with minimum area (less padding) - if candidates.len() > 1 { - *candidates - .iter() - .min_by_key(|(h, w)| h * w) - .unwrap_or(&candidates[0]) - } else { - candidates[0] - } - } - - /// Build a padded [C, H, W] f32 tensor from a smaller image. - /// - /// The image is placed at top-left, and the remaining canvas is filled with - /// the normalized value of black (0). This fuses pad + tensor conversion - /// into one step, avoiding an intermediate padded `RgbImage` allocation. - fn pad_and_normalize_to_tensor( - &self, - image: &DynamicImage, - canvas_w: usize, - canvas_h: usize, - ) -> Array3 { - let (img_w, img_h, raw) = transforms::rgb_bytes(image); - let canvas_pixels = canvas_h * canvas_w; - - // Precompute fused scale/bias: (pixel/255 - mean) / std - let scale: [f32; 3] = std::array::from_fn(|c| 1.0 / (255.0 * self.std[c] as f32)); - let bias: [f32; 3] = std::array::from_fn(|c| -(self.mean[c] as f32) / (self.std[c] as f32)); - - let mut data = vec![0.0f32; 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 * scale + bias = bias - r_plane.fill(bias[0]); - g_plane.fill(bias[1]); - b_plane.fill(bias[2]); - - // Overwrite image region row-by-row using the shared block-optimized helper - 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") - } - - /// Split image tensor into tiles. - fn split_to_tiles( - &self, - tensor: &Array3, - num_tiles_h: usize, - num_tiles_w: usize, - ) -> Array4 { - let tile = self.tile_size as usize; - let num_tiles = num_tiles_h * num_tiles_w; - - let mut tiles = Array4::::zeros((num_tiles, 3, tile, tile)); - - for h_idx in 0..num_tiles_h { - for w_idx in 0..num_tiles_w { - let tile_idx = h_idx * num_tiles_w + w_idx; - let y_start = h_idx * tile; - let x_start = w_idx * tile; - - let tile_view = - tensor.slice(s![.., y_start..y_start + tile, x_start..x_start + tile]); - tiles.slice_mut(s![tile_idx, .., .., ..]).assign(&tile_view); - } - } - - tiles - } - - /// Create global image by bilinear interpolation to tile size. - fn create_global_image(&self, image: &DynamicImage) -> Array3 { - let tile = self.tile_size; - let resized = transforms::resize(image, tile, tile, FilterType::Triangle); - transforms::to_tensor_and_normalize(&resized, &self.mean, &self.std) - } - - /// Process a single image. - fn process_single_image(&self, image: &DynamicImage) -> (Array4, (usize, usize)) { - let (orig_w, orig_h) = image.dimensions(); - let image_size = (orig_h, orig_w); - - // Step 1: Find best fit resolution (canvas size for padding/tiling) - let target_size = self.get_best_fit(image_size); - let (target_h, target_w) = target_size; - - // Step 2: Compute resize target - limit upscaling if not resize_to_max_canvas - let resize_target = if self.resize_to_max_canvas { - target_size - } else { - let tile = self.tile_size; - let new_target_h = target_h.min(orig_h.max(tile)); - let new_target_w = target_w.min(orig_w.max(tile)); - (new_target_h, new_target_w) - }; - - // Step 3: Resize preserving aspect ratio to fit within resize_target - let new_size = Self::get_max_res_without_distortion(image_size, resize_target); - let (new_h, new_w) = (new_size.0.max(1), new_size.1.max(1)); - - let resized = transforms::resize(image, new_w, new_h, FilterType::Triangle); - - // Fused pad + tensor: build the padded f32 tensor directly from the - // resized RGB bytes, avoiding an intermediate padded RgbImage allocation. - let tensor = if new_w != target_w || new_h != target_h { - self.pad_and_normalize_to_tensor(&resized, target_w as usize, target_h as usize) - } else { - transforms::to_tensor_and_normalize(&resized, &self.mean, &self.std) - }; - - // Step 6: Calculate tile counts based on target_size (canvas size) - let tile = self.tile_size as usize; - let num_tiles_h = target_h as usize / tile; - let num_tiles_w = target_w as usize / tile; - - // Step 7: Split into tiles + global tile - let tiles = self.split_to_tiles(&tensor, num_tiles_h, num_tiles_w); - let num_tiles = num_tiles_h * num_tiles_w; - - let output = if num_tiles > 1 { - let global_tile = self.create_global_image(image); - let mut combined = Array4::::zeros((num_tiles + 1, 3, tile, tile)); - combined - .slice_mut(s![..num_tiles, .., .., ..]) - .assign(&tiles); - combined - .slice_mut(s![num_tiles, .., .., ..]) - .assign(&global_tile); - combined - } else { - tiles - }; - - (output, (num_tiles_h, num_tiles_w)) - } - - /// Calculate number of image tokens for a given aspect ratio. - pub fn calculate_num_tokens_for_aspect_ratio(&self, aspect_ratio: (usize, usize)) -> usize { - let (h_tiles, w_tiles) = aspect_ratio; - let num_tiles = h_tiles * w_tiles; - // Add 1 for global tile if num_tiles > 1 - let total_tiles = if num_tiles > 1 { - num_tiles + 1 - } else { - num_tiles - }; - let tokens_per_tile = (self.tile_size as usize / PATCH_SIZE).pow(2); - total_tiles * tokens_per_tile - } -} - -impl VisionPreProcessor for Llama4VisionProcessor { - fn default_mean(&self) -> [f64; 3] { - self.mean - } - - fn default_std(&self) -> [f64; 3] { - self.std - } - - fn preprocess( - &self, - images: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - if images.is_empty() { - return Err(TransformError::InvalidShape { - expected: "non-empty image batch".to_string(), - actual: vec![0], - }); - } - - let owned_processor; - let processor = if config.max_image_tiles.is_some() - || config.image_mean.is_some() - || config.image_std.is_some() - || config.size.is_some() - { - owned_processor = Self::from_preprocessor_config(config); - &owned_processor - } else { - self - }; - - let mut all_outputs = Vec::new(); - let mut all_aspect_ratios = Vec::new(); - let mut item_sizes = Vec::new(); - let mut feature_token_counts = Vec::new(); - - for image in images { - let (output, aspect_ratio) = processor.process_single_image(image); - let tokens = processor.calculate_num_tokens_for_aspect_ratio(aspect_ratio); - - all_outputs.push(output); - all_aspect_ratios.push(aspect_ratio); - item_sizes.push(image.dimensions()); - feature_token_counts.push(tokens); - } - - // Per-image tile counts (must be computed before remove/concatenate) - let patches_per_image: Vec = all_outputs.iter().map(|o| o.shape()[0] as i64).collect(); - - // Concatenate all tiles from all images into a single 4D tensor - // [total_tiles, C, H, W] — no batch dimension, no zero-padding. - // This matches what sglang and vLLM vision models expect. - let encoder_input = if all_outputs.len() == 1 { - all_outputs.remove(0) - } else { - let tile_views: Vec> = - all_outputs.iter().map(|o| o.view()).collect(); - ndarray::concatenate(ndarray::Axis(0), &tile_views).map_err(|e| { - TransformError::ShapeError(format!("Failed to concatenate tiles: {e}")) - })? - }; - - // Store aspect ratios and patches_per_image as model-specific data - let mut model_specific = std::collections::HashMap::new(); - let batch_size = images.len(); - - let aspect_ratios_flat: Vec = all_aspect_ratios - .iter() - .flat_map(|&(h, w)| vec![h as i64, w as i64]) - .collect(); - model_specific.insert( - "aspect_ratios".to_string(), - ModelSpecificValue::IntTensor { - data: aspect_ratios_flat, - shape: vec![batch_size, 2], - }, - ); - model_specific.insert( - "patches_per_image".to_string(), - ModelSpecificValue::int_1d(patches_per_image), - ); - - Ok(PreprocessedEncoderInputs { - encoder_input: encoder_input.into_dyn(), - feature_token_counts, - item_sizes, - model_specific, - }) - } - - fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { - let processor = Self::from_preprocessor_config(config); - let image_size = (height, width); - // target_size from get_best_fit determines the canvas and tile count - let target_size = processor.get_best_fit(image_size); - - let tile = processor.tile_size as usize; - let num_tiles_h = target_size.0 as usize / tile; - let num_tiles_w = target_size.1 as usize / tile; - - processor.calculate_num_tokens_for_aspect_ratio((num_tiles_h, num_tiles_w)) - } - - fn model_name(&self) -> &'static str { - "llama4-vision" - } - - fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> { - // For LLaMA 4, the size depends on the input image - let _ = config; - None - } -} - -#[cfg(test)] -mod tests { - use image::{Rgb, RgbImage}; - - use super::*; - - fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { - DynamicImage::from(RgbImage::from_pixel(width, height, color)) - } - - #[test] - fn test_llama4_vision_processor_default() { - let processor = Llama4VisionProcessor::new(); - assert_eq!(processor.tile_size(), TILE_SIZE); - assert_eq!(processor.max_patches(), DEFAULT_MAX_PATCHES); - assert_eq!(processor.mean, LLAMA4_MEAN); - assert_eq!(processor.std, LLAMA4_STD); - } - - #[test] - fn test_get_factors() { - let factors = Llama4VisionProcessor::get_factors(12); - assert!(factors.contains(&1)); - assert!(factors.contains(&2)); - assert!(factors.contains(&3)); - assert!(factors.contains(&4)); - assert!(factors.contains(&6)); - assert!(factors.contains(&12)); - assert_eq!(factors.len(), 6); - } - - #[test] - fn test_find_supported_resolutions() { - let processor = Llama4VisionProcessor::with_max_patches(4); - let resolutions = processor.find_supported_resolutions(); - - // Should include 1x1, 1x2, 2x1, 1x3, 3x1, 2x2, 1x4, 4x1 - let expected: Vec<(u32, u32)> = vec![ - (336, 336), // 1x1 - (336, 672), // 1x2 - (672, 336), // 2x1 - (336, 1008), // 1x3 - (1008, 336), // 3x1 - (672, 672), // 2x2 - (336, 1344), // 1x4 - (1344, 336), // 4x1 - ]; - - for exp in expected { - assert!( - resolutions.contains(&exp), - "Expected resolution {exp:?} not found" - ); - } - } - - #[test] - fn test_get_best_fit_square() { - let processor = Llama4VisionProcessor::new(); - let best = processor.get_best_fit((500, 500)); - // Square image should get a square or near-square resolution - assert!(best.0 == best.1 || (best.0 as i32 - best.1 as i32).abs() <= 336); - } - - #[test] - fn test_get_best_fit_wide() { - let processor = Llama4VisionProcessor::new(); - let best = processor.get_best_fit((300, 900)); - // Wide image should get wider resolution - assert!(best.1 >= best.0); - } - - #[test] - fn test_get_best_fit_tall() { - let processor = Llama4VisionProcessor::new(); - let best = processor.get_best_fit((900, 300)); - // Tall image should get taller resolution - assert!(best.0 >= best.1); - } - - #[test] - fn test_preprocess_square_image() { - let processor = Llama4VisionProcessor::new(); - let config = PreProcessorConfig::default(); - - let image = create_test_image(500, 500, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - // 4D output: [total_tiles, C, H, W] - assert_eq!(result.encoder_input.ndim(), 4); - assert_eq!(result.feature_token_counts.len(), 1); - assert!(result.feature_token_counts[0] > 0); - - // Check pixel values are normalized - let flat = result.encoder_input_flat(); - assert!(flat.iter().all(|&v| (-1.5..=1.5).contains(&v))); - } - - #[test] - fn test_preprocess_wide_image() { - let processor = Llama4VisionProcessor::new(); - let config = PreProcessorConfig::default(); - - let image = create_test_image(1000, 300, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - // 4D output: [total_tiles, C, H, W] - assert_eq!(result.encoder_input.ndim(), 4); - assert_eq!(result.feature_token_counts.len(), 1); - // Wide image should have more tiles in width direction - let aspect_ratios = result.model_specific.get("aspect_ratios").unwrap(); - if let ModelSpecificValue::IntTensor { data, .. } = aspect_ratios { - let h_tiles = data[0]; - let w_tiles = data[1]; - assert!(w_tiles >= h_tiles); - } - } - - #[test] - fn test_preprocess_multiple_images() { - let processor = Llama4VisionProcessor::new(); - let config = PreProcessorConfig::default(); - - let images = vec![ - create_test_image(500, 500, Rgb([100, 100, 100])), - create_test_image(800, 400, Rgb([150, 150, 150])), - ]; - - let result = processor.preprocess(&images, &config).unwrap(); - - // 4D output: [total_tiles, C, H, W] — tiles from both images concatenated - assert_eq!(result.encoder_input.ndim(), 4); - assert_eq!(result.feature_token_counts.len(), 2); - assert_eq!(result.item_sizes.len(), 2); - // Total tiles should be > 2 (at least 1 tile per image) - assert!(result.encoder_input.shape()[0] >= 2); - } - - #[test] - fn test_global_tile_added_for_multiple_tiles() { - let processor = Llama4VisionProcessor::new(); - let config = PreProcessorConfig::default(); - - // Large image that will require multiple tiles - let image = create_test_image(1000, 1000, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - let aspect_ratios = result.model_specific.get("aspect_ratios").unwrap(); - if let ModelSpecificValue::IntTensor { data, .. } = aspect_ratios { - let h_tiles = data[0] as usize; - let w_tiles = data[1] as usize; - let num_tiles = h_tiles * w_tiles; - - if num_tiles > 1 { - // 4D output: [total_tiles, C, H, W] - // total_tiles = num_tiles + 1 (global tile) - let shape = result.encoder_input.shape(); - assert_eq!(shape[0], num_tiles + 1); - } - } - } - - #[test] - fn test_model_name() { - let processor = Llama4VisionProcessor::new(); - assert_eq!(processor.model_name(), "llama4-vision"); - } - - #[test] - fn test_normalization_values() { - let processor = Llama4VisionProcessor::new(); - assert_eq!(processor.default_mean(), [0.5, 0.5, 0.5]); - assert_eq!(processor.default_std(), [0.5, 0.5, 0.5]); - } - - #[test] - fn test_token_count_calculation() { - let processor = Llama4VisionProcessor::new(); - // 1x1 tile: 576 tokens - assert_eq!(processor.calculate_num_tokens_for_aspect_ratio((1, 1)), 576); - // 2x2 tiles + 1 global: 5 * 576 = 2880 tokens - assert_eq!( - processor.calculate_num_tokens_for_aspect_ratio((2, 2)), - 2880 - ); - // 1x2 tiles + 1 global: 3 * 576 = 1728 tokens - assert_eq!( - processor.calculate_num_tokens_for_aspect_ratio((1, 2)), - 1728 - ); - } -} diff --git a/crates/multimodal/src/vision/processors/llava.rs b/crates/multimodal/src/vision/processors/llava.rs deleted file mode 100644 index 1d0699135..000000000 --- a/crates/multimodal/src/vision/processors/llava.rs +++ /dev/null @@ -1,957 +0,0 @@ -//! LLaVA family image processors. -//! -//! This module implements preprocessing for: -//! - LLaVA 1.5: CLIP-based preprocessing with configurable aspect ratio handling -//! - LLaVA-NeXT: Multi-crop anyres processing for higher resolution -//! -//! # Image Aspect Ratio Modes -//! -//! The processing behavior depends on the `image_aspect_ratio` config: -//! -//! - **None/Square**: Standard CLIP processing (resize shortest edge, center crop) -//! - **"pad"**: Expand to square with mean color padding, then resize -//! - **"anyres"**: Multi-crop processing for higher resolution (LLaVA-NeXT) -//! -//! # Processing Pipeline -//! -//! ## LLaVA 1.5 (Standard - no expand_to_square) -//! Used for `llava-hf/*` models where `image_aspect_ratio` is not set: -//! 1. Resize so shortest edge = target_size (preserving aspect ratio) -//! 2. Center crop to target_size x target_size -//! 3. Rescale by 1/255 -//! 4. Normalize with CLIP mean/std -//! -//! ## LLaVA 1.5 (Pad mode - with expand_to_square) -//! Used for `liuhaotian/llava-*` models where `image_aspect_ratio = "pad"`: -//! 1. Expand image to square by padding with mean color -//! 2. Resize to target size (typically 336x336) -//! 3. Normalize with CLIP mean/std -//! -//! ## LLaVA-NeXT -//! 1. Select best resolution from grid pinpoints -//! 2. Resize and pad to best resolution -//! 3. Divide into crops -//! 4. Process each crop + original resized image -//! 5. Stack all processed patches - -use image::{DynamicImage, GenericImageView}; -use ndarray::{self, Array3}; - -use crate::vision::{ - preprocessor_config::PreProcessorConfig, - processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::{ - center_crop, expand_to_square, mean_to_rgb, normalize, pil_to_filter, resize, stack_batch, - to_tensor, TransformError, - }, -}; - -/// CLIP normalization mean values used by LLaVA models. -pub const CLIP_MEAN: [f64; 3] = [0.48145466, 0.4578275, 0.40821073]; - -/// CLIP normalization std values used by LLaVA models. -pub const CLIP_STD: [f64; 3] = [0.26862954, 0.26130258, 0.27577711]; - -/// Image aspect ratio handling mode. -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub enum ImageAspectRatio { - /// Standard CLIP processing: resize shortest edge, center crop. - /// Used for llava-hf/* models where image_aspect_ratio is not set. - #[default] - Square, - /// Expand to square with mean color padding, then resize. - /// Used for liuhaotian/llava-* models where image_aspect_ratio = "pad". - Pad, - /// Multi-crop anyres processing (handled by LlavaNextProcessor). - Anyres, -} - -impl std::str::FromStr for ImageAspectRatio { - type Err = std::convert::Infallible; - - /// Parse from string value (e.g., from config). - fn from_str(s: &str) -> Result { - Ok(match s.to_lowercase().as_str() { - "pad" => Self::Pad, - "anyres" => Self::Anyres, - _ if s.contains("anyres") => Self::Anyres, // anyres_max_12, etc. - _ => Self::Square, - }) - } -} - -/// LLaVA 1.5 image processor. -/// -/// Implements CLIP-based preprocessing with configurable aspect ratio handling. -/// This processor is used for LLaVA 1.5 and similar models that expect fixed-size -/// square inputs. -/// -/// # Aspect Ratio Modes -/// -/// - `Square` (default): Standard CLIP processing (llava-hf/*) -/// - `Pad`: Expand to square with mean padding (liuhaotian/llava-*) -/// -/// # Token Calculation -/// -/// For LLaVA 1.5, the number of image tokens is: -/// ```text -/// num_tokens = (image_size / patch_size)² -/// ``` -/// With default settings (336x336, patch_size=14): 576 tokens -#[derive(Debug, Clone)] -pub struct LlavaProcessor { - /// Patch size for token calculation (typically 14) - pub patch_size: u32, - /// Target image size after processing (typically 336) - pub image_size: u32, - /// Image aspect ratio handling mode - pub aspect_ratio: ImageAspectRatio, -} - -impl Default for LlavaProcessor { - fn default() -> Self { - Self::new() - } -} - -impl LlavaProcessor { - /// Create a new LLaVA 1.5 processor with default settings. - /// - /// Default: patch_size=14, image_size=336, aspect_ratio=Square - /// This matches the llava-hf/* model behavior. - pub fn new() -> Self { - Self { - patch_size: 14, - image_size: 336, - aspect_ratio: ImageAspectRatio::Square, - } - } - - /// Create a processor with "pad" aspect ratio mode. - /// - /// This matches the liuhaotian/llava-* model behavior where - /// images are expanded to square before processing. - pub fn new_with_pad() -> Self { - Self { - patch_size: 14, - image_size: 336, - aspect_ratio: ImageAspectRatio::Pad, - } - } - - /// Create a processor with custom settings. - pub fn with_config(patch_size: u32, image_size: u32, aspect_ratio: ImageAspectRatio) -> Self { - Self { - patch_size, - image_size, - aspect_ratio, - } - } - - /// Create a processor from model config JSON. - /// - /// Extracts patch_size, image_size, and image_aspect_ratio from config. - pub fn from_config(config: &serde_json::Value) -> Self { - let patch_size = config - .get("vision_config") - .and_then(|v| v.get("patch_size")) - .and_then(|v| v.as_u64()) - .map(|v| v as u32) - .unwrap_or(14); - - let image_size = config - .get("vision_config") - .and_then(|v| v.get("image_size")) - .and_then(|v| v.as_u64()) - .map(|v| v as u32) - .unwrap_or(336); - - let aspect_ratio = config - .get("image_aspect_ratio") - .and_then(|v| v.as_str()) - .and_then(|s| s.parse().ok()) - .unwrap_or_default(); - - Self { - patch_size, - image_size, - aspect_ratio, - } - } - - /// Process a single image through the LLaVA 1.5 pipeline. - /// - /// The processing flow depends on `self.aspect_ratio`: - /// - `Square`: Standard CLIP (resize shortest edge, center crop) - /// - `Pad`: Expand to square with mean padding, then resize - fn process_one_image(&self, image: &DynamicImage, config: &PreProcessorConfig) -> Array3 { - let mean = config.get_image_mean(); - let std = config.get_image_std(); - let filter = pil_to_filter(config.resampling); - - // Get target size from config or use default - let target_size = config - .get_target_size() - .map(|(h, _w)| h) - .unwrap_or(self.image_size); - - // Get crop size (may be different from target_size) - let crop_size = config - .get_crop_size() - .map(|(h, _w)| h) - .unwrap_or(target_size); - - let processed = match self.aspect_ratio { - ImageAspectRatio::Pad => { - // Pad mode: expand to square with mean color padding, then resize - let (w, h) = image.dimensions(); - let squared = if w == h { - std::borrow::Cow::Borrowed(image) - } else { - let mean_color = mean_to_rgb(&mean); - std::borrow::Cow::Owned(expand_to_square(image, mean_color)) - }; - - // Resize to target size (maintaining square) - if config.do_resize.unwrap_or(true) { - resize(&squared, target_size, target_size, filter) - } else { - squared.into_owned() - } - } - ImageAspectRatio::Square | ImageAspectRatio::Anyres => { - // Square mode: Standard CLIP processing - // 1. Resize so shortest edge = target_size (preserving aspect ratio) - // 2. Center crop to crop_size x crop_size - let resized = if config.do_resize.unwrap_or(true) { - // Resize so shortest edge = target_size - let (w, h) = image.dimensions(); - let scale = if w < h { - target_size as f32 / w as f32 - } else { - target_size as f32 / h as f32 - }; - let new_w = (w as f32 * scale).round() as u32; - let new_h = (h as f32 * scale).round() as u32; - resize(image, new_w, new_h, filter) - } else { - image.clone() - }; - - // Center crop to crop_size (skip if image already fits) - if config.do_center_crop.unwrap_or(true) { - let (rw, rh) = resized.dimensions(); - if crop_size >= rw && crop_size >= rh { - resized - } else { - center_crop(&resized, crop_size, crop_size) - } - } else { - resized - } - } - }; - - // Convert to tensor [C, H, W] normalized to [0, 1] - let mut tensor = to_tensor(&processed); - - // Normalize with mean/std - if config.do_normalize.unwrap_or(true) { - normalize(&mut tensor, &mean, &std); - } - - tensor - } -} - -impl VisionPreProcessor for LlavaProcessor { - fn default_mean(&self) -> [f64; 3] { - CLIP_MEAN - } - - fn default_std(&self) -> [f64; 3] { - CLIP_STD - } - - fn preprocess( - &self, - images: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - if images.is_empty() { - return Err(TransformError::EmptyBatch); - } - - // Store original sizes - let item_sizes: Vec<(u32, u32)> = images.iter().map(|img| img.dimensions()).collect(); - - // Process each image - let tensors: Vec> = images - .iter() - .map(|img| self.process_one_image(img, config)) - .collect(); - - // Stack into batch - let encoder_input = stack_batch(&tensors)?; - - // Calculate token counts - let feature_token_counts: Vec = images - .iter() - .map(|_| self.calculate_num_tokens(self.image_size, self.image_size, config)) - .collect(); - - Ok(PreprocessedEncoderInputs::new( - encoder_input, - feature_token_counts, - item_sizes, - )) - } - - fn calculate_num_tokens( - &self, - _width: u32, - _height: u32, - config: &PreProcessorConfig, - ) -> usize { - // For LLaVA 1.5, token count is based on processed image size and patch size - let patch_size = config.get_patch_size(self.patch_size as usize) as u32; - let image_size = config - .get_target_size() - .map(|(h, _w)| h) - .unwrap_or(self.image_size); - - let patches_per_side = image_size / patch_size; - (patches_per_side * patches_per_side) as usize - } - - fn model_name(&self) -> &'static str { - "llava" - } - - fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> { - let size = config - .get_target_size() - .map(|(h, _w)| h) - .unwrap_or(self.image_size); - Some((size, size)) - } -} - -// ============================================================================ -// LLaVA-NeXT (Anyres) Support -// ============================================================================ - -/// LLaVA-NeXT image processor with anyres (multi-crop) support. -/// -/// LLaVA-NeXT processes high-resolution images by: -/// 1. Selecting the best resolution from predefined grid pinpoints -/// 2. Resizing and padding the image to that resolution -/// 3. Dividing into crops -/// 4. Processing each crop plus the original resized image -/// -/// # Token Calculation -/// -/// For LLaVA-NeXT, the number of tokens depends on the selected resolution: -/// ```text -/// base_tokens = (image_size / patch_size)² -/// grid_shape = (best_width / patch_size, best_height / patch_size) -/// unpad_shape = adjusted for aspect ratio -/// total_tokens = base_tokens + (unpad_w + 1) * unpad_h -/// ``` -#[derive(Debug, Clone)] -pub struct LlavaNextProcessor { - /// Base processor for individual patches - pub base: LlavaProcessor, - /// Grid pinpoints for resolution selection [(width, height), ...] - pub image_grid_pinpoints: Vec<(u32, u32)>, -} - -impl Default for LlavaNextProcessor { - fn default() -> Self { - Self::new() - } -} - -impl LlavaNextProcessor { - /// Create a new LLaVA-NeXT processor with default settings. - /// - /// Default grid pinpoints are common LLaVA-NeXT resolutions. - pub fn new() -> Self { - Self { - base: LlavaProcessor::new(), - // Common LLaVA-NeXT grid pinpoints - image_grid_pinpoints: vec![ - (336, 672), - (672, 336), - (672, 672), - (1008, 336), - (336, 1008), - ], - } - } - - /// Create a processor with custom grid pinpoints. - pub fn with_grid_pinpoints(grid_pinpoints: Vec<(u32, u32)>) -> Self { - Self { - base: LlavaProcessor::new(), - image_grid_pinpoints: grid_pinpoints, - } - } - - /// Create a processor from model config. - pub fn from_config(config: &serde_json::Value) -> Self { - let base = LlavaProcessor::from_config(config); - - let grid_pinpoints = config - .get("image_grid_pinpoints") - .and_then(|v| v.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|p| { - let pair = p.as_array()?; - // HF config stores pinpoints as [height, width] - let h = pair.first()?.as_u64()? as u32; - let w = pair.get(1)?.as_u64()? as u32; - Some((w, h)) - }) - .collect() - }) - .unwrap_or_else(|| vec![(336, 672), (672, 336), (672, 672), (1008, 336), (336, 1008)]); - - Self { - base, - image_grid_pinpoints: grid_pinpoints, - } - } - - /// Select the best resolution from grid pinpoints for the given image. - /// - /// Minimizes wasted pixels while maximizing effective resolution. - pub fn select_best_resolution(&self, original_size: (u32, u32)) -> (u32, u32) { - select_best_resolution(original_size, &self.image_grid_pinpoints) - } - - /// Get the grid shape (in crops) for anyres processing. - /// - /// Returns (grid_width, grid_height) — the number of 336×336 crops along - /// each axis of the best-fit resolution. - pub fn get_anyres_grid_shape(&self, image_size: (u32, u32)) -> (u32, u32) { - let (width, height) = self.select_best_resolution(image_size); - (width / self.base.image_size, height / self.base.image_size) - } - - /// Calculate unpad dimensions based on original aspect ratio. - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - pub fn calculate_unpad(&self, grid_shape: (u32, u32), original_size: (u32, u32)) -> (u32, u32) { - calculate_unpad(grid_shape, original_size) - } - - /// Resize and pad image to target resolution, maintaining aspect ratio. - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - fn resize_and_pad_image(&self, image: &DynamicImage, target: (u32, u32)) -> DynamicImage { - resize_and_pad_image(image, target) - } - - /// Divide image into crops of specified size. - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - fn divide_to_samples(&self, image: &DynamicImage, crop_size: (u32, u32)) -> Vec { - divide_to_samples(image, crop_size) - } - - /// Process a single patch/crop. - fn process_patch(&self, image: &DynamicImage, config: &PreProcessorConfig) -> Array3 { - let mean = config.get_image_mean(); - let std = config.get_image_std(); - let filter = pil_to_filter(config.resampling); - - // Get target size for patches - let target_size = config - .get_target_size() - .map(|(h, _w)| h) - .unwrap_or(self.base.image_size); - - // Resize patch to target size - let resized = if config.do_resize.unwrap_or(true) { - resize(image, target_size, target_size, filter) - } else { - image.clone() - }; - - // Center crop if configured (skip if image already fits) - let cropped = if config.do_center_crop.unwrap_or(true) { - if let Some((crop_h, crop_w)) = config.get_crop_size() { - let (rw, rh) = resized.dimensions(); - if crop_w >= rw && crop_h >= rh { - resized - } else { - center_crop(&resized, crop_w, crop_h) - } - } else { - resized - } - } else { - resized - }; - - // Convert to tensor - let mut tensor = to_tensor(&cropped); - - // Normalize - if config.do_normalize.unwrap_or(true) { - normalize(&mut tensor, &mean, &std); - } - - tensor - } -} - -impl VisionPreProcessor for LlavaNextProcessor { - fn default_mean(&self) -> [f64; 3] { - CLIP_MEAN - } - - fn default_std(&self) -> [f64; 3] { - CLIP_STD - } - - fn preprocess( - &self, - images: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - if images.is_empty() { - return Err(TransformError::EmptyBatch); - } - - let mut patches_per_image: Vec>> = Vec::with_capacity(images.len()); - let mut feature_token_counts = Vec::with_capacity(images.len()); - let mut item_sizes = Vec::with_capacity(images.len()); - - let filter = pil_to_filter(config.resampling); - let target_size = config - .get_target_size() - .map(|(h, _w)| h) - .unwrap_or(self.base.image_size); - let crop_size = config.get_crop_size().unwrap_or((target_size, target_size)); - - for image in images { - let original_size = image.dimensions(); - item_sizes.push(original_size); - - let best_resolution = self.select_best_resolution(original_size); - let image_padded = self.resize_and_pad_image(image, best_resolution); - let image_original_resize = resize(image, target_size, target_size, filter); - - let mut samples = vec![image_original_resize]; - samples.extend(self.divide_to_samples(&image_padded, crop_size)); - - let patches: Vec> = samples - .iter() - .map(|s| self.process_patch(s, config)) - .collect(); - patches_per_image.push(patches); - - feature_token_counts.push(self.calculate_num_tokens( - original_size.0, - original_size.1, - config, - )); - } - - // Build 5D encoder_input [num_images, max_patches, C, H, W] matching the - // HF LlavaNextImageProcessor output that vLLM expects with Batched layout. - let max_patches = patches_per_image.iter().map(|p| p.len()).max().unwrap_or(0); - let (c, h, w) = if let Some(first) = patches_per_image.iter().find_map(|p| p.first()) { - let s = first.shape(); - (s[0], s[1], s[2]) - } else { - return Err(TransformError::EmptyBatch); - }; - - let num_images = images.len(); - let mut encoder_input = ndarray::Array5::::zeros((num_images, max_patches, c, h, w)); - for (i, patches) in patches_per_image.iter().enumerate() { - for (j, patch) in patches.iter().enumerate() { - encoder_input - .slice_mut(ndarray::s![i, j, .., .., ..]) - .assign(patch); - } - } - - // Build model-specific image_sizes tensor as [num_images, 2] in (height, width) order - // to match the HF LlavaNextImageProcessor output format that vLLM expects. - let image_sizes_flat: Vec = item_sizes - .iter() - .flat_map(|&(w, h)| [h as i64, w as i64]) - .collect(); - - let mut result = - PreprocessedEncoderInputs::new(encoder_input, feature_token_counts, item_sizes); - result.model_specific.insert( - "image_sizes".to_string(), - ModelSpecificValue::int_2d(image_sizes_flat, num_images, 2), - ); - Ok(result) - } - - fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { - let original_size = (width, height); - - // Use effective geometry from config, falling back to base defaults. - let image_size = config - .get_target_size() - .map(|(h, _w)| h) - .unwrap_or(self.base.image_size); - let patch_size = self.base.patch_size; - - // Base feature tokens (from original resized image through ViT). - // CLIP ViT produces (image_size/patch_size)^2 + 1 tokens (patches + CLS). - // With vision_feature_select_strategy="default" CLS is removed, - // leaving (image_size/patch_size)^2 = 576 features. - let npatches = image_size / patch_size; // 24 for 336/14 - let base_features = (npatches * npatches) as usize; // 576 - - // Grid shape in crops (e.g. 1x2 for 336x672 best resolution) - let (grid_w, grid_h) = self.get_anyres_grid_shape(original_size); - - // Unpadded feature dimensions — matches vLLM's - // _get_num_unpadded_features which operates in feature space - // (npatches * grid_dim) per axis. - let current_w = (npatches * grid_w) as f32; - let current_h = (npatches * grid_h) as f32; - let aspect_ratio = width as f32 / height as f32; - let current_aspect = current_w / current_h; - - let (feat_h, feat_w) = if aspect_ratio > current_aspect { - // Landscape: height is unpadded - let new_h = (height as f32 * (current_w / width as f32)).round(); - let padding = ((current_h - new_h) / 2.0).floor() as usize; - (current_h as usize - 2 * padding, current_w as usize) - } else { - // Portrait or square: width is unpadded - let new_w = (width as f32 * (current_h / height as f32)).round(); - let padding = ((current_w - new_w) / 2.0).floor() as usize; - (current_h as usize, current_w as usize - 2 * padding) - }; - - let unpadded_features = feat_h * feat_w; - let newline_features = feat_h; // one newline token per row - - unpadded_features + newline_features + base_features - } - - fn model_name(&self) -> &'static str { - "llava-next" - } - - fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> { - // LLaVA-NeXT has variable output size based on crops - // Return the base patch size - let size = config - .get_target_size() - .map(|(h, _w)| h) - .unwrap_or(self.base.image_size); - Some((size, size)) - } -} - -// ============================================================================ -// Helper Functions (ported from mistral.rs) -// ============================================================================ - -/// Select the best resolution from possible resolutions for the given image size. -/// -/// Minimizes wasted pixels while maximizing effective resolution. -fn select_best_resolution( - original_size: (u32, u32), - possible_resolutions: &[(u32, u32)], -) -> (u32, u32) { - let (original_width, original_height) = original_size; - let mut best_fit = (0, 0); - let original_width_f = original_width as f32; - let original_height_f = original_height as f32; - let mut max_effective_resolution = 0_u32; - let mut min_wasted_resolution = u32::MAX; - - for &(width, height) in possible_resolutions { - let width_f = width as f32; - let height_f = height as f32; - let scale = (width_f / original_width_f).min(height_f / original_height_f); - let (downscaled_width, downscaled_height) = ( - (original_width_f * scale) as u32, - (original_height_f * scale) as u32, - ); - let effective_resolution = std::cmp::min( - downscaled_width * downscaled_height, - original_width * original_height, - ); - let wasted_resolution = width * height - effective_resolution; - - if effective_resolution > max_effective_resolution - || (effective_resolution == max_effective_resolution - && wasted_resolution < min_wasted_resolution) - { - best_fit = (width, height); - max_effective_resolution = effective_resolution; - min_wasted_resolution = wasted_resolution; - } - } - best_fit -} - -/// Calculate unpad dimensions based on aspect ratio. -fn calculate_unpad(size: (u32, u32), original_size: (u32, u32)) -> (u32, u32) { - let (original_width, original_height) = original_size; - let (current_width, current_height) = size; - let original_aspect_ratio = original_width as f32 / original_height as f32; - let current_aspect_ratio = current_width as f32 / current_height as f32; - - if original_aspect_ratio > current_aspect_ratio { - let scale_factor = current_width as f32 / original_width as f32; - let new_height = (original_height as f32 * scale_factor).floor() as u32; - let padding = (current_height - new_height) / 2; - (current_width, current_height - 2 * padding) - } else { - let scale_factor = current_height as f32 / original_height as f32; - let new_width = (original_width as f32 * scale_factor).floor() as u32; - let padding = (current_width - new_width) / 2; - (current_width - 2 * padding, current_height) - } -} - -/// Resize and pad image to target resolution, centering the image. -fn resize_and_pad_image(image: &DynamicImage, target: (u32, u32)) -> DynamicImage { - let (original_width, original_height) = image.dimensions(); - let (target_width, target_height) = target; - - let scale_w = target_width as f32 / original_width as f32; - let scale_h = target_height as f32 / original_height as f32; - - let (new_width, new_height) = if scale_w < scale_h { - ( - target_width, - std::cmp::min( - (original_height as f32 * scale_w).ceil() as u32, - target_height, - ), - ) - } else { - ( - std::cmp::min( - (original_width as f32 * scale_h).ceil() as u32, - target_width, - ), - target_height, - ) - }; - - let resized = resize( - image, - new_width, - new_height, - image::imageops::FilterType::CatmullRom, - ); - - let mut new_image = DynamicImage::new_rgb8(target_width, target_height); - let paste_x = (target_width - new_width) as i64 / 2; - let paste_y = (target_height - new_height) as i64 / 2; - - image::imageops::overlay(&mut new_image, &resized, paste_x, paste_y); - new_image -} - -/// Divide image into crops of specified size. -fn divide_to_samples(image: &DynamicImage, crop_size: (u32, u32)) -> Vec { - let (width, height) = image.dimensions(); - let mut samples = Vec::new(); - - for y in (0..height).step_by(crop_size.1 as usize) { - for x in (0..width).step_by(crop_size.0 as usize) { - let patch = image.crop_imm(x, y, crop_size.0, crop_size.1); - samples.push(patch); - } - } - samples -} - -#[cfg(test)] -mod tests { - use image::{Rgb, RgbImage}; - - use super::*; - - fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { - DynamicImage::from(RgbImage::from_pixel(width, height, color)) - } - - #[test] - fn test_llava_processor_default() { - let processor = LlavaProcessor::new(); - assert_eq!(processor.patch_size, 14); - assert_eq!(processor.image_size, 336); - assert_eq!(processor.aspect_ratio, ImageAspectRatio::Square); - } - - #[test] - fn test_llava_processor_with_pad() { - let processor = LlavaProcessor::new_with_pad(); - assert_eq!(processor.patch_size, 14); - assert_eq!(processor.image_size, 336); - assert_eq!(processor.aspect_ratio, ImageAspectRatio::Pad); - } - - #[test] - fn test_llava_token_calculation() { - let processor = LlavaProcessor::new(); - let config = PreProcessorConfig::default(); - - // 336 / 14 = 24, 24 * 24 = 576 - let tokens = processor.calculate_num_tokens(336, 336, &config); - assert_eq!(tokens, 576); - } - - #[test] - fn test_llava_preprocess_square() { - let processor = LlavaProcessor::new(); - let config = PreProcessorConfig { - do_resize: Some(true), - do_center_crop: Some(true), - do_normalize: Some(true), - image_mean: Some(CLIP_MEAN.to_vec()), - image_std: Some(CLIP_STD.to_vec()), - ..Default::default() - }; - - let image = create_test_image(336, 336, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - assert_eq!(result.batch_size(), 1); - assert_eq!(result.height().unwrap(), 336); - assert_eq!(result.width().unwrap(), 336); - assert_eq!(result.feature_token_counts[0], 576); - } - - #[test] - fn test_llava_preprocess_rectangular_square_mode() { - // Square mode (default): resize shortest edge, center crop - let processor = LlavaProcessor::new(); - let config = PreProcessorConfig { - do_resize: Some(true), - do_center_crop: Some(true), - do_normalize: Some(true), - size: Some([("shortest_edge".to_string(), 336)].into_iter().collect()), - crop_size: Some( - [("height".to_string(), 336), ("width".to_string(), 336)] - .into_iter() - .collect(), - ), - ..Default::default() - }; - - // Tall image - should be resized so shortest edge = 336, then center cropped - let image = create_test_image(200, 400, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - assert_eq!(result.batch_size(), 1); - assert_eq!(result.height().unwrap(), 336); - assert_eq!(result.width().unwrap(), 336); - } - - #[test] - fn test_llava_preprocess_rectangular_pad_mode() { - // Pad mode: expand to square with mean padding, then resize - let processor = LlavaProcessor::new_with_pad(); - let config = PreProcessorConfig { - do_resize: Some(true), - do_center_crop: Some(false), - do_normalize: Some(true), - ..Default::default() - }; - - // Tall image should be padded to square first - let image = create_test_image(200, 400, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - assert_eq!(result.batch_size(), 1); - // After expand_to_square: 400x400, then resize to 336x336 - assert_eq!(result.height().unwrap(), 336); - assert_eq!(result.width().unwrap(), 336); - } - - #[test] - fn test_select_best_resolution() { - let pinpoints = vec![(336, 672), (672, 336), (672, 672), (1008, 336), (336, 1008)]; - - // Square image should pick square resolution - let best = select_best_resolution((500, 500), &pinpoints); - assert_eq!(best, (672, 672)); - - // Wide image should pick wide resolution - let best = select_best_resolution((800, 400), &pinpoints); - assert_eq!(best, (672, 336)); - - // Tall image should pick tall resolution - let best = select_best_resolution((400, 800), &pinpoints); - assert_eq!(best, (336, 672)); - } - - #[test] - fn test_calculate_unpad() { - // Square grid, wide original -> should reduce width padding - let unpad = calculate_unpad((24, 24), (800, 400)); - assert!(unpad.0 >= unpad.1); // Width should be >= height - - // Square grid, tall original -> should reduce height padding - let unpad = calculate_unpad((24, 24), (400, 800)); - assert!(unpad.1 >= unpad.0); // Height should be >= width - } - - #[test] - fn test_llava_next_processor_default() { - let processor = LlavaNextProcessor::new(); - assert!(!processor.image_grid_pinpoints.is_empty()); - assert_eq!(processor.base.patch_size, 14); - } - - #[test] - fn test_llava_next_preprocess() { - let processor = LlavaNextProcessor::new(); - let config = PreProcessorConfig { - do_resize: Some(true), - do_center_crop: Some(false), - do_normalize: Some(true), - ..Default::default() - }; - - let image = create_test_image(500, 500, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - // 5D: [num_images=1, num_patches, C, H, W] - assert_eq!(result.batch_size(), 1); - // Should have multiple patches (original + crops) in the second dimension - assert!(result.encoder_input.shape()[1] > 1); - } - - #[test] - fn test_divide_to_samples() { - let image = create_test_image(672, 672, Rgb([128, 128, 128])); - let samples = divide_to_samples(&image, (336, 336)); - - // 672x672 / 336x336 = 2x2 = 4 patches - assert_eq!(samples.len(), 4); - - for sample in &samples { - assert_eq!(sample.width(), 336); - assert_eq!(sample.height(), 336); - } - } -} diff --git a/crates/multimodal/src/vision/processors/mod.rs b/crates/multimodal/src/vision/processors/mod.rs deleted file mode 100644 index 4c8c4135e..000000000 --- a/crates/multimodal/src/vision/processors/mod.rs +++ /dev/null @@ -1,39 +0,0 @@ -//! Model-specific vision processors. -//! -//! This module contains implementations of `VisionPreProcessor` for various -//! vision-language model families. -//! -//! # Supported Models -//! -//! - **LLaVA 1.5** (`llava`): CLIP-based preprocessing with configurable aspect ratio -//! - **LLaVA-NeXT** (`llava`): Multi-crop anyres processing -//! - **Qwen2-VL** (`qwen2_vl`): Dynamic resolution with smart resizing -//! - **Qwen2.5-VL** (`qwen2_vl`): Same processor as Qwen2-VL (identical preprocessing) -//! - **Qwen3-VL** (`qwen3_vl`): Similar to Qwen2-VL but with patch_size=16 and [0.5,0.5,0.5] normalization -//! - **Qwen3-Omni** (`qwen3_omni_vision`): Qwen3 vision preprocessing with Omni video limits and timing metadata -//! - **Kimi-K2.5** (`kimi_k25`): MoonViT resize and zero-padding to patch alignment -//! - **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 -//! - **Pixtral/Mistral3** (`pixtral`): CLIP-based preprocessing with dynamic resolution - -pub mod kimi_k25; -pub mod llama4_vision; -pub mod llava; -pub mod phi3_vision; -pub mod phi4_vision; -pub mod pixtral; -pub mod qwen2_vl; -pub mod qwen3_omni_vision; -pub mod qwen3_vl; -pub mod qwen_vl_base; - -pub use kimi_k25::KimiK25Processor; -pub use llama4_vision::Llama4VisionProcessor; -pub use llava::{ImageAspectRatio, LlavaNextProcessor, LlavaProcessor}; -pub use phi3_vision::Phi3VisionProcessor; -pub use phi4_vision::Phi4VisionProcessor; -pub use pixtral::PixtralProcessor; -pub use qwen2_vl::Qwen2VLProcessor; -pub use qwen3_omni_vision::Qwen3OmniVisionProcessor; -pub use qwen3_vl::Qwen3VLProcessor; diff --git a/crates/multimodal/src/vision/processors/phi3_vision.rs b/crates/multimodal/src/vision/processors/phi3_vision.rs deleted file mode 100644 index a266d4940..000000000 --- a/crates/multimodal/src/vision/processors/phi3_vision.rs +++ /dev/null @@ -1,563 +0,0 @@ -//! Phi3-Vision image processor. -//! -//! This module implements the Phi3-Vision image preprocessing pipeline with -//! Dynamic High Definition (HD) transform, which tiles images into 336x336 crops. -//! -//! # Processing Pipeline -//! -//! 1. **HD Transform**: Resize and pad image to multiples of 336 -//! 2. **Normalize**: Apply CLIP normalization -//! 3. **Create Global Image**: Bicubic interpolate to 336x336 -//! 4. **Tile**: Reshape into (num_tiles, 3, 336, 336) -//! 5. **Concatenate**: [global_image, tiles...] -//! 6. **Pad**: Zero-pad to (num_crops+1, 3, 336, 336) -//! -//! # Key Features -//! -//! - Dynamic resolution via HD transform -//! - Default num_crops: 16 -//! - CLIP normalization: mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711] -//! - Token count formula: `((h//336)*(w//336)+1)*144 + 1 + (h//336+1)*12` - -use image::{imageops::FilterType, DynamicImage, GenericImageView, Rgb, RgbImage}; -use ndarray::{s, Array3, Array4, IxDyn}; - -use crate::vision::{ - preprocessor_config::PreProcessorConfig, - processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::{self, TransformError}, -}; - -/// CLIP normalization mean values. -pub const CLIP_MEAN: [f64; 3] = [0.48145466, 0.4578275, 0.40821073]; - -/// CLIP normalization std values. -pub const CLIP_STD: [f64; 3] = [0.26862954, 0.26130258, 0.27577711]; - -/// Default number of crops for HD transform. -pub const DEFAULT_NUM_CROPS: usize = 16; - -/// Default number of image tokens per crop (144 per tile + base). -pub const DEFAULT_NUM_IMG_TOKENS: usize = 144; - -/// Tile size used in Phi3-Vision (336x336). -pub const TILE_SIZE: u32 = 336; - -/// Phi3-Vision image processor. -/// -/// Implements Dynamic HD transform with tile-based processing. -#[derive(Debug, Clone)] -pub struct Phi3VisionProcessor { - /// Maximum number of HD crops (not including global). - num_crops: usize, - /// Normalization mean. - mean: [f64; 3], - /// Normalization std. - std: [f64; 3], -} - -impl Default for Phi3VisionProcessor { - fn default() -> Self { - Self::new() - } -} - -impl Phi3VisionProcessor { - /// Create a new Phi3-Vision processor with default settings. - pub fn new() -> Self { - Self { - num_crops: DEFAULT_NUM_CROPS, - mean: CLIP_MEAN, - std: CLIP_STD, - } - } - - /// Create a processor with custom settings. - pub fn with_config(num_crops: usize) -> Self { - Self { - num_crops, - mean: CLIP_MEAN, - std: CLIP_STD, - } - } - - /// Create a processor from preprocessor config. - pub fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { - Self { - num_crops: config.num_crops.unwrap_or(DEFAULT_NUM_CROPS), - mean: config - .image_mean - .as_ref() - .map(|v| [v[0], v[1], v[2]]) - .unwrap_or(CLIP_MEAN), - std: config - .image_std - .as_ref() - .map(|v| [v[0], v[1], v[2]]) - .unwrap_or(CLIP_STD), - } - } - - /// Get the number of crops. - pub fn num_crops(&self) -> usize { - self.num_crops - } - - /// HD transform: resize and pad image to multiples of 336. - /// - /// Algorithm: - /// 1. If width < height, transpose (flip over main diagonal) - /// 2. Calculate scale: while scale * ceil(scale/ratio) <= hd_num: scale++ - /// 3. Resize to new_w = scale * 336, new_h = new_w / ratio - /// 4. Pad height to multiple of 336 (centered, white padding) - /// 5. If transposed, transpose back - pub fn hd_transform(&self, image: &DynamicImage) -> DynamicImage { - let (width, height) = image.dimensions(); - - let (img, transposed) = if width < height { - // Transpose (PIL's Image.TRANSPOSE): equivalent to fliph + rotate270 (ccw 90°) - // This swaps x and y coordinates: pixel at (x, y) goes to (y, x) - (image.fliph().rotate270(), true) - } else { - (image.clone(), false) - }; - - let (width, height) = img.dimensions(); - let ratio = width as f64 / height as f64; - - // Calculate scale factor - let mut scale = 1.0f64; - while scale * (scale / ratio).ceil() <= self.num_crops as f64 { - scale += 1.0; - } - scale -= 1.0; - - let new_w = (scale * TILE_SIZE as f64) as u32; - let new_h = (new_w as f64 / ratio) as u32; - - // Resize using bilinear filter (matching torchvision's bilinear+antialias) - // HuggingFace uses torchvision.transforms.functional.resize with - // BILINEAR interpolation and antialias=True. PIL's BILINEAR includes - // implicit antialiasing that closely matches torchvision. - let resized = transforms::resize(&img, new_w, new_h, FilterType::Triangle); - - // Pad height to multiple of 336 - let padded = self.padding_336(&resized); - - // Transpose back if needed (transpose is self-inverse) - if transposed { - padded.fliph().rotate270() - } else { - padded - } - } - - /// Pad image height to multiple of 336 (centered, white padding). - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - fn padding_336(&self, image: &DynamicImage) -> DynamicImage { - let (width, height) = image.dimensions(); - let target_h = ((height as f64 / TILE_SIZE as f64).ceil() * TILE_SIZE as f64) as u32; - - if height == target_h { - return image.clone(); - } - - let top_padding = (target_h - height) / 2; - - // Create white-padded image - let mut new_image = - DynamicImage::from(RgbImage::from_pixel(width, target_h, Rgb([255, 255, 255]))); - - // Copy original image to center - image::imageops::overlay(&mut new_image, image, 0, top_padding as i64); - - new_image - } - - /// Create global image by bicubic interpolation to 336x336. - /// - /// Uses the shared `bicubic_resize` which matches PyTorch's - /// `torch.nn.functional.interpolate(mode='bicubic', align_corners=False)`. - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - fn create_global_image(&self, tensor: &Array3) -> Array3 { - transforms::bicubic_resize(tensor, TILE_SIZE as usize, TILE_SIZE as usize) - } - - /// Reshape HD image into tiles. - /// - /// Transforms [3, H, W] -> [num_tiles, 3, 336, 336] - /// where H and W are multiples of 336. - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - fn reshape_to_tiles(&self, tensor: &Array3) -> Vec> { - let (_c, h, w) = (tensor.shape()[0], tensor.shape()[1], tensor.shape()[2]); - let grid_h = h / TILE_SIZE as usize; - let grid_w = w / TILE_SIZE as usize; - - let mut tiles = Vec::with_capacity(grid_h * grid_w); - - for gh in 0..grid_h { - for gw in 0..grid_w { - let y_start = gh * TILE_SIZE as usize; - let x_start = gw * TILE_SIZE as usize; - let y_end = y_start + TILE_SIZE as usize; - let x_end = x_start + TILE_SIZE as usize; - - let tile_view = tensor.slice(s![.., y_start..y_end, x_start..x_end]); - tiles.push(tile_view.to_owned()); - } - } - - tiles - } - - /// Calculate number of image tokens for given HD size. - /// - /// Formula: `((h//336)*(w//336)+1)*144 + 1 + (h//336+1)*12` - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - pub fn calculate_num_tokens(&self, h: usize, w: usize) -> usize { - let grid_h = h / TILE_SIZE as usize; - let grid_w = w / TILE_SIZE as usize; - - // ((h//336)*(w//336)+1)*144 + 1 + (h//336+1)*12 - (grid_h * grid_w + 1) * 144 + 1 + (grid_h + 1) * 12 - } - - /// Process a single image through the full pipeline. - fn process_single_image( - &self, - image: &DynamicImage, - config: &PreProcessorConfig, - ) -> (Array4, (usize, usize), usize) { - // 1. Convert to RGB - let image = DynamicImage::ImageRgb8(image.to_rgb8()); - - // 2. HD transform - let hd_image = self.hd_transform(&image); - let (hd_w, hd_h) = hd_image.dimensions(); - - // 3. To tensor [0, 1] and normalize - let mut tensor = transforms::to_tensor(&hd_image); - let mean = config - .image_mean - .as_ref() - .map(|v| [v[0], v[1], v[2]]) - .unwrap_or(self.mean); - let std = config - .image_std - .as_ref() - .map(|v| [v[0], v[1], v[2]]) - .unwrap_or(self.std); - transforms::normalize(&mut tensor, &mean, &std); - - // 4. Create global image (336x336) - let global_image = self.create_global_image(&tensor); - - // 5. Reshape HD image into tiles - let tiles = self.reshape_to_tiles(&tensor); - - // 6. Concatenate global + tiles - let max_crops = self.num_crops + 1; // num_crops + 1 for global - - // Create output tensor [max_crops, 3, 336, 336] - let mut output = - Array4::::zeros((max_crops, 3, TILE_SIZE as usize, TILE_SIZE as usize)); - - // Copy global image (first position) - output.slice_mut(s![0, .., .., ..]).assign(&global_image); - - // Copy tiles (positions 1..num_actual_crops) - for (i, tile) in tiles.iter().enumerate() { - if i + 1 < max_crops { - output.slice_mut(s![i + 1, .., .., ..]).assign(tile); - } - } - - // Calculate token count - let num_tokens = self.calculate_num_tokens(hd_h as usize, hd_w as usize); - - // The returned size is the HD-transformed image size. - (output, (hd_h as usize, hd_w as usize), num_tokens) - } -} - -impl VisionPreProcessor for Phi3VisionProcessor { - fn default_mean(&self) -> [f64; 3] { - self.mean - } - - fn default_std(&self) -> [f64; 3] { - self.std - } - - fn preprocess( - &self, - images: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - if images.is_empty() { - return Err(TransformError::InvalidShape { - expected: "at least one image".to_string(), - actual: vec![0], - }); - } - - let mut all_pixel_values = Vec::with_capacity(images.len()); - let mut all_image_sizes = Vec::with_capacity(images.len()); - let mut all_num_tokens = Vec::with_capacity(images.len()); - - for image in images { - let (encoder_input, image_size, num_tokens) = self.process_single_image(image, config); - all_pixel_values.push(encoder_input); - all_image_sizes.push((image_size.1 as u32, image_size.0 as u32)); // (width, height) - all_num_tokens.push(num_tokens); - } - - // Stack into batch [B, num_crops+1, 3, 336, 336] - let max_crops = self.num_crops + 1; - let batch_size = images.len(); - let mut batch_tensor = ndarray::Array5::::zeros(( - batch_size, - max_crops, - 3, - TILE_SIZE as usize, - TILE_SIZE as usize, - )); - - for (i, pv) in all_pixel_values.iter().enumerate() { - batch_tensor.slice_mut(s![i, .., .., .., ..]).assign(pv); - } - - // Convert to dynamic array for storage - let shape = batch_tensor.shape().to_vec(); - let (flat_data, _offset) = batch_tensor.into_raw_vec_and_offset(); - - // Store model-specific image_sizes data. - let mut model_specific = std::collections::HashMap::new(); - - // image_sizes as [batch, 2] tensor (h, w for each image) - let image_sizes_data: Vec = all_image_sizes - .iter() - .flat_map(|(w, h)| [*h as i64, *w as i64]) // [h, w] for each image - .collect(); - model_specific.insert( - "image_sizes".to_string(), - ModelSpecificValue::IntTensor { - data: image_sizes_data, - shape: vec![batch_size, 2], - }, - ); - - // feature_token_counts as list - model_specific.insert( - "num_img_tokens".to_string(), - ModelSpecificValue::IntVec(all_num_tokens.iter().map(|&t| t as i64).collect()), - ); - - // Convert 5D tensor to appropriate format - // Phi3-Vision expects [B, num_crops+1, C, H, W] - let encoder_input = ndarray::ArrayD::::from_shape_vec(IxDyn(&shape), flat_data) - .map_err(|e| TransformError::InvalidShape { - expected: format!("valid 5D shape, but failed with error: {e}"), - actual: shape.clone(), - })?; - - Ok(PreprocessedEncoderInputs { - encoder_input, - feature_token_counts: all_num_tokens, - item_sizes: all_image_sizes, - model_specific, - }) - } - - fn calculate_num_tokens(&self, width: u32, height: u32, _config: &PreProcessorConfig) -> usize { - // First apply HD transform to get the actual size - let image = DynamicImage::new_rgb8(width, height); - let hd_image = self.hd_transform(&image); - let (_, hd_h) = hd_image.dimensions(); - let hd_w = hd_image.width(); - - self.calculate_num_tokens(hd_h as usize, hd_w as usize) - } - - fn model_name(&self) -> &'static str { - "phi3-vision" - } - - fn get_processed_size(&self, _config: &PreProcessorConfig) -> Option<(u32, u32)> { - // Phi3-Vision has dynamic size based on HD transform - None - } -} - -#[cfg(test)] -mod tests { - use image::RgbImage; - - use super::*; - - fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { - DynamicImage::from(RgbImage::from_pixel(width, height, color)) - } - - #[test] - fn test_phi3_vision_processor_default() { - let processor = Phi3VisionProcessor::new(); - assert_eq!(processor.num_crops(), 16); - assert_eq!(processor.default_mean(), CLIP_MEAN); - assert_eq!(processor.default_std(), CLIP_STD); - } - - #[test] - fn test_hd_transform_square() { - let processor = Phi3VisionProcessor::new(); - let image = create_test_image(504, 504, Rgb([128, 128, 128])); - - let hd_image = processor.hd_transform(&image); - let (w, h) = hd_image.dimensions(); - - // Should be multiple of 336 - assert_eq!(h % 336, 0); - assert_eq!(w % 336, 0); - - // Should respect num_crops limit - let num_tiles = (h / 336) * (w / 336); - assert!(num_tiles <= processor.num_crops() as u32); - } - - #[test] - fn test_hd_transform_tall() { - let processor = Phi3VisionProcessor::new(); - let image = create_test_image(400, 600, Rgb([100, 100, 100])); - - let hd_image = processor.hd_transform(&image); - let (w, h) = hd_image.dimensions(); - - // Should be multiple of 336 - assert_eq!(h % 336, 0); - assert_eq!(w % 336, 0); - } - - #[test] - fn test_hd_transform_wide() { - let processor = Phi3VisionProcessor::new(); - let image = create_test_image(600, 400, Rgb([150, 150, 150])); - - let hd_image = processor.hd_transform(&image); - let (w, h) = hd_image.dimensions(); - - // Should be multiple of 336 - assert_eq!(h % 336, 0); - assert_eq!(w % 336, 0); - } - - #[test] - fn test_calculate_num_tokens() { - let processor = Phi3VisionProcessor::new(); - - // 1344x1344 -> 4x4 grid -> (16+1)*144 + 1 + (4+1)*12 = 2448 + 1 + 60 = 2509 - let tokens = processor.calculate_num_tokens(1344, 1344); - assert_eq!(tokens, 2509); - - // 1008x1344 -> 3x4 grid -> (12+1)*144 + 1 + (3+1)*12 = 1872 + 1 + 48 = 1921 - let tokens = processor.calculate_num_tokens(1008, 1344); - assert_eq!(tokens, 1921); - - // 1344x1008 -> 4x3 grid -> (12+1)*144 + 1 + (4+1)*12 = 1872 + 1 + 60 = 1933 - let tokens = processor.calculate_num_tokens(1344, 1008); - assert_eq!(tokens, 1933); - } - - #[test] - fn test_phi3_vision_preprocess() { - let processor = Phi3VisionProcessor::new(); - let config = PreProcessorConfig::default(); - - let image = create_test_image(504, 504, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - assert_eq!(result.batch_size(), 1); - - // Check output shape is [1, num_crops+1, 3, 336, 336] - let shape = result.encoder_input.shape(); - assert_eq!(shape.len(), 5); - assert_eq!(shape[0], 1); // batch - assert_eq!(shape[1], 17); // num_crops + 1 - assert_eq!(shape[2], 3); // channels - assert_eq!(shape[3], 336); // height - assert_eq!(shape[4], 336); // width - - // Check model-specific outputs - assert!(result.model_specific.contains_key("image_sizes")); - assert!(result.model_specific.contains_key("num_img_tokens")); - } - - #[test] - fn test_phi3_vision_preprocess_multiple() { - let processor = Phi3VisionProcessor::new(); - let config = PreProcessorConfig::default(); - - let images = vec![ - create_test_image(504, 504, Rgb([100, 100, 100])), - create_test_image(400, 600, Rgb([150, 150, 150])), - ]; - - let result = processor.preprocess(&images, &config).unwrap(); - - assert_eq!(result.batch_size(), 2); - assert_eq!(result.item_sizes.len(), 2); - assert_eq!(result.feature_token_counts.len(), 2); - } - - #[test] - fn test_model_name() { - let processor = Phi3VisionProcessor::new(); - assert_eq!(processor.model_name(), "phi3-vision"); - } - - #[test] - fn test_from_config() { - let config = PreProcessorConfig { - num_crops: Some(8), - image_mean: Some(vec![0.5, 0.5, 0.5]), - image_std: Some(vec![0.5, 0.5, 0.5]), - ..Default::default() - }; - - let processor = Phi3VisionProcessor::from_preprocessor_config(&config); - assert_eq!(processor.num_crops(), 8); - } - - #[test] - fn test_transpose_equivalence() { - // Test that fliph().rotate270() correctly implements PIL's Image.TRANSPOSE - // TRANSPOSE swaps x and y coordinates: pixel at (x, y) goes to (y, x) - use image::{GenericImageView, Rgb, RgbImage}; - - let mut img = RgbImage::new(100, 200); - img.put_pixel(0, 0, Rgb([255, 0, 0])); // Top-left = red - img.put_pixel(99, 0, Rgb([0, 255, 0])); // Top-right = green - img.put_pixel(0, 199, Rgb([0, 0, 255])); // Bottom-left = blue - img.put_pixel(99, 199, Rgb([255, 255, 0])); // Bottom-right = yellow - - let img = DynamicImage::ImageRgb8(img); - let transposed = img.fliph().rotate270(); - - // After TRANSPOSE: (x, y) -> (y, x) - assert_eq!(transposed.get_pixel(0, 0).0[0..3], [255, 0, 0]); // (0,0) -> (0,0) - assert_eq!(transposed.get_pixel(0, 99).0[0..3], [0, 255, 0]); // (99,0) -> (0,99) - assert_eq!(transposed.get_pixel(199, 0).0[0..3], [0, 0, 255]); // (0,199) -> (199,0) - assert_eq!(transposed.get_pixel(199, 99).0[0..3], [255, 255, 0]); // (99,199) -> (199,99) - } -} diff --git a/crates/multimodal/src/vision/processors/phi4_vision.rs b/crates/multimodal/src/vision/processors/phi4_vision.rs deleted file mode 100644 index eb74d7c10..000000000 --- a/crates/multimodal/src/vision/processors/phi4_vision.rs +++ /dev/null @@ -1,740 +0,0 @@ -//! Phi4-Vision (Phi4-Multimodal) image processor. -//! -//! This module implements the Phi4-Vision image preprocessing pipeline with -//! Dynamic HD transform using a 448x448 base resolution. -//! -//! # Key Differences from Phi3-Vision -//! -//! | Feature | Phi3-Vision | Phi4-Vision | -//! |---------|-------------|-------------| -//! | Base resolution | 336 | 448 | -//! | Normalization | CLIP | [0.5, 0.5, 0.5] | -//! | Default max crops | 16 | 36 | -//! | Aspect ratio selection | Simple scale | Target ratio matching | -//! | Attention mask | No | Yes | -//! -//! # Processing Pipeline -//! -//! 1. **Dynamic Preprocess**: Calculate target resolution based on aspect ratio -//! 2. **Resize**: Scale to target resolution maintaining aspect ratio -//! 3. **Pad**: Add white padding to reach exact target dimensions -//! 4. **Normalize**: Apply [0.5, 0.5, 0.5] mean/std normalization -//! 5. **Create Global Image**: Bilinear interpolate to 448x448 -//! 6. **Tile**: Reshape HD image into (h_crops * w_crops, 3, 448, 448) tiles -//! 7. **Concatenate**: [global_image, tiles...] -//! 8. **Generate Attention Mask**: Track valid (non-padding) regions -//! -//! # Token Count Formula -//! -//! `num_tokens = 256 + 1 + mask_sum + mask_col0_sum + 16` -//! -//! Where: -//! - 256: base global tokens -//! - mask_sum: sum of downsampled attention mask -//! - mask_col0_sum: sum of first column of mask (row separators) - -use std::collections::HashSet; - -use image::{imageops::FilterType, DynamicImage, GenericImageView, Rgb, RgbImage}; -use ndarray::{s, Array2, Array3, Array4, IxDyn}; - -use crate::vision::{ - preprocessor_config::PreProcessorConfig, - processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::{self, TransformError}, -}; - -/// Simple normalization mean for Phi4-Vision. -pub const PHI4_MEAN: [f64; 3] = [0.5, 0.5, 0.5]; - -/// Simple normalization std for Phi4-Vision. -pub const PHI4_STD: [f64; 3] = [0.5, 0.5, 0.5]; - -/// Default dynamic_hd value (max crops) for Phi4-Vision. -pub const DEFAULT_DYNAMIC_HD: usize = 36; - -/// Base resolution for Phi4-Vision (448x448). -pub const BASE_RESOLUTION: u32 = 448; - -/// Mask resolution (base_resolution / patch_size = 448 / 14). -pub const MASK_RESOLUTION: usize = 32; - -/// Patch size used in Phi4-Vision. -pub const PATCH_SIZE: usize = 14; - -/// Result type for single image processing. -/// Contains (encoder_input, attention_mask, (height, width), num_tokens). -type SingleImageResult = (Array4, Array3, (u32, u32), usize); - -/// Phi4-Vision image processor. -/// -/// Implements Dynamic HD transform with aspect ratio matching. -#[derive(Debug, Clone)] -pub struct Phi4VisionProcessor { - /// Maximum number of crops (dynamic_hd parameter). - dynamic_hd: usize, - /// Base resolution for tiles. - base_resolution: u32, - /// Mask resolution (base_resolution / patch_size). - mask_resolution: usize, - /// Normalization mean. - mean: [f64; 3], - /// Normalization std. - std: [f64; 3], -} - -impl Default for Phi4VisionProcessor { - fn default() -> Self { - Self::new() - } -} - -impl Phi4VisionProcessor { - /// Create a new Phi4-Vision processor with default settings. - pub fn new() -> Self { - Self { - dynamic_hd: DEFAULT_DYNAMIC_HD, - base_resolution: BASE_RESOLUTION, - mask_resolution: MASK_RESOLUTION, - mean: PHI4_MEAN, - std: PHI4_STD, - } - } - - /// Create a processor with custom dynamic_hd setting. - pub fn with_dynamic_hd(dynamic_hd: usize) -> Self { - Self { - dynamic_hd, - base_resolution: BASE_RESOLUTION, - mask_resolution: MASK_RESOLUTION, - mean: PHI4_MEAN, - std: PHI4_STD, - } - } - - /// Create a processor from preprocessor config. - pub fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { - Self { - dynamic_hd: config.dynamic_hd.unwrap_or(DEFAULT_DYNAMIC_HD), - base_resolution: BASE_RESOLUTION, - mask_resolution: MASK_RESOLUTION, - mean: config - .image_mean - .as_ref() - .map(|v| [v[0], v[1], v[2]]) - .unwrap_or(PHI4_MEAN), - std: config - .image_std - .as_ref() - .map(|v| [v[0], v[1], v[2]]) - .unwrap_or(PHI4_STD), - } - } - - /// Get the dynamic_hd (max crops) value. - pub fn dynamic_hd(&self) -> usize { - self.dynamic_hd - } - - /// Get the base resolution. - pub fn base_resolution(&self) -> u32 { - self.base_resolution - } - - /// Compute valid target aspect ratios for the given crop range. - /// - /// Returns sorted list of (width_crops, height_crops) tuples where - /// min_num <= w * h <= max_num. - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - fn compute_target_ratios(&self, min_num: usize, max_num: usize) -> Vec<(usize, usize)> { - let mut ratios: HashSet<(usize, usize)> = HashSet::new(); - for n in min_num..=max_num { - // Find factor pairs by iterating up to sqrt(n) - for i in 1..=(n as f64).sqrt() as usize { - if n % i == 0 { - ratios.insert((i, n / i)); - ratios.insert((n / i, i)); - } - } - } - let mut sorted_ratios: Vec<(usize, usize)> = ratios.into_iter().collect(); - sorted_ratios.sort_by_key(|&(i, j)| i * j); - sorted_ratios - } - - /// Find the closest aspect ratio from the target ratios. - /// - /// Selects the ratio that minimizes the difference from the original - /// aspect ratio. When tied, prefers ratios where the resized area - /// exceeds half the base area product. - fn find_closest_aspect_ratio( - &self, - aspect_ratio: f64, - target_ratios: &[(usize, usize)], - width: u32, - height: u32, - ) -> (usize, usize) { - let mut best_ratio_diff = f64::INFINITY; - let mut best_ratio = (1, 1); - let area = (width * height) as f64; - let base_area = (self.base_resolution * self.base_resolution) as f64; - - for &(w_ratio, h_ratio) in target_ratios { - let target_aspect_ratio = w_ratio as f64 / h_ratio as f64; - let ratio_diff = (aspect_ratio - target_aspect_ratio).abs(); - - if ratio_diff < best_ratio_diff { - best_ratio_diff = ratio_diff; - best_ratio = (w_ratio, h_ratio); - } else if (ratio_diff - best_ratio_diff).abs() < 1e-6 { - // Tie-breaker: prefer ratio if area > 0.5 * base_area * w * h - if area > 0.5 * base_area * (w_ratio * h_ratio) as f64 { - best_ratio = (w_ratio, h_ratio); - } - } - } - best_ratio - } - - /// Dynamic preprocess: calculate target dimensions and create attention mask. - /// - /// Returns (processed_image, attention_mask, target_h_crops, target_w_crops) - fn dynamic_preprocess( - &self, - image: &DynamicImage, - ) -> (DynamicImage, Array2, usize, usize) { - let (orig_w, orig_h) = image.dimensions(); - let base_res = self.base_resolution as f64; - - // Calculate natural crop numbers - let w_crop_num = (orig_w as f64 / base_res).ceil() as usize; - let h_crop_num = (orig_h as f64 / base_res).ceil() as usize; - - let (target_w_crops, target_h_crops, target_width, target_height) = - if w_crop_num * h_crop_num > self.dynamic_hd { - // Image exceeds max crops, need to find best aspect ratio - let aspect_ratio = orig_w as f64 / orig_h as f64; - let target_ratios = self.compute_target_ratios(1, self.dynamic_hd); - let (w_ratio, h_ratio) = - self.find_closest_aspect_ratio(aspect_ratio, &target_ratios, orig_w, orig_h); - - let target_width = self.base_resolution * w_ratio as u32; - let target_height = self.base_resolution * h_ratio as u32; - (w_ratio, h_ratio, target_width, target_height) - } else { - // Image fits within max crops - let target_width = self.base_resolution * w_crop_num as u32; - let target_height = self.base_resolution * h_crop_num as u32; - (w_crop_num, h_crop_num, target_width, target_height) - }; - - // Calculate resize ratios - let ratio_width = target_width as f64 / orig_w as f64; - let ratio_height = target_height as f64 / orig_h as f64; - - let (new_w, new_h, padding_width, padding_height) = if ratio_width < ratio_height { - // Width is the limiting factor - let new_w = target_width; - let new_h = (orig_h as f64 * ratio_width) as u32; - (new_w, new_h, 0u32, target_height - new_h) - } else { - // Height is the limiting factor - let new_h = target_height; - let new_w = (orig_w as f64 * ratio_height) as u32; - (new_w, new_h, target_width - new_w, 0u32) - }; - - // Create attention mask (tracks valid regions) - let mask_h = self.mask_resolution * target_h_crops; - let mask_w = self.mask_resolution * target_w_crops; - let mut attention_mask = Array2::::ones((mask_h, mask_w)); - - // Mark padding regions as 0 in mask - if padding_width >= PATCH_SIZE as u32 { - let padding_mask_cols = (padding_width as usize) / PATCH_SIZE; - for row in 0..mask_h { - for col in (mask_w - padding_mask_cols)..mask_w { - attention_mask[[row, col]] = 0; - } - } - } - if padding_height >= PATCH_SIZE as u32 { - let padding_mask_rows = (padding_height as usize) / PATCH_SIZE; - for row in (mask_h - padding_mask_rows)..mask_h { - for col in 0..mask_w { - attention_mask[[row, col]] = 0; - } - } - } - - // Resize image with bilinear interpolation (matching HuggingFace torchvision) - // HuggingFace uses torchvision.transforms.functional.resize with BILINEAR + antialias=True. - // FilterType::Triangle (bilinear) closely matches this behavior. - let resized = transforms::resize(image, new_w, new_h, FilterType::Triangle); - - // Pad to target dimensions (white padding on right/bottom) - let padded = self.pad_image(&resized, target_width, target_height); - - (padded, attention_mask, target_h_crops, target_w_crops) - } - - /// Pad image to target dimensions with white padding. - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - fn pad_image(&self, image: &DynamicImage, target_w: u32, target_h: u32) -> DynamicImage { - let (w, h) = image.dimensions(); - if w == target_w && h == target_h { - return image.clone(); - } - - // Create white background - let white = Rgb([255u8, 255, 255]); - let mut padded = RgbImage::from_pixel(target_w, target_h, white); - - // Copy image to top-left using efficient overlay - image::imageops::overlay(&mut padded, &image.to_rgb8(), 0, 0); - - DynamicImage::ImageRgb8(padded) - } - - /// Create global image by bicubic interpolation to base resolution. - /// - /// Uses the shared `bicubic_resize` which matches PyTorch's - /// `torch.nn.functional.interpolate(mode='bicubic', align_corners=False)`. - fn create_global_image(&self, tensor: &Array3) -> Array3 { - let target = self.base_resolution as usize; - transforms::bicubic_resize(tensor, target, target) - } - - /// Tile the HD image into crops of base_resolution x base_resolution. - fn tile_image(&self, tensor: &Array3, h_crops: usize, w_crops: usize) -> Array4 { - let base = self.base_resolution as usize; - let num_tiles = h_crops * w_crops; - - let mut tiles = Array4::::zeros((num_tiles, 3, base, base)); - - for h_idx in 0..h_crops { - for w_idx in 0..w_crops { - let tile_idx = h_idx * w_crops + w_idx; - let y_start = h_idx * base; - let x_start = w_idx * base; - - for c in 0..3 { - for y in 0..base { - for x in 0..base { - tiles[[tile_idx, c, y, x]] = tensor[[c, y_start + y, x_start + x]]; - } - } - } - } - } - - tiles - } - - /// Downsample attention mask by factor of 2. - fn downsample_mask(&self, mask: &Array2, h_crops: usize, w_crops: usize) -> Array2 { - let half_res = self.mask_resolution / 2; - let out_h = h_crops * half_res; - let out_w = w_crops * half_res; - - let mut downsampled = Array2::::zeros((out_h, out_w)); - - for y in 0..out_h { - for x in 0..out_w { - // Sample every other pixel - let src_y = y * 2; - let src_x = x * 2; - if src_y < mask.shape()[0] && src_x < mask.shape()[1] { - downsampled[[y, x]] = mask[[src_y, src_x]]; - } - } - } - - downsampled - } - - /// Calculate number of image tokens. - /// - /// Formula: 256 + 1 + mask_sum + mask_col0_sum + 16 - /// - 256: global image tokens - /// - 1: separator - /// - mask_sum: sum of downsampled attention mask (valid HD tokens) - /// - mask_col0_sum: sum of first column (row separators) - /// - 16: additional fixed tokens - #[expect( - clippy::unused_self, - reason = "method logically belongs to the processor; keeps API consistent" - )] - fn calculate_num_tokens(&self, downsampled_mask: &Array2) -> usize { - let mask_sum: u32 = downsampled_mask.iter().sum(); - let mask_col0_sum: u32 = downsampled_mask.column(0).iter().sum(); - 256 + 1 + mask_sum as usize + mask_col0_sum as usize + 16 - } - - /// Process a single image. - fn process_single_image(&self, image: &DynamicImage) -> SingleImageResult { - // Step 1: Dynamic preprocess (resize, pad, create attention mask) - let (hd_image, attention_mask, h_crops, w_crops) = self.dynamic_preprocess(image); - - let hd_h = hd_image.height(); - let hd_w = hd_image.width(); - - // Step 2: Convert to tensor and normalize - let mut hd_tensor = transforms::to_tensor(&hd_image); - transforms::normalize(&mut hd_tensor, &self.mean, &self.std); - - // Step 3: Create global image - let global_tensor = self.create_global_image(&hd_tensor); - - // Step 4: Tile HD image - let tiles = self.tile_image(&hd_tensor, h_crops, w_crops); - let num_hd_tiles = h_crops * w_crops; - - // Step 5: Concatenate global + tiles - // Output shape: [num_hd_tiles + 1, 3, base_resolution, base_resolution] - let base = self.base_resolution as usize; - let total_crops = num_hd_tiles + 1; - let mut output = Array4::::zeros((total_crops, 3, base, base)); - - // First slot is global image - output.slice_mut(s![0, .., .., ..]).assign(&global_tensor); - - // Remaining slots are HD tiles - if num_hd_tiles > 0 { - output.slice_mut(s![1.., .., .., ..]).assign(&tiles); - } - - // Step 6: Create combined attention mask [total_crops, mask_resolution, mask_resolution] - let mask_res = self.mask_resolution; - let mut combined_mask = Array3::::zeros((total_crops, mask_res, mask_res)); - - // Global mask is all ones - combined_mask.slice_mut(s![0, .., ..]).fill(1); - - // Tile attention masks - for h_idx in 0..h_crops { - for w_idx in 0..w_crops { - let tile_idx = h_idx * w_crops + w_idx + 1; // +1 for global - let mask_y_start = h_idx * mask_res; - let mask_x_start = w_idx * mask_res; - - let tile_mask = attention_mask.slice(s![ - mask_y_start..mask_y_start + mask_res, - mask_x_start..mask_x_start + mask_res - ]); - combined_mask - .slice_mut(s![tile_idx, .., ..]) - .assign(&tile_mask); - } - } - - // Step 7: Calculate token count - let downsampled = self.downsample_mask(&attention_mask, h_crops, w_crops); - let num_tokens = self.calculate_num_tokens(&downsampled); - - (output, combined_mask, (hd_h, hd_w), num_tokens) - } -} - -impl VisionPreProcessor for Phi4VisionProcessor { - fn default_mean(&self) -> [f64; 3] { - self.mean - } - - fn default_std(&self) -> [f64; 3] { - self.std - } - - fn preprocess( - &self, - images: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - if images.is_empty() { - return Err(TransformError::InvalidShape { - expected: "non-empty image batch".to_string(), - actual: vec![0], - }); - } - - let processor = if config.dynamic_hd.is_some() || config.image_mean.is_some() { - Self::from_preprocessor_config(config) - } else { - self.clone() - }; - - let mut all_outputs = Vec::new(); - let mut all_masks = Vec::new(); - let mut item_sizes = Vec::new(); - let mut feature_token_counts = Vec::new(); - - for image in images { - let (output, mask, size, tokens) = processor.process_single_image(image); - all_outputs.push(output); - all_masks.push(mask); - item_sizes.push(size); - feature_token_counts.push(tokens); - } - - // Find max crops across batch for padding - let max_crops = all_outputs - .iter() - .map(|o| o.shape()[0]) - .max() - .ok_or(TransformError::EmptyBatch)?; - let base = self.base_resolution as usize; - let mask_res = self.mask_resolution; - - // Pad all outputs to max_crops - let batch_size = images.len(); - let mut encoder_input = - ndarray::ArrayD::::zeros(IxDyn(&[batch_size, max_crops, 3, base, base])); - let mut attention_masks = - ndarray::ArrayD::::zeros(IxDyn(&[batch_size, max_crops, mask_res, mask_res])); - - for (b, (output, mask)) in all_outputs.iter().zip(all_masks.iter()).enumerate() { - let num_crops = output.shape()[0]; - for t in 0..num_crops { - for c in 0..3 { - for y in 0..base { - for x in 0..base { - encoder_input[[b, t, c, y, x]] = output[[t, c, y, x]]; - } - } - } - for y in 0..mask_res { - for x in 0..mask_res { - attention_masks[[b, t, y, x]] = mask[[t, y, x]]; - } - } - } - // Remaining crops stay as zeros (padding) - } - - // Convert to standard format - let mut model_specific = std::collections::HashMap::new(); - - // Store attention mask as model-specific data - let mask_flat: Vec = attention_masks.iter().map(|&v| v as i64).collect(); - model_specific.insert( - "pixel_attention_mask".to_string(), - ModelSpecificValue::IntTensor { - data: mask_flat, - shape: vec![batch_size, max_crops, mask_res, mask_res], - }, - ); - - // Store image sizes (H, W after HD transform) - let sizes_flat: Vec = item_sizes - .iter() - .flat_map(|&(h, w)| vec![h as i64, w as i64]) - .collect(); - model_specific.insert( - "image_sizes".to_string(), - ModelSpecificValue::IntTensor { - data: sizes_flat, - shape: vec![batch_size, 2], - }, - ); - - Ok(PreprocessedEncoderInputs { - encoder_input: encoder_input.into_dyn(), - feature_token_counts, - item_sizes, - model_specific, - }) - } - - fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { - let processor = Self::from_preprocessor_config(config); - let base_res = processor.base_resolution as f64; - - let w_crop_num = (width as f64 / base_res).ceil() as usize; - let h_crop_num = (height as f64 / base_res).ceil() as usize; - - let (target_w_crops, target_h_crops) = if w_crop_num * h_crop_num > processor.dynamic_hd { - let aspect_ratio = width as f64 / height as f64; - let target_ratios = processor.compute_target_ratios(1, processor.dynamic_hd); - processor.find_closest_aspect_ratio(aspect_ratio, &target_ratios, width, height) - } else { - (w_crop_num, h_crop_num) - }; - - // Approximate token count (without actual mask) - // Full mask would have target_w_crops * target_h_crops * (mask_res/2)^2 tokens - let half_res = processor.mask_resolution / 2; - let mask_area = target_h_crops * target_w_crops * half_res * half_res; - let mask_col0 = target_h_crops * half_res; - - 256 + 1 + mask_area + mask_col0 + 16 - } - - fn model_name(&self) -> &'static str { - "phi4-vision" - } - - fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> { - // For Phi4, the size depends on the input image - let _ = config; - None - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { - DynamicImage::from(RgbImage::from_pixel(width, height, color)) - } - - #[test] - fn test_phi4_vision_processor_default() { - let processor = Phi4VisionProcessor::new(); - assert_eq!(processor.dynamic_hd(), DEFAULT_DYNAMIC_HD); - assert_eq!(processor.base_resolution(), BASE_RESOLUTION); - assert_eq!(processor.mean, PHI4_MEAN); - assert_eq!(processor.std, PHI4_STD); - } - - #[test] - fn test_compute_target_ratios() { - let processor = Phi4VisionProcessor::new(); - - let ratios = processor.compute_target_ratios(1, 4); - // Should include (1,1), (1,2), (2,1), (1,3), (3,1), (2,2), (1,4), (4,1) - assert!(ratios.contains(&(1, 1))); - assert!(ratios.contains(&(2, 2))); - assert!(ratios.contains(&(1, 4))); - assert!(ratios.contains(&(4, 1))); - } - - #[test] - fn test_find_closest_aspect_ratio_square() { - let processor = Phi4VisionProcessor::new(); - let ratios = processor.compute_target_ratios(1, 36); - - // Square image should get close to (1,1) or similar square ratio - let result = processor.find_closest_aspect_ratio(1.0, &ratios, 500, 500); - assert_eq!(result.0, result.1); // Should be square - } - - #[test] - fn test_find_closest_aspect_ratio_wide() { - let processor = Phi4VisionProcessor::new(); - let ratios = processor.compute_target_ratios(1, 36); - - // Wide image (2:1 aspect ratio) - let result = processor.find_closest_aspect_ratio(2.0, &ratios, 1000, 500); - assert!(result.0 > result.1); // Width crops > height crops - } - - #[test] - fn test_find_closest_aspect_ratio_tall() { - let processor = Phi4VisionProcessor::new(); - let ratios = processor.compute_target_ratios(1, 36); - - // Tall image (1:2 aspect ratio) - let result = processor.find_closest_aspect_ratio(0.5, &ratios, 500, 1000); - assert!(result.0 < result.1); // Width crops < height crops - } - - #[test] - fn test_pad_image() { - let processor = Phi4VisionProcessor::new(); - let image = create_test_image(300, 200, Rgb([100, 100, 100])); - - let padded = processor.pad_image(&image, 448, 448); - assert_eq!(padded.width(), 448); - assert_eq!(padded.height(), 448); - - // Check original content is preserved - let p = padded.get_pixel(100, 100); - assert_eq!(p.0[0], 100); - - // Check padding is white - let p = padded.get_pixel(400, 400); - assert_eq!(p.0[0], 255); - } - - #[test] - fn test_preprocess_square_image() { - let processor = Phi4VisionProcessor::new(); - let config = PreProcessorConfig::default(); - - let image = create_test_image(500, 500, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - assert_eq!(result.batch_size(), 1); - assert!(result.feature_token_counts[0] > 256); // At least global tokens - - // Check pixel values are normalized - let flat = result.encoder_input_flat(); - assert!(flat.iter().all(|&v| (-1.5..=1.5).contains(&v))); - } - - #[test] - fn test_preprocess_wide_image() { - let processor = Phi4VisionProcessor::new(); - let config = PreProcessorConfig::default(); - - let image = create_test_image(1000, 500, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - assert_eq!(result.batch_size(), 1); - // Wide image should have more crops in width direction - assert!(result.item_sizes[0].1 >= result.item_sizes[0].0); - } - - #[test] - fn test_preprocess_multiple_images() { - let processor = Phi4VisionProcessor::new(); - let config = PreProcessorConfig::default(); - - let images = vec![ - create_test_image(500, 500, Rgb([100, 100, 100])), - create_test_image(800, 400, Rgb([150, 150, 150])), - ]; - - let result = processor.preprocess(&images, &config).unwrap(); - - assert_eq!(result.batch_size(), 2); - assert_eq!(result.item_sizes.len(), 2); - assert_eq!(result.feature_token_counts.len(), 2); - } - - #[test] - fn test_model_name() { - let processor = Phi4VisionProcessor::new(); - assert_eq!(processor.model_name(), "phi4-vision"); - } - - #[test] - fn test_normalization_values() { - let processor = Phi4VisionProcessor::new(); - assert_eq!(processor.default_mean(), [0.5, 0.5, 0.5]); - assert_eq!(processor.default_std(), [0.5, 0.5, 0.5]); - } - - #[test] - fn test_phi4_vs_phi3_differences() { - // Verify key differences from Phi3 - let processor = Phi4VisionProcessor::new(); - - // Phi4 uses 448 base resolution (vs 336 in Phi3) - assert_eq!(processor.base_resolution(), 448); - - // Phi4 uses simple 0.5 normalization (vs CLIP in Phi3) - assert_eq!(processor.mean, [0.5, 0.5, 0.5]); - assert_eq!(processor.std, [0.5, 0.5, 0.5]); - - // Phi4 default dynamic_hd is 36 (vs 16 num_crops in Phi3) - assert_eq!(processor.dynamic_hd(), 36); - } -} diff --git a/crates/multimodal/src/vision/processors/pixtral.rs b/crates/multimodal/src/vision/processors/pixtral.rs deleted file mode 100644 index c414f5d91..000000000 --- a/crates/multimodal/src/vision/processors/pixtral.rs +++ /dev/null @@ -1,412 +0,0 @@ -//! Pixtral/Mistral3 Vision image processor implementation. -//! -//! This module implements the image preprocessing for Pixtral/Mistral3 models, -//! matching the behavior of HuggingFace's `PixtralImageProcessor`. -//! -//! Key characteristics: -//! - CLIP normalization: mean [0.48145466, 0.4578275, 0.40821073], std [0.26862954, 0.26130258, 0.27577711] -//! - Bicubic resampling for resize -//! - Images resized to fit within longest_edge (default 1024) -//! - Output dimensions are multiples of patch_size (default 16) -//! - No tiling - single image output per input - -use std::collections::HashMap; - -use image::{imageops::FilterType, DynamicImage}; -use ndarray::{Array4, IxDyn}; - -use crate::vision::{ - preprocessor_config::PreProcessorConfig, - processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::{self, TransformError}, -}; - -/// Default normalization mean values (CLIP) -const DEFAULT_IMAGE_MEAN: [f64; 3] = [0.48145466, 0.4578275, 0.40821073]; - -/// Default normalization std values (CLIP) -const DEFAULT_IMAGE_STD: [f64; 3] = [0.26862954, 0.26130258, 0.27577711]; - -/// Default longest edge for resize -const DEFAULT_LONGEST_EDGE: u32 = 1024; - -/// Default patch size -const DEFAULT_PATCH_SIZE: u32 = 16; - -/// Pixtral/Mistral3 Vision image processor. -/// -/// This processor handles image preprocessing for Pixtral and Mistral3 vision models. -/// Unlike tile-based processors (Phi3, LLaMA4), Pixtral processes images at their -/// natural resolution (up to a maximum), preserving aspect ratio. -#[derive(Debug, Clone)] -pub struct PixtralProcessor { - /// Maximum dimension for the longest edge - longest_edge: u32, - /// Patch size for calculating output dimensions - patch_size: u32, - /// Normalization mean values - image_mean: [f64; 3], - /// Normalization std values - image_std: [f64; 3], -} - -impl Default for PixtralProcessor { - fn default() -> Self { - Self::new() - } -} - -impl PixtralProcessor { - /// Creates a new Pixtral processor with default settings. - pub fn new() -> Self { - Self { - longest_edge: DEFAULT_LONGEST_EDGE, - patch_size: DEFAULT_PATCH_SIZE, - image_mean: DEFAULT_IMAGE_MEAN, - image_std: DEFAULT_IMAGE_STD, - } - } - - /// Creates a processor from a HuggingFace preprocessor config. - pub fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { - let longest_edge = config - .size - .as_ref() - .and_then(|s| s.get("longest_edge").copied()) - .unwrap_or(DEFAULT_LONGEST_EDGE); - - // Patch size uses the new PatchSize type from config - let patch_size = config.get_patch_size(DEFAULT_PATCH_SIZE as usize) as u32; - - let image_mean = config - .image_mean - .as_ref() - .filter(|m| m.len() >= 3) - .map(|m| [m[0], m[1], m[2]]) - .unwrap_or(DEFAULT_IMAGE_MEAN); - - let image_std = config - .image_std - .as_ref() - .filter(|s| s.len() >= 3) - .map(|s| [s[0], s[1], s[2]]) - .unwrap_or(DEFAULT_IMAGE_STD); - - Self { - longest_edge, - patch_size, - image_mean, - image_std, - } - } - - /// Calculates the target output size for an image. - /// - /// The image is resized to fit within `longest_edge` while preserving aspect ratio. - /// The output dimensions are then adjusted to be multiples of `patch_size`. - fn get_resize_output_size(&self, height: u32, width: u32) -> (u32, u32) { - let max_size = self.longest_edge; - let patch_size = self.patch_size; - - // Calculate ratio for scaling down (only if larger than max_size) - let ratio = f64::max( - height as f64 / max_size as f64, - width as f64 / max_size as f64, - ); - - let (new_height, new_width) = if ratio > 1.0 { - // Scale down using floor to ensure we don't exceed max_size - let new_height = (height as f64 / ratio).floor() as u32; - let new_width = (width as f64 / ratio).floor() as u32; - (new_height, new_width) - } else { - (height, width) - }; - - // Calculate number of patches in each dimension - // Using: num_tokens = (dim - 1) / patch_size + 1 (i.e., ceiling division) - let num_height_tokens = (new_height.max(1) - 1) / patch_size + 1; - let num_width_tokens = (new_width.max(1) - 1) / patch_size + 1; - - // Final size is patches * patch_size - ( - num_height_tokens * patch_size, - num_width_tokens * patch_size, - ) - } - - /// Processes a single image through the Pixtral pipeline. - fn process_single_image( - &self, - image: &DynamicImage, - ) -> Result<(Array4, (usize, usize)), TransformError> { - let (orig_width, orig_height) = (image.width(), image.height()); - - // Step 1: Calculate output size - let (target_h, target_w) = self.get_resize_output_size(orig_height, orig_width); - - // Step 2: Resize image using bicubic interpolation - let resized = transforms::resize(image, target_w, target_h, FilterType::CatmullRom); - - // Step 3: Convert to tensor (0-1 range) and normalize - let mut tensor = transforms::to_tensor(&resized); - transforms::normalize(&mut tensor, &self.image_mean, &self.image_std); - - // Step 4: Reshape to (1, C, H, W) - let (c, h, w) = (tensor.shape()[0], tensor.shape()[1], tensor.shape()[2]); - let output = tensor - .into_shape_with_order((1, c, h, w)) - .map_err(|e| TransformError::ShapeError(e.to_string()))?; - - Ok((output, (target_h as usize, target_w as usize))) - } -} - -impl VisionPreProcessor for PixtralProcessor { - fn default_mean(&self) -> [f64; 3] { - self.image_mean - } - - fn default_std(&self) -> [f64; 3] { - self.image_std - } - - fn preprocess( - &self, - images: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - if images.is_empty() { - return Err(TransformError::InvalidShape { - expected: "non-empty image batch".to_string(), - actual: vec![0], - }); - } - - // Apply config overrides if present - let processor = if config.size.is_some() - || config.patch_size.is_some() - || config.image_mean.is_some() - || config.image_std.is_some() - { - Self::from_preprocessor_config(config) - } else { - self.clone() - }; - - let mut all_pixel_values = Vec::new(); - let mut all_image_sizes = Vec::new(); - let mut original_sizes = Vec::new(); - let mut feature_token_counts = Vec::new(); - - for image in images { - let (pixels, size) = processor.process_single_image(image)?; - let tokens = processor.calculate_num_tokens(image.width(), image.height(), config); - - all_pixel_values.push(pixels); - all_image_sizes.push(size); - original_sizes.push((image.height(), image.width())); - feature_token_counts.push(tokens); - } - - // Pad images to the same size for batching - let max_height = all_image_sizes.iter().map(|(h, _)| *h).max().unwrap_or(0); - let max_width = all_image_sizes.iter().map(|(_, w)| *w).max().unwrap_or(0); - - // Create batch tensor with padding - let batch_size = all_pixel_values.len(); - let channels = 3; - let mut batch_tensor = - ndarray::ArrayD::::zeros(IxDyn(&[batch_size, channels, max_height, max_width])); - - for (i, (pixels, (h, w))) in all_pixel_values - .iter() - .zip(all_image_sizes.iter()) - .enumerate() - { - // Copy the image data into the batch (top-left aligned, zero-padded) - for c in 0..channels { - for y in 0..*h { - for x in 0..*w { - batch_tensor[[i, c, y, x]] = pixels[[0, c, y, x]]; - } - } - } - } - - // Store image sizes as model-specific data - let mut model_specific = HashMap::new(); - let image_sizes_flat: Vec = all_image_sizes - .iter() - .flat_map(|&(h, w)| vec![h as i64, w as i64]) - .collect(); - model_specific.insert( - "image_sizes".to_string(), - ModelSpecificValue::IntTensor { - data: image_sizes_flat, - shape: vec![batch_size, 2], - }, - ); - - Ok(PreprocessedEncoderInputs { - encoder_input: batch_tensor, - feature_token_counts, - item_sizes: original_sizes, - model_specific, - }) - } - - fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { - let processor = Self::from_preprocessor_config(config); - let (target_h, target_w) = processor.get_resize_output_size(height, width); - let patch_size = processor.patch_size; - - // Number of tokens = num_patches_h * num_patches_w - let num_patches_h = target_h / patch_size; - let num_patches_w = target_w / patch_size; - (num_patches_h * num_patches_w) as usize - } - - fn model_name(&self) -> &'static str { - "pixtral" - } - - fn get_processed_size(&self, _config: &PreProcessorConfig) -> Option<(u32, u32)> { - // Pixtral has dynamic size based on input - None - } -} - -#[cfg(test)] -mod tests { - use image::{Rgb, RgbImage}; - - use super::*; - use crate::vision::preprocessor_config::PatchSize; - - fn create_test_image(width: u32, height: u32) -> DynamicImage { - let mut img = RgbImage::new(width, height); - for y in 0..height { - for x in 0..width { - let r = ((x * 255) / width.max(1)) as u8; - let g = ((y * 255) / height.max(1)) as u8; - let b = (((x + y) * 128) / (width + height).max(1)) as u8; - img.put_pixel(x, y, Rgb([r, g, b])); - } - } - DynamicImage::ImageRgb8(img) - } - - #[test] - fn test_resize_output_size_small_image() { - let processor = PixtralProcessor::new(); - - // Small image that doesn't need resizing - just pad to patch boundary - // 100x100 -> patches: ceil(100/16) = 7, output: 7*16 = 112 - let (h, w) = processor.get_resize_output_size(100, 100); - assert_eq!((h, w), (112, 112)); - } - - #[test] - fn test_resize_output_size_large_image() { - let processor = PixtralProcessor::new(); - - // Large image that needs resizing - // 2048x1024: ratio = 2048/1024 = 2.0 - // scaled: 2048/2 = 1024, 1024/2 = 512 - // patches h: ceil(1024/16) = 64, patches w: ceil(512/16) = 32 - // output: 64*16 = 1024, 32*16 = 512 - let (h, w) = processor.get_resize_output_size(2048, 1024); - assert_eq!((h, w), (1024, 512)); - } - - #[test] - fn test_resize_output_size_at_limit() { - let processor = PixtralProcessor::new(); - - // Image exactly at limit - // 1024x768: ratio = max(1024/1024, 768/1024) = 1.0 - // No resize needed - // patches h: ceil(1024/16) = 64, patches w: ceil(768/16) = 48 - // output: 64*16 = 1024, 48*16 = 768 - let (h, w) = processor.get_resize_output_size(1024, 768); - assert_eq!((h, w), (1024, 768)); - } - - #[test] - fn test_process_single_image() { - let processor = PixtralProcessor::new(); - let image = create_test_image(200, 150); - - let (tensor, size) = processor.process_single_image(&image).unwrap(); - - // 200x150 -> patches h: ceil(150/16) = 10, patches w: ceil(200/16) = 13 - // output: 10*16 = 160, 13*16 = 208 - assert_eq!(size, (160, 208)); - assert_eq!(tensor.shape(), &[1, 3, 160, 208]); - } - - #[test] - fn test_preprocess_batch() { - let processor = PixtralProcessor::new(); - let config = PreProcessorConfig::default(); - - let images = vec![create_test_image(200, 150), create_test_image(300, 100)]; - - let result = processor.preprocess(&images, &config).unwrap(); - - // First image: 150x200 -> 160x208 - // Second image: 100x300 -> 112x304 (ceil(100/16)=7, ceil(300/16)=19) - // Batch padded to max: 160x304 - assert_eq!(result.encoder_input.shape()[0], 2); // batch size - assert_eq!(result.encoder_input.shape()[1], 3); // channels - } - - #[test] - fn test_normalization_values() { - let processor = PixtralProcessor::new(); - - // Verify CLIP normalization values - assert!((processor.image_mean[0] - 0.48145466).abs() < 1e-6); - assert!((processor.image_mean[1] - 0.4578275).abs() < 1e-6); - assert!((processor.image_mean[2] - 0.40821073).abs() < 1e-6); - - assert!((processor.image_std[0] - 0.26862954).abs() < 1e-6); - assert!((processor.image_std[1] - 0.26130258).abs() < 1e-6); - assert!((processor.image_std[2] - 0.27577711).abs() < 1e-6); - } - - #[test] - fn test_from_config() { - let mut size = HashMap::new(); - size.insert("longest_edge".to_string(), 2048u32); - - let config = PreProcessorConfig { - size: Some(size), - patch_size: Some(PatchSize { - height: Some(14), - width: Some(14), - }), - image_mean: Some(vec![0.5, 0.5, 0.5]), - image_std: Some(vec![0.5, 0.5, 0.5]), - ..Default::default() - }; - - let processor = PixtralProcessor::from_preprocessor_config(&config); - - assert_eq!(processor.longest_edge, 2048); - assert_eq!(processor.patch_size, 14); - assert_eq!(processor.image_mean, [0.5, 0.5, 0.5]); - assert_eq!(processor.image_std, [0.5, 0.5, 0.5]); - } - - #[test] - fn test_calculate_num_tokens() { - let processor = PixtralProcessor::new(); - let config = PreProcessorConfig::default(); - - // 200x150 -> 208x160 -> 13*10 = 130 patches - let tokens = processor.calculate_num_tokens(200, 150, &config); - assert_eq!(tokens, 130); - } -} diff --git a/crates/multimodal/src/vision/processors/qwen2_vl.rs b/crates/multimodal/src/vision/processors/qwen2_vl.rs deleted file mode 100644 index 7e2b5dca8..000000000 --- a/crates/multimodal/src/vision/processors/qwen2_vl.rs +++ /dev/null @@ -1,529 +0,0 @@ -//! Qwen2-VL family image processors. -//! -//! This module provides the Qwen2-VL processor which wraps the shared -//! `QwenVLProcessorBase` with Qwen2-VL specific default parameters. -//! -//! # Key Features -//! -//! - **Smart Resize**: Resizes images to fit within min/max pixel bounds while -//! preserving aspect ratio and aligning to patch boundaries -//! - **Dynamic Token Count**: Token count depends on actual image dimensions -//! - **image_grid_thw**: Returns (T, H, W) grid dimensions for position encoding -//! -//! # Qwen2-VL Parameters -//! -//! - patch_size: 14 -//! - merge_size: 2 -//! - factor: 28 (patch_size * merge_size) -//! - normalization: CLIP mean/std - -use std::ops::Deref; - -use image::DynamicImage; - -use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase, QwenVideoResizeMode}; -use crate::vision::{ - preprocessor_config::PreProcessorConfig, - processor::{PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::TransformError, -}; - -/// CLIP normalization mean values used by Qwen2-VL models. -pub const CLIP_MEAN: [f64; 3] = [0.48145466, 0.4578275, 0.40821073]; - -/// CLIP normalization std values used by Qwen2-VL models. -pub const CLIP_STD: [f64; 3] = [0.26862954, 0.26130258, 0.27577711]; - -/// Default minimum pixels (256 * 28 * 28 = 200,704) -pub const DEFAULT_MIN_PIXELS: usize = 256 * 28 * 28; - -/// Default maximum pixels (1280 * 28 * 28 = 1,003,520) -pub const DEFAULT_MAX_PIXELS: usize = 1280 * 28 * 28; - -/// Default patch size -pub const DEFAULT_PATCH_SIZE: usize = 14; - -/// Default merge size for token reduction -pub const DEFAULT_MERGE_SIZE: usize = 2; - -/// Default temporal patch size (for video frames) -pub const DEFAULT_TEMPORAL_PATCH_SIZE: usize = 2; - -/// Qwen2-VL image processor. -/// -/// This is a thin wrapper around `QwenVLProcessorBase` with Qwen2-VL -/// specific default parameters: -/// - patch_size: 14 -/// - merge_size: 2 -/// - CLIP normalization mean/std -#[derive(Debug, Clone)] -pub struct Qwen2VLProcessor { - inner: QwenVLProcessorBase, -} - -impl Default for Qwen2VLProcessor { - fn default() -> Self { - Self::new() - } -} - -impl Qwen2VLProcessor { - /// Create a new Qwen2-VL processor with default settings. - /// - /// Defaults: - /// - patch_size: 14 - /// - merge_size: 2 - /// - min_pixels: 200,704 (256 * 28 * 28) - /// - max_pixels: 1,003,520 (1280 * 28 * 28) - /// - temporal_patch_size: 2 - /// - normalization: CLIP mean/std - pub fn new() -> Self { - Self { - inner: QwenVLProcessorBase::new(QwenVLConfig { - patch_size: DEFAULT_PATCH_SIZE, - merge_size: DEFAULT_MERGE_SIZE, - min_pixels: DEFAULT_MIN_PIXELS, - max_pixels: DEFAULT_MAX_PIXELS, - video_min_pixels: DEFAULT_MIN_PIXELS, - video_max_pixels: DEFAULT_MAX_PIXELS, - video_resize_mode: QwenVideoResizeMode::TotalVolume, - temporal_patch_size: DEFAULT_TEMPORAL_PATCH_SIZE, - mean: CLIP_MEAN, - std: CLIP_STD, - model_name: "qwen2-vl", - }), - } - } - - /// Create a processor with custom settings. - pub fn with_config( - patch_size: usize, - merge_size: usize, - min_pixels: usize, - max_pixels: usize, - temporal_patch_size: usize, - ) -> Self { - Self { - inner: QwenVLProcessorBase::new(QwenVLConfig { - patch_size, - merge_size, - min_pixels, - max_pixels, - video_min_pixels: min_pixels, - video_max_pixels: max_pixels, - video_resize_mode: QwenVideoResizeMode::TotalVolume, - temporal_patch_size, - mean: CLIP_MEAN, - std: CLIP_STD, - model_name: "qwen2-vl", - }), - } - } - - /// Create a processor from preprocessor config. - pub fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { - Self { - inner: QwenVLProcessorBase::new(QwenVLConfig { - patch_size: config.get_patch_size(DEFAULT_PATCH_SIZE), - merge_size: config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), - min_pixels: config.min_pixels.unwrap_or(DEFAULT_MIN_PIXELS), - max_pixels: config.max_pixels.unwrap_or(DEFAULT_MAX_PIXELS), - video_min_pixels: config.min_pixels.unwrap_or(DEFAULT_MIN_PIXELS), - video_max_pixels: config.max_pixels.unwrap_or(DEFAULT_MAX_PIXELS), - video_resize_mode: QwenVideoResizeMode::TotalVolume, - temporal_patch_size: config - .temporal_patch_size - .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), - mean: CLIP_MEAN, - std: CLIP_STD, - model_name: "qwen2-vl", - }), - } - } - - /// Build the effective processor for a request, applying any structural - /// overrides from `config`; otherwise reuse the existing defaults. - fn with_preprocessor_config(&self, config: &PreProcessorConfig) -> Self { - if config.has_structural_overrides() { - Self::from_preprocessor_config(config) - } else { - self.clone() - } - } - - /// Get the patch size. - pub fn patch_size(&self) -> usize { - self.inner.patch_size() - } - - /// Get the merge size. - pub fn merge_size(&self) -> usize { - self.inner.merge_size() - } - - /// Get the minimum pixels. - pub fn min_pixels(&self) -> usize { - self.inner.min_pixels() - } - - /// Get the maximum pixels. - pub fn max_pixels(&self) -> usize { - self.inner.max_pixels() - } - - /// Get the temporal patch size. - pub fn temporal_patch_size(&self) -> usize { - self.inner.temporal_patch_size() - } - - /// Get the factor for dimension alignment. - #[inline] - pub fn get_factor(&self) -> usize { - self.inner.get_factor() - } - - /// Smart resize algorithm for Qwen2-VL. - pub fn smart_resize( - &self, - height: usize, - width: usize, - ) -> Result<(usize, usize), TransformError> { - self.inner.smart_resize(height, width) - } - - /// Calculate the grid dimensions (T, H, W) for an image. - pub fn calculate_grid_thw( - &self, - height: usize, - width: usize, - num_frames: usize, - ) -> (usize, usize, usize) { - self.inner.calculate_grid_thw(height, width, num_frames) - } - - /// Calculate the number of image tokens after merge. - pub fn calculate_tokens_from_grid(&self, grid_t: usize, grid_h: usize, grid_w: usize) -> usize { - self.inner - .calculate_tokens_from_grid(grid_t, grid_h, grid_w) - } -} - -impl Deref for Qwen2VLProcessor { - type Target = QwenVLProcessorBase; - - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -impl VisionPreProcessor for Qwen2VLProcessor { - 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 { - let processor = self.with_preprocessor_config(config); - processor.inner.preprocess(images, config) - } - - fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { - let processor = self.with_preprocessor_config(config); - processor.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::{preprocessor_config::PatchSize, processor::ModelSpecificValue}; - - fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { - DynamicImage::from(RgbImage::from_pixel(width, height, color)) - } - - #[test] - fn test_qwen2_vl_processor_default() { - let processor = Qwen2VLProcessor::new(); - assert_eq!(processor.patch_size(), 14); - assert_eq!(processor.merge_size(), 2); - assert_eq!(processor.min_pixels(), DEFAULT_MIN_PIXELS); - assert_eq!(processor.max_pixels(), DEFAULT_MAX_PIXELS); - assert_eq!(processor.get_factor(), 28); // 14 * 2 - } - - #[test] - fn test_smart_resize_within_bounds() { - let processor = Qwen2VLProcessor::new(); - - // Image that's already within bounds - let (h, w) = processor.smart_resize(500, 500).unwrap(); - - // Should be aligned to factor (28) - assert_eq!(h % 28, 0); - assert_eq!(w % 28, 0); - - // Should be within bounds - assert!(h * w >= processor.min_pixels()); - assert!(h * w <= processor.max_pixels()); - } - - #[test] - fn test_smart_resize_too_large() { - let processor = Qwen2VLProcessor::new(); - - // Very large image - let (h, w) = processor.smart_resize(3000, 3000).unwrap(); - - // Should be scaled down - assert!(h * w <= processor.max_pixels()); - assert_eq!(h % 28, 0); - assert_eq!(w % 28, 0); - } - - #[test] - fn test_smart_resize_too_small() { - let processor = Qwen2VLProcessor::new(); - - // Small image (but above minimum dimension) - let (h, w) = processor.smart_resize(100, 100).unwrap(); - - // Should be scaled up to min_pixels - assert!(h * w >= processor.min_pixels()); - assert_eq!(h % 28, 0); - assert_eq!(w % 28, 0); - } - - #[test] - fn test_smart_resize_aspect_ratio_preserved() { - let processor = Qwen2VLProcessor::new(); - - // 2:1 aspect ratio - let (h, w) = processor.smart_resize(400, 800).unwrap(); - - // Aspect ratio should be approximately preserved - let original_ratio = 800.0 / 400.0; - let new_ratio = w as f64 / h as f64; - assert!((new_ratio - original_ratio).abs() < 0.5); - } - - #[test] - fn test_smart_resize_extreme_aspect_ratio_error() { - let processor = Qwen2VLProcessor::new(); - - // 300:1 aspect ratio - should fail - let result = processor.smart_resize(100, 30000); - assert!(result.is_err()); - } - - #[test] - fn test_smart_resize_small_dimension_clamps_to_factor() { - let processor = Qwen2VLProcessor::new(); - - // Dimension smaller than factor (28) should be clamped up, not rejected - let (h, w) = processor.smart_resize(10, 100).unwrap(); - assert!(h >= 28); - assert!(w >= 28); - assert_eq!(h % 28, 0); - assert_eq!(w % 28, 0); - } - - #[test] - fn test_calculate_grid_thw_image() { - let processor = Qwen2VLProcessor::new(); - - // 448x448 image (16x16 grid patches) - let (t, h, w) = processor.calculate_grid_thw(448, 448, 1); - - assert_eq!(t, 1); // Single image - assert_eq!(h, 448 / 14); // 32 - assert_eq!(w, 448 / 14); // 32 - } - - #[test] - fn test_calculate_tokens() { - let processor = Qwen2VLProcessor::new(); - - // With merge_size=2, tokens = (t * h * w) / 4 - let tokens = processor.calculate_tokens_from_grid(1, 32, 32); - assert_eq!(tokens, (32 * 32) / 4); // 256 - } - - #[test] - fn test_qwen2_vl_preprocess() { - let processor = Qwen2VLProcessor::new(); - let config = PreProcessorConfig { - do_resize: Some(true), - do_normalize: Some(true), - image_mean: Some(CLIP_MEAN.to_vec()), - image_std: Some(CLIP_STD.to_vec()), - patch_size: Some(PatchSize { - height: Some(14), - width: Some(14), - }), - merge_size: Some(2), - min_pixels: Some(DEFAULT_MIN_PIXELS), - max_pixels: Some(DEFAULT_MAX_PIXELS), - ..Default::default() - }; - - let image = create_test_image(600, 400, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - // encoder_input is patchified: [total_patches, patch_features] - assert_eq!(result.encoder_input.ndim(), 2); - assert!(result.encoder_input.shape()[0] > 0); // total_patches > 0 - - // Check pixel values are normalized - let flat = result.encoder_input_flat(); - // After normalization with CLIP mean/std, gray (0.5) should be near 0 - // (0.5 - 0.48) / 0.27 ≈ 0.07 - assert!(flat.iter().all(|&v| v.abs() < 1.0)); // Should be normalized - - // Check image_grid_thw and patches_per_image are present - assert!(result.model_specific.contains_key("image_grid_thw")); - assert!(result.model_specific.contains_key("patches_per_image")); - - // Verify token count is reasonable - assert!(result.feature_token_counts[0] > 0); - } - - #[test] - fn test_qwen2_vl_preprocess_multiple() { - let processor = Qwen2VLProcessor::new(); - let config = PreProcessorConfig::default(); - - let images = vec![ - create_test_image(600, 400, Rgb([100, 100, 100])), - create_test_image(400, 600, Rgb([150, 150, 150])), - ]; - - let result = processor.preprocess(&images, &config).unwrap(); - - // Both images processed - assert_eq!(result.item_sizes.len(), 2); - assert_eq!(result.feature_token_counts.len(), 2); - - // encoder_input is 2D [total_patches, patch_features] - assert_eq!(result.encoder_input.ndim(), 2); - - // Check grid_thw shape - if let Some(ModelSpecificValue::IntTensor { data, shape }) = - result.model_specific.get("image_grid_thw") - { - assert_eq!(shape, &[2, 3]); // 2 images, 3 values (T, H, W) each - assert_eq!(data.len(), 6); - } else { - panic!("Expected image_grid_thw to be IntTensor"); - } - - // Check patches_per_image - if let Some(ModelSpecificValue::IntTensor { data, shape }) = - result.model_specific.get("patches_per_image") - { - assert_eq!(shape, &[2]); // 2 images - assert_eq!(data.len(), 2); - let total: i64 = data.iter().sum(); - assert_eq!(total as usize, result.encoder_input.shape()[0]); - } else { - panic!("Expected patches_per_image to be IntTensor"); - } - } - - #[test] - fn test_qwen2_vl_from_config() { - let config = PreProcessorConfig { - patch_size: Some(PatchSize { - height: Some(16), - width: Some(16), - }), - merge_size: Some(4), - min_pixels: Some(100000), - max_pixels: Some(500000), - temporal_patch_size: Some(4), - ..Default::default() - }; - - let processor = Qwen2VLProcessor::from_preprocessor_config(&config); - - assert_eq!(processor.patch_size(), 16); - assert_eq!(processor.merge_size(), 4); - assert_eq!(processor.min_pixels(), 100000); - assert_eq!(processor.max_pixels(), 500000); - assert_eq!(processor.temporal_patch_size(), 4); - } - - #[test] - fn test_calculate_num_tokens_honors_config_max_pixels() { - let processor = Qwen2VLProcessor::new(); - - // A 1400x1400 image clamps to max_pixels before the grid is computed, - // so a lower config max_pixels must yield fewer tokens. - let default_tokens = - processor.calculate_num_tokens(1400, 1400, &PreProcessorConfig::default()); - assert_eq!(default_tokens, 1225); // resized to 980x980 -> (70*70)/4 - - let config = PreProcessorConfig { - max_pixels: Some(512 * 28 * 28), // 401,408, below the 1,003,520 default - ..Default::default() - }; - let config_tokens = processor.calculate_num_tokens(1400, 1400, &config); - assert_eq!(config_tokens, 484); // resized to 616x616 -> (44*44)/4 - assert_ne!( - config_tokens, default_tokens, - "config max_pixels override must change the token count" - ); - } - - #[test] - fn test_preprocess_honors_config_max_pixels() { - let processor = Qwen2VLProcessor::new(); - let image = create_test_image(1400, 1400, Rgb([128, 128, 128])); - - let config = PreProcessorConfig { - max_pixels: Some(512 * 28 * 28), - ..Default::default() - }; - let result = processor.preprocess(&[image], &config).unwrap(); - - // 616x616 -> grid (1, 44, 44) -> (44*44)/4 = 484 tokens, vs 1225 at the default. - assert_eq!(result.feature_token_counts[0], 484); - if let Some(ModelSpecificValue::IntTensor { data, .. }) = - result.model_specific.get("image_grid_thw") - { - assert_eq!(data, &[1, 44, 44]); - } else { - panic!("Expected image_grid_thw to be IntTensor"); - } - } - - #[test] - fn test_model_name() { - let processor = Qwen2VLProcessor::new(); - assert_eq!(processor.model_name(), "qwen2-vl"); - } - - #[test] - fn test_default_mean_std() { - let processor = Qwen2VLProcessor::new(); - assert_eq!(processor.default_mean(), CLIP_MEAN); - assert_eq!(processor.default_std(), CLIP_STD); - } -} diff --git a/crates/multimodal/src/vision/processors/qwen3_omni_vision.rs b/crates/multimodal/src/vision/processors/qwen3_omni_vision.rs deleted file mode 100644 index e295494af..000000000 --- a/crates/multimodal/src/vision/processors/qwen3_omni_vision.rs +++ /dev/null @@ -1,295 +0,0 @@ -//! Qwen3-Omni image and video preprocessing. -//! -//! This processor shares patchification and normalization through -//! `QwenVLProcessorBase`, but cannot use `Qwen3VLProcessor` directly. Omni -//! applies video pixel limits per frame rather than across the sampled clip, -//! gives its video preprocessor config precedence for video-specific limits, -//! and emits timing metadata required by mixed-modality M-RoPE. - -use image::DynamicImage; - -use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase, QwenVideoResizeMode}; -use crate::{ - types::RgbFrameRef, - vision::{ - preprocessor_config::PreProcessorConfig, - processor::{PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::TransformError, - }, -}; - -pub const QWEN3_OMNI_MEAN: [f64; 3] = [0.5; 3]; -pub const QWEN3_OMNI_STD: [f64; 3] = [0.5; 3]; -pub const DEFAULT_IMAGE_MIN_PIXELS: usize = 3_136; -pub const DEFAULT_IMAGE_MAX_PIXELS: usize = 12_845_056; -pub const DEFAULT_VIDEO_MIN_PIXELS: usize = 128 * 32 * 32; -pub const DEFAULT_VIDEO_MAX_PIXELS: usize = 768 * 32 * 32; -pub const DEFAULT_PATCH_SIZE: usize = 16; -pub const DEFAULT_MERGE_SIZE: usize = 2; -pub const DEFAULT_TEMPORAL_PATCH_SIZE: usize = 2; - -#[derive(Debug, Clone)] -pub struct Qwen3OmniVisionProcessor { - inner: QwenVLProcessorBase, -} - -impl Default for Qwen3OmniVisionProcessor { - fn default() -> Self { - Self::new() - } -} - -impl Qwen3OmniVisionProcessor { - pub fn new() -> Self { - Self::with_limits( - DEFAULT_IMAGE_MIN_PIXELS, - DEFAULT_IMAGE_MAX_PIXELS, - DEFAULT_VIDEO_MIN_PIXELS, - DEFAULT_VIDEO_MAX_PIXELS, - DEFAULT_PATCH_SIZE, - DEFAULT_MERGE_SIZE, - DEFAULT_TEMPORAL_PATCH_SIZE, - ) - } - - fn with_limits( - image_min_pixels: usize, - image_max_pixels: usize, - video_min_pixels: usize, - video_max_pixels: usize, - patch_size: usize, - merge_size: usize, - temporal_patch_size: usize, - ) -> Self { - Self { - inner: QwenVLProcessorBase::new(QwenVLConfig { - patch_size, - merge_size, - min_pixels: image_min_pixels, - max_pixels: image_max_pixels, - video_min_pixels, - video_max_pixels, - video_resize_mode: QwenVideoResizeMode::PerFrame, - temporal_patch_size, - mean: QWEN3_OMNI_MEAN, - std: QWEN3_OMNI_STD, - model_name: "qwen3-omni", - }), - } - } - - fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { - let configured_min = config.min_pixels.or_else(|| config.get_shortest_edge()); - let configured_max = config.max_pixels.or_else(|| config.get_longest_edge()); - Self::with_limits( - configured_min.unwrap_or(DEFAULT_IMAGE_MIN_PIXELS), - configured_max.unwrap_or(DEFAULT_IMAGE_MAX_PIXELS), - DEFAULT_VIDEO_MIN_PIXELS, - DEFAULT_VIDEO_MAX_PIXELS, - config.get_patch_size(DEFAULT_PATCH_SIZE), - config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), - config - .temporal_patch_size - .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), - ) - } - - fn from_video_preprocessor_config(config: &PreProcessorConfig) -> Self { - let configured_min = config.min_pixels.or_else(|| config.get_shortest_edge()); - let configured_max = config.max_pixels.or_else(|| config.get_longest_edge()); - Self::with_limits( - DEFAULT_IMAGE_MIN_PIXELS, - DEFAULT_IMAGE_MAX_PIXELS, - configured_min.unwrap_or(DEFAULT_VIDEO_MIN_PIXELS), - configured_max.unwrap_or(DEFAULT_VIDEO_MAX_PIXELS), - config.get_patch_size(DEFAULT_PATCH_SIZE), - config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), - config - .temporal_patch_size - .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), - ) - } - - fn with_image_preprocessor_config(&self, config: &PreProcessorConfig) -> Self { - if config.has_structural_overrides() { - Self::from_preprocessor_config(config) - } else { - self.clone() - } - } - - fn with_video_preprocessor_config(&self, config: &PreProcessorConfig) -> Self { - if config.has_structural_overrides() { - if config.is_image_only_processor_type() { - // Qwen3-Omni's shared preprocessor_config.json carries image - // limits. The HF processor supplies separate video defaults at - // call time, so those image limits must not become a per-frame - // video budget here. - Self::from_preprocessor_config(config) - } else { - Self::from_video_preprocessor_config(config) - } - } else { - self.clone() - } - } -} - -impl VisionPreProcessor for Qwen3OmniVisionProcessor { - 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.with_image_preprocessor_config(config) - .inner - .preprocess(images, config) - } - - fn preprocess_video( - &self, - frames: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - self.with_video_preprocessor_config(config) - .inner - .preprocess_video(frames, config) - } - - fn preprocess_video_rgb( - &self, - frames: &[RgbFrameRef<'_>], - config: &PreProcessorConfig, - ) -> Result { - self.with_video_preprocessor_config(config) - .inner - .preprocess_video_rgb(frames, config) - } - - fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { - self.with_image_preprocessor_config(config) - .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::{DynamicImage, RgbImage}; - - use super::*; - use crate::vision::{processor::ModelSpecificValue, processors::Qwen3VLProcessor}; - - #[test] - fn omni_multiframe_resize_differs_from_qwen3_volume_budget() { - let omni = Qwen3OmniVisionProcessor::new(); - let qwen3_vl = Qwen3VLProcessor::with_config( - DEFAULT_PATCH_SIZE, - DEFAULT_MERGE_SIZE, - DEFAULT_VIDEO_MIN_PIXELS, - DEFAULT_VIDEO_MAX_PIXELS, - DEFAULT_TEMPORAL_PATCH_SIZE, - ); - - assert_eq!(omni.inner.min_pixels(), DEFAULT_IMAGE_MIN_PIXELS); - assert_eq!(omni.inner.max_pixels(), DEFAULT_IMAGE_MAX_PIXELS); - assert_eq!(omni.inner.video_min_pixels(), DEFAULT_VIDEO_MIN_PIXELS); - assert_eq!(omni.inner.video_max_pixels(), DEFAULT_VIDEO_MAX_PIXELS); - assert_eq!( - omni.inner.video_resize_mode(), - QwenVideoResizeMode::PerFrame - ); - - let omni_size = omni.inner.smart_resize_video(16, 720, 1280).unwrap(); - let volume_size = qwen3_vl.smart_resize_video(16, 720, 1280).unwrap(); - - assert_eq!(omni_size, (640, 1152)); - assert_eq!(volume_size, (160, 288)); - assert_ne!(omni_size, volume_size); - } - - #[test] - fn video_preprocessor_config_overrides_per_frame_limits() { - let config = PreProcessorConfig::from_json( - r#"{"size":{"shortest_edge":65536,"longest_edge":262144},"temporal_patch_size":2}"#, - ) - .unwrap(); - let processor = Qwen3OmniVisionProcessor::from_video_preprocessor_config(&config); - - assert_eq!(processor.inner.min_pixels(), DEFAULT_IMAGE_MIN_PIXELS); - assert_eq!(processor.inner.max_pixels(), DEFAULT_IMAGE_MAX_PIXELS); - assert_eq!(processor.inner.video_min_pixels(), 65_536); - assert_eq!(processor.inner.video_max_pixels(), 262_144); - assert_eq!( - processor.inner.smart_resize_video(32, 720, 1280).unwrap(), - (384, 672) - ); - } - - #[test] - fn shared_image_config_keeps_omni_video_defaults() { - let config = PreProcessorConfig::from_json( - r#"{"image_processor_type":"Qwen2VLImageProcessor","min_pixels":3136,"max_pixels":12845056,"patch_size":16,"merge_size":2,"temporal_patch_size":2}"#, - ) - .unwrap(); - let processor = Qwen3OmniVisionProcessor::new().with_video_preprocessor_config(&config); - - assert_eq!(processor.inner.video_min_pixels(), DEFAULT_VIDEO_MIN_PIXELS); - assert_eq!(processor.inner.video_max_pixels(), DEFAULT_VIDEO_MAX_PIXELS); - assert_eq!( - processor.inner.smart_resize_video(16, 720, 1280).unwrap(), - (640, 1152) - ); - } - - #[test] - fn empty_config_uses_omni_half_normalization() { - let image = DynamicImage::ImageRgb8(RgbImage::new(32, 32)); - let output = Qwen3OmniVisionProcessor::new() - .preprocess(&[image], &PreProcessorConfig::default()) - .unwrap(); - - assert!(output - .encoder_input - .iter() - .all(|value| (*value + 1.0).abs() < 1e-6)); - } - - #[test] - fn sampled_video_fps_controls_mrope_grid_timing() { - let config = PreProcessorConfig::from_json( - r#"{"size":{"shortest_edge":1024,"longest_edge":65536},"fps":4.0}"#, - ) - .unwrap(); - let frames = vec![ - DynamicImage::ImageRgb8(RgbImage::new(32, 32)), - DynamicImage::ImageRgb8(RgbImage::new(32, 32)), - ]; - - let output = Qwen3OmniVisionProcessor::new() - .preprocess_video(&frames, &config) - .unwrap(); - - assert!(matches!( - output.model_specific.get("video_second_per_grid"), - Some(ModelSpecificValue::Tensor { data, shape }) - if data == &vec![0.5] && shape == &vec![1] - )); - } -} diff --git a/crates/multimodal/src/vision/processors/qwen3_vl.rs b/crates/multimodal/src/vision/processors/qwen3_vl.rs deleted file mode 100644 index 4f62424c6..000000000 --- a/crates/multimodal/src/vision/processors/qwen3_vl.rs +++ /dev/null @@ -1,793 +0,0 @@ -//! Qwen3-VL family image processors. -//! -//! This module provides the Qwen3-VL processor which wraps the shared -//! `QwenVLProcessorBase` with Qwen3-VL specific default parameters. -//! -//! # Key Differences from Qwen2-VL -//! -//! - **Patch Size**: 16 (vs 14 in Qwen2-VL) -//! - **Factor**: 32 (patch_size * merge_size) (vs 28 in Qwen2-VL) -//! - **Normalization**: [0.5, 0.5, 0.5] mean/std (vs CLIP in Qwen2-VL) -//! -//! # Qwen3-VL Parameters -//! -//! - patch_size: 16 -//! - merge_size: 2 -//! - factor: 32 (patch_size * merge_size) -//! - normalization: [0.5, 0.5, 0.5] mean/std - -use std::ops::Deref; - -use image::DynamicImage; - -use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase, QwenVideoResizeMode}; -use crate::{ - types::RgbFrameRef, - vision::{ - preprocessor_config::PreProcessorConfig, - processor::{PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::TransformError, - }, -}; - -/// Qwen3-VL normalization mean values (simple [0.5, 0.5, 0.5]). -pub const QWEN3_MEAN: [f64; 3] = [0.5, 0.5, 0.5]; - -/// Qwen3-VL normalization std values (simple [0.5, 0.5, 0.5]). -pub const QWEN3_STD: [f64; 3] = [0.5, 0.5, 0.5]; - -/// Default minimum pixels for Qwen3-VL -/// This corresponds to shortest_edge = 65536 from HF config -pub const DEFAULT_MIN_PIXELS: usize = 65536; - -/// Default maximum pixels for Qwen3-VL -/// This corresponds to longest_edge = 16777216 from HF config -pub const DEFAULT_MAX_PIXELS: usize = 16777216; - -/// Default minimum pixels for a complete Qwen3-VL video volume. -/// This corresponds to shortest_edge = 4096 from the HF video config. -pub const DEFAULT_VIDEO_MIN_PIXELS: usize = 4096; - -/// Default maximum pixels for a complete Qwen3-VL video volume. -/// This corresponds to longest_edge = 25165824 from the HF video config. -pub const DEFAULT_VIDEO_MAX_PIXELS: usize = 25165824; - -/// Default patch size for Qwen3-VL (16, vs 14 in Qwen2-VL) -pub const DEFAULT_PATCH_SIZE: usize = 16; - -/// Default merge size for token reduction -pub const DEFAULT_MERGE_SIZE: usize = 2; - -/// Default temporal patch size (for video frames) -pub const DEFAULT_TEMPORAL_PATCH_SIZE: usize = 2; - -/// Qwen3-VL image processor. -/// -/// This is a thin wrapper around `QwenVLProcessorBase` with Qwen3-VL -/// specific default parameters: -/// - patch_size: 16 -/// - merge_size: 2 -/// - [0.5, 0.5, 0.5] normalization mean/std -#[derive(Debug, Clone)] -pub struct Qwen3VLProcessor { - inner: QwenVLProcessorBase, -} - -impl Default for Qwen3VLProcessor { - fn default() -> Self { - Self::new() - } -} - -impl Qwen3VLProcessor { - /// Create a new Qwen3-VL processor with default settings. - /// - /// Defaults: - /// - patch_size: 16 - /// - merge_size: 2 - /// - min_pixels: 65,536 - /// - max_pixels: 16,777,216 - /// - temporal_patch_size: 2 - /// - normalization: [0.5, 0.5, 0.5] mean/std - pub fn new() -> Self { - Self::with_limits( - DEFAULT_MIN_PIXELS, - DEFAULT_MAX_PIXELS, - DEFAULT_VIDEO_MIN_PIXELS, - DEFAULT_VIDEO_MAX_PIXELS, - DEFAULT_PATCH_SIZE, - DEFAULT_MERGE_SIZE, - DEFAULT_TEMPORAL_PATCH_SIZE, - ) - } - - fn with_limits( - image_min_pixels: usize, - image_max_pixels: usize, - video_min_pixels: usize, - video_max_pixels: usize, - patch_size: usize, - merge_size: usize, - temporal_patch_size: usize, - ) -> Self { - Self { - inner: QwenVLProcessorBase::new(QwenVLConfig { - patch_size, - merge_size, - min_pixels: image_min_pixels, - max_pixels: image_max_pixels, - video_min_pixels, - video_max_pixels, - video_resize_mode: QwenVideoResizeMode::TotalVolume, - temporal_patch_size, - mean: QWEN3_MEAN, - std: QWEN3_STD, - model_name: "qwen3-vl", - }), - } - } - - /// Create a processor with custom settings. - pub fn with_config( - patch_size: usize, - merge_size: usize, - min_pixels: usize, - max_pixels: usize, - temporal_patch_size: usize, - ) -> Self { - Self::with_limits( - min_pixels, - max_pixels, - min_pixels, - max_pixels, - patch_size, - merge_size, - temporal_patch_size, - ) - } - - /// Create a processor from preprocessor config. - pub fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { - let configured_min = config.min_pixels.or_else(|| config.get_shortest_edge()); - let configured_max = config.max_pixels.or_else(|| config.get_longest_edge()); - let min_pixels = configured_min.unwrap_or(DEFAULT_MIN_PIXELS); - let max_pixels = configured_max.unwrap_or(DEFAULT_MAX_PIXELS); - Self::with_limits( - min_pixels, - max_pixels, - min_pixels, - max_pixels, - config.get_patch_size(DEFAULT_PATCH_SIZE), - config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), - config - .temporal_patch_size - .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), - ) - } - - fn from_image_preprocessor_config(config: &PreProcessorConfig) -> Self { - let configured_min = config.min_pixels.or_else(|| config.get_shortest_edge()); - let configured_max = config.max_pixels.or_else(|| config.get_longest_edge()); - Self::with_limits( - configured_min.unwrap_or(DEFAULT_MIN_PIXELS), - configured_max.unwrap_or(DEFAULT_MAX_PIXELS), - DEFAULT_VIDEO_MIN_PIXELS, - DEFAULT_VIDEO_MAX_PIXELS, - config.get_patch_size(DEFAULT_PATCH_SIZE), - config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), - config - .temporal_patch_size - .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), - ) - } - - fn from_video_preprocessor_config(config: &PreProcessorConfig) -> Self { - let configured_min = config.min_pixels.or_else(|| config.get_shortest_edge()); - let configured_max = config.max_pixels.or_else(|| config.get_longest_edge()); - Self::with_limits( - DEFAULT_MIN_PIXELS, - DEFAULT_MAX_PIXELS, - configured_min.unwrap_or(DEFAULT_VIDEO_MIN_PIXELS), - configured_max.unwrap_or(DEFAULT_VIDEO_MAX_PIXELS), - config.get_patch_size(DEFAULT_PATCH_SIZE), - config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), - config - .temporal_patch_size - .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), - ) - } - - fn with_preprocessor_config(&self, config: &PreProcessorConfig) -> Self { - if config.has_structural_overrides() { - Self::from_image_preprocessor_config(config) - } else { - self.clone() - } - } - - fn with_video_preprocessor_config(&self, config: &PreProcessorConfig) -> Self { - if !config.has_structural_overrides() { - return self.clone(); - } - if config.is_image_only_processor_type() { - Self::from_image_preprocessor_config(config) - } else { - Self::from_video_preprocessor_config(config) - } - } - - /// Get the patch size. - pub fn patch_size(&self) -> usize { - self.inner.patch_size() - } - - /// Get the merge size. - pub fn merge_size(&self) -> usize { - self.inner.merge_size() - } - - /// Get the minimum pixels. - pub fn min_pixels(&self) -> usize { - self.inner.min_pixels() - } - - /// Get the maximum pixels. - pub fn max_pixels(&self) -> usize { - self.inner.max_pixels() - } - - /// Get the temporal patch size. - pub fn temporal_patch_size(&self) -> usize { - self.inner.temporal_patch_size() - } - - /// Get the factor for dimension alignment. - #[inline] - pub fn get_factor(&self) -> usize { - self.inner.get_factor() - } - - /// Smart resize algorithm for Qwen3-VL. - pub fn smart_resize( - &self, - height: usize, - width: usize, - ) -> Result<(usize, usize), TransformError> { - self.inner.smart_resize(height, width) - } - - /// Calculate the grid dimensions (T, H, W) for an image. - pub fn calculate_grid_thw( - &self, - height: usize, - width: usize, - num_frames: usize, - ) -> (usize, usize, usize) { - self.inner.calculate_grid_thw(height, width, num_frames) - } - - /// Calculate the number of image tokens after merge. - pub fn calculate_tokens_from_grid(&self, grid_t: usize, grid_h: usize, grid_w: usize) -> usize { - self.inner - .calculate_tokens_from_grid(grid_t, grid_h, grid_w) - } -} - -impl Deref for Qwen3VLProcessor { - type Target = QwenVLProcessorBase; - - fn deref(&self) -> &Self::Target { - &self.inner - } -} - -impl VisionPreProcessor for Qwen3VLProcessor { - 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 { - let processor = self.with_preprocessor_config(config); - processor.inner.preprocess(images, config) - } - - fn preprocess_video( - &self, - frames: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - let processor = self.with_video_preprocessor_config(config); - processor.inner.preprocess_video(frames, config) - } - - fn preprocess_video_rgb( - &self, - frames: &[RgbFrameRef<'_>], - config: &PreProcessorConfig, - ) -> Result { - let processor = self.with_video_preprocessor_config(config); - processor.inner.preprocess_video_rgb(frames, config) - } - - fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { - let processor = self.with_preprocessor_config(config); - processor.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 std::collections::HashMap; - - use image::{Rgb, RgbImage}; - - use super::*; - 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)) - } - - #[test] - fn test_qwen3_vl_processor_default() { - let processor = Qwen3VLProcessor::new(); - assert_eq!(processor.patch_size(), 16); - assert_eq!(processor.merge_size(), 2); - assert_eq!(processor.min_pixels(), DEFAULT_MIN_PIXELS); - assert_eq!(processor.max_pixels(), DEFAULT_MAX_PIXELS); - assert_eq!(processor.inner.video_min_pixels(), DEFAULT_VIDEO_MIN_PIXELS); - assert_eq!(processor.inner.video_max_pixels(), DEFAULT_VIDEO_MAX_PIXELS); - assert_eq!(processor.get_factor(), 32); // 16 * 2 - } - - #[test] - fn test_with_config_preserves_shared_image_and_video_limits() { - let processor = Qwen3VLProcessor::with_config(16, 2, 8192, 1_048_576, 2); - - assert_eq!(processor.min_pixels(), 8192); - assert_eq!(processor.max_pixels(), 1_048_576); - assert_eq!(processor.inner.video_min_pixels(), 8192); - assert_eq!(processor.inner.video_max_pixels(), 1_048_576); - } - - #[test] - fn test_smart_resize_within_bounds() { - let processor = Qwen3VLProcessor::new(); - - // Image that's within bounds - let (h, w) = processor.smart_resize(500, 500).unwrap(); - - // Should be aligned to factor (32) - assert_eq!(h % 32, 0); - assert_eq!(w % 32, 0); - - // Should be within bounds - assert!(h * w >= processor.min_pixels()); - assert!(h * w <= processor.max_pixels()); - } - - #[test] - fn test_smart_resize_aspect_ratio_preserved() { - let processor = Qwen3VLProcessor::new(); - - // 2:1 aspect ratio - let (h, w) = processor.smart_resize(400, 800).unwrap(); - - // Aspect ratio should be approximately preserved - let original_ratio = 800.0 / 400.0; - let new_ratio = w as f64 / h as f64; - assert!((new_ratio - original_ratio).abs() < 0.5); - } - - #[test] - fn test_smart_resize_extreme_aspect_ratio_error() { - let processor = Qwen3VLProcessor::new(); - - // 300:1 aspect ratio - should fail - let result = processor.smart_resize(100, 30000); - assert!(result.is_err()); - } - - #[test] - fn test_smart_resize_small_dimension_clamps_to_factor() { - let processor = Qwen3VLProcessor::new(); - - // Dimension smaller than factor (32) should be clamped up, not rejected - let (h, w) = processor.smart_resize(10, 100).unwrap(); - assert!(h >= 32); - assert!(w >= 32); - assert_eq!(h % 32, 0); - assert_eq!(w % 32, 0); - } - - #[test] - fn test_calculate_grid_thw_image() { - let processor = Qwen3VLProcessor::new(); - - // 480x640 image - let (t, h, w) = processor.calculate_grid_thw(480, 640, 1); - - assert_eq!(t, 1); // Single image - assert_eq!(h, 480 / 16); // 30 - assert_eq!(w, 640 / 16); // 40 - } - - #[test] - fn test_calculate_tokens() { - let processor = Qwen3VLProcessor::new(); - - // With merge_size=2, tokens = (t * h * w) / 4 - let tokens = processor.calculate_tokens_from_grid(1, 30, 40); - assert_eq!(tokens, (30 * 40) / 4); // 300 - } - - #[test] - fn test_qwen3_vl_preprocess() { - let processor = Qwen3VLProcessor::new(); - let config = PreProcessorConfig { - do_resize: Some(true), - do_normalize: Some(true), - image_mean: Some(QWEN3_MEAN.to_vec()), - image_std: Some(QWEN3_STD.to_vec()), - patch_size: Some(PatchSize { - height: Some(16), - width: Some(16), - }), - merge_size: Some(2), - min_pixels: Some(DEFAULT_MIN_PIXELS), - max_pixels: Some(DEFAULT_MAX_PIXELS), - ..Default::default() - }; - - let image = create_test_image(640, 480, Rgb([128, 128, 128])); - let result = processor.preprocess(&[image], &config).unwrap(); - - // encoder_input is patchified: [total_patches, patch_features] - assert_eq!(result.encoder_input.ndim(), 2); - assert!(result.encoder_input.shape()[0] > 0); // total_patches > 0 - - // Check pixel values are normalized - let flat = result.encoder_input_flat(); - // After normalization with [0.5, 0.5, 0.5] mean/std: - // (0.5 - 0.5) / 0.5 = 0.0 for gray - // Values should be in [-1, 1] range - assert!(flat.iter().all(|&v| (-1.5..=1.5).contains(&v))); - - // Check image_grid_thw and patches_per_image are present - assert!(result.model_specific.contains_key("image_grid_thw")); - assert!(result.model_specific.contains_key("patches_per_image")); - - // Verify token count is reasonable - assert!(result.feature_token_counts[0] > 0); - } - - #[test] - fn test_qwen3_vl_preprocess_multiple() { - let processor = Qwen3VLProcessor::new(); - let config = PreProcessorConfig { - image_mean: Some(QWEN3_MEAN.to_vec()), - image_std: Some(QWEN3_STD.to_vec()), - ..Default::default() - }; - - let images = vec![ - create_test_image(640, 480, Rgb([100, 100, 100])), - create_test_image(480, 640, Rgb([150, 150, 150])), - ]; - - let result = processor.preprocess(&images, &config).unwrap(); - - // Both images processed - assert_eq!(result.item_sizes.len(), 2); - assert_eq!(result.feature_token_counts.len(), 2); - - // encoder_input is 2D [total_patches, patch_features] - assert_eq!(result.encoder_input.ndim(), 2); - - // Check grid_thw shape - if let Some(ModelSpecificValue::IntTensor { data, shape }) = - result.model_specific.get("image_grid_thw") - { - assert_eq!(shape, &[2, 3]); // 2 images, 3 values (T, H, W) each - assert_eq!(data.len(), 6); - } else { - panic!("Expected image_grid_thw to be IntTensor"); - } - - // Check patches_per_image - if let Some(ModelSpecificValue::IntTensor { data, shape }) = - result.model_specific.get("patches_per_image") - { - assert_eq!(shape, &[2]); // 2 images - assert_eq!(data.len(), 2); - // Total patches should match encoder_input first dim - let total: i64 = data.iter().sum(); - assert_eq!(total as usize, result.encoder_input.shape()[0]); - } else { - panic!("Expected patches_per_image to be IntTensor"); - } - } - - // Per-image-independence guard for the smg gateway pixel cache (EPD P1): the - // cache stores one entry per image and reassembles requests from per-image - // payloads, which is only sound if preprocessing image i is independent of the - // other images in the batch. This asserts that a batched preprocess equals the - // concatenation of per-image preprocesses (encoder_input rows, grid_thw rows, - // and feature_token_counts). If a processor ever introduces cross-image work, - // this fails and the cache assumption must be revisited. - #[test] - fn per_image_preprocess_equals_batched_slices() { - use ndarray::{Axis, Slice}; - - let processor = Qwen3VLProcessor::new(); - let config = PreProcessorConfig { - image_mean: Some(QWEN3_MEAN.to_vec()), - image_std: Some(QWEN3_STD.to_vec()), - ..Default::default() - }; - - let image_a = create_test_image(640, 480, Rgb([100, 110, 120])); - let image_b = create_test_image(420, 560, Rgb([10, 200, 90])); - - let batched = processor - .preprocess(&[image_a.clone(), image_b.clone()], &config) - .unwrap(); - let single_a = processor.preprocess(&[image_a], &config).unwrap(); - let single_b = processor.preprocess(&[image_b], &config).unwrap(); - - // Feature token counts line up per image. - assert_eq!( - batched.feature_token_counts, - vec![ - single_a.feature_token_counts[0], - single_b.feature_token_counts[0] - ] - ); - - // encoder_input rows: batch == [single_a rows ++ single_b rows]. - let pa = single_a.encoder_input.shape()[0]; - let pb = single_b.encoder_input.shape()[0]; - assert_eq!(batched.encoder_input.shape()[0], pa + pb); - assert_eq!( - batched - .encoder_input - .slice_axis(Axis(0), Slice::from(0..pa)) - .to_owned(), - single_a.encoder_input - ); - assert_eq!( - batched - .encoder_input - .slice_axis(Axis(0), Slice::from(pa..pa + pb)) - .to_owned(), - single_b.encoder_input - ); - - // image_grid_thw rows: batch == [single_a row ++ single_b row]. - let grid = |inputs: &PreprocessedEncoderInputs| match inputs - .model_specific - .get("image_grid_thw") - { - Some(ModelSpecificValue::IntTensor { data, .. }) => data.clone(), - other => panic!("expected image_grid_thw IntTensor, got {other:?}"), - }; - let mut expected_grid = grid(&single_a); - expected_grid.extend(grid(&single_b)); - assert_eq!(grid(&batched), expected_grid); - } - - #[test] - fn test_qwen3_vl_preprocess_video() { - let processor = Qwen3VLProcessor::new(); - let config = PreProcessorConfig { - image_mean: Some(QWEN3_MEAN.to_vec()), - image_std: Some(QWEN3_STD.to_vec()), - ..Default::default() - }; - - let frames = vec![ - create_test_image(640, 480, Rgb([100, 100, 100])), - create_test_image(640, 480, Rgb([150, 150, 150])), - create_test_image(640, 480, Rgb([200, 200, 200])), - ]; - - let result = processor.preprocess_video(&frames, &config).unwrap(); - assert_eq!(result.encoder_input.ndim(), 2); - assert_eq!(result.feature_token_counts.len(), 1); - assert!(result.model_specific.contains_key("video_grid_thw")); - - if let Some(ModelSpecificValue::IntTensor { data, shape }) = - result.model_specific.get("video_grid_thw") - { - assert_eq!(shape, &[1, 3]); - assert_eq!(data[0], 2); // 3 frames padded to 4, temporal_patch_size=2 - } else { - panic!("Expected video_grid_thw to be IntTensor"); - } - assert!(matches!( - result.model_specific.get("video_second_per_grid"), - Some(ModelSpecificValue::Tensor { data, shape }) - if data == &vec![1.0] && shape == &vec![1] - )); - } - - #[test] - fn test_qwen3_vl_preprocess_video_rgb_applies_config() { - let processor = Qwen3VLProcessor::new(); - let config = PreProcessorConfig { - patch_size: Some(PatchSize { - height: Some(8), - width: Some(8), - }), - merge_size: Some(1), - temporal_patch_size: Some(4), - min_pixels: Some(1), - max_pixels: Some(4096), - ..Default::default() - }; - - let frames = vec![ - create_test_image(32, 32, Rgb([100, 100, 100])), - create_test_image(32, 32, Rgb([150, 150, 150])), - create_test_image(32, 32, Rgb([200, 200, 200])), - ]; - let rgb_images: Vec = frames.iter().map(|frame| frame.to_rgb8()).collect(); - let rgb_frames: Vec> = rgb_images - .iter() - .map(|frame| RgbFrameRef { - width: frame.width(), - height: frame.height(), - data: frame.as_raw(), - }) - .collect(); - - let dynamic = processor.preprocess_video(&frames, &config).unwrap(); - let rgb = processor - .preprocess_video_rgb(&rgb_frames, &config) - .unwrap(); - - assert_eq!(rgb.encoder_input.shape(), dynamic.encoder_input.shape()); - assert_eq!(rgb.feature_token_counts, dynamic.feature_token_counts); - let Some(ModelSpecificValue::IntTensor { - data: rgb_grid, - shape: rgb_shape, - }) = rgb.model_specific.get("video_grid_thw") - else { - panic!("Expected RGB video_grid_thw to be IntTensor"); - }; - let Some(ModelSpecificValue::IntTensor { - data: dynamic_grid, - shape: dynamic_shape, - }) = dynamic.model_specific.get("video_grid_thw") - else { - panic!("Expected dynamic video_grid_thw to be IntTensor"); - }; - assert_eq!(rgb_shape, dynamic_shape); - assert_eq!(rgb_grid, dynamic_grid); - } - - #[test] - fn test_qwen3_vl_from_config() { - let config = PreProcessorConfig { - patch_size: Some(PatchSize { - height: Some(16), - width: Some(16), - }), - merge_size: Some(4), - min_pixels: Some(100000), - max_pixels: Some(500000), - temporal_patch_size: Some(4), - ..Default::default() - }; - - let processor = Qwen3VLProcessor::from_preprocessor_config(&config); - - assert_eq!(processor.patch_size(), 16); - assert_eq!(processor.merge_size(), 4); - assert_eq!(processor.min_pixels(), 100000); - assert_eq!(processor.max_pixels(), 500000); - assert_eq!(processor.temporal_patch_size(), 4); - assert_eq!(processor.inner.video_min_pixels(), 100000); - assert_eq!(processor.inner.video_max_pixels(), 500000); - } - - #[test] - fn test_qwen3_vl_video_config_uses_size_edges() { - let config = PreProcessorConfig { - size: Some(HashMap::from([ - ("shortest_edge".to_string(), 4096), - ("longest_edge".to_string(), 25165824), - ])), - ..Default::default() - }; - - let processor = Qwen3VLProcessor::new().with_video_preprocessor_config(&config); - - assert_eq!(processor.min_pixels(), DEFAULT_MIN_PIXELS); - assert_eq!(processor.max_pixels(), DEFAULT_MAX_PIXELS); - assert_eq!(processor.inner.video_min_pixels(), 4096); - assert_eq!(processor.inner.video_max_pixels(), 25165824); - assert_eq!( - processor.inner.smart_resize_video(239, 720, 1280).unwrap(), - (224, 416) - ); - } - - #[test] - fn test_qwen3_vl_image_config_keeps_video_defaults() { - let config = PreProcessorConfig { - image_processor_type: Some("Qwen3VLImageProcessor".to_string()), - min_pixels: Some(100000), - max_pixels: Some(500000), - ..Default::default() - }; - - let processor = Qwen3VLProcessor::new().with_video_preprocessor_config(&config); - - assert_eq!(processor.min_pixels(), 100000); - assert_eq!(processor.max_pixels(), 500000); - assert_eq!(processor.inner.video_min_pixels(), DEFAULT_VIDEO_MIN_PIXELS); - assert_eq!(processor.inner.video_max_pixels(), DEFAULT_VIDEO_MAX_PIXELS); - } - - #[test] - fn test_model_name() { - let processor = Qwen3VLProcessor::new(); - assert_eq!(processor.model_name(), "qwen3-vl"); - } - - #[test] - fn test_default_mean_std() { - let processor = Qwen3VLProcessor::new(); - assert_eq!(processor.default_mean(), QWEN3_MEAN); - assert_eq!(processor.default_std(), QWEN3_STD); - } - - #[test] - fn test_qwen3_vs_qwen2_differences() { - // Verify the key differences from Qwen2-VL - let processor = Qwen3VLProcessor::new(); - - // Qwen3-VL uses patch_size=16 (vs 14 in Qwen2) - assert_eq!(processor.patch_size(), 16); - - // Factor is 32 (vs 28 in Qwen2) - assert_eq!(processor.get_factor(), 32); - - // Mean/std are [0.5, 0.5, 0.5] (vs CLIP values in Qwen2) - assert_eq!(processor.default_mean(), [0.5, 0.5, 0.5]); - assert_eq!(processor.default_std(), [0.5, 0.5, 0.5]); - } - - #[test] - fn test_smart_resize_grayscale_400x300() { - // grayscale.jpg is 400x300 - // 400/32 = 12.5 -> rounds to 12 (banker's rounding) -> 384 - // 300/32 = 9.375 -> rounds to 9 -> 288 - // Expected: 384x288, giving grid [1, 18, 24] - let processor = Qwen3VLProcessor::new(); - - // smart_resize takes (height, width) - let (h, w) = processor.smart_resize(300, 400).unwrap(); - - // Expected from HuggingFace: 288x384 -> grid [1, 18, 24] - assert_eq!(h, 288, "Height should be 288"); - assert_eq!(w, 384, "Width should be 384"); - } -} diff --git a/crates/multimodal/src/vision/processors/qwen_vl_base.rs b/crates/multimodal/src/vision/processors/qwen_vl_base.rs deleted file mode 100644 index c76336af5..000000000 --- a/crates/multimodal/src/vision/processors/qwen_vl_base.rs +++ /dev/null @@ -1,1918 +0,0 @@ -//! Shared base implementation for Qwen VL family image processors. -//! -//! This module provides a generic processor that handles the common logic -//! for Qwen2-VL, Qwen2.5-VL, and Qwen3-VL models. The specific variants -//! differ only in their default parameters (patch_size, normalization values). -//! -//! # Processing Pipeline -//! -//! 1. Validate aspect ratio (must be < 200:1) -//! 2. Smart resize to fit within min/max pixel bounds -//! 3. Align dimensions to (patch_size * merge_size) boundary -//! 4. Convert to tensor and normalize -//! 5. Reshape into patches for the vision encoder -//! -//! # Token Calculation -//! -//! ```text -//! grid_t = 1 (for images, temporal dimension is 1) -//! grid_h = resized_height / patch_size -//! grid_w = resized_width / patch_size -//! num_tokens = (grid_t * grid_h * grid_w) / merge_size² -//! ``` - -use std::borrow::Cow; - -use image::{imageops::FilterType, DynamicImage, GenericImageView}; -use ndarray::{Array2, Array3}; - -use crate::{ - types::RgbFrameRef, - vision::{ - execution::{scope as parallel_scope, task_count}, - preprocessor_config::PreProcessorConfig, - processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::{ - par_threads, pil_to_filter, resize, resize_bicubic_pil, resize_bicubic_pil_rgb, - resize_rgb_bytes, rgb_bytes, TransformError, - }, - }, -}; - -/// Python-compatible rounding (banker's rounding / round half to even). -/// -/// This matches Python's `round()` behavior where 0.5 is rounded to the nearest -/// even number (e.g. `12.5 -> 12`, `13.5 -> 14`), unlike Rust's `f64::round()` -/// which rounds half away from zero. -#[inline] -fn round_half_to_even(x: f64) -> f64 { - let rounded = x.round(); - // Check if we're exactly at a .5 case - if (x - x.floor() - 0.5).abs() < 1e-9 { - // Round to nearest even - if rounded as i64 % 2 != 0 { - return rounded - 1.0; - } - } - rounded -} - -/// Configuration for a Qwen VL processor variant. -#[derive(Debug, Clone)] -pub struct QwenVLConfig { - /// Vision encoder patch size - pub patch_size: usize, - /// Merge size for token reduction - pub merge_size: usize, - /// Minimum total pixels allowed - pub min_pixels: usize, - /// Maximum total pixels allowed - pub max_pixels: usize, - /// Minimum video pixels, interpreted according to `video_resize_mode`. - pub video_min_pixels: usize, - /// Maximum video pixels, interpreted according to `video_resize_mode`. - pub video_max_pixels: usize, - /// Whether the video budget applies per frame or to the sampled volume. - pub video_resize_mode: QwenVideoResizeMode, - /// Temporal patch size for video - pub temporal_patch_size: usize, - /// Normalization mean values - pub mean: [f64; 3], - /// Normalization std values - pub std: [f64; 3], - /// Model name for identification - pub model_name: &'static str, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum QwenVideoResizeMode { - TotalVolume, - PerFrame, -} - -#[derive(Clone)] -struct VideoFrameRgb<'a> { - width: usize, - height: usize, - data: Cow<'a, [u8]>, -} - -struct QwenImagePlan { - target_width: u32, - target_height: u32, - needs_resize: bool, - grid_t: usize, - grid_h: usize, - grid_w: usize, - num_patches: usize, - patch_values: usize, - tokens: usize, -} - -struct QwenVideoPlan { - original_size: (u32, u32), - target_width: u32, - target_height: u32, - grid_t: usize, - grid_h: usize, - grid_w: usize, - patch_features: usize, - num_patches: usize, - output_values: usize, - tokens: usize, - second_per_grid: f32, - filter: FilterType, - do_resize: bool, - lut: [[f32; 256]; 3], -} - -fn normalization_lut( - config: &PreProcessorConfig, - default_mean: [f64; 3], - default_std: [f64; 3], -) -> [[f32; 256]; 3] { - let mean = config - .image_mean - .as_ref() - .filter(|values| values.len() >= 3) - .map(|values| [values[0], values[1], values[2]]) - .unwrap_or(default_mean); - let std = config - .image_std - .as_ref() - .filter(|values| values.len() >= 3) - .map(|values| [values[0], values[1], values[2]]) - .unwrap_or(default_std); - let do_normalize = config.do_normalize.unwrap_or(true); - let scale: [f32; 3] = if do_normalize { - std::array::from_fn(|channel| 1.0 / (255.0 * std[channel] as f32)) - } else { - [1.0 / 255.0; 3] - }; - let bias: [f32; 3] = if do_normalize { - std::array::from_fn(|channel| -(mean[channel] as f32) / std[channel] as f32) - } else { - [0.0; 3] - }; - std::array::from_fn(|channel| { - std::array::from_fn(|value| value as f32 * scale[channel] + bias[channel]) - }) -} - -fn dispatch_patch_blocks( - region: &mut [f32], - n_blocks: usize, - block_out: usize, - write_blocks: impl Fn(usize, &mut [f32]) + Sync, -) { - let nthreads = par_threads(size_of_val(region), n_blocks); - if nthreads <= 1 { - write_blocks(0, region); - return; - } - - let chunk_blocks = n_blocks.div_ceil(nthreads); - parallel_scope(|scope| { - let write_blocks = &write_blocks; - let mut rest = region; - let mut block_start = 0usize; - while block_start < n_blocks { - let blocks = chunk_blocks.min(n_blocks - block_start); - let (band, tail) = rest.split_at_mut(blocks * block_out); - rest = tail; - let start = block_start; - scope.spawn(move |_| write_blocks(start, band)); - block_start += blocks; - } - }); -} - -fn resize_rgb_frame_to_raw( - frame: RgbFrameRef<'_>, - target_width: u32, - target_height: u32, - filter: FilterType, -) -> Result, TransformError> { - // BICUBIC (Qwen default) uses the PIL-compatible path, same as the image - // path; other filters keep the SIMD resizer. - let resized = if filter == FilterType::CatmullRom { - resize_bicubic_pil_rgb( - frame.data, - frame.width, - frame.height, - target_width, - target_height, - )? - } else { - resize_rgb_bytes( - frame.data, - frame.width, - frame.height, - target_width, - target_height, - filter, - )? - }; - Ok(resized.into_raw()) -} - -fn prepare_video_rgb_frame<'a>( - frame: RgbFrameRef<'a>, - target_width: u32, - target_height: u32, - filter: FilterType, - do_resize: bool, -) -> Result, TransformError> { - if do_resize && (frame.width != target_width || frame.height != target_height) { - Ok(VideoFrameRgb { - width: target_width as usize, - height: target_height as usize, - data: Cow::Owned(resize_rgb_frame_to_raw( - frame, - target_width, - target_height, - filter, - )?), - }) - } else { - Ok(VideoFrameRgb { - width: frame.width as usize, - height: frame.height as usize, - data: Cow::Borrowed(frame.data), - }) - } -} - -fn prepare_video_rgb_chunk<'a>( - frames: &[RgbFrameRef<'a>], - temporal_index: usize, - temporal_patch_size: usize, - target_width: u32, - target_height: u32, - filter: FilterType, - do_resize: bool, -) -> Result>, TransformError> { - let mut prepared = Vec::with_capacity(temporal_patch_size); - prepare_video_frame_chunk( - frames.len(), - temporal_index, - temporal_patch_size, - &mut prepared, - |frame_index| { - prepare_video_rgb_frame( - frames[frame_index], - target_width, - target_height, - filter, - do_resize, - ) - }, - )?; - Ok(prepared) -} - -fn prepare_video_frame_chunk<'a>( - frame_count: usize, - temporal_index: usize, - temporal_patch_size: usize, - prepared: &mut Vec>, - mut prepare_frame: impl FnMut(usize) -> Result, TransformError>, -) -> Result<(), TransformError> { - prepared.clear(); - let mut previous_frame_index = None; - for temporal_offset in 0..temporal_patch_size { - let frame_index = - (temporal_index * temporal_patch_size + temporal_offset).min(frame_count - 1); - if previous_frame_index == Some(frame_index) { - if let Some(previous) = prepared.last().cloned() { - prepared.push(previous); - continue; - } - } - prepared.push(prepare_frame(frame_index)?); - previous_frame_index = Some(frame_index); - } - Ok(()) -} - -fn resize_dynamic_frame_to_raw( - frame: &DynamicImage, - target_width: u32, - target_height: u32, - filter: FilterType, -) -> (usize, usize, Vec) { - let resized = if filter == FilterType::CatmullRom { - resize_bicubic_pil(frame, target_width, target_height) - } else { - resize(frame, target_width, target_height, filter) - }; - let (width, height, data) = rgb_bytes(&resized); - (width, height, data.into_owned()) -} - -/// Generic Qwen VL image processor. -/// -/// This struct implements the shared preprocessing logic for all Qwen VL -/// model variants. Each variant (Qwen2-VL, Qwen3-VL, etc.) uses this with -/// different configuration values. -#[derive(Debug, Clone)] -pub struct QwenVLProcessorBase { - config: QwenVLConfig, -} - -impl QwenVLProcessorBase { - /// Create a new processor with the given configuration. - pub fn new(config: QwenVLConfig) -> Self { - Self { config } - } - - /// Get the patch size. - pub fn patch_size(&self) -> usize { - self.config.patch_size - } - - /// Get the merge size. - pub fn merge_size(&self) -> usize { - self.config.merge_size - } - - /// Get the minimum pixels. - pub fn min_pixels(&self) -> usize { - self.config.min_pixels - } - - /// Get the maximum pixels. - pub fn max_pixels(&self) -> usize { - self.config.max_pixels - } - - pub fn video_min_pixels(&self) -> usize { - self.config.video_min_pixels - } - - pub fn video_max_pixels(&self) -> usize { - self.config.video_max_pixels - } - - pub fn video_resize_mode(&self) -> QwenVideoResizeMode { - self.config.video_resize_mode - } - - /// Get the temporal patch size. - pub fn temporal_patch_size(&self) -> usize { - self.config.temporal_patch_size - } - - fn plan_video( - &self, - frame_count: usize, - width: u32, - height: u32, - config: &PreProcessorConfig, - ) -> Result { - let temporal_patch_size = self.config.temporal_patch_size; - let padded_frames = frame_count.div_ceil(temporal_patch_size) * temporal_patch_size; - let (target_height, target_width) = - self.smart_resize_video(frame_count, height as usize, width as usize)?; - let (grid_t, grid_h, grid_w) = - self.calculate_grid_thw(target_height, target_width, padded_frames); - let patch_features = - 3 * temporal_patch_size * self.config.patch_size * self.config.patch_size; - let num_patches = grid_t - .checked_mul(grid_h) - .and_then(|value| value.checked_mul(grid_w)) - .ok_or_else(|| { - TransformError::ShapeError(format!( - "Qwen video patch count overflow: grid=({grid_t}, {grid_h}, {grid_w})" - )) - })?; - let output_values = num_patches.checked_mul(patch_features).ok_or_else(|| { - TransformError::ShapeError(format!( - "Qwen video patch buffer size overflow: patches={num_patches}, features={patch_features}" - )) - })?; - // MediaConnector samples video at 2 fps by default. A checkpoint may - // override that value in video_preprocessor_config.json. - let sample_fps = config.get_extra::("fps").unwrap_or(2.0); - if !sample_fps.is_finite() || sample_fps <= 0.0 { - return Err(TransformError::ShapeError(format!( - "Qwen video fps must be finite and positive, got {sample_fps}" - ))); - } - - Ok(QwenVideoPlan { - original_size: (width, height), - target_width: target_width as u32, - target_height: target_height as u32, - grid_t, - grid_h, - grid_w, - patch_features, - num_patches, - output_values, - tokens: self.calculate_tokens_from_grid(grid_t, grid_h, grid_w), - second_per_grid: temporal_patch_size as f32 / sample_fps, - filter: pil_to_filter(config.resampling.or(Some(3))), - do_resize: config.do_resize.unwrap_or(true), - lut: normalization_lut(config, self.config.mean, self.config.std), - }) - } - - fn finish_video( - plan: QwenVideoPlan, - patches: Vec, - ) -> Result { - let encoder_input = - Array2::from_shape_vec((plan.num_patches, plan.patch_features), patches).map_err( - |error| { - TransformError::ShapeError(format!( - "Failed to create video encoder_input [{}, {}]: {error}", - plan.num_patches, plan.patch_features - )) - }, - )?; - - Ok(PreprocessedEncoderInputs::new( - encoder_input, - vec![plan.tokens], - vec![plan.original_size], - ) - .with_extra( - "video_grid_thw", - ModelSpecificValue::int_2d( - vec![plan.grid_t as i64, plan.grid_h as i64, plan.grid_w as i64], - 1, - 3, - ), - ) - .with_extra( - "patches_per_video", - ModelSpecificValue::int_1d(vec![plan.num_patches as i64]), - ) - .with_extra( - "patches_per_image", - ModelSpecificValue::int_1d(vec![plan.num_patches as i64]), - ) - .with_extra( - "video_second_per_grid", - ModelSpecificValue::Tensor { - data: vec![plan.second_per_grid], - shape: vec![1], - }, - )) - } - - /// Get the factor for dimension alignment. - /// - /// Dimensions must be divisible by (patch_size * merge_size). - #[inline] - pub fn get_factor(&self) -> usize { - self.config.patch_size * self.config.merge_size - } - - /// Smart resize algorithm for Qwen VL models. - /// - /// Resizes image dimensions to fit within min/max pixel bounds while: - /// - Preserving aspect ratio - /// - Aligning to (patch_size * merge_size) boundaries - /// - /// # Arguments - /// * `height` - Original image height - /// * `width` - Original image width - /// - /// # Returns - /// (new_height, new_width) or error if aspect ratio is too extreme - /// - /// # Errors - /// - If height or width is zero - /// - If aspect ratio exceeds 200:1 - pub fn smart_resize( - &self, - height: usize, - width: usize, - ) -> Result<(usize, usize), TransformError> { - let factor = self.get_factor(); - - // Validate non-zero dimensions - if height == 0 || width == 0 { - return Err(TransformError::InvalidShape { - expected: "non-zero dimensions".to_string(), - actual: vec![height, width], - }); - } - - // Validate aspect ratio - let max_dim = height.max(width) as f64; - let min_dim = height.min(width) as f64; - let aspect_ratio = max_dim / min_dim; - if aspect_ratio > 200.0 { - return Err(TransformError::InvalidShape { - expected: "aspect ratio < 200:1".to_string(), - actual: vec![height, width], - }); - } - - // Round to nearest factor multiple using Python-compatible rounding - // Python uses banker's rounding (round half to even), which affects - // edge cases like 400/32 = 12.5 -> 12 (not 13) - let mut h_bar = round_half_to_even(height as f64 / factor as f64) as usize * factor; - let mut w_bar = round_half_to_even(width as f64 / factor as f64) as usize * factor; - - // Ensure minimum size - h_bar = h_bar.max(factor); - w_bar = w_bar.max(factor); - - // Scale down if exceeding max_pixels - if h_bar * w_bar > self.config.max_pixels { - let beta = ((height * width) as f64 / self.config.max_pixels as f64).sqrt(); - h_bar = ((height as f64 / beta / factor as f64).floor() as usize) * factor; - w_bar = ((width as f64 / beta / factor as f64).floor() as usize) * factor; - // Ensure minimum size after scaling down - h_bar = h_bar.max(factor); - w_bar = w_bar.max(factor); - } - // Scale up if below min_pixels - else if h_bar * w_bar < self.config.min_pixels { - let beta = (self.config.min_pixels as f64 / (height * width) as f64).sqrt(); - h_bar = ((height as f64 * beta / factor as f64).ceil() as usize) * factor; - w_bar = ((width as f64 * beta / factor as f64).ceil() as usize) * factor; - } - - Ok((h_bar, w_bar)) - } - - /// Smart resize for Qwen3-style video processors. - /// - /// `TotalVolume` applies the pixel budget to the padded sampled video - /// volume (`T * H * W`), while `PerFrame` applies it to each frame's - /// spatial area (`H * W`). - pub fn smart_resize_video( - &self, - num_frames: usize, - height: usize, - width: usize, - ) -> Result<(usize, usize), TransformError> { - let factor = self.get_factor(); - - if num_frames == 0 { - return Err(TransformError::InvalidShape { - expected: "num_frames > 0".to_string(), - actual: vec![num_frames], - }); - } - - if height < factor || width < factor { - return Err(TransformError::InvalidShape { - expected: format!("height and width >= factor ({factor})"), - actual: vec![height, width], - }); - } - - let max_dim = height.max(width) as f64; - let min_dim = height.min(width) as f64; - let aspect_ratio = max_dim / min_dim; - if aspect_ratio > 200.0 { - return Err(TransformError::InvalidShape { - expected: "aspect ratio < 200:1".to_string(), - actual: vec![height, width], - }); - } - - let mut h_bar = round_half_to_even(height as f64 / factor as f64) as usize * factor; - let mut w_bar = round_half_to_even(width as f64 / factor as f64) as usize * factor; - h_bar = h_bar.max(factor); - w_bar = w_bar.max(factor); - - let (budget_scale, resized_pixels) = match self.config.video_resize_mode { - QwenVideoResizeMode::TotalVolume => { - let padded_frames = num_frames.div_ceil(self.config.temporal_patch_size) - * self.config.temporal_patch_size; - ( - num_frames as f64, - padded_frames as f64 * h_bar as f64 * w_bar as f64, - ) - } - QwenVideoResizeMode::PerFrame => (1.0, h_bar as f64 * w_bar as f64), - }; - let source_pixels = budget_scale * height as f64 * width as f64; - if resized_pixels > self.config.video_max_pixels as f64 { - let beta = (source_pixels / self.config.video_max_pixels as f64).sqrt(); - h_bar = ((height as f64 / beta / factor as f64).floor() as usize) * factor; - w_bar = ((width as f64 / beta / factor as f64).floor() as usize) * factor; - h_bar = h_bar.max(factor); - w_bar = w_bar.max(factor); - } else if resized_pixels < self.config.video_min_pixels as f64 { - let beta = (self.config.video_min_pixels as f64 / source_pixels).sqrt(); - h_bar = ((height as f64 * beta / factor as f64).ceil() as usize) * factor; - w_bar = ((width as f64 * beta / factor as f64).ceil() as usize) * factor; - } - - Ok((h_bar, w_bar)) - } - - /// Calculate the grid dimensions (T, H, W) for an image. - /// - /// For single images, T=1. For video, T = num_frames / temporal_patch_size. - /// - /// # Arguments - /// * `height` - Resized image height - /// * `width` - Resized image width - /// * `num_frames` - Number of frames (1 for images) - /// - /// # Returns - /// (grid_t, grid_h, grid_w) - pub fn calculate_grid_thw( - &self, - height: usize, - width: usize, - num_frames: usize, - ) -> (usize, usize, usize) { - let grid_t = - num_frames.max(self.config.temporal_patch_size) / self.config.temporal_patch_size; - let grid_h = height / self.config.patch_size; - let grid_w = width / self.config.patch_size; - (grid_t, grid_h, grid_w) - } - - /// Calculate the number of image tokens after merge. - /// - /// tokens = (grid_t * grid_h * grid_w) / merge_size² - pub fn calculate_tokens_from_grid(&self, grid_t: usize, grid_h: usize, grid_w: usize) -> usize { - (grid_t * grid_h * grid_w) / (self.config.merge_size * self.config.merge_size) - } - - /// Patchify tensor directly into an output buffer (avoids intermediate Vec allocation). - /// Patchify a [C, H, W] tensor and append the patches to `output`. - /// - /// Output layout per image: - /// `[grid_t, patch_rows, patch_cols, merge_h, merge_w, C, temporal, patch_h, patch_w]` - /// - /// Each "merged patch" covers a `(merge_size * patch_size)²` spatial region. - /// Within it, `merge_size²` sub-patches are emitted, each containing all channels. - pub fn patchify_into( - &self, - tensor: &Array3, - grid_t: usize, - grid_h: usize, - grid_w: usize, - output: &mut Vec, - ) -> Result<(), TransformError> { - let channel = tensor.shape()[0]; - let height = tensor.shape()[1]; - let width = tensor.shape()[2]; - let patch_size = self.config.patch_size; - let merge_size = self.config.merge_size; - let temporal_patch_size = self.config.temporal_patch_size; - - debug_assert_eq!( - height, - grid_h * patch_size, - "Height must match grid_h * patch_size" - ); - debug_assert_eq!( - width, - grid_w * patch_size, - "Width must match grid_w * patch_size" - ); - - let num_patches = grid_t * grid_h * grid_w; - let patch_features = channel * temporal_patch_size * patch_size * patch_size; - let base_idx = output.len(); - output.resize(base_idx + num_patches * patch_features, 0.0); - - let data = tensor.as_standard_layout(); - let flat = data.as_slice().ok_or_else(|| { - TransformError::ShapeError("tensor not contiguous after as_standard_layout".to_string()) - })?; - let planes: Vec<&[f32]> = (0..channel) - .map(|c| &flat[c * height * width..(c + 1) * height * width]) - .collect(); - - let merged_patch = merge_size * patch_size; - let pr_blocks = grid_h / merge_size; - let pc_blocks = grid_w / merge_size; - let n_blocks = grid_t * pr_blocks * pc_blocks; - let block_out = merge_size * merge_size * patch_features; - // Each (gt,pr,pc) block writes a contiguous, deterministic output - // region of pure copies, so banding blocks across threads is - // BIT-IDENTICAL. Small grids stay serial. - let region = &mut output[base_idx..base_idx + n_blocks * block_out]; - dispatch_patch_blocks(region, n_blocks, block_out, |block_start, band| { - Self::patchify_block_band( - &planes, - width, - patch_size, - merge_size, - temporal_patch_size, - merged_patch, - pr_blocks, - pc_blocks, - block_start, - band, - ); - }); - - Ok(()) - } - - /// Fill `band` with the patchified output for blocks - /// `[block_start, block_start + band.len()/block_out)` in (gt, pr, pc) - /// row-major order. Pure gather/copy from `planes`; deterministic and - /// independent per block (safe to call concurrently on disjoint bands). - #[expect( - clippy::too_many_arguments, - reason = "block-band patchifier: planes + grid dims + output band" - )] - fn patchify_block_band( - planes: &[&[f32]], - width: usize, - patch_size: usize, - merge_size: usize, - temporal_patch_size: usize, - merged_patch: usize, - pr_blocks: usize, - pc_blocks: usize, - block_start: usize, - band: &mut [f32], - ) { - let block_out = - merge_size * merge_size * planes.len() * temporal_patch_size * patch_size * patch_size; - let per_t = pr_blocks * pc_blocks; - for (bi, chunk) in band.chunks_mut(block_out).enumerate() { - let blk = block_start + bi; - let rem = blk % per_t; - let pr = rem / pc_blocks; - let pc = rem % pc_blocks; - let y0 = pr * merged_patch; - let x0 = pc * merged_patch; - let mut o = 0usize; - for mh in 0..merge_size { - for mw in 0..merge_size { - for plane in planes { - for _tp in 0..temporal_patch_size { - for py in 0..patch_size { - let row = - (y0 + mh * patch_size + py) * width + x0 + mw * patch_size; - chunk[o..o + patch_size] - .copy_from_slice(&plane[row..row + patch_size]); - o += patch_size; - } - } - } - } - } - } - } - - /// Patchify a sequence of frame tensors into Qwen's video patch layout. - /// - /// `tensors` must already be resized and normalized to the same spatial - /// shape. If the frame count is not divisible by `temporal_patch_size`, - /// the caller should pad by repeating the final frame before calling this. - pub fn patchify_video_into( - &self, - tensors: &[Array3], - grid_t: usize, - grid_h: usize, - grid_w: usize, - output: &mut Vec, - ) -> Result<(), TransformError> { - if tensors.is_empty() { - return Err(TransformError::EmptyBatch); - } - - let channel = tensors[0].shape()[0]; - let height = tensors[0].shape()[1]; - let width = tensors[0].shape()[2]; - let patch_size = self.config.patch_size; - let merge_size = self.config.merge_size; - let temporal_patch_size = self.config.temporal_patch_size; - - debug_assert_eq!(height, grid_h * patch_size); - debug_assert_eq!(width, grid_w * patch_size); - debug_assert_eq!(tensors.len(), grid_t * temporal_patch_size); - - let num_patches = grid_t * grid_h * grid_w; - let patch_features = channel * temporal_patch_size * patch_size * patch_size; - let base_idx = output.len(); - output.resize(base_idx + num_patches * patch_features, 0.0); - - let frame_planes: Vec> = tensors - .iter() - .map(|tensor| { - let flat = tensor.as_slice().ok_or_else(|| { - TransformError::ShapeError("video frame tensor is not contiguous".to_string()) - })?; - Ok((0..channel) - .map(|c| &flat[c * height * width..(c + 1) * height * width]) - .collect::>()) - }) - .collect::>()?; - - let merged_patch = merge_size * patch_size; - let mut out_idx = base_idx; - - for gt in 0..grid_t { - let frame_start = gt * temporal_patch_size; - for pr in 0..grid_h / merge_size { - for pc in 0..grid_w / merge_size { - let y0 = pr * merged_patch; - let x0 = pc * merged_patch; - - for mh in 0..merge_size { - for mw in 0..merge_size { - let frame_window = - &frame_planes[frame_start..frame_start + temporal_patch_size]; - for channel_frames in (0..channel).map(|channel_idx| { - frame_window.iter().map(move |planes| planes[channel_idx]) - }) { - for plane in channel_frames { - for py in 0..patch_size { - let row = (y0 + mh * patch_size + py) * width - + x0 - + mw * patch_size; - output[out_idx..out_idx + patch_size] - .copy_from_slice(&plane[row..row + patch_size]); - out_idx += patch_size; - } - } - } - } - } - } - } - } - - Ok(()) - } - - fn patchify_video_rgb_chunk_into( - &self, - frames: &[VideoFrameRgb<'_>], - grid_h: usize, - grid_w: usize, - output: &mut [f32], - out_idx: &mut usize, - lut: &[[f32; 256]; 3], - ) -> Result<(), TransformError> { - let patch_size = self.config.patch_size; - let merge_size = self.config.merge_size; - let temporal_patch_size = self.config.temporal_patch_size; - if frames.len() != temporal_patch_size { - return Err(TransformError::InvalidShape { - expected: format!("{temporal_patch_size} video frames in temporal patch"), - actual: vec![frames.len()], - }); - } - - let height = grid_h * patch_size; - let width = grid_w * patch_size; - for frame in frames { - if frame.height != height || frame.width != width { - return Err(TransformError::InvalidShape { - expected: format!("video frame size {width}x{height}"), - actual: vec![frame.width, frame.height], - }); - } - } - - let merged_patch = merge_size * patch_size; - let pr_blocks = grid_h / merge_size; - let pc_blocks = grid_w / merge_size; - let n_blocks = pr_blocks.checked_mul(pc_blocks).ok_or_else(|| { - TransformError::ShapeError("Qwen video patch block count overflow".to_string()) - })?; - let block_out = merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; - let base_idx = *out_idx; - let patch_values = n_blocks.checked_mul(block_out).ok_or_else(|| { - TransformError::ShapeError("Qwen video patch output size overflow".to_string()) - })?; - let end_idx = base_idx.checked_add(patch_values).ok_or_else(|| { - TransformError::ShapeError("Qwen video patch output range overflow".to_string()) - })?; - let region = output.get_mut(base_idx..end_idx).ok_or_else(|| { - TransformError::ShapeError("Qwen video patch output range out of bounds".to_string()) - })?; - dispatch_patch_blocks(region, n_blocks, block_out, |block_start, band| { - Self::patchify_video_rgb_block_band( - frames, - width, - patch_size, - merge_size, - merged_patch, - pc_blocks, - block_start, - band, - lut, - ); - }); - *out_idx = end_idx; - - Ok(()) - } - - #[expect( - clippy::too_many_arguments, - reason = "RGB video patchifier: frame window + grid dims + output band" - )] - fn patchify_video_rgb_block_band( - frames: &[VideoFrameRgb<'_>], - width: usize, - patch_size: usize, - merge_size: usize, - merged_patch: usize, - pc_blocks: usize, - block_start: usize, - band: &mut [f32], - lut: &[[f32; 256]; 3], - ) { - let block_out = merge_size * merge_size * 3 * frames.len() * patch_size * patch_size; - for (bi, chunk) in band.chunks_mut(block_out).enumerate() { - let blk = block_start + bi; - let pr = blk / pc_blocks; - let pc = blk % pc_blocks; - let y0 = pr * merged_patch; - let x0 = pc * merged_patch; - let mut o = 0usize; - - for mh in 0..merge_size { - for mw in 0..merge_size { - for (c, lut_c) in lut.iter().enumerate().take(3) { - for frame in frames { - let raw = frame.data.as_ref(); - for py in 0..patch_size { - let row = - (y0 + mh * patch_size + py) * width + x0 + mw * patch_size; - let source_start = row * 3; - let source_end = (row + patch_size) * 3; - for (dst, pixel) in chunk[o..o + patch_size] - .iter_mut() - .zip(raw[source_start..source_end].chunks_exact(3)) - { - *dst = lut_c[pixel[c] as usize]; - } - o += patch_size; - } - } - } - } - } - } - } - - fn patchify_image_rgb_into( - &self, - image: &DynamicImage, - grid_h: usize, - grid_w: usize, - output: &mut [f32], - out_idx: &mut usize, - lut: &[[f32; 256]; 3], - ) -> Result<(), TransformError> { - let (width, height, data) = rgb_bytes(image); - let patch_size = self.config.patch_size; - let merge_size = self.config.merge_size; - let temporal_patch_size = self.config.temporal_patch_size; - let expected_height = grid_h * patch_size; - let expected_width = grid_w * patch_size; - if height != expected_height || width != expected_width { - return Err(TransformError::InvalidShape { - expected: format!("image size {expected_width}x{expected_height}"), - actual: vec![width, height], - }); - } - - let raw = data.as_ref(); - let merged_patch = merge_size * patch_size; - let pr_blocks = grid_h / merge_size; - let pc_blocks = grid_w / merge_size; - let n_blocks = pr_blocks.checked_mul(pc_blocks).ok_or_else(|| { - TransformError::ShapeError("Qwen image patch block count overflow".to_string()) - })?; - let block_out = merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; - let base_idx = *out_idx; - let patch_values = n_blocks.checked_mul(block_out).ok_or_else(|| { - TransformError::ShapeError("Qwen image patch output size overflow".to_string()) - })?; - let end_idx = base_idx.checked_add(patch_values).ok_or_else(|| { - TransformError::ShapeError("Qwen image patch output range overflow".to_string()) - })?; - let region = output.get_mut(base_idx..end_idx).ok_or_else(|| { - TransformError::ShapeError("Qwen image patch output range out of bounds".to_string()) - })?; - dispatch_patch_blocks(region, n_blocks, block_out, |block_start, band| { - Self::patchify_image_rgb_block_band( - raw, - width, - patch_size, - merge_size, - temporal_patch_size, - merged_patch, - pc_blocks, - block_start, - band, - lut, - ); - }); - *out_idx = end_idx; - - Ok(()) - } - - #[expect( - clippy::too_many_arguments, - reason = "RGB image patchifier: raw bytes + grid dims + output band" - )] - fn patchify_image_rgb_block_band( - raw: &[u8], - width: usize, - patch_size: usize, - merge_size: usize, - temporal_patch_size: usize, - merged_patch: usize, - pc_blocks: usize, - block_start: usize, - band: &mut [f32], - lut: &[[f32; 256]; 3], - ) { - let block_out = merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; - for (bi, chunk) in band.chunks_mut(block_out).enumerate() { - let blk = block_start + bi; - let pr = blk / pc_blocks; - let pc = blk % pc_blocks; - let y0 = pr * merged_patch; - let x0 = pc * merged_patch; - let mut o = 0usize; - - for mh in 0..merge_size { - for mw in 0..merge_size { - for (c, lut_c) in lut.iter().enumerate().take(3) { - for _tp in 0..temporal_patch_size { - for py in 0..patch_size { - let row = - (y0 + mh * patch_size + py) * width + x0 + mw * patch_size; - let mut src_idx = row * 3 + c; - for dst in &mut chunk[o..o + patch_size] { - *dst = lut_c[raw[src_idx] as usize]; - src_idx += 3; - } - o += patch_size; - } - } - } - } - } - } - } -} - -impl VisionPreProcessor for QwenVLProcessorBase { - 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); - } - - // Qwen2VL/Qwen3VL image processors default to BICUBIC (PIL resample=3) - // when the preprocessor config omits `resample`. The global pil_to_filter - // fallback is bilinear, which yields smoother features and measurably - // degrades VLM accuracy, so pin the HF-correct default here. - let filter = pil_to_filter(config.resampling.or(Some(3))); - - let patch_size = self.config.patch_size; - let temporal_patch_size = self.config.temporal_patch_size; - let patch_features = 3 * temporal_patch_size * patch_size * patch_size; - let do_resize = config.do_resize.unwrap_or(true); - let lut = normalization_lut(config, self.config.mean, self.config.std); - - let mut image_plans = Vec::with_capacity(images.len()); - let mut item_sizes = Vec::with_capacity(images.len()); - let mut total_patch_values = 0usize; - let mut total_patches = 0usize; - for image in images { - let (w, h) = image.dimensions(); - item_sizes.push((w, h)); - let (target_h, target_w) = self.smart_resize(h as usize, w as usize)?; - let (tw32, th32) = (target_w as u32, target_h as u32); - let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(target_h, target_w, 1); - let num_patches = grid_t - .checked_mul(grid_h) - .and_then(|value| value.checked_mul(grid_w)) - .ok_or_else(|| { - TransformError::ShapeError(format!( - "Qwen image patch count overflow: grid=({grid_t}, {grid_h}, {grid_w})" - )) - })?; - total_patches = total_patches.checked_add(num_patches).ok_or_else(|| { - TransformError::ShapeError("Qwen image total patch count overflow".to_string()) - })?; - let patch_values = num_patches.checked_mul(patch_features).ok_or_else(|| { - TransformError::ShapeError(format!( - "Qwen image patch buffer size overflow: patches={num_patches}, features={patch_features}" - )) - })?; - total_patch_values = total_patch_values - .checked_add(patch_values) - .ok_or_else(|| { - TransformError::ShapeError( - "Qwen image patch buffer total size overflow".to_string(), - ) - })?; - image_plans.push(QwenImagePlan { - target_width: tw32, - target_height: th32, - needs_resize: do_resize && (w != tw32 || h != th32), - grid_t, - grid_h, - grid_w, - num_patches, - patch_values, - tokens: self.calculate_tokens_from_grid(grid_t, grid_h, grid_w), - }); - } - - let mut all_patches: Vec = Vec::with_capacity(total_patch_values); - 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, plan) in images.iter().zip(image_plans) { - // Resize to the image's own target size (skip if dimensions match) - let resized; - let img_ref = if plan.needs_resize { - // BICUBIC (Qwen default) uses the PIL-compatible path; other - // filters keep the SIMD path. - resized = if filter == FilterType::CatmullRom { - resize_bicubic_pil(image, plan.target_width, plan.target_height) - } else { - resize(image, plan.target_width, plan.target_height, filter) - }; - &resized - } else { - image - }; - - grid_thw_data.push(plan.grid_t as i64); - grid_thw_data.push(plan.grid_h as i64); - grid_thw_data.push(plan.grid_w as i64); - - feature_token_counts.push(plan.tokens); - - // Patchify directly from RGB bytes to avoid the intermediate - // [C,H,W] tensor allocation. This matches the tensor path's - // channel/temporal/spatial order. - let base_idx = all_patches.len(); - all_patches.resize(base_idx + plan.patch_values, 0.0); - let mut out_idx = base_idx; - self.patchify_image_rgb_into( - img_ref, - plan.grid_h, - plan.grid_w, - &mut all_patches, - &mut out_idx, - &lut, - )?; - debug_assert_eq!(out_idx, all_patches.len()); - patches_per_image.push(plan.num_patches as i64); - } - - let encoder_input = - Array2::from_shape_vec((total_patches, patch_features), all_patches).map_err(|e| { - TransformError::ShapeError(format!( - "Failed to create patchified encoder_input [{total_patches}, {patch_features}]: {e}" - )) - })?; - - let result = - PreprocessedEncoderInputs::new(encoder_input, feature_token_counts, item_sizes) - .with_extra( - "image_grid_thw", - ModelSpecificValue::int_2d(grid_thw_data, images.len(), 3), - ) - .with_extra( - "patches_per_image", - ModelSpecificValue::int_1d(patches_per_image), - ); - - Ok(result) - } - - fn preprocess_video( - &self, - frames: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - if frames.is_empty() { - return Err(TransformError::EmptyBatch); - } - - let (width, height) = frames[0].dimensions(); - let plan = self.plan_video(frames.len(), width, height, config)?; - let temporal_patch_size = self.config.temporal_patch_size; - let mut all_patches = vec![0.0; plan.output_values]; - let mut out_idx = 0; - let mut frame_rgbs = Vec::with_capacity(temporal_patch_size); - for gt in 0..plan.grid_t { - prepare_video_frame_chunk( - frames.len(), - gt, - temporal_patch_size, - &mut frame_rgbs, - |frame_index| { - let frame = &frames[frame_index]; - let needs_resize = plan.do_resize - && (frame.width() != plan.target_width - || frame.height() != plan.target_height); - if needs_resize { - let (width, height, data) = resize_dynamic_frame_to_raw( - frame, - plan.target_width, - plan.target_height, - plan.filter, - ); - Ok(VideoFrameRgb { - width, - height, - data: Cow::Owned(data), - }) - } else { - let (width, height, data) = rgb_bytes(frame); - Ok(VideoFrameRgb { - width, - height, - data, - }) - } - }, - )?; - - self.patchify_video_rgb_chunk_into( - &frame_rgbs, - plan.grid_h, - plan.grid_w, - &mut all_patches, - &mut out_idx, - &plan.lut, - )?; - } - debug_assert_eq!(out_idx, all_patches.len()); - Self::finish_video(plan, all_patches) - } - - fn preprocess_video_rgb( - &self, - frames: &[RgbFrameRef<'_>], - config: &PreProcessorConfig, - ) -> Result { - if frames.is_empty() { - return Err(TransformError::EmptyBatch); - } - - let plan = self.plan_video(frames.len(), frames[0].width, frames[0].height, config)?; - let temporal_patch_size = self.config.temporal_patch_size; - let mut all_patches = vec![0.0; plan.output_values]; - for frame in frames { - let expected_len = (frame.width as usize) - .checked_mul(frame.height as usize) - .and_then(|pixels| pixels.checked_mul(3)) - .ok_or_else(|| { - TransformError::ShapeError(format!( - "video frame dimensions are too large: {}x{}", - frame.width, frame.height - )) - })?; - if frame.data.len() != expected_len { - return Err(TransformError::InvalidShape { - expected: format!( - "RGB frame byte length {expected_len} for {}x{}", - frame.width, frame.height - ), - actual: vec![frame.data.len()], - }); - } - } - - let values_per_group = plan.grid_h * plan.grid_w * plan.patch_features; - let parallel_tasks = task_count(all_patches.len() * size_of::(), plan.grid_t, 1); - let groups_per_task = plan.grid_t.div_ceil(parallel_tasks); - let mut errors = (0..parallel_tasks).map(|_| None).collect::>(); - parallel_scope(|scope| { - for (task_index, (output_band, error_slot)) in all_patches - .chunks_mut(groups_per_task * values_per_group) - .zip(errors.iter_mut()) - .enumerate() - { - let first_group = task_index * groups_per_task; - scope.spawn(move |_| { - let outcome = (|| { - for (group_offset, group_output) in - output_band.chunks_mut(values_per_group).enumerate() - { - let temporal_index = first_group + group_offset; - let prepared = prepare_video_rgb_chunk( - frames, - temporal_index, - temporal_patch_size, - plan.target_width, - plan.target_height, - plan.filter, - plan.do_resize, - )?; - let mut output_index = 0; - self.patchify_video_rgb_chunk_into( - &prepared, - plan.grid_h, - plan.grid_w, - group_output, - &mut output_index, - &plan.lut, - )?; - debug_assert_eq!(output_index, group_output.len()); - } - Ok::<_, TransformError>(()) - })(); - *error_slot = outcome.err(); - }); - } - }); - if let Some(error) = errors.into_iter().flatten().next() { - return Err(error); - } - - Self::finish_video(plan, all_patches) - } - - fn calculate_num_tokens(&self, width: u32, height: u32, _config: &PreProcessorConfig) -> usize { - // Calculate resized dimensions - let (new_height, new_width) = match self.smart_resize(height as usize, width as usize) { - Ok((h, w)) => (h, w), - Err(_) => { - // Fallback: use minimum size - let factor = self.get_factor(); - (factor, factor) - } - }; - - // Calculate grid and tokens - let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(new_height, new_width, 1); - self.calculate_tokens_from_grid(grid_t, grid_h, grid_w) - } - - fn model_name(&self) -> &'static str { - self.config.model_name - } - - fn get_processed_size(&self, _config: &PreProcessorConfig) -> Option<(u32, u32)> { - // Qwen VL models have dynamic sizing, no fixed output size - None - } -} - -#[cfg(test)] -mod tests { - use image::RgbImage; - - use super::*; - use crate::vision::transforms::to_tensor_and_normalize; - - fn create_test_config() -> QwenVLConfig { - QwenVLConfig { - patch_size: 14, - merge_size: 2, - min_pixels: 256 * 28 * 28, - max_pixels: 1280 * 28 * 28, - video_min_pixels: 256 * 28 * 28, - video_max_pixels: 1280 * 28 * 28, - video_resize_mode: QwenVideoResizeMode::TotalVolume, - temporal_patch_size: 2, - mean: [0.5, 0.5, 0.5], - std: [0.5, 0.5, 0.5], - model_name: "test-qwen-vl", - } - } - - fn create_video_test_config() -> QwenVLConfig { - QwenVLConfig { - patch_size: 2, - merge_size: 1, - min_pixels: 1, - max_pixels: 1024 * 1024, - video_min_pixels: 1, - video_max_pixels: 1024 * 1024, - video_resize_mode: QwenVideoResizeMode::TotalVolume, - temporal_patch_size: 2, - mean: [0.5, 0.25, 0.75], - std: [0.5, 0.25, 0.5], - model_name: "test-qwen-vl-video", - } - } - - fn create_pattern_frame(seed: u8) -> DynamicImage { - let mut image = RgbImage::new(4, 4); - for y in 0..4 { - for x in 0..4 { - image.put_pixel( - x, - y, - image::Rgb([ - seed.wrapping_add((x * 3 + y) as u8), - seed.wrapping_add((x + y * 5) as u8), - seed.wrapping_add((x * 7 + y * 11) as u8), - ]), - ); - } - } - DynamicImage::ImageRgb8(image) - } - - fn create_sized_pattern_frame(w: u32, h: u32, seed: u8) -> DynamicImage { - let mut image = RgbImage::new(w, h); - for y in 0..h { - for x in 0..w { - image.put_pixel( - x, - y, - image::Rgb([ - seed.wrapping_add((x * 3 + y) as u8), - seed.wrapping_add((x + y * 5) as u8), - seed.wrapping_add((x * 7 + y * 11) as u8), - ]), - ); - } - } - DynamicImage::ImageRgb8(image) - } - - #[test] - fn test_prepare_video_frame_chunk_reuses_temporal_padding() { - let frames = [[1_u8, 2, 3], [4, 5, 6], [7, 8, 9]]; - let prepare_calls = std::cell::Cell::new(0); - let mut prepared = Vec::new(); - - prepare_video_frame_chunk(3, 1, 2, &mut prepared, |frame_index| { - prepare_calls.set(prepare_calls.get() + 1); - Ok(VideoFrameRgb { - width: 1, - height: 1, - data: Cow::Borrowed(&frames[frame_index]), - }) - }) - .unwrap(); - - assert_eq!(prepare_calls.get(), 1); - assert_eq!(prepared.len(), 2); - assert_eq!(prepared[0].data.as_ref(), prepared[1].data.as_ref()); - } - - #[test] - fn test_preprocess_image_matches_tensor_patchify_with_resize() { - let processor = QwenVLProcessorBase::new(create_video_test_config()); - let config = PreProcessorConfig { - image_mean: Some(processor.default_mean().to_vec()), - image_std: Some(processor.default_std().to_vec()), - ..Default::default() - }; - let image = create_sized_pattern_frame(7, 9, 3); - let (target_h, target_w) = processor.smart_resize(9, 7).unwrap(); - assert!( - (target_w as u32, target_h as u32) != (7u32, 9u32), - "test must force a resize; target {target_w}x{target_h} should differ from 7x9" - ); - - let result = processor - .preprocess(std::slice::from_ref(&image), &config) - .unwrap(); - let actual = result.encoder_input.as_slice_memory_order().unwrap(); - - let resized = resize_bicubic_pil(&image, target_w as u32, target_h as u32); - let tensor = to_tensor_and_normalize( - &resized, - &processor.default_mean(), - &processor.default_std(), - ); - let (grid_t, grid_h, grid_w) = processor.calculate_grid_thw(target_h, target_w, 1); - let mut expected = Vec::new(); - processor - .patchify_into(&tensor, grid_t, grid_h, grid_w, &mut expected) - .unwrap(); - - assert_eq!(actual.len(), expected.len()); - for (idx, (&got, &want)) in actual.iter().zip(expected.iter()).enumerate() { - assert_eq!( - got.to_bits(), - want.to_bits(), - "image patch value differs at index {idx}: got {got}, want {want}" - ); - } - } - - #[test] - fn test_patchify_image_rgb_block_band_matches_tensor_patchify_with_resize() { - let processor = QwenVLProcessorBase::new(create_video_test_config()); - let image = create_sized_pattern_frame(7, 9, 3); - let (target_h, target_w) = processor.smart_resize(9, 7).unwrap(); - let resized = resize_bicubic_pil(&image, target_w as u32, target_h as u32); - let tensor = to_tensor_and_normalize( - &resized, - &processor.default_mean(), - &processor.default_std(), - ); - let (grid_t, grid_h, grid_w) = processor.calculate_grid_thw(target_h, target_w, 1); - let mut expected = Vec::new(); - processor - .patchify_into(&tensor, grid_t, grid_h, grid_w, &mut expected) - .unwrap(); - - let (width, height, raw) = rgb_bytes(&resized); - assert_eq!((width, height), (target_w, target_h)); - let patch_size = processor.config.patch_size; - let merge_size = processor.config.merge_size; - let temporal_patch_size = processor.config.temporal_patch_size; - let merged_patch = merge_size * patch_size; - let pr_blocks = grid_h / merge_size; - let pc_blocks = grid_w / merge_size; - let n_blocks = pr_blocks * pc_blocks; - assert!( - n_blocks > 1, - "test must exercise multiple image patch blocks" - ); - let block_out = merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; - let mean = processor.default_mean(); - let std = processor.default_std(); - 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)); - let lut: [[f32; 256]; 3] = - std::array::from_fn(|c| std::array::from_fn(|v| v as f32 * scale[c] + bias[c])); - - let mut actual = vec![0.0; expected.len()]; - let split_blocks = n_blocks / 2; - let split_at = split_blocks * block_out; - let (first, second) = actual.split_at_mut(split_at); - QwenVLProcessorBase::patchify_image_rgb_block_band( - raw.as_ref(), - width, - patch_size, - merge_size, - temporal_patch_size, - merged_patch, - pc_blocks, - 0, - first, - &lut, - ); - QwenVLProcessorBase::patchify_image_rgb_block_band( - raw.as_ref(), - width, - patch_size, - merge_size, - temporal_patch_size, - merged_patch, - pc_blocks, - split_blocks, - second, - &lut, - ); - - assert_eq!(actual.len(), expected.len()); - for (idx, (&got, &want)) in actual.iter().zip(expected.iter()).enumerate() { - assert_eq!( - got.to_bits(), - want.to_bits(), - "image block-band patch value differs at index {idx}: got {got}, want {want}" - ); - } - } - - /// When a video frame actually needs resizing, the DynamicImage path - /// (`preprocess_video` → `resize_bicubic_pil`) and the raw-RGB path - /// (`preprocess_video_rgb` → `resize_bicubic_pil_rgb`) must produce - /// byte-for-byte identical encoder inputs. The other video tests use 4x4 - /// frames that need no resize, so this is the one exercising the - /// default-bicubic resize branch. - #[test] - fn test_preprocess_video_rgb_matches_dynamic_with_resize() { - let processor = QwenVLProcessorBase::new(create_video_test_config()); - let config = PreProcessorConfig { - image_mean: Some(processor.default_mean().to_vec()), - image_std: Some(processor.default_std().to_vec()), - ..Default::default() - }; - // Odd dimensions force smart_resize_video to a different factor-aligned - // target, guaranteeing the resize branch runs. - let frames = vec![ - create_sized_pattern_frame(7, 9, 3), - create_sized_pattern_frame(7, 9, 101), - ]; - let (target_h, target_w) = processor.smart_resize_video(frames.len(), 9, 7).unwrap(); - assert!( - (target_w as u32, target_h as u32) != (7u32, 9u32), - "test must force a resize; target {target_w}x{target_h} should differ from 7x9" - ); - - let rgb_frames = frames - .iter() - .map(|frame| { - let DynamicImage::ImageRgb8(rgb) = frame else { - panic!("test frame is not RGB8"); - }; - RgbFrameRef { - width: rgb.width(), - height: rgb.height(), - data: rgb.as_raw(), - } - }) - .collect::>(); - - let dynamic = processor.preprocess_video(&frames, &config).unwrap(); - let rgb = processor - .preprocess_video_rgb(&rgb_frames, &config) - .unwrap(); - - let a = dynamic.encoder_input.as_slice_memory_order().unwrap(); - let b = rgb.encoder_input.as_slice_memory_order().unwrap(); - assert_eq!( - a.len(), - b.len(), - "resized video encoder input length differs" - ); - for (idx, (&got, &want)) in a.iter().zip(b.iter()).enumerate() { - assert_eq!( - got.to_bits(), - want.to_bits(), - "resized video path diverges at index {idx}: dynamic {got} vs rgb {want}" - ); - } - } - - #[test] - fn test_preprocess_video_rgb_matches_dynamic_with_resize_and_padding() { - let processor = QwenVLProcessorBase::new(create_video_test_config()); - let config = PreProcessorConfig { - image_mean: Some(processor.default_mean().to_vec()), - image_std: Some(processor.default_std().to_vec()), - ..Default::default() - }; - let frames = vec![ - create_sized_pattern_frame(7, 9, 3), - create_sized_pattern_frame(7, 9, 101), - create_sized_pattern_frame(7, 9, 177), - ]; - assert_ne!( - frames.len() % processor.temporal_patch_size(), - 0, - "test must force temporal padding" - ); - let (target_h, target_w) = processor.smart_resize_video(frames.len(), 9, 7).unwrap(); - assert!( - (target_w as u32, target_h as u32) != (7u32, 9u32), - "test must force a resize; target {target_w}x{target_h} should differ from 7x9" - ); - - let rgb_frames = frames - .iter() - .map(|frame| { - let DynamicImage::ImageRgb8(rgb) = frame else { - panic!("test frame is not RGB8"); - }; - RgbFrameRef { - width: rgb.width(), - height: rgb.height(), - data: rgb.as_raw(), - } - }) - .collect::>(); - - let dynamic = processor.preprocess_video(&frames, &config).unwrap(); - let rgb = processor - .preprocess_video_rgb(&rgb_frames, &config) - .unwrap(); - - let a = dynamic.encoder_input.as_slice_memory_order().unwrap(); - let b = rgb.encoder_input.as_slice_memory_order().unwrap(); - assert_eq!(a.len(), b.len()); - for (idx, (&got, &want)) in a.iter().zip(b.iter()).enumerate() { - assert_eq!( - got.to_bits(), - want.to_bits(), - "resized padded video path diverges at index {idx}: dynamic {got} vs rgb {want}" - ); - } - } - - #[test] - fn test_qwen_vl_base_factor() { - let processor = QwenVLProcessorBase::new(create_test_config()); - assert_eq!(processor.get_factor(), 28); // 14 * 2 - } - - #[test] - fn test_smart_resize_within_bounds() { - let processor = QwenVLProcessorBase::new(create_test_config()); - let (h, w) = processor.smart_resize(500, 500).unwrap(); - - assert_eq!(h % 28, 0); - assert_eq!(w % 28, 0); - assert!(h * w >= processor.min_pixels()); - assert!(h * w <= processor.max_pixels()); - } - - #[test] - fn test_smart_resize_video_matches_hf_actual_frame_beta() { - let processor = QwenVLProcessorBase::new(QwenVLConfig { - patch_size: 16, - merge_size: 2, - min_pixels: 1, - max_pixels: 16_777_216, - video_min_pixels: 1, - video_max_pixels: 16_777_216, - video_resize_mode: QwenVideoResizeMode::TotalVolume, - temporal_patch_size: 2, - mean: [0.5; 3], - std: [0.5; 3], - model_name: "test-qwen-vl-video-budget", - }); - - let (height, width) = processor.smart_resize_video(1, 3000, 3000).unwrap(); - - // Hugging Face uses the padded frame count for the threshold, but the - // actual frame count for beta. Preserve that behavior for parity. - assert_eq!((height, width), (4096, 4096)); - } - - #[test] - fn test_smart_resize_extreme_aspect_ratio_error() { - let processor = QwenVLProcessorBase::new(create_test_config()); - let result = processor.smart_resize(100, 30000); - assert!(result.is_err()); - } - - #[test] - fn test_calculate_grid_thw() { - let processor = QwenVLProcessorBase::new(create_test_config()); - let (t, h, w) = processor.calculate_grid_thw(448, 448, 1); - - assert_eq!(t, 1); - assert_eq!(h, 448 / 14); - assert_eq!(w, 448 / 14); - } - - #[test] - fn test_calculate_tokens() { - let processor = QwenVLProcessorBase::new(create_test_config()); - let tokens = processor.calculate_tokens_from_grid(1, 32, 32); - assert_eq!(tokens, (32 * 32) / 4); - } - - #[test] - fn test_preprocess_video_matches_tensor_patchify() { - let processor = QwenVLProcessorBase::new(create_video_test_config()); - let config = PreProcessorConfig { - image_mean: Some(processor.default_mean().to_vec()), - image_std: Some(processor.default_std().to_vec()), - ..Default::default() - }; - let frames = vec![create_pattern_frame(3), create_pattern_frame(101)]; - - let result = processor.preprocess_video(&frames, &config).unwrap(); - let actual = result.encoder_input.as_slice_memory_order().unwrap(); - - let tensors = frames - .iter() - .map(|frame| { - to_tensor_and_normalize(frame, &processor.default_mean(), &processor.default_std()) - }) - .collect::>(); - let (grid_t, grid_h, grid_w) = processor.calculate_grid_thw(4, 4, frames.len()); - let mut expected = Vec::new(); - processor - .patchify_video_into(&tensors, grid_t, grid_h, grid_w, &mut expected) - .unwrap(); - - assert_eq!(actual.len(), expected.len()); - for (idx, (&got, &want)) in actual.iter().zip(expected.iter()).enumerate() { - assert_eq!( - got.to_bits(), - want.to_bits(), - "video patch value differs at index {idx}: got {got}, want {want}" - ); - } - } - - #[test] - fn test_preprocess_video_rgb_matches_dynamic_video() { - let processor = QwenVLProcessorBase::new(create_video_test_config()); - let config = PreProcessorConfig { - image_mean: Some(processor.default_mean().to_vec()), - image_std: Some(processor.default_std().to_vec()), - ..Default::default() - }; - let frames = vec![ - create_pattern_frame(3), - create_pattern_frame(101), - create_pattern_frame(177), - ]; - - let rgb_frames = frames - .iter() - .map(|frame| { - let DynamicImage::ImageRgb8(rgb) = frame else { - panic!("test frame is not RGB8"); - }; - RgbFrameRef { - width: rgb.width(), - height: rgb.height(), - data: rgb.as_raw(), - } - }) - .collect::>(); - - let dynamic_result = processor.preprocess_video(&frames, &config).unwrap(); - let rgb_result = processor - .preprocess_video_rgb(&rgb_frames, &config) - .unwrap(); - - assert_eq!( - dynamic_result.encoder_input.shape(), - rgb_result.encoder_input.shape() - ); - let mut dynamic_keys = dynamic_result.model_specific.keys().collect::>(); - let mut rgb_keys = rgb_result.model_specific.keys().collect::>(); - dynamic_keys.sort(); - rgb_keys.sort(); - assert_eq!(dynamic_keys, rgb_keys); - - let dynamic_values = dynamic_result - .encoder_input - .as_slice_memory_order() - .unwrap(); - let rgb_values = rgb_result.encoder_input.as_slice_memory_order().unwrap(); - for (idx, (&got, &want)) in rgb_values.iter().zip(dynamic_values.iter()).enumerate() { - assert_eq!( - got.to_bits(), - want.to_bits(), - "RGB video patch value differs at index {idx}: got {got}, want {want}" - ); - } - } - - #[test] - fn test_preprocess_video_rgb_matches_dynamic_video_parallel_blocks() { - let processor = QwenVLProcessorBase::new(create_video_test_config()); - let config = PreProcessorConfig { - image_mean: Some(processor.default_mean().to_vec()), - image_std: Some(processor.default_std().to_vec()), - ..Default::default() - }; - let frames = vec![ - create_sized_pattern_frame(280, 280, 3), - create_sized_pattern_frame(280, 280, 101), - ]; - - let rgb_frames = frames - .iter() - .map(|frame| { - let DynamicImage::ImageRgb8(rgb) = frame else { - panic!("test frame is not RGB8"); - }; - RgbFrameRef { - width: rgb.width(), - height: rgb.height(), - data: rgb.as_raw(), - } - }) - .collect::>(); - - let dynamic_result = processor.preprocess_video(&frames, &config).unwrap(); - let rgb_result = processor - .preprocess_video_rgb(&rgb_frames, &config) - .unwrap(); - - assert_eq!( - dynamic_result.encoder_input.shape(), - rgb_result.encoder_input.shape() - ); - let dynamic_values = dynamic_result - .encoder_input - .as_slice_memory_order() - .unwrap(); - let rgb_values = rgb_result.encoder_input.as_slice_memory_order().unwrap(); - for (idx, (&got, &want)) in rgb_values.iter().zip(dynamic_values.iter()).enumerate() { - assert_eq!( - got.to_bits(), - want.to_bits(), - "parallel RGB video patch value differs at index {idx}: got {got}, want {want}" - ); - } - } -} diff --git a/crates/multimodal/src/vision/scratch.rs b/crates/multimodal/src/vision/scratch.rs deleted file mode 100644 index bf22234c8..000000000 --- a/crates/multimodal/src/vision/scratch.rs +++ /dev/null @@ -1,58 +0,0 @@ -//! Recycling pools for the large per-image vision buffers. -//! -//! The preprocess pipeline allocates tens of MB per image (the [C, H, W] f32 -//! tensor and the batched patch buffer are each large). -//! Freshly-allocated buffers of this size bypass the allocator's reuse paths -//! (glibc caps non-main-arena chunks at 64 MB and mmaps anything larger or -//! colder), so every image pays tens of thousands of minor page faults; the -//! fault path serializes process-wide and caps the data plane's effective -//! parallelism. Recycling keeps the pages mapped and hot. -//! -//! A lock-free thread-local pool serves same-thread take/give (preprocess -//! internals run on blocking-pool threads). The pool is capped to bound -//! residency; buffers beyond the cap are dropped. - -use std::cell::RefCell; - -/// Max recycled buffers kept per thread per class; excess is dropped. The -/// vision path holds at most a couple of live tensors per request, so a small -/// cap captures same-thread reuse. -const MAX_THREAD_POOLED: usize = 2; - -thread_local! { - static F32_LOCAL: RefCell>> = const { RefCell::new(Vec::new()) }; -} - -macro_rules! pool_impl { - ($take_cap:ident, $give:ident, $ty:ty, $local:ident) => { - /// Take an empty `Vec` with at least `cap` capacity, reusing pooled storage. - pub fn $take_cap(cap: usize) -> Vec<$ty> { - let mut v = $local.with(|p| p.borrow_mut().pop()).unwrap_or_default(); - v.clear(); - v.reserve(cap); - v - } - - /// Return a buffer for reuse by a later same-thread take. - pub fn $give(v: Vec<$ty>) { - if v.capacity() == 0 { - return; - } - $local.with(|p| { - let mut p = p.borrow_mut(); - if p.len() < MAX_THREAD_POOLED { - p.push(v); - } - }); - } - }; -} - -pool_impl!(take_f32_cap, give_f32, f32, F32_LOCAL); - -/// Take a zero-filled `Vec` of exactly `len`, reusing pooled storage. -pub fn take_f32(len: usize) -> Vec { - let mut v = take_f32_cap(len); - v.resize(len, 0.0); - v -} diff --git a/crates/multimodal/src/vision/transforms.rs b/crates/multimodal/src/vision/transforms.rs deleted file mode 100644 index ee0217497..000000000 --- a/crates/multimodal/src/vision/transforms.rs +++ /dev/null @@ -1,1245 +0,0 @@ -//! Image transformation functions for vision preprocessing. -//! -//! This module provides composable transforms that match HuggingFace image processor -//! behavior, enabling pure Rust preprocessing without Python dependencies. - -use std::cell::RefCell; - -use fast_image_resize::{ - images::{Image as FirImage, ImageRef as FirImageRef}, - IntoImageView, PixelType, ResizeAlg, ResizeOptions, Resizer, -}; -use image::{imageops::FilterType, DynamicImage, GenericImageView, Rgb, RgbImage}; -use ndarray::{s, Array3, Array4}; - -use super::{ - execution::{scope as parallel_scope, task_count}, - scratch, -}; -pub use crate::error::TransformError; - -pub type Result = std::result::Result; - -/// Extract RGB pixel data from a DynamicImage, avoiding a copy when already RGB8. -/// Returns (width, height, raw_bytes) where raw_bytes is interleaved R,G,B,R,G,B,... -pub fn rgb_bytes(image: &DynamicImage) -> (usize, usize, std::borrow::Cow<'_, [u8]>) { - match image { - DynamicImage::ImageRgb8(rgb) => ( - rgb.width() as usize, - rgb.height() as usize, - std::borrow::Cow::Borrowed(rgb.as_raw()), - ), - _ => { - let rgb = image.to_rgb8(); - let w = rgb.width() as usize; - let h = rgb.height() as usize; - (w, h, std::borrow::Cow::Owned(rgb.into_raw())) - } - } -} - -/// 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]`. -/// -/// Processes 8 pixels at a time so the compiler can unroll and auto-vectorize -/// the stride-3 gather pattern. -pub fn deinterleave_rgb_to_planes( - rgb: &[u8], - r_plane: &mut [f32], - g_plane: &mut [f32], - b_plane: &mut [f32], - scale: [f32; 3], - bias: [f32; 3], -) { - let pixels = r_plane.len(); - debug_assert_eq!(pixels, g_plane.len()); - debug_assert_eq!(pixels, b_plane.len()); - debug_assert!(rgb.len() >= pixels * 3); - - // Each output element depends only on its own input byte, so banding the - // pixel range across threads is BIT-IDENTICAL (elementwise f32, no - // reduction). Small images stay serial. - let nthreads = par_threads(pixels * 3 * 4, pixels); - if nthreads <= 1 { - deinterleave_contiguous(rgb, r_plane, g_plane, b_plane, scale, bias); - return; - } - let chunk = pixels.div_ceil(nthreads); - let (mut rr, mut gg, mut bb) = (r_plane, g_plane, b_plane); - parallel_scope(|s| { - let mut p0 = 0usize; - while p0 < pixels { - let n = chunk.min(pixels - p0); - let (rb, rt) = rr.split_at_mut(n); - let (gb, gt) = gg.split_at_mut(n); - let (bbnd, bt) = bb.split_at_mut(n); - rr = rt; - gg = gt; - bb = bt; - let rgb_band = &rgb[p0 * 3..(p0 + n) * 3]; - s.spawn(move |_| deinterleave_contiguous(rgb_band, rb, gb, bbnd, scale, bias)); - p0 += n; - } - }); -} - -/// Deinterleave a contiguous pixel range (planes/rgb already sliced to the band). -fn deinterleave_contiguous( - rgb: &[u8], - r_plane: &mut [f32], - g_plane: &mut [f32], - b_plane: &mut [f32], - scale: [f32; 3], - bias: [f32; 3], -) { - let pixels = r_plane.len(); - let full_blocks = pixels / 8; - let remainder = pixels % 8; - - for block in 0..full_blocks { - let dst = block * 8; - let src_base = dst * 3; - let src = &rgb[src_base..src_base + 24]; - let rd = &mut r_plane[dst..dst + 8]; - let gd = &mut g_plane[dst..dst + 8]; - let bd = &mut b_plane[dst..dst + 8]; - - for i in 0..8 { - let s = i * 3; - rd[i] = src[s] as f32 * scale[0] + bias[0]; - gd[i] = src[s + 1] as f32 * scale[1] + bias[1]; - bd[i] = src[s + 2] as f32 * scale[2] + bias[2]; - } - } - - let tail_dst = full_blocks * 8; - let tail_src = tail_dst * 3; - for i in 0..remainder { - let s = tail_src + i * 3; - r_plane[tail_dst + i] = rgb[s] as f32 * scale[0] + bias[0]; - g_plane[tail_dst + i] = rgb[s + 1] as f32 * scale[1] + bias[1]; - b_plane[tail_dst + i] = rgb[s + 2] as f32 * scale[2] + bias[2]; - } -} - -/// Build a [C, H, W] f32 tensor from interleaved RGB bytes with per-channel -/// `scale` and `bias`: `output[c][i] = raw[i*3 + c] * scale[c] + bias[c]`. -fn build_planar_tensor( - raw: &[u8], - w: usize, - h: usize, - scale: [f32; 3], - bias: [f32; 3], -) -> Array3 { - let pixels = h * w; - // Pooled: this large per-image buffer is the data plane's hottest allocation. - let mut data = scratch::take_f32(3 * pixels); - let (r_plane, rest) = data.split_at_mut(pixels); - let (g_plane, b_plane) = rest.split_at_mut(pixels); - - deinterleave_rgb_to_planes(raw, r_plane, g_plane, b_plane, scale, bias); - - #[expect( - clippy::expect_used, - reason = "data has exactly 3*h*w elements by construction" - )] - Array3::from_shape_vec((3, h, w), data).expect("shape matches pre-allocated buffer") -} - -/// Convert image to tensor [C, H, W] normalized to [0, 1]. -/// -/// This matches the default behavior of `torchvision.transforms.ToTensor()`. -pub fn to_tensor(image: &DynamicImage) -> Array3 { - let (w, h, raw) = rgb_bytes(image); - let s = 1.0 / 255.0; - build_planar_tensor(&raw, w, h, [s, s, s], [0.0, 0.0, 0.0]) -} - -/// Convert image to tensor [C, H, W] without normalization (keeps [0, 255]). -#[cfg(test)] -pub fn to_tensor_no_norm(image: &DynamicImage) -> Array3 { - let (w, h, raw) = rgb_bytes(image); - build_planar_tensor(&raw, w, h, [1.0, 1.0, 1.0], [0.0, 0.0, 0.0]) -} - -/// Normalize tensor per channel: (x - mean) / std. -/// -/// This matches `torchvision.transforms.Normalize(mean, std)`. -/// -/// # Arguments -/// * `tensor` - Input tensor of shape [C, H, W] -/// * `mean` - Per-channel mean values -/// * `std` - Per-channel standard deviation values -pub fn normalize(tensor: &mut Array3, mean: &[f64; 3], std: &[f64; 3]) { - let [h, w] = [tensor.shape()[1], tensor.shape()[2]]; - let pixels = h * w; - - if let Some(flat) = tensor.as_slice_mut() { - // Fast path: contiguous memory, process channel planes directly - for c in 0..3 { - let mean_c = mean[c] as f32; - let inv_std_c = 1.0 / std[c] as f32; - let plane = &mut flat[c * pixels..(c + 1) * pixels]; - for v in plane.iter_mut() { - *v = (*v - mean_c) * inv_std_c; - } - } - } else { - for c in 0..3 { - let mean_c = mean[c] as f32; - let std_c = std[c] as f32; - tensor - .slice_mut(s![c, .., ..]) - .mapv_inplace(|v| (v - mean_c) / std_c); - } - } -} - -/// Convert image to tensor and normalize in a single pass. -/// -/// Fuses `to_tensor` (u8→f32 with /255) and `normalize` ((x-mean)/std) -/// into one loop to avoid an extra pass over the data. -pub fn to_tensor_and_normalize( - image: &DynamicImage, - mean: &[f64; 3], - std: &[f64; 3], -) -> Array3 { - let (w, h, raw) = rgb_bytes(image); - // Fused: (pixel/255 - mean) / std = pixel * (1/(255*std)) - mean/std - 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)); - build_planar_tensor(&raw, w, h, scale, bias) -} - -/// Rescale tensor by a constant factor. -/// -/// Used when `do_rescale=True` in HuggingFace configs (typically 1/255). -pub fn rescale(tensor: &mut Array3, factor: f64) { - let factor = factor as f32; - tensor.mapv_inplace(|v| v * factor); -} - -/// Map `image` crate filter types to `fast_image_resize` algorithm. -fn to_fir_algorithm(filter: FilterType) -> ResizeAlg { - use fast_image_resize::FilterType as FirFilter; - match filter { - FilterType::Nearest => ResizeAlg::Nearest, - FilterType::Triangle => ResizeAlg::Convolution(FirFilter::Bilinear), - FilterType::CatmullRom => ResizeAlg::Convolution(FirFilter::CatmullRom), - FilterType::Gaussian => ResizeAlg::Convolution(FirFilter::Gaussian), - FilterType::Lanczos3 => ResizeAlg::Convolution(FirFilter::Lanczos3), - } -} - -thread_local! { - static RESIZER: RefCell = RefCell::new(Resizer::new()); -} - -/// Resize image to exact dimensions using SIMD-accelerated resizer. -/// -/// # Arguments -/// * `image` - Input image -/// * `width` - Target width -/// * `height` - Target height -/// * `filter` - Interpolation filter (Nearest, Triangle/Bilinear, CatmullRom/Bicubic, Lanczos3) -pub fn resize(image: &DynamicImage, width: u32, height: u32, filter: FilterType) -> 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 ok = RESIZER.with(|r| r.borrow_mut().resize(image, &mut dst, &options).is_ok()); - if !ok { - return image.resize_exact(width, height, filter); - } - fir_image_to_dynamic(dst, width, height, image, filter) -} - -/// Resize borrowed interleaved RGB bytes without first materializing an -/// `image::RgbImage` over an owned input buffer. -pub fn resize_rgb_bytes( - data: &[u8], - width: u32, - height: u32, - target_width: u32, - target_height: u32, - filter: FilterType, -) -> Result { - let src = FirImageRef::new(width, height, data, PixelType::U8x3) - .map_err(|e| TransformError::ShapeError(format!("invalid RGB source image: {e}")))?; - let mut dst = FirImage::new(target_width, target_height, PixelType::U8x3); - let options = ResizeOptions::new().resize_alg(to_fir_algorithm(filter)); - RESIZER - .with(|r| r.borrow_mut().resize(&src, &mut dst, &options)) - .map_err(|e| TransformError::ShapeError(format!("RGB resize failed: {e}")))?; - - RgbImage::from_raw(target_width, target_height, dst.into_vec()).ok_or_else(|| { - TransformError::ShapeError(format!( - "failed to build resized RGB image for {target_width}x{target_height}" - )) - }) -} - -/// Convert a `fast_image_resize::Image` back to a `DynamicImage`. -/// -/// Falls back to the `image` crate resize for unhandled pixel formats. -fn fir_image_to_dynamic( - img: FirImage<'_>, - width: u32, - height: u32, - source: &DynamicImage, - filter: FilterType, -) -> DynamicImage { - let buf = img.into_vec(); - match source { - DynamicImage::ImageRgb8(_) => { - RgbImage::from_raw(width, height, buf).map(DynamicImage::ImageRgb8) - } - DynamicImage::ImageRgba8(_) => { - image::RgbaImage::from_raw(width, height, buf).map(DynamicImage::ImageRgba8) - } - DynamicImage::ImageLuma8(_) => { - image::GrayImage::from_raw(width, height, buf).map(DynamicImage::ImageLuma8) - } - _ => None, - } - .unwrap_or_else(|| source.resize_exact(width, height, filter)) -} - -// --------------------------------------------------------------------------- -// Pillow-exact bicubic resize. -// -// Qwen image processors resize via `PIL.Image.resize(size, BICUBIC)` on the -// uint8 image. The SIMD `fast_image_resize` path above is the same filter -// *family* (Catmull-Rom, a=-0.5) but diverges bit-wise on non-integer ratios -// (support scaling + fixed-point details), which the vision encoder amplifies -// into a large embedding shift. This routine replicates Pillow's `Resample.c` -// algorithm exactly, validated against Pillow. -const PIL_PRECISION_BITS: i64 = 32 - 8 - 2; -const PIL_BICUBIC_SUPPORT: f64 = 2.0; - -#[inline] -fn pil_cubic(x: f64) -> f64 { - // Keys cubic with a = -0.5 (Pillow's BICUBIC). - const A: f64 = -0.5; - let x = x.abs(); - if x < 1.0 { - ((A + 2.0) * x - (A + 3.0)) * x * x + 1.0 - } else if x < 2.0 { - (((x - 5.0) * x + 8.0) * x - 4.0) * A - } else { - 0.0 - } -} - -/// Pillow `precompute_coeffs` for one axis: integer (fixed-point) kernels plus -/// per-output bounds `(start, count)`. -fn pil_precompute_coeffs(in_size: usize, out_size: usize) -> (Vec<(usize, usize)>, Vec>) { - let scale = in_size as f64 / out_size as f64; - let filterscale = if scale >= 1.0 { scale } else { 1.0 }; - let support = PIL_BICUBIC_SUPPORT * filterscale; - let inv = 1.0 / filterscale; - let coeff_scale = (1_i64 << PIL_PRECISION_BITS) as f64; - - let mut bounds = Vec::with_capacity(out_size); - let mut kernels = Vec::with_capacity(out_size); - for xx in 0..out_size { - let center = (xx as f64 + 0.5) * scale; - let mut xmin = (center - support + 0.5) as i64; - if xmin < 0 { - xmin = 0; - } - let mut xmax = (center + support + 0.5) as i64; - if xmax > in_size as i64 { - xmax = in_size as i64; - } - let xmin = xmin as usize; - let xmax = (xmax as usize).saturating_sub(xmin); - - let mut w = vec![0.0_f64; xmax]; - let mut tot = 0.0; - for (x, wx) in w.iter_mut().enumerate() { - let v = pil_cubic(((x + xmin) as f64 - center + 0.5) * inv); - *wx = v; - tot += v; - } - if tot != 0.0 { - for wx in &mut w { - *wx /= tot; - } - } - // Pillow normalize_coeffs_8bpc: round half away from zero into fixed point. - let k: Vec = w - .iter() - .map(|&c| { - if c < 0.0 { - (-0.5 + c * coeff_scale) as i64 - } else { - (0.5 + c * coeff_scale) as i64 - } - }) - .collect(); - bounds.push((xmin, xmax)); - kernels.push(k); - } - (bounds, kernels) -} - -#[inline] -fn pil_clip8(v: i64) -> u8 { - let v = v >> PIL_PRECISION_BITS; - if v < 0 { - 0 - } else if v > 255 { - 255 - } else { - v as u8 - } -} - -/// Number of threads to split an elementwise or row-banded preprocessing pass -/// across. Each output row/element is independent, so banding work over threads -/// yields BIT-IDENTICAL output: no shared accumulation and no inner-loop order -/// changes. Small images run serial to avoid thread-spawn overhead. -pub(crate) fn par_threads(out_bytes: usize, out_rows: usize) -> usize { - task_count(out_bytes, out_rows, 32) -} - -/// Process output rows `[oy0, oy0 + out_band.len()/row_out)` of the horizontal -/// pass into `out_band`. Horizontal pass preserves row count, so output row i -/// reads input row `oy0 + i`. -#[expect( - clippy::too_many_arguments, - reason = "row-band resampler: precomputed coeffs + dims + output band" -)] -fn pil_h_band( - src: &[u8], - bounds: &[(usize, usize)], - kernels: &[Vec], - half: i64, - in_w: usize, - out_w: usize, - channels: usize, - oy0: usize, - out_band: &mut [u8], -) { - let row_out = out_w * channels; - for (i, orow) in out_band.chunks_mut(row_out).enumerate() { - let y = oy0 + i; - let row = &src[y * in_w * channels..(y + 1) * in_w * channels]; - for xx in 0..out_w { - let (xmin, xmax) = bounds[xx]; - let k = &kernels[xx]; - for c in 0..channels { - let mut ss = half; - for x in 0..xmax { - ss += row[(xmin + x) * channels + c] as i64 * k[x]; - } - orow[xx * channels + c] = pil_clip8(ss); - } - } - } -} - -#[expect( - clippy::too_many_arguments, - reason = "RGB row-band resampler: precomputed coeffs + dims + output band" -)] -fn pil_h_band_rgb( - src: &[u8], - bounds: &[(usize, usize)], - kernels: &[Vec], - half: i64, - in_w: usize, - out_w: usize, - oy0: usize, - out_band: &mut [u8], -) { - let row_out = out_w * 3; - for (i, output_row) in out_band.chunks_mut(row_out).enumerate() { - let y = oy0 + i; - let row = &src[y * in_w * 3..(y + 1) * in_w * 3]; - for output_x in 0..out_w { - let (source_x, source_columns) = bounds[output_x]; - let kernel = &kernels[output_x]; - let mut red = half; - let mut green = half; - let mut blue = half; - let source_start = source_x * 3; - let source_end = (source_x + source_columns) * 3; - for (pixel, &coefficient) in row[source_start..source_end].chunks_exact(3).zip(kernel) { - red += pixel[0] as i64 * coefficient; - green += pixel[1] as i64 * coefficient; - blue += pixel[2] as i64 * coefficient; - } - let output = output_x * 3; - output_row[output] = pil_clip8(red); - output_row[output + 1] = pil_clip8(green); - output_row[output + 2] = pil_clip8(blue); - } - } -} - -/// Resample interleaved `channels`-channel u8 data along the width axis. -/// `src` is `rows * in_w * channels`; returns `rows * out_w * channels`. -fn pil_resample_horizontal( - src: &[u8], - rows: usize, - in_w: usize, - out_w: usize, - channels: usize, -) -> Vec { - let (bounds, kernels) = pil_precompute_coeffs(in_w, out_w); - let half = 1_i64 << (PIL_PRECISION_BITS - 1); - let row_out = out_w * channels; - let mut out = vec![0_u8; rows * row_out]; - let nthreads = par_threads(out.len(), rows); - if nthreads <= 1 { - pil_h_band( - src, &bounds, &kernels, half, in_w, out_w, channels, 0, &mut out, - ); - } else { - let chunk_rows = rows.div_ceil(nthreads); - parallel_scope(|s| { - let (b, k) = (&bounds, &kernels); - let mut rest = out.as_mut_slice(); - let mut oy0 = 0usize; - while oy0 < rows { - let n = chunk_rows.min(rows - oy0); - let (band, tail) = rest.split_at_mut(n * row_out); - rest = tail; - let start = oy0; - s.spawn(move |_| { - pil_h_band(src, b, k, half, in_w, out_w, channels, start, band); - }); - oy0 += n; - } - }); - } - out -} - -fn pil_resample_horizontal_rgb(src: &[u8], rows: usize, in_w: usize, out_w: usize) -> Vec { - let (bounds, kernels) = pil_precompute_coeffs(in_w, out_w); - let half = 1_i64 << (PIL_PRECISION_BITS - 1); - let row_out = out_w * 3; - let mut out = vec![0_u8; rows * row_out]; - let nthreads = par_threads(out.len(), rows); - if nthreads <= 1 { - pil_h_band_rgb(src, &bounds, &kernels, half, in_w, out_w, 0, &mut out); - } else { - let chunk_rows = rows.div_ceil(nthreads); - parallel_scope(|scope| { - let (bounds, kernels) = (&bounds, &kernels); - let mut rest = out.as_mut_slice(); - let mut output_y = 0; - while output_y < rows { - let band_rows = chunk_rows.min(rows - output_y); - let (band, tail) = rest.split_at_mut(band_rows * row_out); - rest = tail; - let start = output_y; - scope.spawn(move |_| { - pil_h_band_rgb(src, bounds, kernels, half, in_w, out_w, start, band); - }); - output_y += band_rows; - } - }); - } - out -} - -/// Process output rows `[oy0, oy0 + out_band.len()/row_out)` of the vertical -/// pass into `out_band`. -#[expect( - clippy::too_many_arguments, - reason = "row-band resampler: precomputed coeffs + dims + output band" -)] -fn pil_v_band( - src: &[u8], - bounds: &[(usize, usize)], - kernels: &[Vec], - half: i64, - width: usize, - channels: usize, - oy0: usize, - out_band: &mut [u8], -) { - let row_out = width * channels; - for (i, orow) in out_band.chunks_mut(row_out).enumerate() { - let yy = oy0 + i; - let (ymin, ymax) = bounds[yy]; - let k = &kernels[yy]; - for x in 0..width { - for c in 0..channels { - let mut ss = half; - for y in 0..ymax { - ss += src[((ymin + y) * width + x) * channels + c] as i64 * k[y]; - } - orow[x * channels + c] = pil_clip8(ss); - } - } - } -} - -fn pil_v_band_rgb( - src: &[u8], - bounds: &[(usize, usize)], - kernels: &[Vec], - half: i64, - width: usize, - oy0: usize, - out_band: &mut [u8], -) { - let row_out = width * 3; - for (i, output_row) in out_band.chunks_mut(row_out).enumerate() { - let output_y = oy0 + i; - let (source_y, source_rows) = bounds[output_y]; - let kernel = &kernels[output_y]; - let blocked_width = width / 4 * 4; - for x in (0..blocked_width).step_by(4) { - let mut sums = [[half; 3]; 4]; - for (y, &coefficient) in kernel.iter().take(source_rows).enumerate() { - let source = ((source_y + y) * width + x) * 3; - for (pixel, sums) in sums.iter_mut().enumerate() { - let input = source + pixel * 3; - sums[0] += src[input] as i64 * coefficient; - sums[1] += src[input + 1] as i64 * coefficient; - sums[2] += src[input + 2] as i64 * coefficient; - } - } - let output = x * 3; - for (pixel, sums) in sums.iter().enumerate() { - let target = output + pixel * 3; - output_row[target] = pil_clip8(sums[0]); - output_row[target + 1] = pil_clip8(sums[1]); - output_row[target + 2] = pil_clip8(sums[2]); - } - } - for x in blocked_width..width { - let mut red = half; - let mut green = half; - let mut blue = half; - for (y, &coefficient) in kernel.iter().take(source_rows).enumerate() { - let source = ((source_y + y) * width + x) * 3; - red += src[source] as i64 * coefficient; - green += src[source + 1] as i64 * coefficient; - blue += src[source + 2] as i64 * coefficient; - } - let output = x * 3; - output_row[output] = pil_clip8(red); - output_row[output + 1] = pil_clip8(green); - output_row[output + 2] = pil_clip8(blue); - } - } -} - -/// Resample interleaved `channels`-channel u8 data along the height axis. -fn pil_resample_vertical( - src: &[u8], - in_h: usize, - width: usize, - out_h: usize, - channels: usize, -) -> Vec { - let (bounds, kernels) = pil_precompute_coeffs(in_h, out_h); - let half = 1_i64 << (PIL_PRECISION_BITS - 1); - let row_out = width * channels; - let mut out = vec![0_u8; out_h * row_out]; - let nthreads = par_threads(out.len(), out_h); - if nthreads <= 1 { - pil_v_band(src, &bounds, &kernels, half, width, channels, 0, &mut out); - } else { - let chunk_rows = out_h.div_ceil(nthreads); - parallel_scope(|s| { - let (b, k) = (&bounds, &kernels); - let mut rest = out.as_mut_slice(); - let mut oy0 = 0usize; - while oy0 < out_h { - let n = chunk_rows.min(out_h - oy0); - let (band, tail) = rest.split_at_mut(n * row_out); - rest = tail; - let start = oy0; - s.spawn(move |_| pil_v_band(src, b, k, half, width, channels, start, band)); - oy0 += n; - } - }); - } - out -} - -fn pil_resample_vertical_rgb(src: &[u8], in_h: usize, width: usize, out_h: usize) -> Vec { - let (bounds, kernels) = pil_precompute_coeffs(in_h, out_h); - let half = 1_i64 << (PIL_PRECISION_BITS - 1); - let row_out = width * 3; - let mut out = vec![0_u8; out_h * row_out]; - let nthreads = par_threads(out.len(), out_h); - if nthreads <= 1 { - pil_v_band_rgb(src, &bounds, &kernels, half, width, 0, &mut out); - } else { - let chunk_rows = out_h.div_ceil(nthreads); - parallel_scope(|scope| { - let (bounds, kernels) = (&bounds, &kernels); - let mut rest = out.as_mut_slice(); - let mut output_y = 0; - while output_y < out_h { - let rows = chunk_rows.min(out_h - output_y); - let (band, tail) = rest.split_at_mut(rows * row_out); - rest = tail; - let start = output_y; - scope.spawn(move |_| { - pil_v_band_rgb(src, bounds, kernels, half, width, start, band); - }); - output_y += rows; - } - }); - } - out -} - -/// Pillow-exact BICUBIC resize (RGB8), matching -/// `PIL.Image.resize(.., BICUBIC)`. -pub fn resize_bicubic_pil(image: &DynamicImage, out_w: u32, out_h: u32) -> DynamicImage { - let rgb = image.to_rgb8(); - let (in_w, in_h) = rgb.dimensions(); - let output = resize_bicubic_pil_bytes(rgb.as_raw(), in_w, in_h, out_w, out_h, false); - #[expect( - clippy::expect_used, - reason = "output is exactly out_w*out_h*3 bytes by construction" - )] - DynamicImage::ImageRgb8( - RgbImage::from_raw(out_w, out_h, output).expect("pil resize buffer size"), - ) -} - -/// PIL-exact bicubic resize over borrowed interleaved RGB bytes. -/// -/// Byte-for-byte equivalent of [`resize_bicubic_pil`] but for the raw-RGB video -/// frame path (`preprocess_video_rgb`). Returns an `RgbImage` to drop straight -/// into the existing [`resize_rgb_bytes`] call sites. -pub fn resize_bicubic_pil_rgb( - data: &[u8], - width: u32, - height: u32, - out_w: u32, - out_h: u32, -) -> Result { - let (in_w, in_h) = (width as usize, height as usize); - let expected = in_w.saturating_mul(in_h).saturating_mul(3); - if data.len() != expected { - return Err(TransformError::ShapeError(format!( - "PIL bicubic RGB source has {} bytes, expected {expected} for {width}x{height}", - data.len() - ))); - } - let output = resize_bicubic_pil_bytes(data, width, height, out_w, out_h, true); - RgbImage::from_raw(out_w, out_h, output).ok_or_else(|| { - TransformError::ShapeError(format!( - "failed to build PIL bicubic RGB image for {out_w}x{out_h}" - )) - }) -} - -fn resize_bicubic_pil_bytes( - data: &[u8], - in_w: u32, - in_h: u32, - out_w: u32, - out_h: u32, - joint_rgb: bool, -) -> Vec { - let (in_w, in_h, out_w, out_h) = (in_w as usize, in_h as usize, out_w as usize, out_h as usize); - if in_w == out_w && in_h == out_h { - data.to_vec() - } else if in_w == out_w { - if joint_rgb { - pil_resample_vertical_rgb(data, in_h, in_w, out_h) - } else { - pil_resample_vertical(data, in_h, in_w, out_h, 3) - } - } else { - let horiz = if joint_rgb { - pil_resample_horizontal_rgb(data, in_h, in_w, out_w) - } else { - pil_resample_horizontal(data, in_h, in_w, out_w, 3) - }; - if in_h == out_h { - horiz - } else if joint_rgb { - pil_resample_vertical_rgb(&horiz, in_h, out_w, out_h) - } else { - pil_resample_vertical(&horiz, in_h, out_w, out_h, 3) - } - } -} - -/// Resize image preserving aspect ratio, fitting within max dimensions. -pub fn resize_to_fit( - image: &DynamicImage, - max_width: u32, - max_height: u32, - filter: FilterType, -) -> DynamicImage { - let (w, h) = image.dimensions(); - let ratio = (max_width as f64 / w as f64).min(max_height as f64 / h as f64); - if ratio >= 1.0 { - return image.clone(); - } - let new_w = ((w as f64 * ratio).round() as u32).max(1); - let new_h = ((h as f64 * ratio).round() as u32).max(1); - resize(image, new_w, new_h, filter) -} - -/// Center crop image to specified dimensions. -/// -/// If the crop size is larger than the image, the image is returned unchanged. -pub fn center_crop(image: &DynamicImage, crop_w: u32, crop_h: u32) -> DynamicImage { - let (w, h) = image.dimensions(); - if crop_w >= w && crop_h >= h { - return image.clone(); - } - let left = (w.saturating_sub(crop_w)) / 2; - let top = (h.saturating_sub(crop_h)) / 2; - let actual_w = crop_w.min(w); - let actual_h = crop_h.min(h); - image.crop_imm(left, top, actual_w, actual_h) -} - -/// Expand image to square by padding with background color. -/// -/// This is used by LLaVA models which expect square inputs. The image is -/// centered and padded with the mean color on the shorter dimension. -pub fn expand_to_square(image: &DynamicImage, background: Rgb) -> DynamicImage { - let (w, h) = image.dimensions(); - match w.cmp(&h) { - std::cmp::Ordering::Equal => image.clone(), - std::cmp::Ordering::Less => { - // Height > Width: pad horizontally - let mut new_image = DynamicImage::from(RgbImage::from_pixel(h, h, background)); - image::imageops::overlay(&mut new_image, image, ((h - w) / 2) as i64, 0); - new_image - } - std::cmp::Ordering::Greater => { - // Width > Height: pad vertically - let mut new_image = DynamicImage::from(RgbImage::from_pixel(w, w, background)); - image::imageops::overlay(&mut new_image, image, 0, ((w - h) / 2) as i64); - new_image - } - } -} - -/// Stack multiple [C, H, W] tensors into [B, C, H, W]. -/// -/// All tensors must have the same shape. -pub fn stack_batch(tensors: &[Array3]) -> Result> { - if tensors.is_empty() { - return Err(TransformError::EmptyBatch); - } - - let shape = tensors[0].shape(); - let (c, h, w) = (shape[0], shape[1], shape[2]); - - // Verify all tensors have the same shape - for tensor in tensors.iter().skip(1) { - if tensor.shape() != shape { - return Err(TransformError::InvalidShape { - expected: format!("[{c}, {h}, {w}]"), - actual: tensor.shape().to_vec(), - }); - } - } - - let mut batch = Array4::::zeros((tensors.len(), c, h, w)); - for (i, tensor) in tensors.iter().enumerate() { - batch.slice_mut(s![i, .., .., ..]).assign(tensor); - } - - Ok(batch) -} - -/// Convert PIL/HuggingFace resampling enum to image crate filter. -/// -/// PIL resampling constants: -/// - 0: NEAREST -/// - 1: LANCZOS (also ANTIALIAS) -/// - 2: BILINEAR -/// - 3: BICUBIC -/// - 4: BOX -/// - 5: HAMMING -pub fn pil_to_filter(resampling: Option) -> FilterType { - match resampling { - Some(0) => FilterType::Nearest, - Some(1) => FilterType::Lanczos3, - Some(2) | None => FilterType::Triangle, // Bilinear (default) - Some(3) => FilterType::CatmullRom, // Bicubic - // Box and Hamming don't have direct equivalents, use Triangle - Some(4) | Some(5) => FilterType::Triangle, - _ => FilterType::Triangle, - } -} - -/// Calculate mean color of an image as RGB. -pub fn calculate_mean_color(image: &DynamicImage) -> Rgb { - let rgb = image.to_rgb8(); - let (w, h) = (rgb.width() as u64, rgb.height() as u64); - let total_pixels = w * h; - - if total_pixels == 0 { - return Rgb([128, 128, 128]); - } - - let (mut r_sum, mut g_sum, mut b_sum) = (0u64, 0u64, 0u64); - for pixel in rgb.pixels() { - r_sum += pixel[0] as u64; - g_sum += pixel[1] as u64; - b_sum += pixel[2] as u64; - } - - Rgb([ - (r_sum / total_pixels) as u8, - (g_sum / total_pixels) as u8, - (b_sum / total_pixels) as u8, - ]) -} - -/// Convert normalized mean values [0, 1] to RGB bytes. -pub fn mean_to_rgb(mean: &[f64; 3]) -> Rgb { - Rgb([ - (mean[0] * 255.0).round() as u8, - (mean[1] * 255.0).round() as u8, - (mean[2] * 255.0).round() as u8, - ]) -} - -/// Cubic interpolation weight function (Keys bicubic kernel with a=-0.5). -/// -/// This matches PyTorch's bicubic interpolation used in -/// `torch.nn.functional.interpolate(mode='bicubic')`. -#[inline] -pub fn cubic_weight(x: f32) -> f32 { - let x = x.abs(); - if x < 1.0 { - (1.5 * x - 2.5) * x * x + 1.0 - } else if x < 2.0 { - ((-0.5 * x + 2.5) * x - 4.0) * x + 2.0 - } else { - 0.0 - } -} - -/// Perform bicubic interpolation at a single point in a tensor. -/// -/// Uses a 4x4 kernel with Keys bicubic weights (a=-0.5) to match PyTorch's -/// `torch.nn.functional.interpolate(mode='bicubic')`. -/// -/// # Arguments -/// * `tensor` - Input tensor of shape [C, H, W] -/// * `c` - Channel index -/// * `src_y` - Source Y coordinate (can be fractional) -/// * `src_x` - Source X coordinate (can be fractional) -/// * `h` - Height of the tensor -/// * `w` - Width of the tensor -/// -/// # Returns -/// The interpolated value at the specified position. -pub fn bicubic_interpolate( - tensor: &Array3, - c: usize, - src_y: f32, - src_x: f32, - h: usize, - w: usize, -) -> f32 { - let y_int = src_y.floor() as i32; - let x_int = src_x.floor() as i32; - let y_frac = src_y - y_int as f32; - let x_frac = src_x - x_int as f32; - - let mut result = 0.0f32; - - // Sample 4x4 neighborhood - for dy in -1..=2 { - let y_idx = (y_int + dy).clamp(0, h as i32 - 1) as usize; - let y_weight = cubic_weight(y_frac - dy as f32); - - for dx in -1..=2 { - let x_idx = (x_int + dx).clamp(0, w as i32 - 1) as usize; - let x_weight = cubic_weight(x_frac - dx as f32); - - result += tensor[[c, y_idx, x_idx]] * y_weight * x_weight; - } - } - - result -} - -/// Resize a tensor using bicubic interpolation. -/// -/// This matches PyTorch's `torch.nn.functional.interpolate(mode='bicubic', align_corners=False)`. -/// -/// # Arguments -/// * `tensor` - Input tensor of shape [C, H, W] -/// * `target_h` - Target height -/// * `target_w` - Target width -/// -/// # Returns -/// Resized tensor of shape [C, target_h, target_w]. -pub fn bicubic_resize(tensor: &Array3, target_h: usize, target_w: usize) -> Array3 { - let (c, h, w) = (tensor.shape()[0], tensor.shape()[1], tensor.shape()[2]); - - if h == target_h && w == target_w { - return tensor.clone(); - } - - let mut result = Array3::::zeros((c, target_h, target_w)); - - // PyTorch align_corners=False coordinate mapping - let scale_h = h as f32 / target_h as f32; - let scale_w = w as f32 / target_w as f32; - - for ch in 0..c { - for y in 0..target_h { - for x in 0..target_w { - // PyTorch align_corners=False: src = (dst + 0.5) * scale - 0.5 - let src_y = (y as f32 + 0.5) * scale_h - 0.5; - let src_x = (x as f32 + 0.5) * scale_w - 0.5; - - result[[ch, y, x]] = bicubic_interpolate(tensor, ch, src_y, src_x, h, w); - } - } - } - - result -} - -#[cfg(test)] -mod tests { - use super::*; - - fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { - DynamicImage::from(RgbImage::from_pixel(width, height, color)) - } - - /// The raw-RGB video resizer must be byte-for-byte identical to the - /// DynamicImage PIL-bicubic resizer used for images. Guards the video resize - /// path used by `preprocess_video_rgb`. - #[test] - fn resize_bicubic_pil_rgb_matches_dynamic_path() { - let (src_w, src_h) = (37u32, 23u32); // non-aligned source, non-trivial ratios - let (out_w, out_h) = (16u32, 28u32); // downscale width, upscale height - let mut img = RgbImage::new(src_w, src_h); - for y in 0..src_h { - for x in 0..src_w { - img.put_pixel( - x, - y, - Rgb([ - ((x * 7) ^ (y * 13)) as u8, - (x * 3 + y * 5) as u8, - (x + y * y) as u8, - ]), - ); - } - } - let via_dynamic = resize_bicubic_pil(&DynamicImage::ImageRgb8(img.clone()), out_w, out_h); - let via_bytes = resize_bicubic_pil_rgb(img.as_raw(), src_w, src_h, out_w, out_h).unwrap(); - assert_eq!( - via_dynamic.to_rgb8().into_raw(), - via_bytes.into_raw(), - "raw-RGB PIL bicubic must equal DynamicImage PIL bicubic byte-for-byte" - ); - } - - #[test] - fn resize_bicubic_pil_rgb_skips_identity_axes_bit_exactly() { - let (src_w, src_h) = (31u32, 23u32); - let mut data = vec![0u8; src_w as usize * src_h as usize * 3]; - for (index, value) in data.iter_mut().enumerate() { - *value = (index as u8).wrapping_mul(37).wrapping_add(11); - } - - for (out_w, out_h) in [(src_w, 17), (19, src_h), (src_w, src_h)] { - let horizontal = - pil_resample_horizontal(&data, src_h as usize, src_w as usize, out_w as usize, 3); - let expected = pil_resample_vertical( - &horizontal, - src_h as usize, - out_w as usize, - out_h as usize, - 3, - ); - let actual = resize_bicubic_pil_rgb(&data, src_w, src_h, out_w, out_h) - .unwrap() - .into_raw(); - - assert_eq!(actual, expected, "identity-axis fast path changed pixels"); - } - } - - /// `resize_bicubic_pil_rgb` rejects a buffer whose length doesn't match the - /// declared dimensions rather than reading out of bounds. - #[test] - fn resize_bicubic_pil_rgb_rejects_wrong_length() { - assert!( - resize_bicubic_pil_rgb(&[0u8; 10], 4, 4, 2, 2).is_err(), - "wrong-length RGB buffer must error, not panic" - ); - } - - #[test] - fn test_to_tensor_shape() { - let img = create_test_image(10, 20, Rgb([255, 128, 0])); - let tensor = to_tensor(&img); - assert_eq!(tensor.shape(), &[3, 20, 10]); // [C, H, W] - } - - #[test] - fn test_to_tensor_values() { - let img = create_test_image(2, 2, Rgb([255, 128, 0])); - let tensor = to_tensor(&img); - - // Check normalization to [0, 1] - assert!((tensor[[0, 0, 0]] - 1.0).abs() < 1e-6); // R=255 -> 1.0 - assert!((tensor[[1, 0, 0]] - 0.502).abs() < 0.01); // G=128 -> ~0.5 - assert!((tensor[[2, 0, 0]] - 0.0).abs() < 1e-6); // B=0 -> 0.0 - } - - #[test] - fn test_to_tensor_no_norm() { - let img = create_test_image(2, 2, Rgb([255, 128, 64])); - let tensor = to_tensor_no_norm(&img); - - assert!((tensor[[0, 0, 0]] - 255.0).abs() < 1e-6); - assert!((tensor[[1, 0, 0]] - 128.0).abs() < 1e-6); - assert!((tensor[[2, 0, 0]] - 64.0).abs() < 1e-6); - } - - #[test] - fn test_normalize() { - let mut tensor = Array3::::from_elem((3, 2, 2), 0.5); - let mean = [0.5, 0.5, 0.5]; - let std = [0.5, 0.5, 0.5]; - - normalize(&mut tensor, &mean, &std); - - // (0.5 - 0.5) / 0.5 = 0.0 - for val in &tensor { - assert!(val.abs() < 1e-6); - } - } - - #[test] - fn test_rescale() { - let mut tensor = Array3::::from_elem((3, 2, 2), 255.0); - rescale(&mut tensor, 1.0 / 255.0); - - for val in &tensor { - assert!((val - 1.0).abs() < 1e-6); - } - } - - #[test] - fn test_resize() { - let img = create_test_image(100, 50, Rgb([128, 128, 128])); - let resized = resize(&img, 50, 25, FilterType::Triangle); - - assert_eq!(resized.width(), 50); - assert_eq!(resized.height(), 25); - } - - #[test] - fn test_resize_rgb_bytes_matches_resize() { - let mut rgb = RgbImage::new(3, 2); - for y in 0..2 { - for x in 0..3 { - rgb.put_pixel(x, y, Rgb([(x * 40) as u8, (y * 90) as u8, 128])); - } - } - let img = DynamicImage::ImageRgb8(rgb.clone()); - let expected = resize(&img, 2, 2, FilterType::Triangle).to_rgb8(); - let actual = resize_rgb_bytes(rgb.as_raw(), 3, 2, 2, 2, FilterType::Triangle) - .expect("resize_rgb_bytes should resize valid RGB input"); - - assert_eq!(actual.as_raw(), expected.as_raw()); - } - - #[test] - fn test_resize_rgb_bytes_rejects_invalid_length() { - let result = resize_rgb_bytes(&[1, 2, 3, 4, 5], 2, 1, 1, 1, FilterType::Triangle); - assert!(matches!(result, Err(TransformError::ShapeError(_)))); - } - - #[test] - fn test_center_crop() { - let img = create_test_image(100, 100, Rgb([128, 128, 128])); - let cropped = center_crop(&img, 50, 50); - - assert_eq!(cropped.width(), 50); - assert_eq!(cropped.height(), 50); - } - - #[test] - fn test_expand_to_square_horizontal() { - let img = create_test_image(100, 50, Rgb([255, 0, 0])); - let background = Rgb([0, 0, 0]); - let squared = expand_to_square(&img, background); - - assert_eq!(squared.width(), 100); - assert_eq!(squared.height(), 100); - } - - #[test] - fn test_expand_to_square_vertical() { - let img = create_test_image(50, 100, Rgb([255, 0, 0])); - let background = Rgb([0, 0, 0]); - let squared = expand_to_square(&img, background); - - assert_eq!(squared.width(), 100); - assert_eq!(squared.height(), 100); - } - - #[test] - fn test_expand_to_square_already_square() { - let img = create_test_image(100, 100, Rgb([255, 0, 0])); - let background = Rgb([0, 0, 0]); - let squared = expand_to_square(&img, background); - - assert_eq!(squared.width(), 100); - assert_eq!(squared.height(), 100); - } - - #[test] - fn test_stack_batch() { - let t1 = Array3::::zeros((3, 10, 10)); - let t2 = Array3::::ones((3, 10, 10)); - - let batch = stack_batch(&[t1, t2]).unwrap(); - - assert_eq!(batch.shape(), &[2, 3, 10, 10]); - } - - #[test] - fn test_stack_batch_empty() { - let result = stack_batch(&[]); - assert!(matches!(result, Err(TransformError::EmptyBatch))); - } - - #[test] - fn test_pil_to_filter() { - assert!(matches!(pil_to_filter(Some(0)), FilterType::Nearest)); - assert!(matches!(pil_to_filter(Some(1)), FilterType::Lanczos3)); - assert!(matches!(pil_to_filter(Some(2)), FilterType::Triangle)); - assert!(matches!(pil_to_filter(Some(3)), FilterType::CatmullRom)); - assert!(matches!(pil_to_filter(None), FilterType::Triangle)); - } - - #[test] - fn test_mean_to_rgb() { - let mean = [0.5, 0.25, 1.0]; - let rgb = mean_to_rgb(&mean); - - assert_eq!(rgb[0], 128); - assert_eq!(rgb[1], 64); - assert_eq!(rgb[2], 255); - } -} diff --git a/crates/multimodal/tests/decode_preprocess_bench.rs b/crates/multimodal/tests/decode_preprocess_bench.rs deleted file mode 100644 index 64e8fda6b..000000000 --- a/crates/multimodal/tests/decode_preprocess_bench.rs +++ /dev/null @@ -1,69 +0,0 @@ -//! Microbench: SMG (Rust) JPEG decode (libjpeg-turbo) + Qwen3-VL preprocess, -//! on a real image + the real model preprocessor config. Compare against the -//! HF/PIL path that vLLM uses (scripts: bench_hf_preprocess.py). -//! -//! Run: -//! REAL_JPEG=/path/x.jpg PP_CONFIG=/path/preprocessor_config.json \ -//! cargo test -p llm-multimodal --test decode_preprocess_bench -- --ignored --nocapture -#![allow( - clippy::unwrap_used, - clippy::expect_used, - clippy::print_stdout, - clippy::print_stderr -)] -use std::time::Instant; - -use llm_multimodal::{ - jpeg_turbo, - vision::{ - preprocessor_config::PreProcessorConfig, processors::Qwen3VLProcessor, VisionPreProcessor, - }, -}; - -#[test] -#[ignore = "perf microbench; needs REAL_JPEG + PP_CONFIG"] -fn bench_decode_preprocess() { - let jpeg_path = std::env::var("REAL_JPEG").expect("set REAL_JPEG to a JPEG path"); - let cfg_path = std::env::var("PP_CONFIG").expect("set PP_CONFIG to preprocessor_config.json"); - - let bytes = std::fs::read(&jpeg_path).expect("read jpeg"); - let config = - PreProcessorConfig::from_json(&std::fs::read_to_string(&cfg_path).expect("read config")) - .expect("parse preprocessor config"); - let proc = Qwen3VLProcessor::new(); - - // warmup - let img = jpeg_turbo::decode_jpeg_rgb(&bytes).expect("turbojpeg decode"); - let _ = proc - .preprocess(std::slice::from_ref(&img), &config) - .expect("preprocess"); - - let n_dec = 300usize; - let t0 = Instant::now(); - for _ in 0..n_dec { - let _ = jpeg_turbo::decode_jpeg_rgb(&bytes).unwrap(); - } - let dec_ms = t0.elapsed().as_secs_f64() * 1000.0 / n_dec as f64; - - let n_pp = 200usize; - let t1 = Instant::now(); - for _ in 0..n_pp { - let _ = proc - .preprocess(std::slice::from_ref(&img), &config) - .unwrap(); - } - let pp_ms = t1.elapsed().as_secs_f64() * 1000.0 / n_pp as f64; - - eprintln!( - "image: {}x{} ({} bytes jpeg)", - img.width(), - img.height(), - bytes.len() - ); - eprintln!("SMG(Rust) decode (libjpeg-turbo): {dec_ms:.3} ms/img [{n_dec} iters]"); - eprintln!("SMG(Rust) preprocess (Qwen3-VL) : {pp_ms:.3} ms/img [{n_pp} iters]"); - eprintln!( - "SMG(Rust) decode+preprocess total: {:.3} ms/img", - dec_ms + pp_ms - ); -} diff --git a/crates/multimodal/tests/fixtures/golden/qwen_preprocess_fingerprints.json b/crates/multimodal/tests/fixtures/golden/qwen_preprocess_fingerprints.json deleted file mode 100644 index 2f88859fc..000000000 --- a/crates/multimodal/tests/fixtures/golden/qwen_preprocess_fingerprints.json +++ /dev/null @@ -1,50 +0,0 @@ -{ - "generator": "generate_qwen_preprocess_fingerprints.py", - "pillow": "12.2.0", - "transformers": "4.57.1", - "cases": [ - { - "model": "qwen2_vl", - "width": 37, - "height": 23, - "shape": [1092, 1176], - "grid_thw": [1, 26, 42], - "fnv1a_patch_u8": "783a9142f0cbefb9" - }, - { - "model": "qwen2_vl", - "width": 259, - "height": 194, - "shape": [1064, 1176], - "grid_thw": [1, 28, 38], - "fnv1a_patch_u8": "6081ce380d38da91" - }, - { - "model": "qwen3_vl", - "width": 37, - "height": 23, - "shape": [308, 1536], - "grid_thw": [1, 14, 22], - "fnv1a_patch_u8": "ea0821a9c7700bad" - }, - { - "model": "qwen3_vl", - "width": 259, - "height": 194, - "shape": [280, 1536], - "grid_thw": [1, 14, 20], - "fnv1a_patch_u8": "45f7361bfce1b071" - } - ], - "video_cases": [ - { - "model": "qwen3_vl", - "width": 37, - "height": 35, - "frame_count": 3, - "shape": [8, 1536], - "grid_thw": [2, 2, 2], - "fnv1a_patch_u8": "c862984777479f6d" - } - ] -} diff --git a/crates/multimodal/tests/fixtures/images/grayscale.jpg b/crates/multimodal/tests/fixtures/images/grayscale.jpg deleted file mode 100644 index 1ffe2a23f..000000000 Binary files a/crates/multimodal/tests/fixtures/images/grayscale.jpg and /dev/null differ diff --git a/crates/multimodal/tests/fixtures/images/large.jpg b/crates/multimodal/tests/fixtures/images/large.jpg deleted file mode 100644 index eac222d21..000000000 Binary files a/crates/multimodal/tests/fixtures/images/large.jpg and /dev/null differ diff --git a/crates/multimodal/tests/fixtures/images/odd_dims.jpg b/crates/multimodal/tests/fixtures/images/odd_dims.jpg deleted file mode 100644 index 580d57493..000000000 Binary files a/crates/multimodal/tests/fixtures/images/odd_dims.jpg and /dev/null differ diff --git a/crates/multimodal/tests/fixtures/images/small.jpg b/crates/multimodal/tests/fixtures/images/small.jpg deleted file mode 100644 index f859d3c52..000000000 Binary files a/crates/multimodal/tests/fixtures/images/small.jpg and /dev/null differ diff --git a/crates/multimodal/tests/fixtures/images/square.jpg b/crates/multimodal/tests/fixtures/images/square.jpg deleted file mode 100644 index 705344c22..000000000 Binary files a/crates/multimodal/tests/fixtures/images/square.jpg and /dev/null differ diff --git a/crates/multimodal/tests/fixtures/images/tall.jpg b/crates/multimodal/tests/fixtures/images/tall.jpg deleted file mode 100644 index 9d798d9bc..000000000 Binary files a/crates/multimodal/tests/fixtures/images/tall.jpg and /dev/null differ diff --git a/crates/multimodal/tests/fixtures/images/tiny.jpg b/crates/multimodal/tests/fixtures/images/tiny.jpg deleted file mode 100644 index 9f60b2d37..000000000 Binary files a/crates/multimodal/tests/fixtures/images/tiny.jpg and /dev/null differ diff --git a/crates/multimodal/tests/fixtures/images/very_tall.jpg b/crates/multimodal/tests/fixtures/images/very_tall.jpg deleted file mode 100644 index 419726ac2..000000000 Binary files a/crates/multimodal/tests/fixtures/images/very_tall.jpg and /dev/null differ diff --git a/crates/multimodal/tests/fixtures/images/very_wide.jpg b/crates/multimodal/tests/fixtures/images/very_wide.jpg deleted file mode 100644 index 87a0e28f7..000000000 Binary files a/crates/multimodal/tests/fixtures/images/very_wide.jpg and /dev/null differ diff --git a/crates/multimodal/tests/fixtures/images/wide.jpg b/crates/multimodal/tests/fixtures/images/wide.jpg deleted file mode 100644 index 0f644cc6d..000000000 Binary files a/crates/multimodal/tests/fixtures/images/wide.jpg and /dev/null differ diff --git a/crates/multimodal/tests/multimodal_tracker_test.rs b/crates/multimodal/tests/multimodal_tracker_test.rs deleted file mode 100644 index 1efd82365..000000000 --- a/crates/multimodal/tests/multimodal_tracker_test.rs +++ /dev/null @@ -1,171 +0,0 @@ -use std::{path::PathBuf, sync::Arc, time::Duration}; - -use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; -use llm_multimodal::{ - AsyncMultiModalTracker, AudioSource, ImageFetchConfig, ImageSource, MediaConnector, - MediaConnectorConfig, MediaContentPart, MediaSource, Modality, -}; -use reqwest::Client; -use tempfile::tempdir; - -const TINY_PNG_BASE64: &str = - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNgYAAAAAMAASsJTYQAAAAASUVORK5CYII="; - -#[expect( - clippy::expect_used, - reason = "test helper: panic on failure is intentional" -)] -fn tiny_png_bytes() -> Vec { - BASE64_STANDARD - .decode(TINY_PNG_BASE64) - .expect("decode tiny png fixture") -} - -fn wav_i16_mono(sample_rate: u32, samples: &[i16]) -> Vec { - let data_bytes = samples.len() as u32 * 2; - let mut bytes = Vec::new(); - bytes.extend_from_slice(b"RIFF"); - bytes.extend_from_slice(&(36 + data_bytes).to_le_bytes()); - bytes.extend_from_slice(b"WAVEfmt "); - bytes.extend_from_slice(&16_u32.to_le_bytes()); - bytes.extend_from_slice(&1_u16.to_le_bytes()); - bytes.extend_from_slice(&1_u16.to_le_bytes()); - bytes.extend_from_slice(&sample_rate.to_le_bytes()); - bytes.extend_from_slice(&(sample_rate * 2).to_le_bytes()); - bytes.extend_from_slice(&2_u16.to_le_bytes()); - bytes.extend_from_slice(&16_u16.to_le_bytes()); - bytes.extend_from_slice(b"data"); - bytes.extend_from_slice(&data_bytes.to_le_bytes()); - for sample in samples { - bytes.extend_from_slice(&sample.to_le_bytes()); - } - bytes -} - -#[expect( - clippy::expect_used, - reason = "test helper: panic on failure is intentional" -)] -fn test_connector(allowed_path: Option) -> MediaConnector { - let client = Client::builder() - .timeout(Duration::from_secs(5)) - .no_proxy() - .build() - .expect("client"); - MediaConnector::new( - client, - MediaConnectorConfig { - allowed_domains: None, - allowed_local_media_path: allowed_path, - fetch_timeout: Duration::from_secs(5), - }, - ) - .expect("media connector") -} - -#[tokio::test] -async fn fetch_image_from_inline_bytes() { - let connector = test_connector(None); - let bytes = tiny_png_bytes(); - let frame = connector - .fetch_image( - MediaSource::InlineBytes(bytes.clone()), - ImageFetchConfig::default(), - ) - .await - .expect("inline image"); - assert_eq!(frame.data().width(), 1); - assert_eq!(frame.data().height(), 1); - assert_eq!(frame.raw_bytes(), bytes.as_slice()); -} - -#[tokio::test] -async fn fetch_image_from_data_url() { - let connector = test_connector(None); - let bytes = tiny_png_bytes(); - let data_url = format!( - "data:image/png;base64,{}", - BASE64_STANDARD.encode(bytes.clone()) - ); - - let frame = connector - .fetch_image(MediaSource::DataUrl(data_url), ImageFetchConfig::default()) - .await - .expect("data url"); - assert_eq!(frame.data().width(), 1); - assert_eq!(frame.raw_bytes(), bytes.as_slice()); -} - -#[tokio::test] -async fn fetch_image_from_file() { - let tmp = tempdir().expect("tempdir"); - let allowed_root = std::fs::canonicalize(tmp.path()).expect("canonical tmp path"); - let file_path = allowed_root.join("tiny.png"); - std::fs::write(&file_path, tiny_png_bytes()).expect("write png"); - - let connector = test_connector(Some(allowed_root)); - let frame = connector - .fetch_image( - MediaSource::File(file_path.clone()), - ImageFetchConfig::default(), - ) - .await - .expect("file png"); - assert_eq!(frame.data().width(), 1); - let expected = std::fs::canonicalize(&file_path).expect("canonical path"); - match frame.source() { - ImageSource::File { path } => assert_eq!(path, &expected), - other => panic!("expected file source, got {other:?}"), - } -} - -#[tokio::test] -async fn fetch_audio_from_inline_bytes_decodes_samples() { - let connector = test_connector(None); - let bytes = wav_i16_mono(16_000, &[0, 16_384, -16_384]); - let clip = connector - .fetch_audio(MediaSource::InlineBytes(bytes.clone())) - .await - .expect("inline audio"); - - assert_eq!(clip.raw_bytes(), bytes.as_slice()); - assert_eq!(clip.decoded().sample_rate, 16_000); - assert_eq!(clip.decoded().samples.len(), 3); - assert!(clip.decoded().samples[0].abs() < 1e-6); - assert!((clip.decoded().samples[1] - 0.5).abs() < 1e-4); - assert!((clip.decoded().samples[2] + 0.5).abs() < 1e-4); - assert!(matches!(clip.source(), AudioSource::InlineBytes)); -} - -#[tokio::test] -async fn tracker_fetches_images_and_records_uuids() { - let connector = Arc::new(test_connector(None)); - let mut tracker = AsyncMultiModalTracker::new(connector); - - tracker - .push_part(MediaContentPart::Text { - text: "before".into(), - }) - .expect("text part"); - tracker - .push_part(MediaContentPart::ImageData { - data: tiny_png_bytes(), - mime_type: Some("image/png".into()), - uuid: Some("img-1".into()), - detail: None, - }) - .expect("image part"); - tracker - .push_part(MediaContentPart::Text { - text: "after".into(), - }) - .expect("text part"); - - let output = tracker.finalize().await.expect("tracker finalize"); - - let images = output.data.get(&Modality::Image).expect("image entry"); - assert_eq!(images.len(), 1); - - let uuids = output.uuids.get(&Modality::Image).expect("uuid entry"); - assert_eq!(uuids, &vec![Some("img-1".into())]); -} diff --git a/crates/multimodal/tests/preprocess_fingerprint.rs b/crates/multimodal/tests/preprocess_fingerprint.rs deleted file mode 100644 index 85c84b7f7..000000000 --- a/crates/multimodal/tests/preprocess_fingerprint.rs +++ /dev/null @@ -1,80 +0,0 @@ -//! Bit-identity guard for the full Qwen3-VL preprocess (resize + normalize + -//! patchify). Pins the EXACT f32 encoder_input bytes. Any perf change to those -//! stages (parallelization) MUST keep these identical to preserve vLLM/PIL -//! parity (accuracy). -#![allow( - clippy::unwrap_used, - clippy::expect_used, - clippy::print_stdout, - clippy::print_stderr -)] -use image::{DynamicImage, RgbImage}; -use llm_multimodal::vision::{ - preprocessor_config::PreProcessorConfig, processors::Qwen3VLProcessor, VisionPreProcessor, -}; - -fn make(w: u32, h: u32) -> DynamicImage { - let img = RgbImage::from_fn(w, h, |x, y| { - image::Rgb([ - ((x * 7 + y * 3) % 256) as u8, - ((x * 5 + y * 11) % 256) as u8, - ((x + y * 2) % 256) as u8, - ]) - }); - DynamicImage::ImageRgb8(img) -} - -fn config() -> PreProcessorConfig { - PreProcessorConfig::from_json( - r#"{"do_resize":true,"do_normalize":true, - "image_mean":[0.48145466,0.4578275,0.40821073], - "image_std":[0.26862954,0.26130258,0.27577711], - "size":{"shortest_edge":3136,"longest_edge":12845056}, - "resample":3}"#, - ) - .unwrap() -} - -fn fnv1a_f32(data: &[f32]) -> u64 { - let mut h: u64 = 0xcbf2_9ce4_8422_2325; - for &v in data { - for b in v.to_le_bytes() { - h ^= b as u64; - h = h.wrapping_mul(0x0000_0100_0000_01b3); - } - } - h -} - -const CASES: &[(u32, u32)] = &[(560, 420), (840, 560), (1280, 960)]; - -// Captured under serial normalize/patchify; PARALLELIZATION MUST NOT CHANGE THESE. -const EXPECTED: &[u64] = &[0x391ca5deba1ff255, 0x5bde4728a72eba9d, 0x617d3e39f58f1c45]; - -fn fingerprint(w: u32, h: u32) -> (u64, usize) { - let proc = Qwen3VLProcessor::new(); - let res = proc.preprocess(&[make(w, h)], &config()).unwrap(); - let flat = res.encoder_input_flat(); - (fnv1a_f32(flat.as_ref()), flat.len()) -} - -#[test] -#[ignore = "capture mode"] -fn capture_preprocess_fingerprints() { - for (w, h) in CASES { - let (fp, n) = fingerprint(*w, *h); - println!("{w}x{h}: 0x{fp:016x} (len {n})"); - } -} - -#[test] -fn preprocess_bit_identity() { - if EXPECTED.iter().all(|&v| v == 0) { - eprintln!("EXPECTED not filled; run capture_preprocess_fingerprints"); - return; - } - for ((w, h), &exp) in CASES.iter().zip(EXPECTED) { - let (fp, _) = fingerprint(*w, *h); - assert_eq!(fp, exp, "preprocess fingerprint changed for {w}x{h}"); - } -} diff --git a/crates/multimodal/tests/qwen_preprocess_golden.rs b/crates/multimodal/tests/qwen_preprocess_golden.rs deleted file mode 100644 index 269caf3a5..000000000 --- a/crates/multimodal/tests/qwen_preprocess_golden.rs +++ /dev/null @@ -1,191 +0,0 @@ -//! Mandatory HuggingFace golden checks for Qwen image preprocessing. -#![allow(clippy::expect_used, clippy::panic)] - -use image::{DynamicImage, RgbImage}; -use llm_multimodal::vision::{ - processor::ModelSpecificValue, PreProcessorConfig, Qwen2VLProcessor, Qwen3VLProcessor, - VisionPreProcessor, -}; -use serde::Deserialize; - -#[derive(Deserialize)] -struct GoldenDocument { - generator: String, - pillow: String, - transformers: String, - cases: Vec, - video_cases: Vec, -} - -#[derive(Deserialize)] -struct GoldenCase { - model: String, - width: u32, - height: u32, - shape: Vec, - grid_thw: Vec, - fnv1a_patch_u8: String, -} - -#[derive(Deserialize)] -struct GoldenVideoCase { - model: String, - width: u32, - height: u32, - frame_count: usize, - shape: Vec, - grid_thw: Vec, - fnv1a_patch_u8: String, -} - -fn make_image(width: u32, height: u32) -> DynamicImage { - make_seeded_image(width, height, 0) -} - -fn make_seeded_image(width: u32, height: u32, seed: u8) -> DynamicImage { - DynamicImage::ImageRgb8(RgbImage::from_fn(width, height, |x, y| { - image::Rgb([ - seed.wrapping_add(((x * 7 + y * 3) % 256) as u8), - seed.wrapping_add(((x * 5 + y * 11) % 256) as u8), - seed.wrapping_add(((x + y * 2) % 256) as u8), - ]) - })) -} - -fn image_grid(result: &llm_multimodal::vision::PreprocessedEncoderInputs) -> &[i64] { - match result.model_specific.get("image_grid_thw") { - Some(ModelSpecificValue::IntTensor { data, shape }) => { - assert_eq!(shape, &[1, 3]); - data - } - value => panic!("expected image_grid_thw IntTensor, got {value:?}"), - } -} - -fn video_grid(result: &llm_multimodal::vision::PreprocessedEncoderInputs) -> &[i64] { - match result.model_specific.get("video_grid_thw") { - Some(ModelSpecificValue::IntTensor { data, shape }) => { - assert_eq!(shape, &[1, 3]); - data - } - value => panic!("expected video_grid_thw IntTensor, got {value:?}"), - } -} - -fn fnv1a_patch_u8(values: &[f32]) -> u64 { - let mut hash = 0xcbf2_9ce4_8422_2325_u64; - for value in values { - let byte = (value * 255.0).round_ties_even().clamp(0.0, 255.0) as u8; - hash ^= u64::from(byte); - hash = hash.wrapping_mul(0x0000_0100_0000_01b3); - } - hash -} - -fn config(mean: [f64; 3], std: [f64; 3]) -> PreProcessorConfig { - PreProcessorConfig { - do_resize: Some(true), - do_normalize: Some(false), - image_mean: Some(mean.to_vec()), - image_std: Some(std.to_vec()), - resampling: Some(3), - ..Default::default() - } -} - -fn check_case(processor: &dyn VisionPreProcessor, config: &PreProcessorConfig, case: &GoldenCase) { - let result = processor - .preprocess(&[make_image(case.width, case.height)], config) - .expect("Qwen golden preprocessing failed"); - assert_eq!(result.encoder_input.shape(), case.shape); - - assert_eq!(image_grid(&result), &case.grid_thw); - - let expected = - u64::from_str_radix(&case.fnv1a_patch_u8, 16).expect("invalid golden FNV-1a fingerprint"); - let values = result - .encoder_input - .as_slice_memory_order() - .expect("Qwen encoder input must be contiguous"); - assert_eq!( - fnv1a_patch_u8(values), - expected, - "{} {}x{} patchified pixels differ from HuggingFace; first values: {:?}", - case.model, - case.width, - case.height, - // Keep a small sample in failure output to distinguish resize/layout - // regressions from one-bit normalization differences. - &values[..values.len().min(24)] - ); -} - -fn check_video_case( - processor: &Qwen3VLProcessor, - config: &PreProcessorConfig, - case: &GoldenVideoCase, -) { - assert_eq!(case.model, "qwen3_vl"); - let seeds = [3, 101, 177]; - assert_eq!(case.frame_count, seeds.len()); - let frames = seeds - .into_iter() - .map(|seed| make_seeded_image(case.width, case.height, seed)) - .collect::>(); - let result = processor - .preprocess_video(&frames, config) - .expect("Qwen video golden preprocessing failed"); - assert_eq!(result.encoder_input.shape(), case.shape); - assert_eq!(video_grid(&result), &case.grid_thw); - - let expected = u64::from_str_radix(&case.fnv1a_patch_u8, 16) - .expect("invalid video golden FNV-1a fingerprint"); - let values = result - .encoder_input - .as_slice_memory_order() - .expect("Qwen video encoder input must be contiguous"); - assert_eq!( - fnv1a_patch_u8(values), - expected, - "Qwen3-VL {}x{}x{} patchified video differs from HF/Pillow", - case.width, - case.height, - case.frame_count - ); -} - -#[test] -fn qwen_preprocessing_matches_huggingface_golden() { - let golden: GoldenDocument = serde_json::from_str(include_str!( - "fixtures/golden/qwen_preprocess_fingerprints.json" - )) - .expect("invalid checked-in Qwen golden fixture"); - assert_eq!(golden.generator, "generate_qwen_preprocess_fingerprints.py"); - assert!(!golden.transformers.is_empty()); - assert!(!golden.pillow.is_empty()); - assert_eq!(golden.cases.len(), 4, "Qwen golden coverage changed"); - assert_eq!( - golden.video_cases.len(), - 1, - "Qwen video golden coverage changed" - ); - - let qwen2 = Qwen2VLProcessor::new(); - let qwen2_config = config( - [0.48145466, 0.4578275, 0.40821073], - [0.26862954, 0.26130258, 0.27577711], - ); - let qwen3 = Qwen3VLProcessor::new(); - let qwen3_config = config([0.5; 3], [0.5; 3]); - - for case in &golden.cases { - match case.model.as_str() { - "qwen2_vl" => check_case(&qwen2, &qwen2_config, case), - "qwen3_vl" => check_case(&qwen3, &qwen3_config, case), - model => panic!("unknown Qwen golden model {model}"), - } - } - for case in &golden.video_cases { - check_video_case(&qwen3, &qwen3_config, case); - } -} diff --git a/crates/multimodal/tests/resize_fingerprint.rs b/crates/multimodal/tests/resize_fingerprint.rs deleted file mode 100644 index f1ea416f7..000000000 --- a/crates/multimodal/tests/resize_fingerprint.rs +++ /dev/null @@ -1,77 +0,0 @@ -//! Bit-identity guard for resize_bicubic_pil. The fingerprints below pin the -//! EXACT byte output of the Pillow-exact BICUBIC resize. Any change to the -//! resize (e.g. parallelization for speed) MUST keep these identical — the -//! resize feeds vision-encoder input, so its output must stay bit-for-bit -//! stable to preserve vLLM/PIL parity (accuracy). -#![allow( - clippy::unwrap_used, - clippy::expect_used, - clippy::print_stdout, - clippy::print_stderr -)] -use image::{DynamicImage, RgbImage}; -use llm_multimodal::vision::transforms::resize_bicubic_pil; - -fn make(w: u32, h: u32) -> DynamicImage { - // deterministic, non-trivial structure across all 3 channels - let img = RgbImage::from_fn(w, h, |x, y| { - image::Rgb([ - ((x * 7 + y * 3) % 256) as u8, - ((x * 5 + y * 11) % 256) as u8, - ((x + y * 2) % 256) as u8, - ]) - }); - DynamicImage::ImageRgb8(img) -} - -fn fnv1a(bytes: &[u8]) -> u64 { - let mut h: u64 = 0xcbf2_9ce4_8422_2325; - for &b in bytes { - h ^= b as u64; - h = h.wrapping_mul(0x0000_0100_0000_01b3); - } - h -} - -const CASES: &[(u32, u32, u32, u32)] = &[ - (800, 600, 200, 150), // downscale - (640, 480, 336, 336), // downscale to square - (1280, 960, 512, 384), // downscale large - (259, 194, 280, 196), // slight upscale (real MMBench-ish) - (200, 200, 700, 700), // upscale -]; - -// Captured under the serial implementation; PARALLELIZATION MUST NOT CHANGE THESE. -const EXPECTED: &[u64] = &[ - 0xac15afa8701536c4, - 0x9b76033374b3e1a2, - 0xf21c4a6ac5c20c83, - 0xb170a19da1087feb, - 0xf09a480918a7e2ad, -]; - -#[test] -#[ignore = "capture mode: prints fingerprints"] -fn capture_resize_fingerprints() { - for (iw, ih, ow, oh) in CASES { - let out = resize_bicubic_pil(&make(*iw, *ih), *ow, *oh); - let h = fnv1a(out.to_rgb8().as_raw()); - println!("{iw}x{ih}->{ow}x{oh}: 0x{h:016x}"); - } -} - -#[test] -fn resize_bicubic_pil_bit_identity() { - if EXPECTED.iter().all(|&v| v == 0) { - eprintln!("EXPECTED not yet filled; run capture_resize_fingerprints"); - return; - } - for ((iw, ih, ow, oh), &exp) in CASES.iter().zip(EXPECTED) { - let out = resize_bicubic_pil(&make(*iw, *ih), *ow, *oh); - let got = fnv1a(out.to_rgb8().as_raw()); - assert_eq!( - got, exp, - "resize fingerprint changed for {iw}x{ih}->{ow}x{oh}" - ); - } -} diff --git a/crates/multimodal/tests/vision_golden_tests.rs b/crates/multimodal/tests/vision_golden_tests.rs deleted file mode 100644 index 5baa57993..000000000 --- a/crates/multimodal/tests/vision_golden_tests.rs +++ /dev/null @@ -1,1493 +0,0 @@ -//! Golden tests for vision processors. -//! -//! These tests compare Rust preprocessor output against golden outputs -//! generated by HuggingFace transformers to ensure pixel-perfect compatibility. -//! -//! Modes tested: -//! - `llava/` - Standard CLIP processing (llava-hf/* models, no expand-to-square) -//! - `llava_pad/` - Expand-to-square mode (liuhaotian/llava-* models, image_aspect_ratio=pad) -//! - `qwen2_vl/` - Dynamic resolution with smart resize (Qwen/Qwen2-VL-* models) -//! - `qwen3_vl/` - Dynamic resolution with patch_size=16 and [0.5,0.5,0.5] norm (Qwen/Qwen3-VL-* models) -//! -//! To regenerate golden outputs: -//! ```bash -//! python crates/multimodal/scripts/generate_vision_golden.py -//! ``` - -#![expect( - clippy::expect_used, - reason = "integration test helpers: panic on failure is intentional" -)] -#![expect(clippy::print_stdout, reason = "integration tests: diagnostic output")] -#![expect( - clippy::print_stderr, - reason = "integration tests: diagnostic output for missing fixtures" -)] -#![expect( - clippy::panic, - reason = "integration test helpers: assertion-like panics" -)] - -use std::{fs::File, io::Read, path::Path}; - -use llm_multimodal::vision::{ - Llama4VisionProcessor, LlavaProcessor, ModelSpecificValue, Phi3VisionProcessor, - Phi4VisionProcessor, PixtralProcessor, PreProcessorConfig, Qwen2VLProcessor, Qwen3VLProcessor, - VisionPreProcessor, -}; -use ndarray::{Array4, Array5}; - -/// Load a numpy .npz file and extract pixel_values -fn load_golden_npz(path: &Path) -> Array4 { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - // Read pixel_values array (npz stores arrays without .npy extension in the lookup) - let reader = npz - .by_name("pixel_values") - .expect("Failed to read npz") - .expect("No pixel_values"); - - // Get shape from npy header - let shape = reader.shape().to_vec(); - assert_eq!(shape.len(), 4, "Expected 4D tensor [B, C, H, W]"); - - // Read data as f32 vec - let data: Vec = reader.into_vec().expect("Failed to read array"); - - // Convert to Array4 - Array4::from_shape_vec( - ( - shape[0] as usize, - shape[1] as usize, - shape[2] as usize, - shape[3] as usize, - ), - data, - ) - .expect("Shape conversion failed") -} - -/// Load preprocessor config from JSON -fn load_config(path: &Path) -> PreProcessorConfig { - let mut file = File::open(path).expect("Failed to open config"); - let mut contents = String::new(); - file.read_to_string(&mut contents) - .expect("Failed to read config"); - PreProcessorConfig::from_json(&contents).expect("Failed to parse config") -} - -/// Compare two 4D tensors and return max absolute difference -fn max_diff(a: &Array4, b: &ndarray::ArrayD) -> f32 { - assert_eq!(a.shape(), b.shape(), "Shape mismatch"); - // Convert ArrayD to Array4 for comparison - let b_4d = b - .clone() - .into_dimensionality::() - .expect("Expected 4D tensor"); - (a - &b_4d) - .mapv(|v| v.abs()) - .fold(0.0f32, |acc, &v| acc.max(v)) -} - -/// Load image_grid_thw from npz file -fn load_golden_grid_thw(path: &Path) -> Vec { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("image_grid_thw") - .expect("Failed to read npz") - .expect("No image_grid_thw"); - - // Shape not needed, data is flat - let _shape = reader.shape(); - - // Read data as i64 vec (numpy default for int) - reader.into_vec().expect("Failed to read array") -} - -/// Load num_tokens from npz file -fn load_golden_num_tokens(path: &Path) -> usize { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("num_tokens") - .expect("Failed to read npz") - .expect("No num_tokens"); - - // Read single value as i64 - let data: Vec = reader.into_vec().expect("Failed to read array"); - data[0] as usize -} - -/// Run a golden test for a specific mode and image. -/// -/// # Arguments -/// * `mode` - Either "llava" (standard CLIP) or "llava_pad" (expand-to-square mode) -/// * `image_name` - Name of the test image (e.g., "square", "tall", "wide", "small") -fn run_golden_test(mode: &str, image_name: &str) { - let golden_dir = Path::new("crates/multimodal/tests/fixtures/golden").join(mode); - let image_path = - Path::new("crates/multimodal/tests/fixtures/images").join(format!("{image_name}.jpg")); - - if !golden_dir.exists() || !image_path.exists() { - eprintln!("Golden test fixtures for {mode}/{image_name} not found, skipping test"); - eprintln!("Run: python crates/multimodal/scripts/generate_vision_golden.py"); - return; - } - - let golden = load_golden_npz(&golden_dir.join(format!("golden_{image_name}.npz"))); - let config = load_config(&golden_dir.join("preprocessor_config.json")); - - let image = image::open(&image_path).expect("Failed to open image"); - - let processor: Box = match mode { - "llava" => Box::new(LlavaProcessor::new()), - "llava_pad" => Box::new(LlavaProcessor::new_with_pad()), - _ => panic!("Unknown test mode: {mode}"), - }; - - let result = processor - .preprocess(&[image], &config) - .expect("Processing failed"); - - let diff = max_diff(&golden, &result.encoder_input); - println!("{mode} - {image_name} image - Max difference: {diff:.6}"); - println!("Golden shape: {:?}", golden.shape()); - println!("Rust shape: {:?}", result.encoder_input.shape()); - - // Allow tolerance for floating point and interpolation algorithm differences - // Different interpolation implementations (Rust vs Python/PIL) can produce - // small numerical differences, especially for edge cases like tiny or extreme - // aspect ratio images - assert!(diff < 0.1, "Max difference {diff} exceeds tolerance 0.1"); -} - -// ============================================================================ -// Standard CLIP mode tests (llava-hf/* models, no expand-to-square) -// ============================================================================ - -#[test] -fn test_llava_golden_square() { - run_golden_test("llava", "square"); -} - -#[test] -fn test_llava_golden_tall() { - run_golden_test("llava", "tall"); -} - -#[test] -fn test_llava_golden_wide() { - run_golden_test("llava", "wide"); -} - -#[test] -fn test_llava_golden_small() { - run_golden_test("llava", "small"); -} - -#[test] -fn test_llava_golden_tiny() { - run_golden_test("llava", "tiny"); -} - -#[test] -fn test_llava_golden_very_tall() { - run_golden_test("llava", "very_tall"); -} - -#[test] -fn test_llava_golden_very_wide() { - run_golden_test("llava", "very_wide"); -} - -#[test] -fn test_llava_golden_large() { - run_golden_test("llava", "large"); -} - -#[test] -fn test_llava_golden_odd_dims() { - run_golden_test("llava", "odd_dims"); -} - -#[test] -fn test_llava_golden_grayscale() { - run_golden_test("llava", "grayscale"); -} - -// ============================================================================ -// Pad mode tests (liuhaotian/llava-* models, image_aspect_ratio=pad) -// ============================================================================ - -#[test] -fn test_llava_pad_golden_square() { - run_golden_test("llava_pad", "square"); -} - -#[test] -fn test_llava_pad_golden_tall() { - run_golden_test("llava_pad", "tall"); -} - -#[test] -fn test_llava_pad_golden_wide() { - run_golden_test("llava_pad", "wide"); -} - -#[test] -fn test_llava_pad_golden_small() { - run_golden_test("llava_pad", "small"); -} - -#[test] -fn test_llava_pad_golden_tiny() { - run_golden_test("llava_pad", "tiny"); -} - -#[test] -fn test_llava_pad_golden_very_tall() { - run_golden_test("llava_pad", "very_tall"); -} - -#[test] -fn test_llava_pad_golden_very_wide() { - run_golden_test("llava_pad", "very_wide"); -} - -#[test] -fn test_llava_pad_golden_large() { - run_golden_test("llava_pad", "large"); -} - -#[test] -fn test_llava_pad_golden_odd_dims() { - run_golden_test("llava_pad", "odd_dims"); -} - -#[test] -fn test_llava_pad_golden_grayscale() { - run_golden_test("llava_pad", "grayscale"); -} - -// ============================================================================ -// Token count tests -// ============================================================================ - -#[test] -fn test_llava_token_count() { - let golden_dir = Path::new("crates/multimodal/tests/fixtures/golden/llava"); - - if !golden_dir.exists() { - eprintln!("Golden test fixtures not found, skipping test"); - return; - } - - let config = load_config(&golden_dir.join("preprocessor_config.json")); - let processor = LlavaProcessor::new(); - - // LLaVA 1.5 with 336x336 and patch_size=14: (336/14)^2 = 576 tokens - let tokens = processor.calculate_num_tokens(336, 336, &config); - assert_eq!( - tokens, 576, - "Expected 576 tokens for 336x336 with patch_size=14" - ); -} - -// ============================================================================ -// Qwen2-VL tests -// ============================================================================ - -/// Load flattened pixel values from Qwen2-VL npz file. -/// Returns (data, shape) where shape is (num_patches, patch_features). -fn load_golden_qwen2_vl_pixels(path: &Path) -> (Vec, (usize, usize)) { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("pixel_values") - .expect("Failed to read npz") - .expect("No pixel_values"); - - let shape = reader.shape().to_vec(); - assert_eq!(shape.len(), 2, "Expected 2D tensor for Qwen2-VL patches"); - - let data: Vec = reader.into_vec().expect("Failed to read array"); - (data, (shape[0] as usize, shape[1] as usize)) -} - -/// Run a Qwen2-VL golden test for a specific image. -/// -/// This test validates: -/// 1. image_grid_thw matches the HuggingFace output -/// 2. num_tokens calculation is correct -/// 3. Pixel values match after reshaping to patch format -fn run_qwen2_vl_golden_test(image_name: &str) { - let golden_dir = Path::new("crates/multimodal/tests/fixtures/golden/qwen2_vl"); - let image_path = - Path::new("crates/multimodal/tests/fixtures/images").join(format!("{image_name}.jpg")); - - if !golden_dir.exists() || !image_path.exists() { - eprintln!("Golden test fixtures for qwen2_vl/{image_name} not found, skipping test"); - eprintln!( - "Run: python crates/multimodal/scripts/generate_vision_golden.py --model qwen2_vl" - ); - return; - } - - let npz_path = golden_dir.join(format!("golden_{image_name}.npz")); - let config = load_config(&golden_dir.join("preprocessor_config.json")); - - // Load golden values - let golden_grid_thw = load_golden_grid_thw(&npz_path); - let golden_num_tokens = load_golden_num_tokens(&npz_path); - let (golden_pixels, golden_shape) = load_golden_qwen2_vl_pixels(&npz_path); - - // Process image with our Rust processor - let image = image::open(&image_path).expect("Failed to open image"); - let processor = Qwen2VLProcessor::from_preprocessor_config(&config); - let result = processor - .preprocess(&[image], &config) - .expect("Processing failed"); - - // Extract image_grid_thw from result - let rust_grid_thw = match result.model_specific.get("image_grid_thw") { - Some(ModelSpecificValue::IntTensor { data, shape }) => { - assert_eq!(shape, &[1, 3], "Expected shape [1, 3] for single image"); - data.clone() - } - _ => panic!("Expected image_grid_thw in model_specific"), - }; - - // Compare grid dimensions - println!( - "qwen2_vl - {image_name} image - Grid T H W: golden={golden_grid_thw:?}, rust={rust_grid_thw:?}" - ); - assert_eq!( - golden_grid_thw, rust_grid_thw, - "image_grid_thw mismatch for {image_name}" - ); - - // Compare token counts - let rust_num_tokens = result.feature_token_counts[0]; - println!( - "qwen2_vl - {image_name} image - Tokens: golden={golden_num_tokens}, rust={rust_num_tokens}" - ); - assert_eq!( - golden_num_tokens, rust_num_tokens, - "num_tokens mismatch for {image_name}" - ); - - // pixel_values is now already patchified: [total_patches, patch_features] - let rust_patches = result.encoder_input_flat(); - let rust_shape = ( - result.encoder_input.shape()[0], - result.encoder_input.shape()[1], - ); - - println!( - "qwen2_vl - {image_name} image - Patch shape: golden={golden_shape:?}, rust={rust_shape:?}" - ); - assert_eq!(golden_shape, rust_shape, "Patch shape mismatch"); - - // Compare pixel values - let max_diff = rust_patches - .iter() - .zip(golden_pixels.iter()) - .map(|(r, g)| (r - g).abs()) - .fold(0.0f32, f32::max); - - println!("qwen2_vl - {image_name} image - Max pixel diff: {max_diff:.6}"); - - // Allow tolerance for floating point and interpolation differences - // Different interpolation implementations (Rust vs Python/PIL) can produce - // small numerical differences, especially for edge cases - assert!( - max_diff < 0.1, - "Max pixel difference {max_diff} exceeds tolerance 0.1 for {image_name}" - ); -} - -#[test] -fn test_qwen2_vl_golden_square() { - run_qwen2_vl_golden_test("square"); -} - -#[test] -fn test_qwen2_vl_golden_tall() { - run_qwen2_vl_golden_test("tall"); -} - -#[test] -fn test_qwen2_vl_golden_wide() { - run_qwen2_vl_golden_test("wide"); -} - -#[test] -fn test_qwen2_vl_golden_small() { - run_qwen2_vl_golden_test("small"); -} - -#[test] -fn test_qwen2_vl_golden_tiny() { - run_qwen2_vl_golden_test("tiny"); -} - -#[test] -fn test_qwen2_vl_golden_very_tall() { - run_qwen2_vl_golden_test("very_tall"); -} - -#[test] -fn test_qwen2_vl_golden_very_wide() { - run_qwen2_vl_golden_test("very_wide"); -} - -#[test] -fn test_qwen2_vl_golden_large() { - run_qwen2_vl_golden_test("large"); -} - -#[test] -fn test_qwen2_vl_golden_odd_dims() { - run_qwen2_vl_golden_test("odd_dims"); -} - -#[test] -fn test_qwen2_vl_golden_grayscale() { - run_qwen2_vl_golden_test("grayscale"); -} - -// ============================================================================ -// Qwen3-VL tests -// ============================================================================ - -/// Run a Qwen3-VL golden test for a specific image. -/// -/// This test validates: -/// 1. image_grid_thw matches the HuggingFace output -/// 2. num_tokens calculation is correct -/// 3. Pixel values match after reshaping to patch format -/// -/// Key differences from Qwen2-VL: -/// - patch_size: 16 (vs 14) -/// - factor: 32 (vs 28) -/// - normalization: [0.5, 0.5, 0.5] (vs CLIP) -fn run_qwen3_vl_golden_test(image_name: &str) { - let golden_dir = Path::new("crates/multimodal/tests/fixtures/golden/qwen3_vl"); - let image_path = - Path::new("crates/multimodal/tests/fixtures/images").join(format!("{image_name}.jpg")); - - if !golden_dir.exists() || !image_path.exists() { - eprintln!("Golden test fixtures for qwen3_vl/{image_name} not found, skipping test"); - eprintln!( - "Run: python crates/multimodal/scripts/generate_vision_golden.py --model qwen3_vl" - ); - return; - } - - let npz_path = golden_dir.join(format!("golden_{image_name}.npz")); - let config = load_config(&golden_dir.join("preprocessor_config.json")); - - // Load golden values - let golden_grid_thw = load_golden_grid_thw(&npz_path); - let golden_num_tokens = load_golden_num_tokens(&npz_path); - let (golden_pixels, golden_shape) = load_golden_qwen2_vl_pixels(&npz_path); - - // Process image with our Rust processor - let image = image::open(&image_path).expect("Failed to open image"); - let processor = Qwen3VLProcessor::from_preprocessor_config(&config); - let result = processor - .preprocess(&[image], &config) - .expect("Processing failed"); - - // Extract image_grid_thw from result - let rust_grid_thw = match result.model_specific.get("image_grid_thw") { - Some(ModelSpecificValue::IntTensor { data, shape }) => { - assert_eq!(shape, &[1, 3], "Expected shape [1, 3] for single image"); - data.clone() - } - _ => panic!("Expected image_grid_thw in model_specific"), - }; - - // Compare grid dimensions - println!( - "qwen3_vl - {image_name} image - Grid T H W: golden={golden_grid_thw:?}, rust={rust_grid_thw:?}" - ); - assert_eq!( - golden_grid_thw, rust_grid_thw, - "image_grid_thw mismatch for {image_name}" - ); - - // Compare token counts - let rust_num_tokens = result.feature_token_counts[0]; - println!( - "qwen3_vl - {image_name} image - Tokens: golden={golden_num_tokens}, rust={rust_num_tokens}" - ); - assert_eq!( - golden_num_tokens, rust_num_tokens, - "num_tokens mismatch for {image_name}" - ); - - // pixel_values is now already patchified: [total_patches, patch_features] - let rust_patches = result.encoder_input_flat(); - let rust_shape = ( - result.encoder_input.shape()[0], - result.encoder_input.shape()[1], - ); - - println!( - "qwen3_vl - {image_name} image - Patch shape: golden={golden_shape:?}, rust={rust_shape:?}" - ); - assert_eq!(golden_shape, rust_shape, "Patch shape mismatch"); - - // Compare pixel values - let max_diff = rust_patches - .iter() - .zip(golden_pixels.iter()) - .map(|(r, g)| (r - g).abs()) - .fold(0.0f32, f32::max); - - println!("qwen3_vl - {image_name} image - Max pixel diff: {max_diff:.6}"); - - // Allow tolerance for floating point and interpolation differences - // Max diff is ~0.03 due to resize interpolation differences between Rust and HuggingFace - assert!( - max_diff < 0.05, - "Max pixel difference {max_diff} exceeds tolerance 0.05 for {image_name}" - ); -} - -#[test] -fn test_qwen3_vl_golden_square() { - run_qwen3_vl_golden_test("square"); -} - -#[test] -fn test_qwen3_vl_golden_tall() { - run_qwen3_vl_golden_test("tall"); -} - -#[test] -fn test_qwen3_vl_golden_wide() { - run_qwen3_vl_golden_test("wide"); -} - -#[test] -fn test_qwen3_vl_golden_small() { - run_qwen3_vl_golden_test("small"); -} - -#[test] -fn test_qwen3_vl_golden_tiny() { - run_qwen3_vl_golden_test("tiny"); -} - -#[test] -fn test_qwen3_vl_golden_very_tall() { - run_qwen3_vl_golden_test("very_tall"); -} - -#[test] -fn test_qwen3_vl_golden_very_wide() { - run_qwen3_vl_golden_test("very_wide"); -} - -#[test] -fn test_qwen3_vl_golden_large() { - run_qwen3_vl_golden_test("large"); -} - -#[test] -fn test_qwen3_vl_golden_odd_dims() { - run_qwen3_vl_golden_test("odd_dims"); -} - -#[test] -fn test_qwen3_vl_golden_grayscale() { - run_qwen3_vl_golden_test("grayscale"); -} - -// ============================================================================ -// Phi3-Vision tests -// ============================================================================ - -/// Load a 5D numpy .npz file for Phi3-Vision (batch, num_crops+1, C, H, W) -fn load_golden_npz_5d(path: &Path) -> Array5 { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("pixel_values") - .expect("Failed to read npz") - .expect("No pixel_values"); - - let shape = reader.shape().to_vec(); - assert_eq!(shape.len(), 5, "Expected 5D tensor [B, N, C, H, W]"); - - let data: Vec = reader.into_vec().expect("Failed to read array"); - - Array5::from_shape_vec( - ( - shape[0] as usize, - shape[1] as usize, - shape[2] as usize, - shape[3] as usize, - shape[4] as usize, - ), - data, - ) - .expect("Shape conversion failed") -} - -/// Load image_sizes from Phi3-Vision npz file (2D tensor [batch, 2]) -fn load_phi3_image_sizes(path: &Path) -> Vec<(i64, i64)> { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("image_sizes") - .expect("Failed to read npz") - .expect("No image_sizes"); - - let shape = reader.shape().to_vec(); - let data: Vec = reader.into_vec().expect("Failed to read array"); - - // Reshape to pairs - let num_images = shape[0] as usize; - (0..num_images) - .map(|i| (data[i * 2], data[i * 2 + 1])) - .collect() -} - -/// Load num_img_tokens from Phi3-Vision npz file -fn load_phi3_num_img_tokens(path: &Path) -> Vec { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("num_img_tokens") - .expect("Failed to read npz") - .expect("No num_img_tokens"); - - let data: Vec = reader.into_vec().expect("Failed to read array"); - data.into_iter().map(|v| v as usize).collect() -} - -/// Compare two 5D tensors and return max absolute difference -fn max_diff_5d(a: &Array5, b: &Array5) -> f32 { - assert_eq!(a.shape(), b.shape(), "Shape mismatch"); - (a - b).mapv(|v| v.abs()).fold(0.0f32, |acc, &v| acc.max(v)) -} - -/// Find the location and value of max difference between two 5D tensors -fn find_max_diff_location_5d( - golden: &Array5, - rust: &Array5, - image_name: &str, -) -> (f32, (usize, usize, usize, usize, usize)) { - assert_eq!(golden.shape(), rust.shape(), "Shape mismatch"); - let diff = (golden - rust).mapv(|v| v.abs()); - let mut max_diff = 0.0f32; - let mut max_pos = (0, 0, 0, 0, 0); - - // Find per-tile max differences - for b in 0..golden.shape()[0] { - for t in 0..golden.shape()[1] { - let tile_diff = diff.slice(ndarray::s![b, t, .., .., ..]); - let tile_max = tile_diff.fold(0.0f32, |acc, &v| acc.max(v)); - - if tile_max > 0.1 { - let golden_tile = golden.slice(ndarray::s![b, t, .., .., ..]); - let rust_tile = rust.slice(ndarray::s![b, t, .., .., ..]); - println!( - " {image_name} tile {t}: diff={tile_max:.4}, golden_range=[{:.4}, {:.4}], rust_range=[{:.4}, {:.4}]", - golden_tile.fold(f32::MAX, |a, &v| a.min(v)), - golden_tile.fold(f32::MIN, |a, &v| a.max(v)), - rust_tile.fold(f32::MAX, |a, &v| a.min(v)), - rust_tile.fold(f32::MIN, |a, &v| a.max(v)) - ); - } - - if tile_max > max_diff { - max_diff = tile_max; - // Find exact position - for c in 0..golden.shape()[2] { - for h in 0..golden.shape()[3] { - for w in 0..golden.shape()[4] { - if diff[[b, t, c, h, w]] == max_diff { - max_pos = (b, t, c, h, w); - } - } - } - } - } - } - } - - (max_diff, max_pos) -} - -/// Run a Phi3-Vision golden test for a specific image. -/// -/// This test validates: -/// 1. Output shape is [1, num_crops+1, 3, 336, 336] -/// 2. image_sizes matches HuggingFace output -/// 3. num_img_tokens matches HuggingFace output -/// 4. Pixel values match within tolerance -fn run_phi3_vision_golden_test(image_name: &str) { - let golden_dir = Path::new("crates/multimodal/tests/fixtures/golden/phi3_vision"); - let image_path = - Path::new("crates/multimodal/tests/fixtures/images").join(format!("{image_name}.jpg")); - - if !golden_dir.exists() || !image_path.exists() { - eprintln!("Golden test fixtures for phi3_vision/{image_name} not found, skipping test"); - eprintln!( - "Run: python crates/multimodal/scripts/generate_vision_golden.py --model phi3_vision" - ); - return; - } - - let npz_path = golden_dir.join(format!("golden_{image_name}.npz")); - let config = load_config(&golden_dir.join("preprocessor_config.json")); - - // Load golden values - let golden_pixels = load_golden_npz_5d(&npz_path); - let golden_image_sizes = load_phi3_image_sizes(&npz_path); - let golden_num_tokens = load_phi3_num_img_tokens(&npz_path); - - // Process image with our Rust processor - let image = image::open(&image_path).expect("Failed to open image"); - let processor = Phi3VisionProcessor::from_preprocessor_config(&config); - let result = processor - .preprocess(&[image], &config) - .expect("Processing failed"); - - // Check output shape - let rust_shape = result.encoder_input.shape(); - let golden_shape = golden_pixels.shape(); - println!( - "phi3_vision - {image_name} image - Shape: golden={golden_shape:?}, rust={rust_shape:?}" - ); - assert_eq!( - rust_shape, golden_shape, - "Shape mismatch for phi3_vision/{image_name}" - ); - - // Check image_sizes - // Note: HuggingFace returns [h, w], we store as (w, h) but model_specific stores (h, w) - let rust_image_sizes: Vec<(i64, i64)> = match result.model_specific.get("image_sizes") { - Some(ModelSpecificValue::IntTensor { data, shape }) => { - let num_images = shape[0]; - (0..num_images) - .map(|i| (data[i * 2], data[i * 2 + 1])) - .collect() - } - _ => panic!("Expected image_sizes in model_specific"), - }; - - println!( - "phi3_vision - {image_name} image - Image sizes (h, w): golden={golden_image_sizes:?}, rust={rust_image_sizes:?}" - ); - assert_eq!( - golden_image_sizes, rust_image_sizes, - "image_sizes mismatch for {image_name}" - ); - - // Check num_img_tokens - println!( - "phi3_vision - {image_name} image - Num tokens: golden={golden_num_tokens:?}, rust={:?}", - result.feature_token_counts - ); - assert_eq!( - golden_num_tokens, result.feature_token_counts, - "num_img_tokens mismatch for {image_name}" - ); - - // Compare pixel values - // Convert rust ArrayD to Array5 for comparison - let rust_pixels = result - .encoder_input - .clone() - .into_dimensionality::() - .expect("Failed to convert to Ix5"); - - let pixel_diff = max_diff_5d(&golden_pixels, &rust_pixels); - println!("phi3_vision - {image_name} image - Max pixel diff: {pixel_diff:.6}"); - - // If there's a large difference, print detailed info - if pixel_diff > 0.1 { - let (max_diff, max_pos) = - find_max_diff_location_5d(&golden_pixels, &rust_pixels, image_name); - println!( - "phi3_vision - {image_name} image - Max diff {max_diff:.4} at position {max_pos:?}" - ); - let (b, t, c, h, w) = max_pos; - println!( - " golden value: {:.4}, rust value: {:.4}", - golden_pixels[[b, t, c, h, w]], - rust_pixels[[b, t, c, h, w]] - ); - } - - // Allow tolerance for floating point and interpolation differences - // Using bicubic for global image and bilinear for HD resize to match HuggingFace. - assert!( - pixel_diff < 0.08, - "Max pixel difference {pixel_diff} exceeds tolerance 0.08 for {image_name}" - ); -} - -#[test] -fn test_phi3_vision_golden_square() { - run_phi3_vision_golden_test("square"); -} - -#[test] -fn test_phi3_vision_golden_tall() { - run_phi3_vision_golden_test("tall"); -} - -#[test] -fn test_phi3_vision_golden_wide() { - run_phi3_vision_golden_test("wide"); -} - -#[test] -fn test_phi3_vision_golden_small() { - run_phi3_vision_golden_test("small"); -} - -#[test] -fn test_phi3_vision_golden_tiny() { - run_phi3_vision_golden_test("tiny"); -} - -#[test] -fn test_phi3_vision_golden_very_tall() { - run_phi3_vision_golden_test("very_tall"); -} - -#[test] -fn test_phi3_vision_golden_very_wide() { - run_phi3_vision_golden_test("very_wide"); -} - -#[test] -fn test_phi3_vision_golden_large() { - run_phi3_vision_golden_test("large"); -} - -#[test] -fn test_phi3_vision_golden_odd_dims() { - run_phi3_vision_golden_test("odd_dims"); -} - -#[test] -fn test_phi3_vision_golden_grayscale() { - run_phi3_vision_golden_test("grayscale"); -} - -// ============================================================================ -// Phi4-Vision tests -// ============================================================================ - -/// Load num_img_tokens from Phi4-Vision npz file -fn load_phi4_num_img_tokens(path: &Path) -> Vec { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("num_img_tokens") - .expect("Failed to read npz") - .expect("No num_img_tokens"); - - let data: Vec = reader.into_vec().expect("Failed to read array"); - data.into_iter().map(|v| v as usize).collect() -} - -/// Load image_sizes from Phi4-Vision npz file (2D tensor [batch, 2]) -fn load_phi4_image_sizes(path: &Path) -> Vec<(i64, i64)> { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("image_sizes") - .expect("Failed to read npz") - .expect("No image_sizes"); - - let shape = reader.shape().to_vec(); - let data: Vec = reader.into_vec().expect("Failed to read array"); - - // Reshape to pairs - let num_images = shape[0] as usize; - (0..num_images) - .map(|i| (data[i * 2], data[i * 2 + 1])) - .collect() -} - -/// Run a Phi4-Vision golden test for a specific image. -/// -/// This test validates: -/// 1. Output shape is [1, num_crops+1, 3, 448, 448] (note: 448 base resolution) -/// 2. image_sizes matches HuggingFace output -/// 3. num_img_tokens matches HuggingFace output -/// 4. Pixel values match within tolerance -/// -/// Key differences from Phi3-Vision: -/// - Base resolution: 448 (vs 336) -/// - Normalization: [0.5, 0.5, 0.5] (vs CLIP) -/// - Default dynamic_hd: 36 (vs 16) -fn run_phi4_vision_golden_test(image_name: &str) { - let golden_dir = Path::new("crates/multimodal/tests/fixtures/golden/phi4_vision"); - let image_path = - Path::new("crates/multimodal/tests/fixtures/images").join(format!("{image_name}.jpg")); - - if !golden_dir.exists() || !image_path.exists() { - eprintln!("Golden test fixtures for phi4_vision/{image_name} not found, skipping test"); - eprintln!( - "Run: python crates/multimodal/scripts/generate_vision_golden.py --model phi4_vision" - ); - return; - } - - let npz_path = golden_dir.join(format!("golden_{image_name}.npz")); - let config = load_config(&golden_dir.join("preprocessor_config.json")); - - // Load golden values - let golden_pixels = load_golden_npz_5d(&npz_path); - let golden_image_sizes = load_phi4_image_sizes(&npz_path); - let golden_num_tokens = load_phi4_num_img_tokens(&npz_path); - - // Process image with our Rust processor - let image = image::open(&image_path).expect("Failed to open image"); - let processor = Phi4VisionProcessor::from_preprocessor_config(&config); - let result = processor - .preprocess(&[image], &config) - .expect("Processing failed"); - - // Check output shape - let rust_shape = result.encoder_input.shape(); - let golden_shape = golden_pixels.shape(); - println!( - "phi4_vision - {image_name} image - Shape: golden={golden_shape:?}, rust={rust_shape:?}" - ); - assert_eq!( - rust_shape, golden_shape, - "Shape mismatch for phi4_vision/{image_name}" - ); - - // Check image_sizes - let rust_image_sizes: Vec<(i64, i64)> = match result.model_specific.get("image_sizes") { - Some(ModelSpecificValue::IntTensor { data, shape }) => { - let num_images = shape[0]; - (0..num_images) - .map(|i| (data[i * 2], data[i * 2 + 1])) - .collect() - } - _ => panic!("Expected image_sizes in model_specific"), - }; - - println!( - "phi4_vision - {image_name} image - Image sizes (h, w): golden={golden_image_sizes:?}, rust={rust_image_sizes:?}" - ); - assert_eq!( - golden_image_sizes, rust_image_sizes, - "image_sizes mismatch for {image_name}" - ); - - // Check num_img_tokens - println!( - "phi4_vision - {image_name} image - Num tokens: golden={golden_num_tokens:?}, rust={:?}", - result.feature_token_counts - ); - assert_eq!( - golden_num_tokens, result.feature_token_counts, - "num_img_tokens mismatch for {image_name}" - ); - - // Compare pixel values - let rust_pixels = result - .encoder_input - .clone() - .into_dimensionality::() - .expect("Failed to convert to Ix5"); - - let pixel_diff = max_diff_5d(&golden_pixels, &rust_pixels); - println!("phi4_vision - {image_name} image - Max pixel diff: {pixel_diff:.6}"); - - // If there's a large difference, print detailed info - if pixel_diff > 0.1 { - let (max_diff, max_pos) = - find_max_diff_location_5d(&golden_pixels, &rust_pixels, image_name); - println!( - "phi4_vision - {image_name} image - Max diff {max_diff:.4} at position {max_pos:?}" - ); - let (b, t, c, h, w) = max_pos; - println!( - " golden value: {:.4}, rust value: {:.4}", - golden_pixels[[b, t, c, h, w]], - rust_pixels[[b, t, c, h, w]] - ); - } - - // Allow tolerance for floating point and interpolation differences - // Using bilinear for HD resize and bicubic for global image to match HuggingFace. - assert!( - pixel_diff < 0.05, - "Max pixel difference {pixel_diff} exceeds tolerance 0.05 for {image_name}" - ); -} - -#[test] -fn test_phi4_vision_golden_square() { - run_phi4_vision_golden_test("square"); -} - -#[test] -fn test_phi4_vision_golden_tall() { - run_phi4_vision_golden_test("tall"); -} - -#[test] -fn test_phi4_vision_golden_wide() { - run_phi4_vision_golden_test("wide"); -} - -#[test] -fn test_phi4_vision_golden_small() { - run_phi4_vision_golden_test("small"); -} - -#[test] -fn test_phi4_vision_golden_tiny() { - run_phi4_vision_golden_test("tiny"); -} - -#[test] -fn test_phi4_vision_golden_very_tall() { - run_phi4_vision_golden_test("very_tall"); -} - -#[test] -fn test_phi4_vision_golden_very_wide() { - run_phi4_vision_golden_test("very_wide"); -} - -#[test] -fn test_phi4_vision_golden_large() { - run_phi4_vision_golden_test("large"); -} - -#[test] -fn test_phi4_vision_golden_odd_dims() { - run_phi4_vision_golden_test("odd_dims"); -} - -#[test] -fn test_phi4_vision_golden_grayscale() { - run_phi4_vision_golden_test("grayscale"); -} - -// ============================================================================ -// LLaMA 4 Vision tests -// ============================================================================ - -/// Load aspect_ratios from npz file for LLaMA 4 -fn load_llama4_aspect_ratios(path: &Path) -> Vec<(i64, i64)> { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("aspect_ratios") - .expect("Failed to read npz") - .expect("No aspect_ratios"); - - let shape = reader.shape().to_vec(); - - // Read data as i64 vec (numpy default for int) - let data: Vec = reader.into_vec().expect("Failed to read array"); - - // Convert to Vec<(i64, i64)> - let num_images = shape[0] as usize; - (0..num_images) - .map(|i| (data[i * 2], data[i * 2 + 1])) - .collect() -} - -/// Load pixel_values for LLaMA 4 Vision (3D: [num_tiles, C, H, W]) -fn load_llama4_pixels(path: &Path) -> (Vec, Vec) { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("pixel_values") - .expect("Failed to read npz") - .expect("No pixel_values"); - - let shape: Vec = reader.shape().iter().map(|&s| s as usize).collect(); - let data: Vec = reader.into_vec().expect("Failed to read array"); - - (data, shape) -} - -/// Run a LLaMA 4 Vision golden test for a specific image. -/// -/// This test validates: -/// 1. Output shape matches (batch, num_tiles, 3, 336, 336) -/// 2. aspect_ratios match (h_tiles, w_tiles) -/// 3. Pixel values match HuggingFace output -/// 4. Token count is correct -/// -/// LLaMA 4 Vision processing: -/// - Tile size: 336x336 -/// - Max patches: 16 (default) -/// - Normalization: [0.5, 0.5, 0.5] mean/std -/// - Global tile added when num_tiles > 1 -fn run_llama4_vision_golden_test(image_name: &str) { - let golden_dir = Path::new("crates/multimodal/tests/fixtures/golden/llama4_vision"); - let image_path = - Path::new("crates/multimodal/tests/fixtures/images").join(format!("{image_name}.jpg")); - - if !golden_dir.exists() || !image_path.exists() { - eprintln!("Golden test fixtures for llama4_vision/{image_name} not found, skipping test"); - eprintln!( - "Run: python crates/multimodal/scripts/generate_vision_golden.py --model llama4_vision" - ); - return; - } - - let npz_path = golden_dir.join(format!("golden_{image_name}.npz")); - let config = load_config(&golden_dir.join("preprocessor_config.json")); - - // Load golden values - let (golden_pixels, golden_shape) = load_llama4_pixels(&npz_path); - let golden_aspect_ratios = load_llama4_aspect_ratios(&npz_path); - let golden_num_tokens = load_golden_num_tokens(&npz_path); - - // Process image with our Rust processor - let image = image::open(&image_path).expect("Failed to open image"); - let processor = Llama4VisionProcessor::from_preprocessor_config(&config); - let result = processor - .preprocess(&[image], &config) - .expect("Processing failed"); - - // Check aspect_ratios - let rust_aspect_ratios: Vec<(i64, i64)> = match result.model_specific.get("aspect_ratios") { - Some(ModelSpecificValue::IntTensor { data, shape }) => { - let num_images = shape[0]; - (0..num_images) - .map(|i| (data[i * 2], data[i * 2 + 1])) - .collect() - } - _ => panic!("Expected aspect_ratios in model_specific"), - }; - - println!( - "llama4_vision - {image_name} image - Aspect ratios: golden={golden_aspect_ratios:?}, rust={rust_aspect_ratios:?}" - ); - assert_eq!( - golden_aspect_ratios, rust_aspect_ratios, - "aspect_ratios mismatch for {image_name}" - ); - - // Check num_tokens - let rust_num_tokens = result.feature_token_counts[0]; - println!( - "llama4_vision - {image_name} image - Tokens: golden={golden_num_tokens}, rust={rust_num_tokens}" - ); - assert_eq!( - golden_num_tokens, rust_num_tokens, - "num_tokens mismatch for {image_name}" - ); - - // Check output shape - both HuggingFace and Rust output 4D (total_tiles, C, H, W) - let rust_shape = result.encoder_input.shape(); - println!( - "llama4_vision - {image_name} image - Shape: golden={golden_shape:?}, rust={rust_shape:?}" - ); - - assert_eq!( - rust_shape.len(), - 4, - "Expected 4D tensor [total_tiles, C, H, W], got {}D", - rust_shape.len() - ); - assert_eq!( - rust_shape[0], golden_shape[0], - "Expected {} tiles, got {}", - golden_shape[0], rust_shape[0] - ); - - // Compare pixel values - let rust_pixels = result.encoder_input_flat(); - let num_golden_elements: usize = golden_shape.iter().product(); - - // Find the max difference for the actual tiles (not padding) - let mut max_diff = 0.0f32; - for i in 0..num_golden_elements { - let diff = (rust_pixels[i] - golden_pixels[i]).abs(); - max_diff = max_diff.max(diff); - } - - println!("llama4_vision - {image_name} image - Max pixel diff: {max_diff:.6}"); - - // Allow tolerance for floating point and interpolation differences - // LLaMA 4 uses bfloat16 internally which may cause small differences - assert!( - max_diff < 0.03, - "Max pixel difference {max_diff} exceeds tolerance 0.03 for {image_name}" - ); -} - -#[test] -fn test_llama4_vision_golden_square() { - run_llama4_vision_golden_test("square"); -} - -#[test] -fn test_llama4_vision_golden_tall() { - run_llama4_vision_golden_test("tall"); -} - -#[test] -fn test_llama4_vision_golden_wide() { - run_llama4_vision_golden_test("wide"); -} - -#[test] -fn test_llama4_vision_golden_small() { - run_llama4_vision_golden_test("small"); -} - -#[test] -fn test_llama4_vision_golden_tiny() { - run_llama4_vision_golden_test("tiny"); -} - -#[test] -fn test_llama4_vision_golden_very_tall() { - run_llama4_vision_golden_test("very_tall"); -} - -#[test] -fn test_llama4_vision_golden_very_wide() { - run_llama4_vision_golden_test("very_wide"); -} - -#[test] -fn test_llama4_vision_golden_large() { - run_llama4_vision_golden_test("large"); -} - -#[test] -fn test_llama4_vision_golden_odd_dims() { - run_llama4_vision_golden_test("odd_dims"); -} - -#[test] -fn test_llama4_vision_golden_grayscale() { - run_llama4_vision_golden_test("grayscale"); -} - -// ============================================================================ -// Pixtral/Mistral3 Vision tests -// ============================================================================ - -/// Load image_sizes from npz file for Pixtral -fn load_pixtral_image_sizes(path: &Path) -> Vec<(usize, usize)> { - let file = File::open(path).expect("Failed to open golden file"); - let mut npz = npyz::npz::NpzArchive::new(file).expect("Failed to parse npz"); - - let reader = npz - .by_name("image_sizes") - .expect("Failed to read npz") - .expect("No image_sizes"); - - let shape = reader.shape().to_vec(); - - // Read data as i64 vec (numpy default for int) - let data: Vec = reader.into_vec().expect("Failed to read array"); - - // Convert to Vec<(usize, usize)> - let num_images = shape[0] as usize; - (0..num_images) - .map(|i| (data[i * 2] as usize, data[i * 2 + 1] as usize)) - .collect() -} - -/// Run a Pixtral golden test for a specific image. -/// -/// This test validates: -/// 1. Output shape matches (batch, 3, H, W) -/// 2. image_sizes match -/// 3. Pixel values match HuggingFace output -/// 4. Token count is correct -/// -/// Pixtral processing: -/// - Longest edge: 1024 (default) -/// - Patch size: 16 -/// - Normalization: CLIP mean/std -/// - No tiling - single output per image -fn run_pixtral_golden_test(image_name: &str) { - let golden_dir = Path::new("crates/multimodal/tests/fixtures/golden/pixtral"); - let image_path = - Path::new("crates/multimodal/tests/fixtures/images").join(format!("{image_name}.jpg")); - - if !golden_dir.exists() || !image_path.exists() { - eprintln!("Golden test fixtures for pixtral/{image_name} not found, skipping test"); - eprintln!( - "Run: python crates/multimodal/scripts/generate_vision_golden.py --model pixtral" - ); - return; - } - - let npz_path = golden_dir.join(format!("golden_{image_name}.npz")); - let config = load_config(&golden_dir.join("preprocessor_config.json")); - - // Load golden values - let golden_pixels = load_golden_npz(&npz_path); - let golden_shape: Vec = golden_pixels.shape().to_vec(); - let golden_image_sizes = load_pixtral_image_sizes(&npz_path); - let golden_num_tokens = load_golden_num_tokens(&npz_path); - - // Process image with our Rust processor - let image = image::open(&image_path).expect("Failed to open image"); - let processor = PixtralProcessor::from_preprocessor_config(&config); - let result = processor - .preprocess(&[image], &config) - .expect("Processing failed"); - - // Check image_sizes from model_specific - let rust_image_sizes: Vec<(usize, usize)> = match result.model_specific.get("image_sizes") { - Some(ModelSpecificValue::IntTensor { data, shape }) => { - let num_images = shape[0]; - (0..num_images) - .map(|i| (data[i * 2] as usize, data[i * 2 + 1] as usize)) - .collect() - } - _ => panic!("Expected image_sizes in model_specific"), - }; - - println!( - "pixtral - {image_name} image - Image sizes: golden={golden_image_sizes:?}, rust={rust_image_sizes:?}" - ); - assert_eq!( - golden_image_sizes, rust_image_sizes, - "image_sizes mismatch for {image_name}" - ); - - // Check num_tokens - let rust_num_tokens = result.feature_token_counts[0]; - println!( - "pixtral - {image_name} image - Tokens: golden={golden_num_tokens}, rust={rust_num_tokens}" - ); - assert_eq!( - golden_num_tokens, rust_num_tokens, - "num_tokens mismatch for {image_name}" - ); - - // Check output shape - let rust_shape = result.encoder_input.shape(); - println!("pixtral - {image_name} image - Shape: golden={golden_shape:?}, rust={rust_shape:?}"); - - // Pixtral outputs [batch, C, H, W] with padding to max size in batch - // Single image should match golden shape exactly - assert_eq!(rust_shape[0], 1, "Expected batch dim to be 1"); - assert_eq!(rust_shape[1], golden_shape[1], "Channel mismatch"); - assert!( - rust_shape[2] >= golden_shape[2], - "Height {} < golden height {}", - rust_shape[2], - golden_shape[2] - ); - assert!( - rust_shape[3] >= golden_shape[3], - "Width {} < golden width {}", - rust_shape[3], - golden_shape[3] - ); - - // Compare pixel values - only compare the actual image region, not padding - let rust_pixels = result.encoder_input_flat(); - let golden_pixels_flat: Vec = golden_pixels.iter().copied().collect(); - - // Calculate indices for the actual image region (not padding) - let h = golden_shape[2]; - let w = golden_shape[3]; - let rust_w = rust_shape[3]; - - let mut max_diff = 0.0f32; - for c in 0..3 { - for y in 0..h { - for x in 0..w { - let golden_idx = c * h * w + y * w + x; - let rust_idx = c * rust_shape[2] * rust_w + y * rust_w + x; - let diff = (rust_pixels[rust_idx] - golden_pixels_flat[golden_idx]).abs(); - max_diff = max_diff.max(diff); - } - } - } - - println!("pixtral - {image_name} image - Max pixel diff: {max_diff:.6}"); - - // Allow tolerance for bicubic interpolation differences between PIL and Rust image library - // Pixtral uses bicubic which has larger differences than bilinear - assert!( - max_diff < 0.06, - "Max pixel difference {max_diff} exceeds tolerance 0.06 for {image_name}" - ); -} - -#[test] -fn test_pixtral_golden_square() { - run_pixtral_golden_test("square"); -} - -#[test] -fn test_pixtral_golden_tall() { - run_pixtral_golden_test("tall"); -} - -#[test] -fn test_pixtral_golden_wide() { - run_pixtral_golden_test("wide"); -} - -#[test] -fn test_pixtral_golden_small() { - run_pixtral_golden_test("small"); -} - -#[test] -fn test_pixtral_golden_tiny() { - run_pixtral_golden_test("tiny"); -} - -#[test] -fn test_pixtral_golden_very_tall() { - run_pixtral_golden_test("very_tall"); -} - -#[test] -fn test_pixtral_golden_very_wide() { - run_pixtral_golden_test("very_wide"); -} - -#[test] -fn test_pixtral_golden_large() { - run_pixtral_golden_test("large"); -} - -#[test] -fn test_pixtral_golden_odd_dims() { - run_pixtral_golden_test("odd_dims"); -} - -#[test] -fn test_pixtral_golden_grayscale() { - run_pixtral_golden_test("grayscale"); -} diff --git a/model_gateway/Cargo.toml b/model_gateway/Cargo.toml index 9fff00eeb..84ccedb27 100644 --- a/model_gateway/Cargo.toml +++ b/model_gateway/Cargo.toml @@ -75,7 +75,7 @@ smg-auth.workspace = true smg-mcp.workspace = true kv-index.workspace = true smg-data-connector.workspace = true -llm-multimodal = { workspace = true, default-features = false } +llm-multimodal.workspace = true smg-mm-rdma.workspace = true smg-wasm = { workspace = true, features = ["storage-hooks"] } smg-mesh.workspace = true diff --git a/model_gateway/src/routers/grpc/multimodal/mod.rs b/model_gateway/src/routers/grpc/multimodal/mod.rs index 5aa2d8035..fd37cb1b1 100644 --- a/model_gateway/src/routers/grpc/multimodal/mod.rs +++ b/model_gateway/src/routers/grpc/multimodal/mod.rs @@ -22,6 +22,7 @@ use llm_multimodal::{ AudioClip, EncoderFieldLayouts, ImageFrame, Modality, PlaceholderRange, PreprocessedEncoderInputs, VideoClip, }; +use llm_tokenizer::TokenizerTrait; mod assemble; mod capability; @@ -49,6 +50,34 @@ pub(crate) use plan::{ pub(crate) use process::process_multimodal_plan; pub(crate) use transport::{init_mm_transport_defaults, mm_rdma_exporter}; +/// Adapts SMG's tokenizer to the narrow interface used by llm-multimodal. +struct MultimodalTokenizer<'a> { + inner: &'a dyn TokenizerTrait, +} + +impl<'a> MultimodalTokenizer<'a> { + fn new(inner: &'a dyn TokenizerTrait) -> Self { + Self { inner } + } +} + +impl llm_multimodal::Tokenizer for MultimodalTokenizer<'_> { + fn token_to_id(&self, token: &str) -> Option { + self.inner.token_to_id(token) + } + + fn id_to_token(&self, id: u32) -> Option { + self.inner.id_to_token(id) + } + + fn encode_text(&self, text: &str) -> Option> { + self.inner + .encode(text, false) + .ok() + .map(|encoding| encoding.token_ids().to_vec()) + } +} + /// Whether verbose multimodal timing logs are enabled via `SMG_LOG_MM_TIMING`. /// Read from the environment once and cached; the flag is not expected to change /// at runtime, and this is called on every multimodal request. diff --git a/model_gateway/src/routers/grpc/multimodal/plan.rs b/model_gateway/src/routers/grpc/multimodal/plan.rs index 88965808d..e50c60b3d 100644 --- a/model_gateway/src/routers/grpc/multimodal/plan.rs +++ b/model_gateway/src/routers/grpc/multimodal/plan.rs @@ -7,7 +7,7 @@ use anyhow::{Context, Result}; use llm_multimodal::{MediaContentPart, Modality, ModelMetadata}; use llm_tokenizer::TokenizerTrait; -use super::config::MultimodalComponents; +use super::{config::MultimodalComponents, MultimodalTokenizer}; /// Ordered media extracted from an API request. /// @@ -107,7 +107,7 @@ pub(crate) async fn prepare_placeholder_tokens( .await?; let metadata = ModelMetadata { model_id, - tokenizer, + tokenizer: &MultimodalTokenizer::new(tokenizer), config: &model_config.config, }; let spec = components diff --git a/model_gateway/src/routers/grpc/multimodal/process.rs b/model_gateway/src/routers/grpc/multimodal/process.rs index f41b7c490..e9540716b 100644 --- a/model_gateway/src/routers/grpc/multimodal/process.rs +++ b/model_gateway/src/routers/grpc/multimodal/process.rs @@ -21,8 +21,8 @@ use super::{ log_mm_timing_enabled, pixel_cache::{config_fingerprint, CachedPreprocessedItem, PixelCache, PixelCacheKey}, plan::MediaPlan, - MediaBatch, MultimodalIntermediate, MultimodalOutput, PrecomputedMultimodalIntermediate, - PromptBinding, + MediaBatch, MultimodalIntermediate, MultimodalOutput, MultimodalTokenizer, + PrecomputedMultimodalIntermediate, PromptBinding, }; struct PreparedMultimodalPart { @@ -173,7 +173,7 @@ pub(crate) async fn process_multimodal_plan( .and_then(|v| v.as_str()); let metadata = ModelMetadata { model_id, - tokenizer, + tokenizer: &MultimodalTokenizer::new(tokenizer), config: &model_config.config, }; let spec = components diff --git a/scripts/check_release_versions.sh b/scripts/check_release_versions.sh index 4e82d147c..928fcc042 100755 --- a/scripts/check_release_versions.sh +++ b/scripts/check_release_versions.sh @@ -67,7 +67,6 @@ CRATES=( "smg-mcp|crates/mcp|smg-mcp" "kv-index|crates/kv_index|kv-index" "data-connector|crates/data_connector|smg-data-connector" - "llm-multimodal|crates/multimodal|llm-multimodal" "smg-wasm|crates/wasm|smg-wasm" "smg-mesh|crates/mesh|smg-mesh" "smg-grpc-client|crates/grpc_client|smg-grpc-client"