[model-gateway] add llama4 vision image processor (#14438)

sgl-router/tests/vision_golden_tests.rs

@@ -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");
+}