[model-gateway] add llama4 vision image processor (#14438)
This commit is contained in:
2
.github/workflows/pr-test-rust.yml
vendored
2
.github/workflows/pr-test-rust.yml
vendored
@@ -118,7 +118,7 @@ jobs:
|
||||
- name: Generate vision golden fixtures
|
||||
run: |
|
||||
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
|
||||
pip install transformers pillow numpy
|
||||
pip install transformers pillow numpy scipy
|
||||
cd sgl-router/
|
||||
python scripts/generate_vision_golden.py
|
||||
|
||||
|
||||
@@ -61,6 +61,11 @@ MODELS = {
|
||||
"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",
|
||||
},
|
||||
}
|
||||
|
||||
# Default test images
|
||||
@@ -480,6 +485,68 @@ def generate_golden_phi4_vision(image_path: str, output_dir: str) -> dict:
|
||||
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_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}...")
|
||||
@@ -492,6 +559,7 @@ def generate_for_model(model_key: str, image_paths: list, output_dir: str):
|
||||
"qwen3_vl": generate_golden_qwen3_vl,
|
||||
"phi3_vision": generate_golden_phi3_vision,
|
||||
"phi4_vision": generate_golden_phi4_vision,
|
||||
"llama4_vision": generate_golden_llama4_vision,
|
||||
}.get(model_key)
|
||||
|
||||
if generator_fn is None:
|
||||
|
||||
@@ -40,7 +40,7 @@ pub use image_processor::{
|
||||
};
|
||||
pub use preprocessor_config::PreProcessorConfig;
|
||||
pub use processors::{
|
||||
LlavaNextProcessor, LlavaProcessor, Phi3VisionProcessor, Phi4VisionProcessor, Qwen2VLProcessor,
|
||||
Qwen3VLProcessor,
|
||||
Llama4VisionProcessor, LlavaNextProcessor, LlavaProcessor, Phi3VisionProcessor,
|
||||
Phi4VisionProcessor, Qwen2VLProcessor, Qwen3VLProcessor,
|
||||
};
|
||||
pub use transforms::TransformError;
|
||||
|
||||
688
sgl-router/src/multimodal/vision/processors/llama4_vision.rs
Normal file
688
sgl-router/src/multimodal/vision/processors/llama4_vision.rs
Normal file
@@ -0,0 +1,688 @@
|
||||
//! 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, Rgb, RgbImage};
|
||||
use ndarray::{s, Array3, Array4, IxDyn};
|
||||
|
||||
use crate::multimodal::vision::{
|
||||
image_processor::{ImagePreProcessor, ModelSpecificValue, PreprocessedImages},
|
||||
preprocessor_config::PreProcessorConfig,
|
||||
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<usize> {
|
||||
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)
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
let selected_scale = if !upscaling.is_empty() {
|
||||
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)
|
||||
}
|
||||
} else {
|
||||
// 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)
|
||||
};
|
||||
|
||||
// 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]
|
||||
}
|
||||
}
|
||||
|
||||
/// Pad image to target dimensions with black padding.
|
||||
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 black background (LLaMA 4 uses 0 for padding)
|
||||
let black = Rgb([0u8, 0, 0]);
|
||||
let mut padded = RgbImage::from_pixel(target_w, target_h, black);
|
||||
|
||||
// Copy image to top-left using efficient overlay
|
||||
image::imageops::overlay(&mut padded, &image.to_rgb8(), 0, 0);
|
||||
|
||||
DynamicImage::ImageRgb8(padded)
|
||||
}
|
||||
|
||||
/// Split image tensor into tiles.
|
||||
fn split_to_tiles(
|
||||
&self,
|
||||
tensor: &Array3<f32>,
|
||||
num_tiles_h: usize,
|
||||
num_tiles_w: usize,
|
||||
) -> Array4<f32> {
|
||||
let tile = self.tile_size as usize;
|
||||
let num_tiles = num_tiles_h * num_tiles_w;
|
||||
|
||||
let mut tiles = Array4::<f32>::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<f32> {
|
||||
let tile = self.tile_size;
|
||||
let resized = image.resize_exact(tile, tile, FilterType::Triangle);
|
||||
let mut tensor = transforms::to_tensor(&resized);
|
||||
transforms::normalize(&mut tensor, &self.mean, &self.std);
|
||||
tensor
|
||||
}
|
||||
|
||||
/// Process a single image.
|
||||
fn process_single_image(
|
||||
&self,
|
||||
image: &DynamicImage,
|
||||
) -> Result<(Array4<f32>, (usize, usize)), TransformError> {
|
||||
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
|
||||
// This limits how much we resize the image, but we still pad to target_size
|
||||
let resize_target = if !self.resize_to_max_canvas {
|
||||
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)
|
||||
} else {
|
||||
target_size
|
||||
};
|
||||
|
||||
// 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 = image.resize_exact(new_w, new_h, FilterType::Triangle);
|
||||
|
||||
// Step 4: Pad to target_size (the canvas from get_best_fit, not resize_target)
|
||||
let padded = self.pad_image(&resized, target_w, target_h);
|
||||
|
||||
// Step 5: Convert to tensor and normalize
|
||||
let mut tensor = transforms::to_tensor(&padded);
|
||||
transforms::normalize(&mut tensor, &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
|
||||
let tiles = self.split_to_tiles(&tensor, num_tiles_h, num_tiles_w);
|
||||
let num_tiles = num_tiles_h * num_tiles_w;
|
||||
|
||||
// Step 8: Add global tile if there are multiple tiles
|
||||
let output = if num_tiles > 1 {
|
||||
let global_tile = self.create_global_image(image);
|
||||
let mut combined = Array4::<f32>::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
|
||||
};
|
||||
|
||||
Ok((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 ImagePreProcessor 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<PreprocessedImages, TransformError> {
|
||||
if images.is_empty() {
|
||||
return Err(TransformError::InvalidShape {
|
||||
expected: "non-empty image batch".to_string(),
|
||||
actual: vec![0],
|
||||
});
|
||||
}
|
||||
|
||||
let processor = if config.max_image_tiles.is_some()
|
||||
|| config.image_mean.is_some()
|
||||
|| config.image_std.is_some()
|
||||
|| config.size.is_some()
|
||||
{
|
||||
Self::from_preprocessor_config(config)
|
||||
} else {
|
||||
self.clone()
|
||||
};
|
||||
|
||||
let mut all_outputs = Vec::new();
|
||||
let mut all_aspect_ratios = Vec::new();
|
||||
let mut image_sizes = Vec::new();
|
||||
let mut num_img_tokens = 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);
|
||||
image_sizes.push((image.height(), image.width()));
|
||||
num_img_tokens.push(tokens);
|
||||
}
|
||||
|
||||
// Find max tiles across batch for padding
|
||||
let max_tiles = all_outputs.iter().map(|o| o.shape()[0]).max().unwrap();
|
||||
let tile = self.tile_size as usize;
|
||||
|
||||
// Pad all outputs to max_tiles
|
||||
let batch_size = images.len();
|
||||
let mut pixel_values =
|
||||
ndarray::ArrayD::<f32>::zeros(IxDyn(&[batch_size, max_tiles, 3, tile, tile]));
|
||||
|
||||
for (b, output) in all_outputs.iter().enumerate() {
|
||||
let num_tiles = output.shape()[0];
|
||||
for t in 0..num_tiles {
|
||||
pixel_values
|
||||
.slice_mut(s![b, t, .., .., ..])
|
||||
.assign(&output.slice(s![t, .., .., ..]));
|
||||
}
|
||||
// Remaining tiles stay as zeros (padding)
|
||||
}
|
||||
|
||||
// Store aspect ratios as model-specific data
|
||||
let mut model_specific = std::collections::HashMap::new();
|
||||
|
||||
let aspect_ratios_flat: Vec<u32> = all_aspect_ratios
|
||||
.iter()
|
||||
.flat_map(|&(h, w)| vec![h as u32, w as u32])
|
||||
.collect();
|
||||
model_specific.insert(
|
||||
"aspect_ratios".to_string(),
|
||||
ModelSpecificValue::UintTensor {
|
||||
data: aspect_ratios_flat,
|
||||
shape: vec![batch_size, 2],
|
||||
},
|
||||
);
|
||||
|
||||
Ok(PreprocessedImages {
|
||||
pixel_values: pixel_values.into_dyn(),
|
||||
num_img_tokens,
|
||||
image_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 super::*;
|
||||
|
||||
fn create_test_image(width: u32, height: u32, color: Rgb<u8>) -> 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 {:?} not found",
|
||||
exp
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[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();
|
||||
|
||||
assert_eq!(result.batch_size(), 1);
|
||||
assert!(result.num_img_tokens[0] > 0);
|
||||
|
||||
// Check pixel values are normalized
|
||||
let flat = result.pixel_values_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();
|
||||
|
||||
assert_eq!(result.batch_size(), 1);
|
||||
// Wide image should have more tiles in width direction
|
||||
let aspect_ratios = result.model_specific.get("aspect_ratios").unwrap();
|
||||
if let ModelSpecificValue::UintTensor { 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();
|
||||
|
||||
assert_eq!(result.batch_size(), 2);
|
||||
assert_eq!(result.image_sizes.len(), 2);
|
||||
assert_eq!(result.num_img_tokens.len(), 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::UintTensor { 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 {
|
||||
// Output should have num_tiles + 1 (for global tile)
|
||||
let shape = result.pixel_values.shape();
|
||||
assert_eq!(shape[1], 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
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -12,7 +12,9 @@
|
||||
//! - **Qwen3-VL** (`qwen3_vl`): Similar to Qwen2-VL but with patch_size=16 and [0.5,0.5,0.5] normalization
|
||||
//! - **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
|
||||
|
||||
pub mod llama4_vision;
|
||||
pub mod llava;
|
||||
pub mod phi3_vision;
|
||||
pub mod phi4_vision;
|
||||
@@ -20,6 +22,7 @@ pub mod qwen2_vl;
|
||||
pub mod qwen3_vl;
|
||||
pub mod qwen_vl_base;
|
||||
|
||||
pub use llama4_vision::Llama4VisionProcessor;
|
||||
pub use llava::{ImageAspectRatio, LlavaNextProcessor, LlavaProcessor};
|
||||
pub use phi3_vision::Phi3VisionProcessor;
|
||||
pub use phi4_vision::Phi4VisionProcessor;
|
||||
|
||||
@@ -18,8 +18,9 @@ use std::{fs::File, io::Read, path::Path};
|
||||
|
||||
use ndarray::{Array4, Array5};
|
||||
use sgl_model_gateway::multimodal::vision::{
|
||||
image_processor::ModelSpecificValue, ImagePreProcessor, LlavaProcessor, Phi3VisionProcessor,
|
||||
Phi4VisionProcessor, PreProcessorConfig, Qwen2VLProcessor, Qwen3VLProcessor,
|
||||
image_processor::ModelSpecificValue, ImagePreProcessor, Llama4VisionProcessor, LlavaProcessor,
|
||||
Phi3VisionProcessor, Phi4VisionProcessor, PreProcessorConfig, Qwen2VLProcessor,
|
||||
Qwen3VLProcessor,
|
||||
};
|
||||
|
||||
/// Load a numpy .npz file and extract pixel_values
|
||||
@@ -1174,3 +1175,216 @@ fn test_phi4_vision_golden_odd_dims() {
|
||||
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<(u32, u32)> {
|
||||
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<i64> = reader.into_vec().expect("Failed to read array");
|
||||
|
||||
// Convert to Vec<(u32, u32)>
|
||||
let num_images = shape[0] as usize;
|
||||
(0..num_images)
|
||||
.map(|i| (data[i * 2] as u32, data[i * 2 + 1] as u32))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Load pixel_values for LLaMA 4 Vision (3D: [num_tiles, C, H, W])
|
||||
fn load_llama4_pixels(path: &Path) -> (Vec<f32>, Vec<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: Vec<usize> = reader.shape().iter().map(|&s| s as usize).collect();
|
||||
let data: Vec<f32> = 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("tests/fixtures/golden/llama4_vision");
|
||||
let image_path = Path::new("tests/fixtures/images").join(format!("{}.jpg", image_name));
|
||||
|
||||
if !golden_dir.exists() || !image_path.exists() {
|
||||
eprintln!(
|
||||
"Golden test fixtures for llama4_vision/{} not found, skipping test",
|
||||
image_name
|
||||
);
|
||||
eprintln!("Run: python scripts/generate_vision_golden.py --model llama4_vision");
|
||||
return;
|
||||
}
|
||||
|
||||
let npz_path = golden_dir.join(format!("golden_{}.npz", image_name));
|
||||
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<(u32, u32)> = match result.model_specific.get("aspect_ratios") {
|
||||
Some(ModelSpecificValue::UintTensor { 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 - Aspect ratios: golden={:?}, rust={:?}",
|
||||
image_name, golden_aspect_ratios, 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.num_img_tokens[0];
|
||||
println!(
|
||||
"llama4_vision - {} image - Tokens: golden={}, rust={}",
|
||||
image_name, golden_num_tokens, rust_num_tokens
|
||||
);
|
||||
assert_eq!(
|
||||
golden_num_tokens, rust_num_tokens,
|
||||
"num_tokens mismatch for {}",
|
||||
image_name
|
||||
);
|
||||
|
||||
// Check output shape - HuggingFace outputs (num_tiles, 3, 336, 336) without batch
|
||||
// Our Rust outputs (batch, num_tiles, 3, 336, 336) with batch dimension
|
||||
let rust_shape = result.pixel_values.shape();
|
||||
println!(
|
||||
"llama4_vision - {} image - Shape: golden={:?}, rust={:?}",
|
||||
image_name, golden_shape, rust_shape
|
||||
);
|
||||
|
||||
// HuggingFace returns without batch dim, we add batch=1
|
||||
assert!(
|
||||
rust_shape[0] == 1,
|
||||
"Expected batch dim to be 1, got {}",
|
||||
rust_shape[0]
|
||||
);
|
||||
assert!(
|
||||
rust_shape[1] >= golden_shape[0],
|
||||
"Expected at least {} tiles, got {}",
|
||||
golden_shape[0],
|
||||
rust_shape[1]
|
||||
);
|
||||
|
||||
// Compare pixel values
|
||||
let rust_pixels = result.pixel_values_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 - Max pixel diff: {:.6}",
|
||||
image_name, max_diff
|
||||
);
|
||||
|
||||
// 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 {} exceeds tolerance 0.03 for {}",
|
||||
max_diff,
|
||||
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");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user