feat: Image decoder in the frontend

Signed-off-by: Alexandre Milesi <[email protected]>

Author: milesial

Cargo.lock

@@ -21,7 +21,7 @@ testing-full = ["testing-cuda", "testing-nixl"]
 testing-cuda = ["dep:cudarc"]
 testing-nixl = ["dep:nixl-sys"]
 testing-etcd = []
-block-manager = ["dep:nixl-sys", "dep:cudarc", "dep:ndarray", "dep:nix", "dep:aligned-vec"]
+block-manager = ["dep:nixl-sys", "dep:cudarc", "dep:nix", "dep:aligned-vec"]
 cuda = ["dep:cudarc"]
 integration = ["dynamo-runtime/integration"]
 
@@ -97,7 +97,6 @@ dialoguer = { version = "0.11", default-features = false, features = [
 aligned-vec = { version = "0.6.4", optional = true }
 nixl-sys = { version = "=0.7.0", optional = true }
 cudarc = { workspace = true, optional = true }
-ndarray = { version = "0.16", optional = true }
 nix = { version = "0.26", optional = true }
 
 
@@ -143,6 +142,9 @@ json-five = { version = "0.3" }
 # media loading in the preprocessor
 reqwest = { workspace = true }
 base64 = { version = "0.22" }
+image = { version = "0.25" }
+tokio-rayon = {version = "2" }
+ndarray = { version = "0.16" }
 
 # Publishers
 zeromq = "0.4.1"

lib/llm/Cargo.toml

@@ -330,7 +330,7 @@ impl OpenAIPreprocessor {
             let _results = futures::future::join_all(
                 fetch_tasks
                     .iter()
-                    .map(|(_, content_part)| loader.fetch_media_part(content_part)),
+                    .map(|(_, content_part)| loader.fetch_and_decode_media_part(content_part)),
             )
             .await;
 

lib/llm/src/preprocessor.rs

@@ -2,7 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 mod common;
+mod decoders;
 mod loader;
 
 pub use common::EncodedMediaData;
+pub use decoders::{Decoder, ImageDecoder, MediaDecoder};
 pub use loader::MediaLoader;

lib/llm/src/preprocessor/media.rs

@@ -0,0 +1,191 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::Result;
+use image::GenericImageView;
+use ndarray::Array3;
+
+use super::super::common::EncodedMediaData;
+use super::super::decoders::DecodedMediaData;
+use super::Decoder;
+
+#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct ImageDecoder {
+    // maximum total size of the image in pixels
+    #[serde(default)]
+    pub max_pixels: Option<usize>,
+}
+
+impl Decoder for ImageDecoder {
+    fn decode(&self, data: EncodedMediaData) -> Result<DecodedMediaData> {
+        let bytes = data.into_bytes()?;
+        let img = image::load_from_memory(&bytes)?;
+        let (width, height) = img.dimensions();
+        let n_channels = img.color().channel_count();
+
+        let max_pixels = self.max_pixels.unwrap_or(usize::MAX);
+        anyhow::ensure!(
+            (width as usize) * (height as usize) <= max_pixels,
+            "Image dimensions {width}x{height} exceed max pixels {max_pixels}"
+        );
+        let data = match n_channels {
+            1 => img.to_luma8().into_raw(),
+            2 => img.to_luma_alpha8().into_raw(),
+            3 => img.to_rgb8().into_raw(),
+            4 => img.to_rgba8().into_raw(),
+            other => anyhow::bail!("Unsupported channel count {other}"),
+        };
+        let shape = (height as usize, width as usize, n_channels as usize);
+        let array = Array3::from_shape_vec(shape, data)?;
+        Ok(array.into())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use image::{DynamicImage, ImageBuffer};
+    use rstest::rstest;
+    use std::io::Cursor;
+
+    fn create_encoded_media_data(bytes: Vec<u8>) -> EncodedMediaData {
+        EncodedMediaData {
+            bytes,
+            b64_encoded: false,
+        }
+    }
+
+    fn create_test_image(
+        width: u32,
+        height: u32,
+        channels: u32,
+        format: image::ImageFormat,
+    ) -> Vec<u8> {
+        // Create dynamic image based on number of channels with constant values
+        let pixels = vec![128u8; channels as usize].repeat((width * height) as usize);
+        let dynamic_image = match channels {
+            1 => DynamicImage::ImageLuma8(
+                ImageBuffer::from_vec(width, height, pixels).expect("Failed to create image"),
+            ),
+            3 => DynamicImage::ImageRgb8(
+                ImageBuffer::from_vec(width, height, pixels).expect("Failed to create image"),
+            ),
+            4 => DynamicImage::ImageRgba8(
+                ImageBuffer::from_vec(width, height, pixels).expect("Failed to create image"),
+            ),
+            _ => unreachable!("Already validated channel count above"),
+        };
+
+        // Encode to bytes
+        let mut bytes = Vec::new();
+        dynamic_image
+            .write_to(&mut Cursor::new(&mut bytes), format)
+            .expect("Failed to encode test image");
+        bytes
+    }
+
+    #[rstest]
+    #[case(3, image::ImageFormat::Png, 10, 10, 3, "RGB PNG")]
+    #[case(4, image::ImageFormat::Png, 25, 30, 4, "RGBA PNG")]
+    #[case(1, image::ImageFormat::Png, 8, 12, 1, "Grayscale PNG")]
+    #[case(3, image::ImageFormat::Jpeg, 15, 20, 3, "RGB JPEG")]
+    #[case(3, image::ImageFormat::Bmp, 12, 18, 3, "RGB BMP")]
+    #[case(3, image::ImageFormat::WebP, 8, 8, 3, "RGB WebP")]
+    fn test_image_decode(
+        #[case] input_channels: u32,
+        #[case] format: image::ImageFormat,
+        #[case] width: u32,
+        #[case] height: u32,
+        #[case] expected_channels: u32,
+        #[case] description: &str,
+    ) {
+        let decoder = ImageDecoder::default();
+        let image_bytes = create_test_image(width, height, input_channels, format);
+        let encoded_data = create_encoded_media_data(image_bytes);
+
+        let result = decoder.decode(encoded_data);
+        assert!(result.is_ok(), "Failed to decode {}", description);
+
+        let decoded = result.unwrap();
+        assert_eq!(
+            decoded.shape,
+            vec![height as usize, width as usize, expected_channels as usize]
+        );
+        assert_eq!(decoded.dtype, "uint8");
+    }
+
+    #[rstest]
+    #[case(Some(200), 10, 10, image::ImageFormat::Png, true, "within limit")]
+    #[case(Some(50), 10, 10, image::ImageFormat::Jpeg, false, "exceeds limit")]
+    #[case(None, 200, 300, image::ImageFormat::Png, true, "no limit")]
+    fn test_pixel_limits(
+        #[case] max_pixels: Option<usize>,
+        #[case] width: u32,
+        #[case] height: u32,
+        #[case] format: image::ImageFormat,
+        #[case] should_succeed: bool,
+        #[case] test_case: &str,
+    ) {
+        let decoder = ImageDecoder { max_pixels };
+        let image_bytes = create_test_image(width, height, 3, format); // RGB
+        let encoded_data = create_encoded_media_data(image_bytes);
+
+        let result = decoder.decode(encoded_data);
+
+        if should_succeed {
+            assert!(
+                result.is_ok(),
+                "Should decode successfully for case: {} with format {:?}",
+                test_case,
+                format
+            );
+            let decoded = result.unwrap();
+            assert_eq!(decoded.shape, vec![height as usize, width as usize, 3]);
+            assert_eq!(
+                decoded.dtype, "uint8",
+                "dtype should be uint8 for case: {}",
+                test_case
+            );
+        } else {
+            assert!(
+                result.is_err(),
+                "Should fail for case: {} with format {:?}",
+                test_case,
+                format
+            );
+            let error_msg = result.unwrap_err().to_string();
+            assert!(
+                error_msg.contains("exceed max pixels"),
+                "Error should mention exceeding max pixels for case: {}",
+                test_case
+            );
+        }
+    }
+
+    #[rstest]
+    #[case(3, image::ImageFormat::Png)]
+    fn test_decode_1x1_image(#[case] input_channels: u32, #[case] format: image::ImageFormat) {
+        let decoder = ImageDecoder::default();
+        let image_bytes = create_test_image(1, 1, input_channels, format);
+        let encoded_data = create_encoded_media_data(image_bytes);
+
+        let result = decoder.decode(encoded_data);
+        assert!(
+            result.is_ok(),
+            "Should decode 1x1 image with {} channels in {:?} format successfully",
+            input_channels,
+            format
+        );
+
+        let decoded = result.unwrap();
+        assert_eq!(decoded.shape.len(), 3, "Should have 3 dimensions");
+        assert_eq!(decoded.shape[0], 1, "Height should be 1");
+        assert_eq!(decoded.shape[1], 1, "Width should be 1");
+        assert_eq!(
+            decoded.dtype, "uint8",
+            "dtype should be uint8 for {} channels {:?}",
+            input_channels, format
+        );
+    }
+}

lib/llm/src/preprocessor/media/decoders/image.rs

@@ -0,0 +1,52 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::Result;
+
+use super::common::EncodedMediaData;
+use ndarray::{ArrayBase, Dimension, OwnedRepr};
+mod image;
+
+pub use image::ImageDecoder;
+
+// Decoded media data (image RGB, video frames pixels, ...)
+#[derive(Debug)]
+pub struct DecodedMediaData {
+    pub(crate) data: Vec<u8>,
+    pub(crate) shape: Vec<usize>,
+    pub(crate) dtype: String,
+}
+
+// convert Array{N}<u8> to DecodedMediaData
+// TODO: Array1<f32> for audio
+impl<D: Dimension> From<ArrayBase<OwnedRepr<u8>, D>> for DecodedMediaData {
+    fn from(array: ArrayBase<OwnedRepr<u8>, D>) -> Self {
+        let shape = array.shape().to_vec();
+        let (data, _) = array.into_raw_vec_and_offset();
+        Self {
+            data,
+            shape,
+            dtype: "uint8".to_string(),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+pub trait Decoder: Clone + Send + 'static {
+    fn decode(&self, data: EncodedMediaData) -> Result<DecodedMediaData>;
+
+    async fn decode_async(&self, data: EncodedMediaData) -> Result<DecodedMediaData> {
+        // light clone (only config params)
+        let decoder = self.clone();
+        // compute heavy -> rayon
+        let result = tokio_rayon::spawn(move || decoder.decode(data)).await?;
+        Ok(result)
+    }
+}
+
+#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
+pub struct MediaDecoder {
+    #[serde(default)]
+    pub image_decoder: ImageDecoder,
+    // TODO: video, audio decoders
+}