remove multimodal as this is completely dead code (#17750)
This commit is contained in:
@@ -40,7 +40,6 @@ serde_json = { version = "1.0", default-features = false, features = [
|
||||
"std",
|
||||
"preserve_order",
|
||||
] }
|
||||
serde_bytes = "0.11"
|
||||
bytes = "1.8.0"
|
||||
http-body = "1.0"
|
||||
rand = "0.9.2"
|
||||
@@ -48,7 +47,6 @@ reqwest = { version = "0.12.8", features = ["stream", "blocking", "json", "rustl
|
||||
futures-util = "0.3"
|
||||
futures = "0.3"
|
||||
dashmap = "6.1.0"
|
||||
lru = "0.16.2"
|
||||
blake3 = "1.5"
|
||||
xxhash-rust = { version = "0.8", features = ["xxh3"] }
|
||||
bytemuck = { version = "1.21", features = ["derive"] }
|
||||
@@ -98,13 +96,10 @@ rmcp = { version = "0.8.3", features = ["client", "server",
|
||||
"reqwest",
|
||||
"auth"] }
|
||||
serde_yaml = "0.9"
|
||||
oracle = { version = "0.6.3", features = ["chrono"] }
|
||||
subtle = "2.6"
|
||||
jsonwebtoken = { version = "9.3", default-features = false, features = ["use_pem"] }
|
||||
num-traits = "0.2"
|
||||
num-bigint = "0.4"
|
||||
image = { version = "0.25.4", default-features = false, features = ["png", "jpeg", "gif", "bmp", "ico", "tiff", "webp"] }
|
||||
ndarray = "0.16"
|
||||
base64 = "0.22"
|
||||
openai-harmony = { git = "https://github.com/openai/harmony", tag = "v0.0.4" }
|
||||
openmetrics-parser = "0.4.4"
|
||||
@@ -115,17 +110,12 @@ tonic = { version = "0.14.2", features = ["gzip", "transport"] }
|
||||
prost = "0.14.1"
|
||||
prost-types = "0.14.1"
|
||||
tonic-prost = "0.14.2"
|
||||
deadpool = { version = "0.12", features = ["managed", "rt_tokio_1"] }
|
||||
backoff = { version = "0.4", features = ["tokio"] }
|
||||
strum = { version = "0.26", features = ["derive"] }
|
||||
bitflags = "2.10.0"
|
||||
once_cell = "1.21.3"
|
||||
tokio-postgres = { version = "0.7.15", features = ["runtime","with-chrono-0_4","with-serde_json-1","array-impls"] }
|
||||
deadpool-postgres = "0.14.1"
|
||||
|
||||
# CRDT for Mesh state synchronization
|
||||
crdts = "7.3"
|
||||
redis = { version = "0.27.6", features = ["tokio-comp", "json", "connection-manager"] }
|
||||
deadpool-redis = "0.18.0"
|
||||
|
||||
|
||||
# wasm dependencies
|
||||
|
||||
@@ -7,7 +7,6 @@ pub mod grpc_client;
|
||||
pub use smg_mcp as mcp;
|
||||
pub mod mesh;
|
||||
pub mod middleware;
|
||||
pub mod multimodal;
|
||||
pub mod observability;
|
||||
pub mod policies;
|
||||
pub use openai_protocol as protocols;
|
||||
|
||||
@@ -1,47 +0,0 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
use super::types::Modality;
|
||||
|
||||
pub type MultiModalResult<T> = Result<T, MultiModalError>;
|
||||
|
||||
#[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 decode task failed: {0}")]
|
||||
Blocking(#[from] tokio::task::JoinError),
|
||||
#[error("image decode error: {0}")]
|
||||
Image(#[from] image::ImageError),
|
||||
#[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("too many {modality:?} items provided. limit={limit}")]
|
||||
ModalityLimit { modality: Modality, limit: usize },
|
||||
#[error("tracker task join error: {0}")]
|
||||
Join(#[from] tokio::task::JoinError),
|
||||
#[error("tracker validation error: {0}")]
|
||||
Validation(String),
|
||||
}
|
||||
@@ -1,206 +0,0 @@
|
||||
use std::{collections::HashSet, path::PathBuf, sync::Arc, time::Duration};
|
||||
|
||||
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
|
||||
use bytes::Bytes;
|
||||
use image::DynamicImage;
|
||||
use reqwest::Client;
|
||||
use tokio::{fs, task};
|
||||
use url::Url;
|
||||
|
||||
use super::{
|
||||
error::MediaConnectorError,
|
||||
types::{ImageDetail, ImageFrame, ImageSource},
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MediaConnectorConfig {
|
||||
pub allowed_domains: Option<Vec<String>>,
|
||||
pub allowed_local_media_path: Option<PathBuf>,
|
||||
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(Debug, Clone)]
|
||||
pub enum MediaSource {
|
||||
Url(String),
|
||||
DataUrl(String),
|
||||
InlineBytes(Vec<u8>),
|
||||
File(PathBuf),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MediaConnector {
|
||||
client: Client,
|
||||
allowed_domains: Option<HashSet<String>>,
|
||||
allowed_local_media_path: Option<PathBuf>,
|
||||
fetch_timeout: Duration,
|
||||
}
|
||||
|
||||
impl MediaConnector {
|
||||
pub fn new(client: Client, config: MediaConnectorConfig) -> Result<Self, MediaConnectorError> {
|
||||
let allowed_domains = config.allowed_domains.map(|domains| {
|
||||
domains
|
||||
.into_iter()
|
||||
.map(|d| d.to_ascii_lowercase())
|
||||
.collect::<HashSet<_>>()
|
||||
});
|
||||
|
||||
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<Arc<ImageFrame>, 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,
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_http_image(
|
||||
&self,
|
||||
url: String,
|
||||
cfg: ImageFetchConfig,
|
||||
) -> Result<Arc<ImageFrame>, 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 = resp.bytes().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<Arc<ImageFrame>, 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 = BASE64_STANDARD.decode(data)?;
|
||||
self.decode_image(decoded.into(), cfg.detail, ImageSource::DataUrl)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn fetch_file(
|
||||
&self,
|
||||
path: PathBuf,
|
||||
cfg: ImageFetchConfig,
|
||||
) -> Result<Arc<ImageFrame>, 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 = fs::read(&canonical).await?;
|
||||
self.decode_image(
|
||||
bytes.into(),
|
||||
cfg.detail,
|
||||
ImageSource::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<Arc<ImageFrame>, MediaConnectorError> {
|
||||
let raw: Arc<Vec<u8>> = Arc::new(bytes.to_vec());
|
||||
let raw_clone = raw.clone();
|
||||
let image: DynamicImage =
|
||||
task::spawn_blocking(move || image::load_from_memory(&raw_clone)).await??;
|
||||
|
||||
Ok(Arc::new(ImageFrame::new(image, raw, detail, source)))
|
||||
}
|
||||
}
|
||||
@@ -1,17 +0,0 @@
|
||||
pub mod error;
|
||||
pub mod media;
|
||||
pub mod registry;
|
||||
pub mod tracker;
|
||||
pub mod types;
|
||||
pub mod vision;
|
||||
|
||||
pub use error::{MediaConnectorError, MultiModalError, MultiModalResult};
|
||||
pub use media::{ImageFetchConfig, MediaConnector, MediaConnectorConfig, MediaSource};
|
||||
pub use registry::{ModelProcessorSpec, ModelRegistry};
|
||||
pub use tracker::{AsyncMultiModalTracker, TrackerConfig, TrackerOutput};
|
||||
pub use types::*;
|
||||
// Re-export vision processing components
|
||||
pub use vision::{
|
||||
ImagePreProcessor, ImageProcessorRegistry, LlavaNextProcessor, LlavaProcessor,
|
||||
PreProcessorConfig, PreprocessedImages, TransformError,
|
||||
};
|
||||
@@ -1,468 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use once_cell::sync::Lazy;
|
||||
use serde_json::{json, Value};
|
||||
use thiserror::Error;
|
||||
|
||||
use super::types::{ImageSize, Modality, PromptReplacement, TokenId};
|
||||
use crate::tokenizer::traits::Tokenizer as TokenizerTrait;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
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 },
|
||||
}
|
||||
|
||||
pub type RegistryResult<T> = Result<T, ModelRegistryError>;
|
||||
|
||||
/// 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<TokenId> {
|
||||
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<u32> {
|
||||
Self::find_value(self.config, path).and_then(|value| value.as_u64().map(|v| v as u32))
|
||||
}
|
||||
|
||||
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<String>;
|
||||
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId>;
|
||||
fn modality_limits(&self, metadata: &ModelMetadata)
|
||||
-> RegistryResult<HashMap<Modality, usize>>;
|
||||
fn processor_kwargs(&self, metadata: &ModelMetadata) -> RegistryResult<Value>;
|
||||
fn prompt_replacements(
|
||||
&self,
|
||||
metadata: &ModelMetadata,
|
||||
image_sizes: &[ImageSize],
|
||||
) -> RegistryResult<Vec<PromptReplacement>>;
|
||||
}
|
||||
|
||||
pub struct ModelRegistry {
|
||||
specs: Vec<LazySpec>,
|
||||
}
|
||||
|
||||
impl ModelRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
specs: vec![
|
||||
LazySpec::new("llava", || Box::new(LlavaSpec)),
|
||||
LazySpec::new("qwen_vl", || Box::new(QwenVLVisionSpec)),
|
||||
LazySpec::new("phi3_v", || 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<Box<dyn ModelProcessorSpec>>,
|
||||
}
|
||||
|
||||
impl LazySpec {
|
||||
fn new(_id: &'static str, factory: fn() -> Box<dyn ModelProcessorSpec>) -> Self {
|
||||
Self {
|
||||
inner: Lazy::new(factory),
|
||||
}
|
||||
}
|
||||
|
||||
fn get(&self) -> &dyn ModelProcessorSpec {
|
||||
self.inner.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
struct LlavaSpec;
|
||||
|
||||
impl LlavaSpec {
|
||||
fn patch_size(metadata: &ModelMetadata) -> u32 {
|
||||
metadata
|
||||
.config_u32(&["vision_config", "patch_size"])
|
||||
.unwrap_or(14)
|
||||
}
|
||||
|
||||
fn tokens_per_image(metadata: &ModelMetadata, size: &ImageSize) -> usize {
|
||||
let patch = Self::patch_size(metadata);
|
||||
let cols = size.width.div_ceil(patch) as usize;
|
||||
let rows = size.height.div_ceil(patch) as usize;
|
||||
cols * rows
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelProcessorSpec for LlavaSpec {
|
||||
fn name(&self) -> &'static str {
|
||||
"llava"
|
||||
}
|
||||
|
||||
fn matches(&self, metadata: &ModelMetadata) -> bool {
|
||||
metadata.model_id.to_ascii_lowercase().contains("llava")
|
||||
}
|
||||
|
||||
fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult<String> {
|
||||
Ok("<image>".to_string())
|
||||
}
|
||||
|
||||
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
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<HashMap<Modality, usize>> {
|
||||
Ok(HashMap::from([(Modality::Image, 4)]))
|
||||
}
|
||||
|
||||
fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult<Value> {
|
||||
Ok(json!({}))
|
||||
}
|
||||
|
||||
fn prompt_replacements(
|
||||
&self,
|
||||
metadata: &ModelMetadata,
|
||||
image_sizes: &[ImageSize],
|
||||
) -> RegistryResult<Vec<PromptReplacement>> {
|
||||
let token_id = self.placeholder_token_id(metadata)?;
|
||||
let token = self.placeholder_token(metadata)?;
|
||||
Ok(image_sizes
|
||||
.iter()
|
||||
.map(|size| {
|
||||
let count = Self::tokens_per_image(metadata, size);
|
||||
PromptReplacement::repeated(Modality::Image, &token, token_id, count)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
struct QwenVLVisionSpec;
|
||||
|
||||
impl QwenVLVisionSpec {
|
||||
fn pad_token_id(metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
metadata
|
||||
.config_u32(&["vision_token_id"])
|
||||
.map(|v| v as TokenId)
|
||||
.ok_or_else(|| ModelRegistryError::MissingConfigField {
|
||||
field: "vision_token_id".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn start_token_id(metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
metadata
|
||||
.config_u32(&["vision_start_token_id"])
|
||||
.map(|v| v as TokenId)
|
||||
.ok_or_else(|| ModelRegistryError::MissingConfigField {
|
||||
field: "vision_start_token_id".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn patch_grid(metadata: &ModelMetadata, size: &ImageSize) -> (usize, usize) {
|
||||
let patch = metadata
|
||||
.config_u32(&["vision_config", "patch_size"])
|
||||
.unwrap_or(14);
|
||||
let cols = size.width.div_ceil(patch) as usize;
|
||||
let rows = size.height.div_ceil(patch) as usize;
|
||||
(rows, cols)
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelProcessorSpec for QwenVLVisionSpec {
|
||||
fn name(&self) -> &'static str {
|
||||
"qwen_vl"
|
||||
}
|
||||
|
||||
fn matches(&self, metadata: &ModelMetadata) -> bool {
|
||||
metadata.model_id.to_ascii_lowercase().contains("qwen")
|
||||
&& metadata.model_id.to_ascii_lowercase().contains("vl")
|
||||
}
|
||||
|
||||
fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult<String> {
|
||||
Ok("<image>".to_string())
|
||||
}
|
||||
|
||||
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
metadata
|
||||
.config_u32(&["image_token_id"])
|
||||
.map(|v| v as TokenId)
|
||||
.ok_or_else(|| ModelRegistryError::MissingConfigField {
|
||||
field: "image_token_id".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn modality_limits(
|
||||
&self,
|
||||
_metadata: &ModelMetadata,
|
||||
) -> RegistryResult<HashMap<Modality, usize>> {
|
||||
Ok(HashMap::from([(Modality::Image, 10)]))
|
||||
}
|
||||
|
||||
fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult<Value> {
|
||||
Ok(json!({}))
|
||||
}
|
||||
|
||||
fn prompt_replacements(
|
||||
&self,
|
||||
metadata: &ModelMetadata,
|
||||
image_sizes: &[ImageSize],
|
||||
) -> RegistryResult<Vec<PromptReplacement>> {
|
||||
let start_token_id = Self::start_token_id(metadata)?;
|
||||
let pad_token_id = Self::pad_token_id(metadata)?;
|
||||
let placeholder_token = self.placeholder_token(metadata)?;
|
||||
Ok(image_sizes
|
||||
.iter()
|
||||
.map(|size| {
|
||||
let (rows, cols) = Self::patch_grid(metadata, size);
|
||||
let pad_len = rows * cols;
|
||||
let mut tokens = Vec::with_capacity(pad_len + 1);
|
||||
tokens.push(start_token_id);
|
||||
tokens.extend(std::iter::repeat_n(pad_token_id, pad_len));
|
||||
PromptReplacement::sequence(Modality::Image, &placeholder_token, tokens)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
struct Phi3VisionSpec;
|
||||
|
||||
impl Phi3VisionSpec {
|
||||
fn tokens_per_image(metadata: &ModelMetadata) -> usize {
|
||||
metadata
|
||||
.config_u32(&["img_processor", "num_img_tokens"])
|
||||
.unwrap_or(256) as usize
|
||||
}
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
|
||||
fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult<String> {
|
||||
Ok("<image>".to_owned())
|
||||
}
|
||||
|
||||
fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult<TokenId> {
|
||||
metadata.token_id("<image>")
|
||||
}
|
||||
|
||||
fn modality_limits(
|
||||
&self,
|
||||
_metadata: &ModelMetadata,
|
||||
) -> RegistryResult<HashMap<Modality, usize>> {
|
||||
Ok(HashMap::from([(Modality::Image, 4)]))
|
||||
}
|
||||
|
||||
fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult<Value> {
|
||||
Ok(json!({}))
|
||||
}
|
||||
|
||||
fn prompt_replacements(
|
||||
&self,
|
||||
metadata: &ModelMetadata,
|
||||
image_sizes: &[ImageSize],
|
||||
) -> RegistryResult<Vec<PromptReplacement>> {
|
||||
let token_id = self.placeholder_token_id(metadata)?;
|
||||
let token = self.placeholder_token(metadata)?;
|
||||
let count = Self::tokens_per_image(metadata);
|
||||
Ok(image_sizes
|
||||
.iter()
|
||||
.map(|_| PromptReplacement::repeated(Modality::Image, &token, token_id, count))
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::tokenizer::traits::{
|
||||
Decoder, Encoder, Encoding, SpecialTokens, Tokenizer as TokenizerTrait,
|
||||
};
|
||||
|
||||
struct TestTokenizer {
|
||||
vocab: HashMap<String, u32>,
|
||||
}
|
||||
|
||||
impl TestTokenizer {
|
||||
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<Encoding> {
|
||||
Ok(Encoding::Sp(Vec::new()))
|
||||
}
|
||||
|
||||
fn encode_batch(
|
||||
&self,
|
||||
inputs: &[&str],
|
||||
add_special_tokens: bool,
|
||||
) -> anyhow::Result<Vec<Encoding>> {
|
||||
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<String> {
|
||||
Ok(String::new())
|
||||
}
|
||||
}
|
||||
|
||||
impl TokenizerTrait for TestTokenizer {
|
||||
fn vocab_size(&self) -> usize {
|
||||
self.vocab.len()
|
||||
}
|
||||
|
||||
fn get_special_tokens(&self) -> &SpecialTokens {
|
||||
static TOKENS: Lazy<SpecialTokens> = 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<u32> {
|
||||
self.vocab.get(token).copied()
|
||||
}
|
||||
|
||||
fn id_to_token(&self, _id: u32) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn std::any::Any {
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llava_prompt_replacement_uses_config_ids() {
|
||||
let tokenizer = TestTokenizer::new(&[("<image>", 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 replacements = spec
|
||||
.prompt_replacements(&metadata, &[ImageSize::new(336, 336)])
|
||||
.unwrap();
|
||||
assert_eq!(replacements[0].tokens.len(), 576);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen_vision_uses_config_token_ids() {
|
||||
let tokenizer = TestTokenizer::new(&[("<image>", 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");
|
||||
let replacements = spec
|
||||
.prompt_replacements(&metadata, &[ImageSize::new(448, 448)])
|
||||
.unwrap();
|
||||
// 448/14 = 32 patches => 1024 pad tokens + 1 start token
|
||||
assert_eq!(replacements[0].tokens.len(), 1025);
|
||||
assert_eq!(replacements[0].tokens[0], 151652);
|
||||
assert_eq!(replacements[0].tokens[1], 151654);
|
||||
}
|
||||
|
||||
#[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, &[ImageSize::new(336, 336)])
|
||||
.unwrap();
|
||||
assert_eq!(replacements[0].tokens.len(), 144);
|
||||
assert_eq!(replacements[0].tokens[0], 555);
|
||||
}
|
||||
}
|
||||
@@ -1,176 +0,0 @@
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
use super::{
|
||||
error::{MultiModalError, MultiModalResult},
|
||||
media::{ImageFetchConfig, MediaConnector, MediaSource},
|
||||
types::{
|
||||
ChatContentPart, ConversationSegment, ImageDetail, Modality, MultiModalData,
|
||||
MultiModalUUIDs, PlaceholderHandle, PlaceholderMap, TrackedMedia, DEFAULT_PLACEHOLDERS,
|
||||
},
|
||||
};
|
||||
|
||||
type PendingTask = JoinHandle<MultiModalResult<TrackedMedia>>;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TrackerConfig {
|
||||
pub placeholder_tokens: HashMap<Modality, String>,
|
||||
pub modality_limits: HashMap<Modality, usize>,
|
||||
}
|
||||
|
||||
impl Default for TrackerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
placeholder_tokens: DEFAULT_PLACEHOLDERS
|
||||
.iter()
|
||||
.map(|(k, v)| (*k, (*v).to_string()))
|
||||
.collect(),
|
||||
modality_limits: HashMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct TrackerOutput {
|
||||
pub conversation: Vec<ConversationSegment>,
|
||||
pub data: MultiModalData,
|
||||
pub uuids: MultiModalUUIDs,
|
||||
pub placeholders: PlaceholderMap,
|
||||
}
|
||||
|
||||
pub struct AsyncMultiModalTracker {
|
||||
media_connector: Arc<MediaConnector>,
|
||||
config: TrackerConfig,
|
||||
pending: HashMap<Modality, Vec<PendingTask>>,
|
||||
placeholders: PlaceholderMap,
|
||||
conversation: Vec<ConversationSegment>,
|
||||
uuids: MultiModalUUIDs,
|
||||
counts: HashMap<Modality, usize>,
|
||||
}
|
||||
|
||||
impl AsyncMultiModalTracker {
|
||||
pub fn new(media_connector: Arc<MediaConnector>, config: TrackerConfig) -> Self {
|
||||
Self {
|
||||
media_connector,
|
||||
config,
|
||||
pending: HashMap::new(),
|
||||
placeholders: PlaceholderMap::new(),
|
||||
conversation: Vec::new(),
|
||||
uuids: HashMap::new(),
|
||||
counts: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn push_part(&mut self, part: ChatContentPart) -> MultiModalResult<()> {
|
||||
match part {
|
||||
ChatContentPart::Text { text } => {
|
||||
if !text.is_empty() {
|
||||
self.conversation.push(ConversationSegment::text(text));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
ChatContentPart::ImageUrl { url, detail, uuid } => {
|
||||
self.enqueue_image(MediaSource::Url(url), detail.unwrap_or_default(), uuid)
|
||||
}
|
||||
ChatContentPart::ImageData {
|
||||
data,
|
||||
mime_type: _,
|
||||
uuid,
|
||||
detail,
|
||||
} => self.enqueue_image(
|
||||
MediaSource::InlineBytes(data),
|
||||
detail.unwrap_or_default(),
|
||||
uuid,
|
||||
),
|
||||
ChatContentPart::ImageEmbeds { .. } => {
|
||||
Err(MultiModalError::UnsupportedContent("image_embeds"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn finalize(mut self) -> MultiModalResult<TrackerOutput> {
|
||||
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 {
|
||||
conversation: self.conversation,
|
||||
data,
|
||||
uuids: self.uuids,
|
||||
placeholders: self.placeholders,
|
||||
})
|
||||
}
|
||||
|
||||
fn placeholder_token(&self, modality: Modality) -> &str {
|
||||
if let Some(token) = self.config.placeholder_tokens.get(&modality) {
|
||||
token.as_str()
|
||||
} else {
|
||||
DEFAULT_PLACEHOLDERS
|
||||
.get(&modality)
|
||||
.copied()
|
||||
.unwrap_or("<mm>")
|
||||
}
|
||||
}
|
||||
|
||||
fn next_index(&mut self, modality: Modality) -> MultiModalResult<usize> {
|
||||
let count = self.counts.entry(modality).or_insert(0);
|
||||
let next = *count;
|
||||
let limit = self.config.modality_limits.get(&modality).copied();
|
||||
if let Some(limit) = limit {
|
||||
if next >= limit {
|
||||
return Err(MultiModalError::ModalityLimit { modality, limit });
|
||||
}
|
||||
}
|
||||
*count += 1;
|
||||
Ok(next)
|
||||
}
|
||||
|
||||
fn record_uuid(&mut self, modality: Modality, uuid: Option<String>) {
|
||||
self.uuids.entry(modality).or_default().push(uuid);
|
||||
}
|
||||
|
||||
fn add_placeholder(&mut self, token: String, handle: PlaceholderHandle) {
|
||||
self.placeholders.entry(token).or_default().push(handle);
|
||||
}
|
||||
|
||||
fn enqueue_image(
|
||||
&mut self,
|
||||
source: MediaSource,
|
||||
detail: ImageDetail,
|
||||
uuid: Option<String>,
|
||||
) -> MultiModalResult<()> {
|
||||
let modality = Modality::Image;
|
||||
let idx = self.next_index(modality)?;
|
||||
let token = self.placeholder_token(modality).to_string();
|
||||
let text_position = self.conversation.len();
|
||||
self.conversation
|
||||
.push(ConversationSegment::placeholder(token.clone()));
|
||||
self.add_placeholder(
|
||||
token.clone(),
|
||||
PlaceholderHandle {
|
||||
modality,
|
||||
item_index: idx,
|
||||
text_position,
|
||||
},
|
||||
);
|
||||
self.record_uuid(modality, uuid);
|
||||
|
||||
let connector = Arc::clone(&self.media_connector);
|
||||
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);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,290 +0,0 @@
|
||||
use std::{
|
||||
collections::{BTreeMap, HashMap},
|
||||
fmt,
|
||||
path::PathBuf,
|
||||
sync::{Arc, LazyLock},
|
||||
};
|
||||
|
||||
use image::DynamicImage;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
/// 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"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple helper used for default placeholder tokens per modality.
|
||||
pub static DEFAULT_PLACEHOLDERS: LazyLock<HashMap<Modality, &'static str>> = LazyLock::new(|| {
|
||||
HashMap::from([
|
||||
(Modality::Image, "<image>"),
|
||||
(Modality::ImageEmbeds, "<image_embeds>"),
|
||||
(Modality::Audio, "<audio>"),
|
||||
(Modality::Video, "<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 ChatContentPart {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
ImageUrl {
|
||||
url: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
detail: Option<ImageDetail>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
uuid: Option<String>,
|
||||
},
|
||||
ImageData {
|
||||
data: Vec<u8>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
mime_type: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
uuid: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
detail: Option<ImageDetail>,
|
||||
},
|
||||
ImageEmbeds {
|
||||
payload: Value,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
uuid: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Represents one segment in the conversation after placeholders are inserted.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ConversationSegment {
|
||||
Text(String),
|
||||
Placeholder { token: String },
|
||||
}
|
||||
|
||||
impl ConversationSegment {
|
||||
pub fn text<S: Into<String>>(text: S) -> Self {
|
||||
ConversationSegment::Text(text.into())
|
||||
}
|
||||
|
||||
pub fn placeholder<S: Into<String>>(token: S) -> Self {
|
||||
ConversationSegment::Placeholder {
|
||||
token: token.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Metadata describing where placeholder tokens were inserted.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PlaceholderHandle {
|
||||
pub modality: Modality,
|
||||
pub item_index: usize,
|
||||
pub text_position: usize,
|
||||
}
|
||||
|
||||
pub type PlaceholderMap = HashMap<String, Vec<PlaceholderHandle>>;
|
||||
|
||||
/// 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 },
|
||||
}
|
||||
|
||||
/// Concrete image payload captured by the media connector.
|
||||
#[derive(Debug)]
|
||||
pub struct ImageFrame {
|
||||
image: DynamicImage,
|
||||
raw_bytes: Arc<Vec<u8>>,
|
||||
pub detail: ImageDetail,
|
||||
pub source: ImageSource,
|
||||
}
|
||||
|
||||
impl ImageFrame {
|
||||
pub fn new(
|
||||
image: DynamicImage,
|
||||
raw_bytes: Arc<Vec<u8>>,
|
||||
detail: ImageDetail,
|
||||
source: ImageSource,
|
||||
) -> Self {
|
||||
Self {
|
||||
image,
|
||||
raw_bytes,
|
||||
detail,
|
||||
source,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn data(&self) -> &DynamicImage {
|
||||
&self.image
|
||||
}
|
||||
|
||||
pub fn raw_bytes(&self) -> &[u8] {
|
||||
self.raw_bytes.as_slice()
|
||||
}
|
||||
|
||||
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<ImageFrame>),
|
||||
/// Placeholder variants for future modalities.
|
||||
Audio,
|
||||
Video,
|
||||
Embeddings,
|
||||
}
|
||||
|
||||
pub type MultiModalData = HashMap<Modality, Vec<TrackedMedia>>;
|
||||
pub type MultiModalUUIDs = HashMap<Modality, Vec<Option<String>>>;
|
||||
|
||||
pub type TokenId = i32;
|
||||
|
||||
#[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, Serialize, Deserialize)]
|
||||
pub struct MultiModalTensor {
|
||||
pub shape: Vec<usize>,
|
||||
pub dtype: String,
|
||||
#[serde(with = "serde_bytes")]
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum MultiModalValue {
|
||||
Tensor(MultiModalTensor),
|
||||
Json(Value),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct MultiModalInputs {
|
||||
pub prompt_token_ids: Vec<u32>,
|
||||
#[serde(default)]
|
||||
pub mm_kwargs: BTreeMap<String, Vec<MultiModalValue>>,
|
||||
#[serde(default)]
|
||||
pub mm_hashes: BTreeMap<String, Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub mm_placeholders: BTreeMap<String, Vec<PlaceholderRange>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_salt: Option<String>,
|
||||
}
|
||||
|
||||
impl MultiModalInputs {
|
||||
pub fn new(prompt_token_ids: Vec<u32>) -> Self {
|
||||
Self {
|
||||
prompt_token_ids,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptReplacement {
|
||||
pub modality: Modality,
|
||||
pub placeholder_token: String,
|
||||
pub tokens: Vec<TokenId>,
|
||||
}
|
||||
|
||||
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],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sequence(modality: Modality, placeholder_token: &str, sequence: Vec<TokenId>) -> Self {
|
||||
Self {
|
||||
modality,
|
||||
placeholder_token: placeholder_token.to_string(),
|
||||
tokens: sequence,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn multimodal_inputs_defaults() {
|
||||
let inputs = MultiModalInputs::new(vec![1, 2, 3]);
|
||||
assert_eq!(inputs.prompt_token_ids, vec![1, 2, 3]);
|
||||
assert!(inputs.mm_kwargs.is_empty());
|
||||
}
|
||||
|
||||
#[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, "<image>", 100, 3);
|
||||
assert_eq!(rep.tokens, vec![100, 100, 100]);
|
||||
}
|
||||
}
|
||||
@@ -1,491 +0,0 @@
|
||||
//! Image processor trait and output types.
|
||||
//!
|
||||
//! This module defines the interface for model-specific image processors
|
||||
//! and the common output format for preprocessed images.
|
||||
|
||||
use std::{borrow::Cow, collections::HashMap};
|
||||
|
||||
use image::DynamicImage;
|
||||
use ndarray::{Array4, ArrayD};
|
||||
|
||||
use super::{preprocessor_config::PreProcessorConfig, transforms::TransformError};
|
||||
|
||||
/// Model-specific output values that vary by architecture.
|
||||
///
|
||||
/// Different vision models require different auxiliary outputs beyond pixel_values.
|
||||
/// This enum captures the common types of such outputs.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ModelSpecificValue {
|
||||
/// A tensor with shape information (data as flat vec, shape as dims)
|
||||
Tensor { data: Vec<f32>, shape: Vec<usize> },
|
||||
|
||||
/// A tensor of integers (e.g., aspect_ratio_ids)
|
||||
IntTensor { data: Vec<i64>, shape: Vec<usize> },
|
||||
|
||||
/// A tensor of unsigned integers (e.g., image_grid_thw)
|
||||
UintTensor { data: Vec<u32>, shape: Vec<usize> },
|
||||
|
||||
/// Simple integer value
|
||||
Int(i64),
|
||||
|
||||
/// Simple float value
|
||||
Float(f64),
|
||||
|
||||
/// List of integers
|
||||
IntVec(Vec<i64>),
|
||||
|
||||
/// List of unsigned integers
|
||||
UintVec(Vec<u32>),
|
||||
|
||||
/// List of floats
|
||||
FloatVec(Vec<f32>),
|
||||
|
||||
/// List of tuples (e.g., image 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<u32>) -> Self {
|
||||
let len = data.len();
|
||||
Self::UintTensor {
|
||||
data,
|
||||
shape: vec![len],
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a 2D uint tensor.
|
||||
pub fn uint_2d(data: Vec<u32>, 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<i64>) -> Self {
|
||||
let len = data.len();
|
||||
Self::IntTensor {
|
||||
data,
|
||||
shape: vec![len],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Preprocessed images ready for model consumption.
|
||||
///
|
||||
/// This struct contains all the outputs needed by the SGLang scheduler
|
||||
/// to construct `MultimodalInputs` for the model.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PreprocessedImages {
|
||||
/// Pixel values as a dynamic-dimensional float32 tensor.
|
||||
///
|
||||
/// This is the primary input to the vision encoder.
|
||||
/// Shape varies by model:
|
||||
/// - Standard: [B, C, H, W] (4D)
|
||||
/// - Phi3-Vision: [B, num_crops+1, C, H, W] (5D)
|
||||
pub pixel_values: ArrayD<f32>,
|
||||
|
||||
/// Number of image tokens per image in the batch.
|
||||
///
|
||||
/// Used to expand placeholder tokens in the text input.
|
||||
/// For example, LLaVA with 336x336 and patch_size=14 produces 576 tokens.
|
||||
pub num_img_tokens: Vec<usize>,
|
||||
|
||||
/// Original image sizes as (width, height) before preprocessing.
|
||||
///
|
||||
/// Some models need this for proper attention masking or position encoding.
|
||||
pub image_sizes: Vec<(u32, u32)>,
|
||||
|
||||
/// Model-specific auxiliary outputs.
|
||||
///
|
||||
/// Examples:
|
||||
/// - Qwen-VL: `image_grid_thw` for rotary position encoding
|
||||
/// - LLaMA-Vision: `aspect_ratio_ids`, `aspect_ratio_mask`
|
||||
/// - Phi3-Vision: `num_img_tokens` per crop
|
||||
pub model_specific: HashMap<String, ModelSpecificValue>,
|
||||
}
|
||||
|
||||
impl PreprocessedImages {
|
||||
/// Create a new PreprocessedImages with required fields (4D pixel values).
|
||||
pub fn new(
|
||||
pixel_values: Array4<f32>,
|
||||
num_img_tokens: Vec<usize>,
|
||||
image_sizes: Vec<(u32, u32)>,
|
||||
) -> Self {
|
||||
Self {
|
||||
pixel_values: pixel_values.into_dyn(),
|
||||
num_img_tokens,
|
||||
image_sizes,
|
||||
model_specific: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new PreprocessedImages with dynamic-dimensional pixel values.
|
||||
///
|
||||
/// Use this for models like Phi3-Vision that have 5D tensors.
|
||||
pub fn new_dynamic(
|
||||
pixel_values: ArrayD<f32>,
|
||||
num_img_tokens: Vec<usize>,
|
||||
image_sizes: Vec<(u32, u32)>,
|
||||
) -> Self {
|
||||
Self {
|
||||
pixel_values,
|
||||
num_img_tokens,
|
||||
image_sizes,
|
||||
model_specific: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a model-specific value.
|
||||
pub fn with_extra(mut self, key: impl Into<String>, value: ModelSpecificValue) -> Self {
|
||||
self.model_specific.insert(key.into(), value);
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the batch size.
|
||||
pub fn batch_size(&self) -> usize {
|
||||
self.pixel_values.shape()[0]
|
||||
}
|
||||
|
||||
/// 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].
|
||||
pub fn channels(&self) -> usize {
|
||||
match self.pixel_values.ndim() {
|
||||
4 => self.pixel_values.shape()[1],
|
||||
5 => self.pixel_values.shape()[2],
|
||||
ndim => panic!(
|
||||
"Unsupported pixel_values dimension: {}, expected 4 or 5",
|
||||
ndim
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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].
|
||||
pub fn height(&self) -> usize {
|
||||
match self.pixel_values.ndim() {
|
||||
4 => self.pixel_values.shape()[2],
|
||||
5 => self.pixel_values.shape()[3],
|
||||
ndim => panic!(
|
||||
"Unsupported pixel_values dimension: {}, expected 4 or 5",
|
||||
ndim
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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].
|
||||
pub fn width(&self) -> usize {
|
||||
match self.pixel_values.ndim() {
|
||||
4 => self.pixel_values.shape()[3],
|
||||
5 => self.pixel_values.shape()[4],
|
||||
ndim => panic!(
|
||||
"Unsupported pixel_values dimension: {}, expected 4 or 5",
|
||||
ndim
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the number of dimensions of pixel_values.
|
||||
pub fn ndim(&self) -> usize {
|
||||
self.pixel_values.ndim()
|
||||
}
|
||||
|
||||
/// Get total number of image tokens across all images.
|
||||
pub fn total_tokens(&self) -> usize {
|
||||
self.num_img_tokens.iter().sum()
|
||||
}
|
||||
|
||||
/// Get pixel values as a flat f32 slice without copying if possible.
|
||||
pub fn pixel_values_flat(&self) -> Cow<'_, [f32]> {
|
||||
match self.pixel_values.as_slice() {
|
||||
Some(slice) => Cow::Borrowed(slice),
|
||||
None => Cow::Owned(self.pixel_values.iter().copied().collect()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the shape of pixel values as a vector.
|
||||
pub fn pixel_values_shape(&self) -> Vec<usize> {
|
||||
self.pixel_values.shape().to_vec()
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for model-specific image preprocessors.
|
||||
///
|
||||
/// Each vision model (LLaVA, Qwen-VL, Phi3-Vision, etc.) implements this trait
|
||||
/// to provide the correct preprocessing pipeline.
|
||||
pub trait ImagePreProcessor: 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 images ready for the model, or an error.
|
||||
fn preprocess(
|
||||
&self,
|
||||
images: &[DynamicImage],
|
||||
config: &PreProcessorConfig,
|
||||
) -> Result<PreprocessedImages, TransformError>;
|
||||
|
||||
/// Calculate the number of image 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 image processors.
|
||||
pub struct ImageProcessorRegistry {
|
||||
processors: HashMap<String, Box<dyn ImagePreProcessor>>,
|
||||
}
|
||||
|
||||
impl ImageProcessorRegistry {
|
||||
/// 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<String>, processor: Box<dyn ImagePreProcessor>) {
|
||||
self.processors.insert(pattern.into(), processor);
|
||||
}
|
||||
|
||||
/// Find a processor for the given model ID.
|
||||
///
|
||||
/// Matches by substring containment (case-insensitive).
|
||||
pub fn find(&self, model_id: &str) -> Option<&dyn ImagePreProcessor> {
|
||||
let model_lower = model_id.to_lowercase();
|
||||
for (pattern, processor) in &self.processors {
|
||||
if model_lower.contains(&pattern.to_lowercase()) {
|
||||
return Some(processor.as_ref());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Check if a model has a registered processor.
|
||||
pub fn has_processor(&self, model_id: &str) -> bool {
|
||||
self.find(model_id).is_some()
|
||||
}
|
||||
|
||||
/// 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 ImageProcessorRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ImageProcessorRegistry {
|
||||
/// Create a registry with all built-in processors registered.
|
||||
///
|
||||
/// Currently registers:
|
||||
/// - `llava-next` -> LlavaNextProcessor
|
||||
/// - `llava` -> LlavaProcessor (also matches llava-1.5, etc.)
|
||||
/// - `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)
|
||||
/// - `phi-3-vision` -> Phi3VisionProcessor (HD transform with 336x336 tiles)
|
||||
pub fn with_defaults() -> Self {
|
||||
let mut registry = Self::new();
|
||||
|
||||
// Register LLaVA-NeXT first (more specific pattern)
|
||||
registry.register(
|
||||
"llava-next",
|
||||
Box::new(super::processors::LlavaNextProcessor::new()),
|
||||
);
|
||||
registry.register(
|
||||
"llava-v1.6",
|
||||
Box::new(super::processors::LlavaNextProcessor::new()),
|
||||
);
|
||||
|
||||
// Register standard LLaVA (matches llava-1.5, llava-v1.5, etc.)
|
||||
registry.register("llava", 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()),
|
||||
);
|
||||
|
||||
// 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
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use ndarray::Array4;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_preprocessed_images_accessors() {
|
||||
let pixel_values = Array4::<f32>::zeros((2, 3, 336, 336));
|
||||
let images =
|
||||
PreprocessedImages::new(pixel_values, vec![576, 576], vec![(640, 480), (800, 600)]);
|
||||
|
||||
assert_eq!(images.batch_size(), 2);
|
||||
assert_eq!(images.channels(), 3);
|
||||
assert_eq!(images.height(), 336);
|
||||
assert_eq!(images.width(), 336);
|
||||
assert_eq!(images.total_tokens(), 1152);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_preprocessed_images_with_extra() {
|
||||
let pixel_values = Array4::<f32>::zeros((1, 3, 224, 224));
|
||||
let images = PreprocessedImages::new(pixel_values, 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!(images.model_specific.contains_key("image_grid_thw"));
|
||||
assert!(images.model_specific.contains_key("aspect_ratio_id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_model_specific_value_constructors() {
|
||||
let uint_1d = ModelSpecificValue::uint_1d(vec![1, 2, 3]);
|
||||
match uint_1d {
|
||||
ModelSpecificValue::UintTensor { data, shape } => {
|
||||
assert_eq!(data, vec![1, 2, 3]);
|
||||
assert_eq!(shape, vec![3]);
|
||||
}
|
||||
_ => panic!("Expected UintTensor"),
|
||||
}
|
||||
|
||||
let uint_2d = ModelSpecificValue::uint_2d(vec![1, 2, 3, 4], 2, 2);
|
||||
match uint_2d {
|
||||
ModelSpecificValue::UintTensor { data, shape } => {
|
||||
assert_eq!(data, vec![1, 2, 3, 4]);
|
||||
assert_eq!(shape, vec![2, 2]);
|
||||
}
|
||||
_ => panic!("Expected UintTensor"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pixel_values_flat() {
|
||||
let mut pixel_values = Array4::<f32>::zeros((1, 1, 2, 2));
|
||||
pixel_values[[0, 0, 0, 0]] = 1.0;
|
||||
pixel_values[[0, 0, 0, 1]] = 2.0;
|
||||
pixel_values[[0, 0, 1, 0]] = 3.0;
|
||||
pixel_values[[0, 0, 1, 1]] = 4.0;
|
||||
|
||||
let images = PreprocessedImages::new(pixel_values, vec![4], vec![(2, 2)]);
|
||||
let flat = images.pixel_values_flat();
|
||||
|
||||
assert_eq!(flat, vec![1.0, 2.0, 3.0, 4.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_with_defaults() {
|
||||
let registry = ImageProcessorRegistry::with_defaults();
|
||||
|
||||
// Should find LLaVA processor
|
||||
assert!(registry.has_processor("llava-hf/llava-1.5-7b-hf"));
|
||||
assert!(registry.has_processor("liuhaotian/llava-v1.5-7b"));
|
||||
|
||||
// Should find LLaVA-NeXT processor
|
||||
assert!(registry.has_processor("llava-hf/llava-v1.6-mistral-7b-hf"));
|
||||
assert!(registry.has_processor("lmms-lab/llava-next-interleave-qwen-7b"));
|
||||
|
||||
// Get the processor and check model name
|
||||
let processor = registry.find("llava-hf/llava-1.5-7b-hf").unwrap();
|
||||
assert_eq!(processor.model_name(), "llava");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_find() {
|
||||
let mut registry = ImageProcessorRegistry::new();
|
||||
|
||||
// Create a mock processor using LlavaProcessor
|
||||
registry.register(
|
||||
"test-model",
|
||||
Box::new(crate::multimodal::vision::processors::LlavaProcessor::new()),
|
||||
);
|
||||
|
||||
assert!(registry.has_processor("test-model-7b"));
|
||||
assert!(registry.has_processor("TEST-MODEL"));
|
||||
assert!(!registry.has_processor("other-model"));
|
||||
}
|
||||
}
|
||||
@@ -1,46 +0,0 @@
|
||||
//! Pure Rust vision processing module for multimodal models.
|
||||
//!
|
||||
//! This module provides image 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
|
||||
//! - `image_processor`: Trait and output types for processors
|
||||
//! - `processors`: Model-specific implementations (LLaVA, Qwen-VL, etc.)
|
||||
//!
|
||||
//! # Usage
|
||||
//!
|
||||
//! ```rust,ignore
|
||||
//! use smg::multimodal::vision::{
|
||||
//! PreProcessorConfig,
|
||||
//! processors::LlavaProcessor,
|
||||
//! ImagePreProcessor,
|
||||
//! };
|
||||
//!
|
||||
//! // 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 mod image_processor;
|
||||
pub mod preprocessor_config;
|
||||
pub mod processors;
|
||||
pub mod transforms;
|
||||
|
||||
// Re-export commonly used types
|
||||
pub use image_processor::{
|
||||
ImagePreProcessor, ImageProcessorRegistry, ModelSpecificValue, PreprocessedImages,
|
||||
};
|
||||
pub use preprocessor_config::PreProcessorConfig;
|
||||
pub use processors::{
|
||||
Llama4VisionProcessor, LlavaNextProcessor, LlavaProcessor, Phi3VisionProcessor,
|
||||
Phi4VisionProcessor, PixtralProcessor, Qwen2VLProcessor, Qwen3VLProcessor,
|
||||
};
|
||||
pub use transforms::TransformError;
|
||||
@@ -1,470 +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<u32>,
|
||||
pub width: Option<u32>,
|
||||
}
|
||||
|
||||
/// 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<Option<PatchSize>, 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<PatchSize>;
|
||||
|
||||
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
|
||||
formatter.write_str("an integer, a dict with height/width, or null")
|
||||
}
|
||||
|
||||
fn visit_none<E>(self) -> Result<Self::Value, E>
|
||||
where
|
||||
E: de::Error,
|
||||
{
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn visit_unit<E>(self) -> Result<Self::Value, E>
|
||||
where
|
||||
E: de::Error,
|
||||
{
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E>
|
||||
where
|
||||
E: de::Error,
|
||||
{
|
||||
let v = value as u32;
|
||||
Ok(Some(PatchSize {
|
||||
height: Some(v),
|
||||
width: Some(v),
|
||||
}))
|
||||
}
|
||||
|
||||
fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E>
|
||||
where
|
||||
E: de::Error,
|
||||
{
|
||||
let v = value as u32;
|
||||
Ok(Some(PatchSize {
|
||||
height: Some(v),
|
||||
width: Some(v),
|
||||
}))
|
||||
}
|
||||
|
||||
fn visit_map<M>(self, mut map: M) -> Result<Self::Value, M::Error>
|
||||
where
|
||||
M: MapAccess<'de>,
|
||||
{
|
||||
let mut height = None;
|
||||
let mut width = None;
|
||||
|
||||
while let Some(key) = map.next_key::<String>()? {
|
||||
match key.as_str() {
|
||||
"height" => height = Some(map.next_value::<u32>()?),
|
||||
"width" => width = Some(map.next_value::<u32>()?),
|
||||
_ => {
|
||||
let _ = map.next_value::<de::IgnoredAny>()?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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<String>,
|
||||
|
||||
/// Whether to convert to RGB
|
||||
#[serde(default)]
|
||||
pub do_convert_rgb: Option<bool>,
|
||||
|
||||
/// Whether to normalize with mean/std
|
||||
#[serde(default)]
|
||||
pub do_normalize: Option<bool>,
|
||||
|
||||
/// Whether to pad images
|
||||
#[serde(default)]
|
||||
pub do_pad: Option<bool>,
|
||||
|
||||
/// Whether to rescale pixel values (typically by 1/255)
|
||||
#[serde(default)]
|
||||
pub do_rescale: Option<bool>,
|
||||
|
||||
/// Whether to resize images
|
||||
#[serde(default)]
|
||||
pub do_resize: Option<bool>,
|
||||
|
||||
/// Whether to center crop after resizing
|
||||
#[serde(default)]
|
||||
pub do_center_crop: Option<bool>,
|
||||
|
||||
/// Per-channel normalization mean
|
||||
#[serde(default, alias = "norm_mean")]
|
||||
pub image_mean: Option<Vec<f64>>,
|
||||
|
||||
/// Per-channel normalization std
|
||||
#[serde(default, alias = "norm_std")]
|
||||
pub image_std: Option<Vec<f64>>,
|
||||
|
||||
/// Rescale factor (typically 1/255 = 0.00392156862745098)
|
||||
#[serde(default)]
|
||||
pub rescale_factor: Option<f64>,
|
||||
|
||||
/// PIL resampling filter enum (0=Nearest, 1=Lanczos, 2=Bilinear, 3=Bicubic)
|
||||
#[serde(default, alias = "resample")]
|
||||
pub resampling: Option<usize>,
|
||||
|
||||
/// Target size for resizing
|
||||
/// Can be {"height": H, "width": W} or {"shortest_edge": S}
|
||||
#[serde(default)]
|
||||
pub size: Option<HashMap<String, u32>>,
|
||||
|
||||
/// Target size for center cropping
|
||||
#[serde(default)]
|
||||
pub crop_size: Option<HashMap<String, u32>>,
|
||||
|
||||
// =====================
|
||||
// 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<PatchSize>,
|
||||
|
||||
/// Qwen-VL: merge size for token reduction
|
||||
#[serde(default)]
|
||||
pub merge_size: Option<usize>,
|
||||
|
||||
/// Qwen-VL: minimum total pixels
|
||||
#[serde(default)]
|
||||
pub min_pixels: Option<usize>,
|
||||
|
||||
/// Qwen-VL: maximum total pixels
|
||||
#[serde(default)]
|
||||
pub max_pixels: Option<usize>,
|
||||
|
||||
/// Qwen-VL: temporal patch size for video
|
||||
#[serde(default)]
|
||||
pub temporal_patch_size: Option<usize>,
|
||||
|
||||
/// Phi3-Vision: number of image crops
|
||||
#[serde(default)]
|
||||
pub num_crops: Option<usize>,
|
||||
|
||||
/// Phi4-Vision: dynamic HD max crops
|
||||
#[serde(default)]
|
||||
pub dynamic_hd: Option<usize>,
|
||||
|
||||
/// LLaMA-Vision: maximum image tiles
|
||||
#[serde(default)]
|
||||
pub max_image_tiles: Option<usize>,
|
||||
|
||||
/// Fixed number of image tokens (some models use this)
|
||||
#[serde(default)]
|
||||
pub num_img_tokens: Option<usize>,
|
||||
|
||||
// =====================
|
||||
// Special tokens
|
||||
// =====================
|
||||
/// Image start token
|
||||
#[serde(default)]
|
||||
pub im_start_token: Option<String>,
|
||||
|
||||
/// Image end token
|
||||
#[serde(default)]
|
||||
pub im_end_token: Option<String>,
|
||||
|
||||
/// Slice start token (for multi-crop)
|
||||
#[serde(default)]
|
||||
pub slice_start_token: Option<String>,
|
||||
|
||||
/// Slice end token
|
||||
#[serde(default)]
|
||||
pub slice_end_token: Option<String>,
|
||||
|
||||
/// Vision start token (alternative naming)
|
||||
#[serde(default)]
|
||||
pub vision_start_token: Option<String>,
|
||||
|
||||
/// Vision end token
|
||||
#[serde(default)]
|
||||
pub vision_end_token: Option<String>,
|
||||
|
||||
/// Catch-all for model-specific fields not explicitly defined
|
||||
#[serde(flatten)]
|
||||
pub extra: HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
impl PreProcessorConfig {
|
||||
/// Parse from JSON string.
|
||||
pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
|
||||
serde_json::from_str(json)
|
||||
}
|
||||
|
||||
/// Parse from JSON value.
|
||||
pub fn from_value(value: serde_json::Value) -> Result<Self, serde_json::Error> {
|
||||
serde_json::from_value(value)
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
|
||||
/// 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 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<T: serde::de::DeserializeOwned>(&self, key: &str) -> Option<T> {
|
||||
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_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<i32> = config.get_extra("custom_field");
|
||||
assert_eq!(custom, Some(42));
|
||||
|
||||
let nested: Option<HashMap<String, String>> = config.get_extra("nested");
|
||||
assert_eq!(
|
||||
nested.as_ref().unwrap().get("foo"),
|
||||
Some(&"bar".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,688 +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, 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
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,869 +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::Array3;
|
||||
|
||||
use crate::multimodal::vision::{
|
||||
image_processor::{ImagePreProcessor, PreprocessedImages},
|
||||
preprocessor_config::PreProcessorConfig,
|
||||
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<Self, Self::Err> {
|
||||
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,
|
||||
) -> Result<Array3<f32>, TransformError> {
|
||||
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 mean_color = mean_to_rgb(&mean);
|
||||
let squared = 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
|
||||
}
|
||||
}
|
||||
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
|
||||
if config.do_center_crop.unwrap_or(true) {
|
||||
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);
|
||||
}
|
||||
|
||||
Ok(tensor)
|
||||
}
|
||||
}
|
||||
|
||||
impl ImagePreProcessor 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<PreprocessedImages, TransformError> {
|
||||
if images.is_empty() {
|
||||
return Err(TransformError::EmptyBatch);
|
||||
}
|
||||
|
||||
// Store original sizes
|
||||
let image_sizes: Vec<(u32, u32)> = images.iter().map(|img| img.dimensions()).collect();
|
||||
|
||||
// Process each image
|
||||
let tensors: Vec<Array3<f32>> = images
|
||||
.iter()
|
||||
.map(|img| self.process_one_image(img, config))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
|
||||
// Stack into batch
|
||||
let pixel_values = stack_batch(&tensors)?;
|
||||
|
||||
// Calculate token counts
|
||||
let num_img_tokens: Vec<usize> = images
|
||||
.iter()
|
||||
.map(|_| self.calculate_num_tokens(self.image_size, self.image_size, config))
|
||||
.collect();
|
||||
|
||||
Ok(PreprocessedImages::new(
|
||||
pixel_values,
|
||||
num_img_tokens,
|
||||
image_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()?;
|
||||
let w = pair.first()?.as_u64()? as u32;
|
||||
let h = 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 patches) for anyres processing.
|
||||
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.patch_size, height / self.base.patch_size)
|
||||
}
|
||||
|
||||
/// Calculate unpad dimensions based on original aspect ratio.
|
||||
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.
|
||||
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.
|
||||
fn divide_to_samples(&self, image: &DynamicImage, crop_size: (u32, u32)) -> Vec<DynamicImage> {
|
||||
divide_to_samples(image, crop_size)
|
||||
}
|
||||
|
||||
/// Process a single patch/crop.
|
||||
fn process_patch(
|
||||
&self,
|
||||
image: &DynamicImage,
|
||||
config: &PreProcessorConfig,
|
||||
) -> Result<Array3<f32>, TransformError> {
|
||||
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
|
||||
let cropped = if config.do_center_crop.unwrap_or(true) {
|
||||
if let Some((crop_h, crop_w)) = config.get_crop_size() {
|
||||
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);
|
||||
}
|
||||
|
||||
Ok(tensor)
|
||||
}
|
||||
}
|
||||
|
||||
impl ImagePreProcessor 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<PreprocessedImages, TransformError> {
|
||||
if images.is_empty() {
|
||||
return Err(TransformError::EmptyBatch);
|
||||
}
|
||||
|
||||
let mut all_patches = Vec::new();
|
||||
let mut num_img_tokens = Vec::with_capacity(images.len());
|
||||
let mut image_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();
|
||||
image_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));
|
||||
|
||||
for sample in samples {
|
||||
all_patches.push(self.process_patch(&sample, config)?);
|
||||
}
|
||||
|
||||
num_img_tokens.push(self.calculate_num_tokens(
|
||||
original_size.0,
|
||||
original_size.1,
|
||||
config,
|
||||
));
|
||||
}
|
||||
|
||||
let pixel_values = stack_batch(&all_patches)?;
|
||||
|
||||
Ok(PreprocessedImages::new(
|
||||
pixel_values,
|
||||
num_img_tokens,
|
||||
image_sizes,
|
||||
))
|
||||
}
|
||||
|
||||
fn calculate_num_tokens(&self, width: u32, height: u32, _config: &PreProcessorConfig) -> usize {
|
||||
let original_size = (width, height);
|
||||
|
||||
// Base tokens (from original resized image)
|
||||
let patches_per_side = self.base.image_size / self.base.patch_size;
|
||||
let base_tokens = (patches_per_side * patches_per_side) as usize;
|
||||
|
||||
// Grid tokens (from crops)
|
||||
let grid_shape = self.get_anyres_grid_shape(original_size);
|
||||
let unpad_shape = self.calculate_unpad(grid_shape, original_size);
|
||||
|
||||
// Total: base + unpadded area
|
||||
base_tokens + (unpad_shape.0 as usize + 1) * unpad_shape.1 as usize
|
||||
}
|
||||
|
||||
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(width * height, downscaled_width * downscaled_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 = image.resize_exact(
|
||||
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<DynamicImage> {
|
||||
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<u8>) -> 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(), 336);
|
||||
assert_eq!(result.width(), 336);
|
||||
assert_eq!(result.num_img_tokens[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(), 336);
|
||||
assert_eq!(result.width(), 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(), 336);
|
||||
assert_eq!(result.width(), 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();
|
||||
|
||||
// Should have multiple patches (original + crops)
|
||||
assert!(result.batch_size() > 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,33 +0,0 @@
|
||||
//! Model-specific image processors.
|
||||
//!
|
||||
//! This module contains implementations of `ImagePreProcessor` 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
|
||||
//! - **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 llama4_vision;
|
||||
pub mod llava;
|
||||
pub mod phi3_vision;
|
||||
pub mod phi4_vision;
|
||||
pub mod pixtral;
|
||||
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;
|
||||
pub use pixtral::PixtralProcessor;
|
||||
pub use qwen2_vl::Qwen2VLProcessor;
|
||||
pub use qwen3_vl::Qwen3VLProcessor;
|
||||
@@ -1,551 +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::multimodal::vision::{
|
||||
image_processor::{ImagePreProcessor, PreprocessedImages},
|
||||
preprocessor_config::PreProcessorConfig,
|
||||
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 = img.resize_exact(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).
|
||||
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)`.
|
||||
fn create_global_image(&self, tensor: &Array3<f32>) -> Array3<f32> {
|
||||
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.
|
||||
fn reshape_to_tiles(&self, tensor: &Array3<f32>) -> Vec<Array3<f32>> {
|
||||
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`
|
||||
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.
|
||||
#[allow(clippy::type_complexity)]
|
||||
fn process_single_image(
|
||||
&self,
|
||||
image: &DynamicImage,
|
||||
config: &PreProcessorConfig,
|
||||
) -> Result<(Array4<f32>, (usize, usize), usize), TransformError> {
|
||||
// 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::<f32>::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);
|
||||
|
||||
// image_sizes is the HD-transformed size
|
||||
Ok((output, (hd_h as usize, hd_w as usize), num_tokens))
|
||||
}
|
||||
}
|
||||
|
||||
impl ImagePreProcessor 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<PreprocessedImages, TransformError> {
|
||||
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 (pixel_values, image_size, num_tokens) =
|
||||
self.process_single_image(image, config)?;
|
||||
all_pixel_values.push(pixel_values);
|
||||
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::<f32>::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 image_sizes as model-specific 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<u32> = all_image_sizes
|
||||
.iter()
|
||||
.flat_map(|(w, h)| [*h, *w]) // [h, w] for each image
|
||||
.collect();
|
||||
model_specific.insert(
|
||||
"image_sizes".to_string(),
|
||||
crate::multimodal::vision::image_processor::ModelSpecificValue::UintTensor {
|
||||
data: image_sizes_data,
|
||||
shape: vec![batch_size, 2],
|
||||
},
|
||||
);
|
||||
|
||||
// num_img_tokens as list
|
||||
model_specific.insert(
|
||||
"num_img_tokens".to_string(),
|
||||
crate::multimodal::vision::image_processor::ModelSpecificValue::UintVec(
|
||||
all_num_tokens.iter().map(|&t| t as u32).collect(),
|
||||
),
|
||||
);
|
||||
|
||||
// Convert 5D tensor to appropriate format
|
||||
// Phi3-Vision expects [B, num_crops+1, C, H, W]
|
||||
let pixel_values = ndarray::ArrayD::<f32>::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(PreprocessedImages {
|
||||
pixel_values,
|
||||
num_img_tokens: all_num_tokens,
|
||||
image_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<u8>) -> 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.pixel_values.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.image_sizes.len(), 2);
|
||||
assert_eq!(result.num_img_tokens.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)
|
||||
}
|
||||
}
|
||||
@@ -1,724 +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::multimodal::vision::{
|
||||
image_processor::{ImagePreProcessor, ModelSpecificValue, PreprocessedImages},
|
||||
preprocessor_config::PreProcessorConfig,
|
||||
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 (pixel_values, attention_mask, (height, width), num_tokens).
|
||||
type SingleImageResult = (Array4<f32>, Array3<u32>, (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.
|
||||
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,
|
||||
) -> Result<(DynamicImage, Array2<u32>, usize, usize), TransformError> {
|
||||
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::<u32>::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 = image.resize_exact(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);
|
||||
|
||||
Ok((padded, attention_mask, target_h_crops, target_w_crops))
|
||||
}
|
||||
|
||||
/// Pad image to target dimensions with white 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 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<f32>) -> Array3<f32> {
|
||||
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<f32>, h_crops: usize, w_crops: usize) -> Array4<f32> {
|
||||
let base = self.base_resolution as usize;
|
||||
let num_tiles = h_crops * w_crops;
|
||||
|
||||
let mut tiles = Array4::<f32>::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<u32>, h_crops: usize, w_crops: usize) -> Array2<u32> {
|
||||
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::<u32>::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
|
||||
fn calculate_num_tokens(&self, downsampled_mask: &Array2<u32>) -> 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,
|
||||
) -> Result<SingleImageResult, TransformError> {
|
||||
// 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::<f32>::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::<u32>::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);
|
||||
|
||||
Ok((output, combined_mask, (hd_h, hd_w), num_tokens))
|
||||
}
|
||||
}
|
||||
|
||||
impl ImagePreProcessor 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<PreprocessedImages, TransformError> {
|
||||
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 image_sizes = Vec::new();
|
||||
let mut num_img_tokens = Vec::new();
|
||||
|
||||
for image in images {
|
||||
let (output, mask, size, tokens) = processor.process_single_image(image)?;
|
||||
all_outputs.push(output);
|
||||
all_masks.push(mask);
|
||||
image_sizes.push(size);
|
||||
num_img_tokens.push(tokens);
|
||||
}
|
||||
|
||||
// Find max crops across batch for padding
|
||||
let max_crops = all_outputs.iter().map(|o| o.shape()[0]).max().unwrap();
|
||||
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 pixel_values =
|
||||
ndarray::ArrayD::<f32>::zeros(IxDyn(&[batch_size, max_crops, 3, base, base]));
|
||||
let mut attention_masks =
|
||||
ndarray::ArrayD::<u32>::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 {
|
||||
pixel_values[[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<u32> = attention_masks.iter().copied().collect();
|
||||
model_specific.insert(
|
||||
"pixel_attention_mask".to_string(),
|
||||
ModelSpecificValue::UintTensor {
|
||||
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<u32> = image_sizes.iter().flat_map(|&(h, w)| vec![h, w]).collect();
|
||||
model_specific.insert(
|
||||
"image_sizes".to_string(),
|
||||
ModelSpecificValue::UintTensor {
|
||||
data: sizes_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 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<u8>) -> 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.num_img_tokens[0] > 256); // At least global tokens
|
||||
|
||||
// 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 = 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.image_sizes[0].1 >= result.image_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.image_sizes.len(), 2);
|
||||
assert_eq!(result.num_img_tokens.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);
|
||||
}
|
||||
}
|
||||
@@ -1,411 +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::multimodal::vision::{
|
||||
image_processor::{ImagePreProcessor, ModelSpecificValue, PreprocessedImages},
|
||||
preprocessor_config::PreProcessorConfig,
|
||||
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<f32>, (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 = image.resize_exact(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 ImagePreProcessor 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<PreprocessedImages, TransformError> {
|
||||
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 num_img_tokens = 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()));
|
||||
num_img_tokens.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::<f32>::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<i64> = 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(PreprocessedImages {
|
||||
pixel_values: batch_tensor,
|
||||
num_img_tokens,
|
||||
image_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::*;
|
||||
|
||||
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.pixel_values.shape()[0], 2); // batch size
|
||||
assert_eq!(result.pixel_values.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(crate::multimodal::vision::preprocessor_config::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);
|
||||
}
|
||||
}
|
||||
@@ -1,458 +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 ndarray::Array3;
|
||||
|
||||
use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase};
|
||||
use crate::multimodal::vision::{
|
||||
image_processor::{ImagePreProcessor, PreprocessedImages},
|
||||
preprocessor_config::PreProcessorConfig,
|
||||
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,
|
||||
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,
|
||||
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),
|
||||
temporal_patch_size: config
|
||||
.temporal_patch_size
|
||||
.unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE),
|
||||
mean: CLIP_MEAN,
|
||||
std: CLIP_STD,
|
||||
model_name: "qwen2-vl",
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
|
||||
/// Reshape pixel values from [C, H, W] to flattened patches format.
|
||||
pub fn reshape_to_patches(
|
||||
&self,
|
||||
tensor: &Array3<f32>,
|
||||
grid_t: usize,
|
||||
grid_h: usize,
|
||||
grid_w: usize,
|
||||
) -> Vec<f32> {
|
||||
self.inner
|
||||
.reshape_to_patches(tensor, grid_t, grid_h, grid_w)
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for Qwen2VLProcessor {
|
||||
type Target = QwenVLProcessorBase;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl ImagePreProcessor 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<PreprocessedImages, TransformError> {
|
||||
self.inner.preprocess(images, config)
|
||||
}
|
||||
|
||||
fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize {
|
||||
self.inner.calculate_num_tokens(width, height, config)
|
||||
}
|
||||
|
||||
fn model_name(&self) -> &'static str {
|
||||
self.inner.model_name()
|
||||
}
|
||||
|
||||
fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> {
|
||||
self.inner.get_processed_size(config)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use image::{Rgb, RgbImage};
|
||||
|
||||
use super::*;
|
||||
use crate::multimodal::vision::{
|
||||
image_processor::ModelSpecificValue, preprocessor_config::PatchSize,
|
||||
};
|
||||
|
||||
fn create_test_image(width: u32, height: u32, color: Rgb<u8>) -> 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_too_small_dimension_error() {
|
||||
let processor = Qwen2VLProcessor::new();
|
||||
|
||||
// Dimension smaller than factor
|
||||
let result = processor.smart_resize(10, 100);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[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();
|
||||
|
||||
assert_eq!(result.batch_size(), 1);
|
||||
|
||||
// Check pixel values are normalized
|
||||
let flat = result.pixel_values_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 is present
|
||||
assert!(result.model_specific.contains_key("image_grid_thw"));
|
||||
|
||||
// Verify token count is reasonable
|
||||
assert!(result.num_img_tokens[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.image_sizes.len(), 2);
|
||||
assert_eq!(result.num_img_tokens.len(), 2);
|
||||
|
||||
// Check grid_thw shape
|
||||
if let Some(ModelSpecificValue::UintTensor { 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 UintTensor");
|
||||
}
|
||||
}
|
||||
|
||||
#[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_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);
|
||||
}
|
||||
}
|
||||
@@ -1,470 +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 ndarray::Array3;
|
||||
|
||||
use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase};
|
||||
use crate::multimodal::vision::{
|
||||
image_processor::{ImagePreProcessor, PreprocessedImages},
|
||||
preprocessor_config::PreProcessorConfig,
|
||||
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 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 {
|
||||
inner: QwenVLProcessorBase::new(QwenVLConfig {
|
||||
patch_size: DEFAULT_PATCH_SIZE,
|
||||
merge_size: DEFAULT_MERGE_SIZE,
|
||||
min_pixels: DEFAULT_MIN_PIXELS,
|
||||
max_pixels: DEFAULT_MAX_PIXELS,
|
||||
temporal_patch_size: DEFAULT_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 {
|
||||
inner: QwenVLProcessorBase::new(QwenVLConfig {
|
||||
patch_size,
|
||||
merge_size,
|
||||
min_pixels,
|
||||
max_pixels,
|
||||
temporal_patch_size,
|
||||
mean: QWEN3_MEAN,
|
||||
std: QWEN3_STD,
|
||||
model_name: "qwen3-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),
|
||||
temporal_patch_size: config
|
||||
.temporal_patch_size
|
||||
.unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE),
|
||||
mean: QWEN3_MEAN,
|
||||
std: QWEN3_STD,
|
||||
model_name: "qwen3-vl",
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
|
||||
/// Reshape pixel values from [C, H, W] to flattened patches format.
|
||||
pub fn reshape_to_patches(
|
||||
&self,
|
||||
tensor: &Array3<f32>,
|
||||
grid_t: usize,
|
||||
grid_h: usize,
|
||||
grid_w: usize,
|
||||
) -> Vec<f32> {
|
||||
self.inner
|
||||
.reshape_to_patches(tensor, grid_t, grid_h, grid_w)
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for Qwen3VLProcessor {
|
||||
type Target = QwenVLProcessorBase;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl ImagePreProcessor 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<PreprocessedImages, TransformError> {
|
||||
self.inner.preprocess(images, config)
|
||||
}
|
||||
|
||||
fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize {
|
||||
self.inner.calculate_num_tokens(width, height, config)
|
||||
}
|
||||
|
||||
fn model_name(&self) -> &'static str {
|
||||
self.inner.model_name()
|
||||
}
|
||||
|
||||
fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> {
|
||||
self.inner.get_processed_size(config)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use image::{Rgb, RgbImage};
|
||||
|
||||
use super::*;
|
||||
use crate::multimodal::vision::{
|
||||
image_processor::ModelSpecificValue, preprocessor_config::PatchSize,
|
||||
};
|
||||
|
||||
fn create_test_image(width: u32, height: u32, color: Rgb<u8>) -> 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.get_factor(), 32); // 16 * 2
|
||||
}
|
||||
|
||||
#[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_too_small_dimension_error() {
|
||||
let processor = Qwen3VLProcessor::new();
|
||||
|
||||
// Dimension smaller than factor (32)
|
||||
let result = processor.smart_resize(10, 100);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[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();
|
||||
|
||||
assert_eq!(result.batch_size(), 1);
|
||||
|
||||
// Check pixel values are normalized
|
||||
let flat = result.pixel_values_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 is present
|
||||
assert!(result.model_specific.contains_key("image_grid_thw"));
|
||||
|
||||
// Verify token count is reasonable
|
||||
assert!(result.num_img_tokens[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.image_sizes.len(), 2);
|
||||
assert_eq!(result.num_img_tokens.len(), 2);
|
||||
|
||||
// Check grid_thw shape
|
||||
if let Some(ModelSpecificValue::UintTensor { 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 UintTensor");
|
||||
}
|
||||
}
|
||||
|
||||
#[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);
|
||||
}
|
||||
|
||||
#[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");
|
||||
}
|
||||
}
|
||||
@@ -1,500 +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 image::{DynamicImage, GenericImageView};
|
||||
use ndarray::Array3;
|
||||
|
||||
use crate::multimodal::vision::{
|
||||
image_processor::{ImagePreProcessor, ModelSpecificValue, PreprocessedImages},
|
||||
preprocessor_config::PreProcessorConfig,
|
||||
transforms::{normalize, pil_to_filter, resize, stack_batch, to_tensor, 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, unlike Rust's `f64::round()` which rounds half away from zero.
|
||||
///
|
||||
/// Examples:
|
||||
/// - round_half_to_even(12.5) = 12 (not 13)
|
||||
/// - round_half_to_even(13.5) = 14 (not 14)
|
||||
/// - round_half_to_even(12.4) = 12
|
||||
/// - round_half_to_even(12.6) = 13
|
||||
#[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,
|
||||
/// 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,
|
||||
}
|
||||
|
||||
/// 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
|
||||
}
|
||||
|
||||
/// Get the temporal patch size.
|
||||
pub fn temporal_patch_size(&self) -> usize {
|
||||
self.config.temporal_patch_size
|
||||
}
|
||||
|
||||
/// 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 smaller than the factor
|
||||
/// - 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 minimum dimensions
|
||||
if height < factor || width < factor {
|
||||
return Err(TransformError::InvalidShape {
|
||||
expected: format!("dimensions >= {} (patch_size * merge_size)", factor),
|
||||
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))
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
|
||||
/// Reshape pixel values from [C, H, W] to flattened patches format.
|
||||
///
|
||||
/// This matches the HuggingFace Qwen2VLImageProcessor output format:
|
||||
/// `(num_patches, patch_features)` where:
|
||||
/// - num_patches = grid_t * grid_h * grid_w
|
||||
/// - patch_features = C * temporal_patch_size * patch_size * patch_size
|
||||
///
|
||||
/// The transformation follows these steps (matching HuggingFace exactly):
|
||||
/// 1. Start with [C, H, W] tensor, expand to [temporal, C, H, W]
|
||||
/// 2. Reshape to [grid_t, temporal, C, grid_h/merge, merge, patch, grid_w/merge, merge, patch]
|
||||
/// 3. Permute to [grid_t, grid_h/merge, grid_w/merge, merge, merge, C, temporal, patch, patch]
|
||||
/// 4. Flatten to [num_patches, patch_features]
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `tensor` - Input tensor of shape [C, H, W]
|
||||
/// * `grid_t` - Temporal grid size (1 for images)
|
||||
/// * `grid_h` - Height grid size (H / patch_size)
|
||||
/// * `grid_w` - Width grid size (W / patch_size)
|
||||
///
|
||||
/// # Returns
|
||||
/// Flattened patches as Vec<f32> with shape semantics (num_patches, patch_features)
|
||||
pub fn reshape_to_patches(
|
||||
&self,
|
||||
tensor: &Array3<f32>,
|
||||
grid_t: usize,
|
||||
grid_h: usize,
|
||||
grid_w: usize,
|
||||
) -> Vec<f32> {
|
||||
use ndarray::IxDyn;
|
||||
|
||||
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;
|
||||
|
||||
// Verify dimensions match expected grid
|
||||
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"
|
||||
);
|
||||
|
||||
// Step 1: Expand temporal dimension by replicating the frame
|
||||
// [C, H, W] -> [temporal_patch_size, C, H, W]
|
||||
let expanded = tensor
|
||||
.view()
|
||||
.insert_axis(ndarray::Axis(0))
|
||||
.broadcast((temporal_patch_size, channel, height, width))
|
||||
.expect("Broadcast failed")
|
||||
.to_owned();
|
||||
|
||||
// Step 2: Reshape to split spatial dimensions into grid and patch components
|
||||
// [temporal, C, H, W] -> [grid_t, temporal, C, grid_h/merge, merge, patch, grid_w/merge, merge, patch]
|
||||
let grid_h_merged = grid_h / merge_size;
|
||||
let grid_w_merged = grid_w / merge_size;
|
||||
|
||||
// Use IxDyn for 9-dimensional reshape (ndarray only supports up to Ix6 for fixed dims)
|
||||
let shape_9d = IxDyn(&[
|
||||
grid_t,
|
||||
temporal_patch_size,
|
||||
channel,
|
||||
grid_h_merged,
|
||||
merge_size,
|
||||
patch_size,
|
||||
grid_w_merged,
|
||||
merge_size,
|
||||
patch_size,
|
||||
]);
|
||||
|
||||
let reshaped = expanded
|
||||
.into_shape_with_order(shape_9d)
|
||||
.expect("Reshape failed");
|
||||
|
||||
// Step 3: Permute axes to match HuggingFace output order
|
||||
// From: [grid_t, temporal, C, grid_h/merge, merge, patch, grid_w/merge, merge, patch]
|
||||
// [ 0 , 1 , 2, 3 , 4 , 5 , 6 , 7 , 8 ]
|
||||
// To: [grid_t, grid_h/merge, grid_w/merge, merge, merge, C, temporal, patch, patch]
|
||||
// [ 0 , 3 , 6 , 4 , 7 , 2, 1 , 5 , 8 ]
|
||||
let permuted = reshaped.permuted_axes(&[0, 3, 6, 4, 7, 2, 1, 5, 8][..]);
|
||||
|
||||
// Step 4: Flatten to [num_patches, patch_features]
|
||||
let num_patches = grid_t * grid_h * grid_w;
|
||||
let patch_features = channel * temporal_patch_size * patch_size * patch_size;
|
||||
|
||||
// Make contiguous and flatten
|
||||
let contiguous = permuted.as_standard_layout().into_owned();
|
||||
let flat = contiguous
|
||||
.into_shape_with_order(IxDyn(&[num_patches, patch_features]))
|
||||
.expect("Final reshape failed");
|
||||
|
||||
let (vec, _offset) = flat.into_raw_vec_and_offset();
|
||||
vec
|
||||
}
|
||||
}
|
||||
|
||||
impl ImagePreProcessor 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<PreprocessedImages, TransformError> {
|
||||
if images.is_empty() {
|
||||
return Err(TransformError::EmptyBatch);
|
||||
}
|
||||
|
||||
// Store original sizes
|
||||
let image_sizes: Vec<(u32, u32)> = images.iter().map(|img| img.dimensions()).collect();
|
||||
|
||||
// First pass: calculate target dimensions for each image
|
||||
let mut target_sizes = Vec::with_capacity(images.len());
|
||||
for image in images {
|
||||
let (w, h) = image.dimensions();
|
||||
let (new_h, new_w) = self.smart_resize(h as usize, w as usize)?;
|
||||
target_sizes.push((new_h, new_w));
|
||||
}
|
||||
|
||||
// Find max height and width across all images
|
||||
let max_height = target_sizes.iter().map(|(h, _)| *h).max().unwrap_or(0);
|
||||
let max_width = target_sizes.iter().map(|(_, w)| *w).max().unwrap_or(0);
|
||||
|
||||
// Process each image with uniform max dimensions
|
||||
let mean = config.get_image_mean();
|
||||
let std = config.get_image_std();
|
||||
let filter = pil_to_filter(config.resampling);
|
||||
|
||||
let mut tensors = Vec::with_capacity(images.len());
|
||||
let mut grid_thw_data = Vec::with_capacity(images.len() * 3);
|
||||
let mut num_img_tokens = Vec::with_capacity(images.len());
|
||||
|
||||
for (i, image) in images.iter().enumerate() {
|
||||
let (target_h, target_w) = target_sizes[i];
|
||||
|
||||
// Resize to the target size for this image
|
||||
let resized = if config.do_resize.unwrap_or(true) {
|
||||
// For batching: resize to max dimensions to enable stacking
|
||||
resize(image, max_width as u32, max_height as u32, filter)
|
||||
} else {
|
||||
image.clone()
|
||||
};
|
||||
|
||||
// Convert to tensor
|
||||
let mut tensor = to_tensor(&resized);
|
||||
|
||||
// Normalize
|
||||
if config.do_normalize.unwrap_or(true) {
|
||||
normalize(&mut tensor, &mean, &std);
|
||||
}
|
||||
|
||||
tensors.push(tensor);
|
||||
|
||||
// Grid dimensions are based on the individual image's target size
|
||||
let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(target_h, target_w, 1);
|
||||
grid_thw_data.push(grid_t as u32);
|
||||
grid_thw_data.push(grid_h as u32);
|
||||
grid_thw_data.push(grid_w as u32);
|
||||
|
||||
// Token count is based on individual grid
|
||||
let tokens = self.calculate_tokens_from_grid(grid_t, grid_h, grid_w);
|
||||
num_img_tokens.push(tokens);
|
||||
}
|
||||
|
||||
// Stack tensors into batch (now all same size)
|
||||
let pixel_values = stack_batch(&tensors)?;
|
||||
|
||||
// Create result with model-specific image_grid_thw
|
||||
let result = PreprocessedImages::new(pixel_values, num_img_tokens, image_sizes).with_extra(
|
||||
"image_grid_thw",
|
||||
ModelSpecificValue::uint_2d(grid_thw_data, images.len(), 3),
|
||||
);
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
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 super::*;
|
||||
|
||||
fn create_test_config() -> QwenVLConfig {
|
||||
QwenVLConfig {
|
||||
patch_size: 14,
|
||||
merge_size: 2,
|
||||
min_pixels: 256 * 28 * 28,
|
||||
max_pixels: 1280 * 28 * 28,
|
||||
temporal_patch_size: 2,
|
||||
mean: [0.5, 0.5, 0.5],
|
||||
std: [0.5, 0.5, 0.5],
|
||||
model_name: "test-qwen-vl",
|
||||
}
|
||||
}
|
||||
|
||||
#[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_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);
|
||||
}
|
||||
}
|
||||
@@ -1,499 +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 image::{imageops::FilterType, DynamicImage, GenericImageView, Rgb, RgbImage};
|
||||
use ndarray::{s, Array3, Array4};
|
||||
use thiserror::Error;
|
||||
|
||||
/// Errors that can occur during image transformations.
|
||||
#[derive(Error, Debug)]
|
||||
pub enum TransformError {
|
||||
#[error("Invalid tensor shape: expected {expected}, got {actual:?}")]
|
||||
InvalidShape {
|
||||
expected: String,
|
||||
actual: Vec<usize>,
|
||||
},
|
||||
|
||||
#[error("Image operation failed: {0}")]
|
||||
ImageError(#[from] image::ImageError),
|
||||
|
||||
#[error("Empty batch: cannot stack zero tensors")]
|
||||
EmptyBatch,
|
||||
|
||||
#[error("Inconsistent tensor shapes in batch")]
|
||||
InconsistentShapes,
|
||||
|
||||
#[error("Shape error: {0}")]
|
||||
ShapeError(String),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, TransformError>;
|
||||
|
||||
/// 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<f32> {
|
||||
let rgb = image.to_rgb8();
|
||||
let (w, h) = (rgb.width() as usize, rgb.height() as usize);
|
||||
let mut arr = Array3::<f32>::zeros((3, h, w));
|
||||
|
||||
for (x, y, pixel) in rgb.enumerate_pixels() {
|
||||
let (x, y) = (x as usize, y as usize);
|
||||
arr[[0, y, x]] = pixel[0] as f32 / 255.0;
|
||||
arr[[1, y, x]] = pixel[1] as f32 / 255.0;
|
||||
arr[[2, y, x]] = pixel[2] as f32 / 255.0;
|
||||
}
|
||||
arr
|
||||
}
|
||||
|
||||
/// Convert image to tensor [C, H, W] without normalization (keeps [0, 255]).
|
||||
///
|
||||
/// Some models expect unnormalized pixel values.
|
||||
pub fn to_tensor_no_norm(image: &DynamicImage) -> Array3<f32> {
|
||||
let rgb = image.to_rgb8();
|
||||
let (w, h) = (rgb.width() as usize, rgb.height() as usize);
|
||||
let mut arr = Array3::<f32>::zeros((3, h, w));
|
||||
|
||||
for (x, y, pixel) in rgb.enumerate_pixels() {
|
||||
let (x, y) = (x as usize, y as usize);
|
||||
arr[[0, y, x]] = pixel[0] as f32;
|
||||
arr[[1, y, x]] = pixel[1] as f32;
|
||||
arr[[2, y, x]] = pixel[2] as f32;
|
||||
}
|
||||
arr
|
||||
}
|
||||
|
||||
/// 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<f32>, mean: &[f64; 3], std: &[f64; 3]) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
/// Rescale tensor by a constant factor.
|
||||
///
|
||||
/// Used when `do_rescale=True` in HuggingFace configs (typically 1/255).
|
||||
pub fn rescale(tensor: &mut Array3<f32>, factor: f64) {
|
||||
let factor = factor as f32;
|
||||
tensor.mapv_inplace(|v| v * factor);
|
||||
}
|
||||
|
||||
/// Resize image to exact dimensions.
|
||||
///
|
||||
/// # 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 {
|
||||
image.resize_exact(width, height, filter)
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
image.resize(max_width, max_height, 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<u8>) -> 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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Pad image to specified dimensions with background color.
|
||||
///
|
||||
/// Image is placed at top-left corner.
|
||||
pub fn pad_to_size(
|
||||
image: &DynamicImage,
|
||||
target_w: u32,
|
||||
target_h: u32,
|
||||
background: Rgb<u8>,
|
||||
) -> DynamicImage {
|
||||
let (w, h) = image.dimensions();
|
||||
if w >= target_w && h >= target_h {
|
||||
return image.clone();
|
||||
}
|
||||
let new_w = w.max(target_w);
|
||||
let new_h = h.max(target_h);
|
||||
let mut new_image = DynamicImage::from(RgbImage::from_pixel(new_w, new_h, background));
|
||||
image::imageops::overlay(&mut new_image, image, 0, 0);
|
||||
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<f32>]) -> Result<Array4<f32>> {
|
||||
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::<f32>::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<usize>) -> 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<u8> {
|
||||
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<u8> {
|
||||
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<f32>,
|
||||
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<f32>, target_h: usize, target_w: usize) -> Array3<f32> {
|
||||
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::<f32>::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<u8>) -> DynamicImage {
|
||||
DynamicImage::from(RgbImage::from_pixel(width, height, color))
|
||||
}
|
||||
|
||||
#[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::<f32>::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.iter() {
|
||||
assert!(val.abs() < 1e-6);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rescale() {
|
||||
let mut tensor = Array3::<f32>::from_elem((3, 2, 2), 255.0);
|
||||
rescale(&mut tensor, 1.0 / 255.0);
|
||||
|
||||
for val in tensor.iter() {
|
||||
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_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::<f32>::zeros((3, 10, 10));
|
||||
let t2 = Array3::<f32>::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);
|
||||
}
|
||||
}
|
||||
@@ -1,4 +0,0 @@
|
||||
//! Multimodal and vision integration tests
|
||||
|
||||
mod multimodal_tracker_test;
|
||||
mod vision_golden_tests;
|
||||
@@ -1,151 +0,0 @@
|
||||
use std::{collections::HashMap, path::PathBuf, sync::Arc, time::Duration};
|
||||
|
||||
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
|
||||
use reqwest::Client;
|
||||
use smg::multimodal::{
|
||||
AsyncMultiModalTracker, ChatContentPart, ConversationSegment, ImageFetchConfig, ImageSource,
|
||||
MediaConnector, MediaConnectorConfig, MediaSource, Modality, TrackerConfig,
|
||||
};
|
||||
use tempfile::tempdir;
|
||||
|
||||
const TINY_PNG_BASE64: &str =
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNgYAAAAAMAASsJTYQAAAAASUVORK5CYII=";
|
||||
|
||||
fn tiny_png_bytes() -> Vec<u8> {
|
||||
BASE64_STANDARD
|
||||
.decode(TINY_PNG_BASE64)
|
||||
.expect("decode tiny png fixture")
|
||||
}
|
||||
|
||||
fn test_connector(allowed_path: Option<PathBuf>) -> 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 tracker_collects_conversation_and_placeholders() {
|
||||
let connector = Arc::new(test_connector(None));
|
||||
let mut tracker = AsyncMultiModalTracker::new(
|
||||
connector,
|
||||
TrackerConfig {
|
||||
placeholder_tokens: Default::default(),
|
||||
modality_limits: HashMap::from([(Modality::Image, 2)]),
|
||||
},
|
||||
);
|
||||
|
||||
tracker
|
||||
.push_part(ChatContentPart::Text {
|
||||
text: "before".into(),
|
||||
})
|
||||
.expect("text part");
|
||||
tracker
|
||||
.push_part(ChatContentPart::ImageData {
|
||||
data: tiny_png_bytes(),
|
||||
mime_type: Some("image/png".into()),
|
||||
uuid: Some("img-1".into()),
|
||||
detail: None,
|
||||
})
|
||||
.expect("image part");
|
||||
tracker
|
||||
.push_part(ChatContentPart::Text {
|
||||
text: "after".into(),
|
||||
})
|
||||
.expect("text part");
|
||||
|
||||
let output = tracker.finalize().await.expect("tracker finalize");
|
||||
assert_eq!(output.conversation.len(), 3);
|
||||
assert!(matches!(
|
||||
&output.conversation[0],
|
||||
ConversationSegment::Text(text) if text == "before"
|
||||
));
|
||||
assert!(matches!(
|
||||
&output.conversation[1],
|
||||
ConversationSegment::Placeholder { token } if token == "<image>"
|
||||
));
|
||||
assert!(matches!(
|
||||
&output.conversation[2],
|
||||
ConversationSegment::Text(text) if text == "after"
|
||||
));
|
||||
|
||||
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())]);
|
||||
|
||||
let placeholders = output
|
||||
.placeholders
|
||||
.get("<image>")
|
||||
.expect("placeholder entry");
|
||||
assert_eq!(placeholders.len(), 1);
|
||||
assert_eq!(placeholders[0].text_position, 1);
|
||||
assert_eq!(placeholders[0].item_index, 0);
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,6 +0,0 @@
|
||||
//! Multimodal integration tests
|
||||
|
||||
#[path = "common/mod.rs"]
|
||||
pub mod common;
|
||||
|
||||
mod multimodal;
|
||||
Reference in New Issue
Block a user