diff --git a/crates/grpc_client/proto/tokenspeed_encoder.proto b/crates/grpc_client/proto/tokenspeed_encoder.proto index 81cb6337c..ea828a382 100644 --- a/crates/grpc_client/proto/tokenspeed_encoder.proto +++ b/crates/grpc_client/proto/tokenspeed_encoder.proto @@ -5,8 +5,8 @@ package tokenspeed.grpc.encoder; import "tokenspeed_scheduler.proto"; // TokenSpeed EPD encode-stage gRPC service. The gateway ships preprocessed -// multimodal tensors to a vision-tower-only encode worker; the worker runs the -// tower and pushes the resulting image embeddings to a prefill worker over +// multimodal tensors to an encoder-only worker; the worker runs the matching +// tower and pushes the resulting embeddings to a prefill worker over // Mooncake, keyed per item by bootstrap_room. The gateway never sees the // embeddings — it only triggers the encode and assigns each item its room. // @@ -24,7 +24,7 @@ message EncodeRequest { string request_id = 1; // This worker's assigned multimodal item — the same tensor shape the scheduler - // Generate path carries. The encode worker runs the vision tower on it; it + // Generate path carries. The encode worker runs the matching tower on it; it // does not fetch or preprocess (that already happened in the gateway). tokenspeed.grpc.scheduler.MultimodalInputs mm_inputs = 2; diff --git a/crates/grpc_client/proto/tokenspeed_scheduler.proto b/crates/grpc_client/proto/tokenspeed_scheduler.proto index f638a0db2..f37d9204b 100644 --- a/crates/grpc_client/proto/tokenspeed_scheduler.proto +++ b/crates/grpc_client/proto/tokenspeed_scheduler.proto @@ -9,9 +9,9 @@ import "common.proto"; // TokenSpeed scheduler gRPC service. Self-contained wire definition apart // from the cross-engine admin messages in smg.grpc.common (flush/profile). // Trimmed to text+multimodal generation (no embed, no LoRA, no hidden-state -// forwarding). Multimodal carries preprocessed tensors only — image fetch + +// forwarding). Multimodal carries preprocessed tensors only — media fetch + // per-model preprocess happen in the gateway (see crates/multimodal). EPD: a -// prefill request may carry `EncodeBootstrapInfo` so its image embeddings +// prefill request may carry `EncodeBootstrapInfo` so its multimodal embeddings // arrive from an encode worker over Mooncake (see tokenspeed_encoder.proto). service TokenSpeedScheduler { rpc Generate(GenerateRequest) returns (stream GenerateResponse); @@ -108,9 +108,9 @@ message GenerateRequest { // Preprocessed multimodal payload. Absent for text-only requests. MultimodalInputs mm_inputs = 9; - // EPD: present when this request's image embeddings arrive from an encode + // EPD: present when this request's multimodal embeddings arrive from an encode // worker over Mooncake instead of being computed here. The worker waits for - // the embedding keyed by `bootstrap_room` rather than running the vision + // the embedding keyed by `bootstrap_room` rather than running the matching // tower. Absent for non-disaggregated (aggregated) requests. optional EncodeBootstrapInfo encode_bootstrap_info = 10; @@ -142,11 +142,11 @@ message EncodeItemBootstrapInfo { // The encode worker's bootstrap server (data-source side), not the prefill's. string bootstrap_host = 2; int32 bootstrap_port = 3; - // 63-bit rendezvous id: minted random per image by the gateway with no + // 63-bit rendezvous id: minted random per item by the gateway with no // in-flight dedup, so the space must be wide enough that the birthday // collision rate is negligible even with many independent gateways and // thousands of concurrent rooms (int32's 2^31 collides every few days under - // load, silently cross-wiring one image's embedding onto another's room). + // load, silently cross-wiring one item's embedding onto another's room). int64 bootstrap_room = 4; } diff --git a/crates/multimodal/Cargo.toml b/crates/multimodal/Cargo.toml index a06e457a4..aaed07c26 100644 --- a/crates/multimodal/Cargo.toml +++ b/crates/multimodal/Cargo.toml @@ -29,6 +29,8 @@ libloading = "0.8" ndarray = "0.17" once_cell = "1.21.4" rayon = "1.12" +rustfft = "6.4" +symphonia = { version = "0.6", default-features = false, features = ["all"] } opencv = { version = "0.99.0", default-features = false, features = ["clang-runtime", "imgproc", "videoio"], optional = true } reqwest = { workspace = true, features = ["stream"] } serde = { workspace = true, features = ["derive"] } diff --git a/crates/multimodal/src/audio/decode.rs b/crates/multimodal/src/audio/decode.rs new file mode 100644 index 000000000..f7118cd61 --- /dev/null +++ b/crates/multimodal/src/audio/decode.rs @@ -0,0 +1,533 @@ +//! Audio decode helpers shared by model-specific audio preprocessors. +//! +//! The default path mirrors SMG video decode: use an in-process decoder first, +//! then fall back to an external FFmpeg binary for difficult containers/codecs. + +use std::{ + io::{Cursor, Write}, + mem::size_of, + process::{Output, Stdio}, + sync::OnceLock, + time::{Duration, Instant}, +}; + +use symphonia::{ + core::{ + codecs::audio::AudioDecoderOptions, + errors::Error as SymphoniaError, + formats::{probe::Hint, FormatOptions, TrackType}, + io::MediaSourceStream, + meta::MetadataOptions, + }, + default::{get_codecs, get_probe}, +}; +use tokio::{process::Command, task, time}; +use tracing::debug; + +use crate::error::TransformError; + +const DEFAULT_AUDIO_PROCESS_TIMEOUT: Duration = Duration::from_secs(30); +const DEFAULT_AUDIO_MAX_DECODED_BYTES: usize = 256 * 1024 * 1024; + +static AUDIO_PROCESS_TIMEOUT: OnceLock = OnceLock::new(); +static AUDIO_MAX_DECODED_BYTES: OnceLock = OnceLock::new(); + +#[derive(Debug, Clone, PartialEq)] +pub struct DecodedAudio { + pub samples: Vec, + pub sample_rate: usize, +} + +pub async fn decode_audio_mono_f32(bytes: &[u8]) -> Result { + match audio_decode_backend_override() { + Some("symphonia") => decode_audio_with_symphonia_blocking(bytes).await, + Some("ffmpeg") => decode_audio_with_ffmpeg(bytes).await, + Some(backend) => Err(TransformError::ShapeError(format!( + "unsupported SMG_AUDIO_DECODE_BACKEND={backend}; expected auto, symphonia, or ffmpeg" + ))), + None => match decode_audio_with_symphonia_blocking(bytes).await { + Ok(decoded) => Ok(decoded), + Err(symphonia_error) => { + debug!( + error = %symphonia_error, + "smg_mm_timing audio_decode_auto_symphonia_fallback" + ); + decode_audio_with_ffmpeg(bytes).await.map_err(|ffmpeg_error| { + TransformError::ShapeError(format!( + "Symphonia audio decode failed: {symphonia_error}; ffmpeg fallback failed: {ffmpeg_error}" + )) + }) + } + }, + } +} + +async fn decode_audio_with_symphonia_blocking( + bytes: &[u8], +) -> Result { + let bytes = bytes.to_vec(); + task::spawn_blocking(move || decode_audio_mono_f32_symphonia(&bytes)) + .await + .map_err(|e| TransformError::ShapeError(format!("Symphonia decode task failed: {e}")))? +} + +pub(crate) fn decode_audio_mono_f32_symphonia( + bytes: &[u8], +) -> Result { + decode_audio_mono_f32_symphonia_with_limits( + bytes, + audio_max_decoded_bytes(), + audio_process_timeout(), + ) +} + +fn decode_audio_mono_f32_symphonia_with_limits( + bytes: &[u8], + max_decoded_bytes: usize, + timeout: Duration, +) -> Result { + let started = Instant::now(); + let mut hint = Hint::new(); + if let Some(ext) = audio_extension_hint(bytes) { + hint.with_extension(ext); + } + + let cursor = Cursor::new(bytes.to_vec()); + let media_source = MediaSourceStream::new(Box::new(cursor), Default::default()); + let mut format = get_probe() + .probe( + &hint, + media_source, + FormatOptions::default(), + MetadataOptions::default(), + ) + .map_err(|e| TransformError::ShapeError(format!("Symphonia probe failed: {e}")))?; + + let track = format.default_track(TrackType::Audio).ok_or_else(|| { + TransformError::ShapeError("Symphonia found no supported audio track".to_string()) + })?; + let track_id = track.id; + let audio_params = track + .codec_params + .as_ref() + .and_then(|params| params.audio()) + .ok_or_else(|| { + TransformError::ShapeError( + "Symphonia audio track is missing codec parameters".to_string(), + ) + })?; + let mut decoder = get_codecs() + .make_audio_decoder(audio_params, &AudioDecoderOptions::default()) + .map_err(|e| TransformError::ShapeError(format!("Symphonia decoder failed: {e}")))?; + + let mut sample_rate = audio_params.sample_rate.map(|rate| rate as usize); + let mut mono = Vec::new(); + let mut interleaved = Vec::new(); + loop { + ensure_symphonia_deadline(started, timeout)?; + let Some(packet) = format.next_packet().map_err(|error| { + TransformError::ShapeError(format!("Symphonia packet read failed: {error}")) + })? + else { + break; + }; + if packet.track_id != track_id { + continue; + } + + let audio_buf = match decoder.decode(&packet) { + Ok(decoded) => decoded, + Err(SymphoniaError::DecodeError(_)) => continue, + Err(SymphoniaError::IoError(error)) + if error.kind() == std::io::ErrorKind::UnexpectedEof => + { + break; + } + Err(error) => { + return Err(TransformError::ShapeError(format!( + "Symphonia packet decode failed: {error}" + ))); + } + }; + ensure_symphonia_deadline(started, timeout)?; + + let spec = audio_buf.spec(); + sample_rate = Some(spec.rate() as usize); + let channels = spec.channels().count(); + if channels == 0 { + return Err(TransformError::ShapeError( + "decoded audio has zero channels".to_string(), + )); + } + + let interleaved_samples = audio_buf.samples_interleaved(); + ensure_decoded_sample_limit(0, interleaved_samples, max_decoded_bytes)?; + let additional_samples = interleaved_samples / channels; + ensure_decoded_sample_limit(mono.len(), additional_samples, max_decoded_bytes)?; + mono.try_reserve(additional_samples).map_err(|error| { + TransformError::ShapeError(format!( + "failed to reserve {additional_samples} decoded audio samples: {error}" + )) + })?; + interleaved.resize(interleaved_samples, 0.0); + audio_buf.copy_to_slice_interleaved(&mut interleaved); + for frame in interleaved.chunks_exact(channels) { + mono.push(frame.iter().copied().sum::() / channels as f32); + } + } + + let sample_rate = sample_rate.ok_or_else(|| { + TransformError::ShapeError("decoded audio is missing sample rate".to_string()) + })?; + finish_decoded_audio(mono, sample_rate) +} + +async fn decode_audio_with_ffmpeg(bytes: &[u8]) -> Result { + let input_file = write_temp_audio_file_async(bytes).await?; + let sample_rate = probe_audio_sample_rate(input_file.path()).await?; + // Ask FFmpeg for one sample beyond our limit so a longer stream is + // distinguishable from a valid stream whose size is exactly the limit. + let output_limit = audio_max_decoded_bytes() + .saturating_add(size_of::()) + .to_string(); + + let mut command = Command::new("ffmpeg"); + command + .args(["-hide_banner", "-loglevel", "error", "-nostdin", "-i"]) + .arg(input_file.path()) + .args([ + "-map", + "0:a:0", + "-vn", + "-ac", + "1", + "-fs", + &output_limit, + "-f", + "f32le", + "-sample_fmt", + "flt", + "pipe:1", + ]); + let output = run_audio_command_output(command, "ffmpeg").await?; + if !output.status.success() { + return Err(TransformError::ShapeError(format!( + "ffmpeg failed: {}", + String::from_utf8_lossy(&output.stderr) + ))); + } + ensure_decoded_byte_limit(output.stdout.len(), audio_max_decoded_bytes())?; + if output.stdout.len() % 4 != 0 { + return Err(TransformError::ShapeError(format!( + "ffmpeg f32le output has trailing partial sample: {} bytes", + output.stdout.len() % 4 + ))); + } + + let samples = output + .stdout + .chunks_exact(4) + .map(|bytes| { + let mut sample = [0_u8; size_of::()]; + sample.copy_from_slice(bytes); + f32::from_le_bytes(sample) + }) + .collect(); + finish_decoded_audio(samples, sample_rate) +} + +fn ensure_symphonia_deadline(started: Instant, timeout: Duration) -> Result<(), TransformError> { + if started.elapsed() >= timeout { + return Err(TransformError::ShapeError(format!( + "Symphonia timed out after {:.3} seconds", + timeout.as_secs_f64() + ))); + } + Ok(()) +} + +fn ensure_decoded_sample_limit( + existing_samples: usize, + additional_samples: usize, + max_decoded_bytes: usize, +) -> Result<(), TransformError> { + let total_samples = existing_samples + .checked_add(additional_samples) + .ok_or_else(|| { + TransformError::ShapeError("decoded audio sample count overflow".to_string()) + })?; + let decoded_bytes = total_samples.checked_mul(size_of::()).ok_or_else(|| { + TransformError::ShapeError("decoded audio byte size overflow".to_string()) + })?; + ensure_decoded_byte_limit(decoded_bytes, max_decoded_bytes) +} + +fn ensure_decoded_byte_limit( + decoded_bytes: usize, + max_decoded_bytes: usize, +) -> Result<(), TransformError> { + if decoded_bytes > max_decoded_bytes { + return Err(TransformError::ShapeError(format!( + "decoded audio payload is {decoded_bytes} bytes, exceeding SMG_AUDIO_MAX_DECODED_BYTES={max_decoded_bytes}" + ))); + } + Ok(()) +} + +fn finish_decoded_audio( + samples: Vec, + sample_rate: usize, +) -> Result { + if samples.is_empty() { + return Err(TransformError::ShapeError( + "decoded audio produced no samples".to_string(), + )); + } + Ok(DecodedAudio { + samples, + sample_rate, + }) +} + +async fn probe_audio_sample_rate(input_path: &std::path::Path) -> Result { + let mut command = Command::new("ffprobe"); + command + .args([ + "-v", + "error", + "-select_streams", + "a:0", + "-show_entries", + "stream=sample_rate", + "-of", + "default=noprint_wrappers=1:nokey=1", + ]) + .arg(input_path); + let output = run_audio_command_output(command, "ffprobe").await?; + if !output.status.success() { + return Err(TransformError::ShapeError(format!( + "ffprobe failed: {}", + String::from_utf8_lossy(&output.stderr) + ))); + } + let stdout = String::from_utf8_lossy(&output.stdout); + stdout + .lines() + .find_map(|line| line.trim().parse::().ok()) + .filter(|rate| *rate > 0) + .ok_or_else(|| { + TransformError::ShapeError(format!("failed to parse ffprobe sample rate: {stdout:?}")) + }) +} + +async fn write_temp_audio_file_async( + bytes: &[u8], +) -> Result { + let bytes = bytes.to_vec(); + task::spawn_blocking(move || write_temp_audio_file(&bytes)) + .await + .map_err(|e| TransformError::ShapeError(format!("audio tempfile task failed: {e}")))? +} + +fn write_temp_audio_file(bytes: &[u8]) -> Result { + let mut input_file = tempfile::Builder::new() + .prefix("smg-audio-") + .suffix(audio_temp_suffix(bytes)) + .tempfile() + .map_err(|e| TransformError::ShapeError(format!("audio tempfile failed: {e}")))?; + input_file + .write_all(bytes) + .map_err(|e| TransformError::ShapeError(format!("audio tempfile write failed: {e}")))?; + input_file + .flush() + .map_err(|e| TransformError::ShapeError(format!("audio tempfile flush failed: {e}")))?; + Ok(input_file) +} + +async fn run_audio_command_output( + mut command: Command, + program: &'static str, +) -> Result { + command + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + let child = command.spawn().map_err(|e| { + if e.kind() == std::io::ErrorKind::NotFound { + TransformError::ShapeError(format!( + "{program} executable not found; install {program} for audio decode fallback" + )) + } else { + TransformError::ShapeError(format!("{program} spawn failed: {e}")) + } + })?; + + let timeout = audio_process_timeout(); + match time::timeout(timeout, child.wait_with_output()).await { + Ok(Ok(output)) => Ok(output), + Ok(Err(error)) => Err(TransformError::ShapeError(format!( + "{program} wait failed: {error}" + ))), + Err(_) => Err(TransformError::ShapeError(format!( + "{program} timed out after {:.3} seconds", + timeout.as_secs_f64() + ))), + } +} + +fn audio_decode_backend_override() -> Option<&'static str> { + static BACKEND: OnceLock> = OnceLock::new(); + BACKEND + .get_or_init(|| { + std::env::var("SMG_AUDIO_DECODE_BACKEND") + .ok() + .map(|value| value.trim().to_ascii_lowercase()) + .filter(|value| !value.is_empty() && value != "auto") + }) + .as_deref() +} + +fn audio_process_timeout() -> Duration { + *AUDIO_PROCESS_TIMEOUT.get_or_init(|| { + std::env::var("SMG_AUDIO_PROCESS_TIMEOUT_SECS") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|seconds| seconds.is_finite() && *seconds > 0.0) + .map(Duration::from_secs_f64) + .unwrap_or(DEFAULT_AUDIO_PROCESS_TIMEOUT) + }) +} + +fn audio_max_decoded_bytes() -> usize { + *AUDIO_MAX_DECODED_BYTES.get_or_init(|| { + std::env::var("SMG_AUDIO_MAX_DECODED_BYTES") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|bytes| *bytes > 0) + .unwrap_or(DEFAULT_AUDIO_MAX_DECODED_BYTES) + }) +} + +fn audio_temp_suffix(bytes: &[u8]) -> &'static str { + if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && bytes.get(8..12) == Some(b"WAVE") { + return ".wav"; + } + if bytes.starts_with(b"fLaC") { + return ".flac"; + } + if bytes.starts_with(b"ID3") || is_mp3_frame_sync(bytes) { + return ".mp3"; + } + if bytes.starts_with(b"OggS") { + return ".ogg"; + } + if bytes.len() >= 12 && bytes.get(4..8) == Some(b"ftyp") { + return ".m4a"; + } + if bytes.starts_with(&[0x1a, 0x45, 0xdf, 0xa3]) { + return ".webm"; + } + if bytes.len() >= 12 + && bytes.starts_with(b"FORM") + && matches!(bytes.get(8..12), Some(b"AIFF" | b"AIFC")) + { + return ".aiff"; + } + if bytes.len() >= 4 && bytes.starts_with(b"caff") { + return ".caf"; + } + ".audio" +} + +fn audio_extension_hint(bytes: &[u8]) -> Option<&'static str> { + match audio_temp_suffix(bytes) { + ".wav" => Some("wav"), + ".flac" => Some("flac"), + ".mp3" => Some("mp3"), + ".ogg" => Some("ogg"), + ".m4a" => Some("m4a"), + ".webm" => Some("webm"), + ".aiff" => Some("aiff"), + ".caf" => Some("caf"), + _ => None, + } +} + +fn is_mp3_frame_sync(bytes: &[u8]) -> bool { + bytes.len() >= 2 && bytes[0] == 0xff && (bytes[1] & 0xe0) == 0xe0 +} + +#[cfg(test)] +mod tests { + use super::*; + + fn wav_i16_mono(sample_rate: u32, samples: &[i16]) -> Vec { + let data_bytes = samples.len() as u32 * 2; + let mut bytes = Vec::new(); + bytes.extend_from_slice(b"RIFF"); + bytes.extend_from_slice(&(36 + data_bytes).to_le_bytes()); + bytes.extend_from_slice(b"WAVEfmt "); + bytes.extend_from_slice(&16_u32.to_le_bytes()); + bytes.extend_from_slice(&1_u16.to_le_bytes()); + bytes.extend_from_slice(&1_u16.to_le_bytes()); + bytes.extend_from_slice(&sample_rate.to_le_bytes()); + bytes.extend_from_slice(&(sample_rate * 2).to_le_bytes()); + bytes.extend_from_slice(&2_u16.to_le_bytes()); + bytes.extend_from_slice(&16_u16.to_le_bytes()); + bytes.extend_from_slice(b"data"); + bytes.extend_from_slice(&data_bytes.to_le_bytes()); + for sample in samples { + bytes.extend_from_slice(&sample.to_le_bytes()); + } + bytes + } + + #[test] + fn symphonia_decodes_wav_to_mono_f32() { + let wav = wav_i16_mono(16_000, &[0, 16_384, -16_384]); + let decoded = decode_audio_mono_f32_symphonia(&wav).unwrap(); + assert_eq!(decoded.sample_rate, 16_000); + assert_eq!(decoded.samples.len(), 3); + assert!(decoded.samples[0].abs() < 1e-6); + assert!((decoded.samples[1] - 0.5).abs() < 1e-4); + assert!((decoded.samples[2] + 0.5).abs() < 1e-4); + } + + #[test] + fn symphonia_enforces_decoded_byte_limit() { + let wav = wav_i16_mono(16_000, &[0, 1, 2]); + let error = decode_audio_mono_f32_symphonia_with_limits( + &wav, + 2 * size_of::(), + Duration::from_secs(1), + ) + .unwrap_err(); + + assert!(error.to_string().contains("SMG_AUDIO_MAX_DECODED_BYTES=8")); + } + + #[test] + fn symphonia_enforces_decode_deadline() { + let wav = wav_i16_mono(16_000, &[0]); + let error = decode_audio_mono_f32_symphonia_with_limits(&wav, usize::MAX, Duration::ZERO) + .unwrap_err(); + + assert!(error.to_string().contains("Symphonia timed out")); + } + + #[test] + fn empty_decoded_audio_is_rejected() { + let error = finish_decoded_audio(Vec::new(), 16_000).unwrap_err(); + assert!(error.to_string().contains("produced no samples")); + } + + #[test] + fn audio_suffixes_cover_common_containers() { + assert_eq!(audio_temp_suffix(b"fLaC..."), ".flac"); + assert_eq!(audio_temp_suffix(b"ID3..."), ".mp3"); + assert_eq!(audio_temp_suffix(b"OggS..."), ".ogg"); + assert_eq!(audio_temp_suffix(b"\x1a\x45\xdf\xa3..."), ".webm"); + assert_eq!(audio_temp_suffix(b"\0\0\0\x18ftypM4A "), ".m4a"); + } +} diff --git a/crates/multimodal/src/audio/mod.rs b/crates/multimodal/src/audio/mod.rs new file mode 100644 index 000000000..665a7d920 --- /dev/null +++ b/crates/multimodal/src/audio/mod.rs @@ -0,0 +1,10 @@ +//! Audio preprocessing implementations. + +pub mod decode; +pub mod processor; +pub mod processors; +pub(crate) mod transforms; + +pub use decode::{decode_audio_mono_f32, DecodedAudio}; +pub use processor::{AudioPreProcessor, AudioProcessorFactory, AudioProcessorRegistry}; +pub use processors::{Qwen3AudioParams, Qwen3AudioProcessor}; diff --git a/crates/multimodal/src/audio/processor.rs b/crates/multimodal/src/audio/processor.rs new file mode 100644 index 000000000..fca17c33e --- /dev/null +++ b/crates/multimodal/src/audio/processor.rs @@ -0,0 +1,112 @@ +use std::{collections::HashMap, sync::Arc}; + +use serde_json::Value; + +use super::Qwen3AudioProcessor; +use crate::{ + encoder_inputs::PreprocessedEncoderInputs, error::TransformError, types::AudioClip, + vision::PreProcessorConfig, +}; + +/// Audio preprocessing contract selected by [`AudioProcessorRegistry`]. +pub trait AudioPreProcessor: Send + Sync { + fn preprocess( + &self, + clips: &[Arc], + ) -> Result; +} + +pub type AudioProcessorFactory = fn(&Value, &PreProcessorConfig) -> Box; + +/// Registry of model-specific audio processor factories. +/// +/// Audio processors are created with the current model config because their +/// feature shapes and quantization parameters can be checkpoint-specific. +/// Model-family detection remains the responsibility of `ModelRegistry`; this +/// registry is keyed by the resolved model spec name so that matching logic is +/// not duplicated across capability and processor registries. +pub struct AudioProcessorRegistry { + factories: HashMap, +} + +impl AudioProcessorRegistry { + pub fn new() -> Self { + Self { + factories: HashMap::new(), + } + } + + pub fn register(&mut self, model_spec: impl Into, factory: AudioProcessorFactory) { + self.factories.insert(model_spec.into(), factory); + } + + pub fn create( + &self, + model_spec: &str, + model_config: &Value, + preprocessor_config: &PreProcessorConfig, + ) -> Option> { + self.factories + .get(model_spec) + .copied() + .map(|factory| factory(model_config, preprocessor_config)) + } + + pub fn with_defaults() -> Self { + fn qwen3_audio( + config: &Value, + preprocessor_config: &PreProcessorConfig, + ) -> Box { + Box::new(Qwen3AudioProcessor::from_configs( + config, + preprocessor_config, + )) + } + + let mut registry = Self::new(); + registry.register("qwen3_asr", qwen3_audio); + registry.register("qwen3_omni", qwen3_audio); + registry + } +} + +impl Default for AudioProcessorRegistry { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use bytes::Bytes; + + use super::*; + use crate::{audio::DecodedAudio, types::AudioSource}; + + fn clip() -> Arc { + Arc::new(AudioClip::new( + Bytes::from_static(b"audio"), + DecodedAudio { + samples: vec![0.0; 800], + sample_rate: 16_000, + }, + AudioSource::InlineBytes, + "audio-hash".to_string(), + )) + } + + #[test] + fn qwen_registry_applies_preprocessor_config() { + let registry = AudioProcessorRegistry::with_defaults(); + let preprocessor_config = PreProcessorConfig::from_json( + r#"{"feature_size": 16, "sampling_rate": 16000, "n_fft": 400, "hop_length": 160}"#, + ) + .unwrap(); + let processor = registry + .create("qwen3_asr", &serde_json::json!({}), &preprocessor_config) + .expect("Qwen audio processor"); + + let result = processor.preprocess(&[clip()]).unwrap(); + assert_eq!(result.encoder_input.shape(), &[1, 16, 5]); + } +} diff --git a/crates/multimodal/src/audio/processors/mod.rs b/crates/multimodal/src/audio/processors/mod.rs new file mode 100644 index 000000000..77d27a3f6 --- /dev/null +++ b/crates/multimodal/src/audio/processors/mod.rs @@ -0,0 +1,5 @@ +//! Model-specific audio preprocessing implementations. + +mod qwen3_audio; + +pub use qwen3_audio::{Qwen3AudioParams, Qwen3AudioProcessor}; diff --git a/crates/multimodal/src/audio/processors/qwen3_audio.rs b/crates/multimodal/src/audio/processors/qwen3_audio.rs new file mode 100644 index 000000000..43cdb4b7a --- /dev/null +++ b/crates/multimodal/src/audio/processors/qwen3_audio.rs @@ -0,0 +1,599 @@ +//! Qwen3 audio preprocessing using a Whisper-compatible log-mel frontend. + +use std::sync::Arc; + +use ndarray::{Array2, Array3}; +use rustfft::{num_complex::Complex32, Fft, FftPlanner}; +use serde_json::Value; + +use crate::{ + audio::{ + transforms::{bandlimited_resample, hann_window, mel_basis}, + AudioPreProcessor, DecodedAudio, + }, + encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}, + error::TransformError, + types::AudioClip, + vision::PreProcessorConfig, +}; + +/// Parameters used by the Qwen3 audio frontend. +#[derive(Debug, Clone, PartialEq)] +pub struct Qwen3AudioParams { + pub sample_rate: usize, + pub n_mels: usize, + pub n_fft: usize, + pub hop_length: usize, + pub n_window: usize, + pub padding_value: f32, + pub max_samples: Option, +} + +impl Default for Qwen3AudioParams { + fn default() -> Self { + Self { + sample_rate: 16_000, + n_mels: 128, + n_fft: 400, + hop_length: 160, + n_window: 50, + padding_value: 0.0, + max_samples: None, + } + } +} + +impl Qwen3AudioParams { + pub fn from_configs(model_config: &Value, preprocessor_config: &PreProcessorConfig) -> Self { + let mut params = Self::default(); + + if let Some(value) = preprocessor_config.get_extra::("sampling_rate") { + params.sample_rate = value; + } + if let Some(value) = preprocessor_config.get_extra::("feature_size") { + params.n_mels = value; + } else if let Some(value) = find_model_usize( + model_config, + &[ + &["thinker_config", "audio_config", "num_mel_bins"], + &["audio_config", "num_mel_bins"], + ], + ) { + params.n_mels = value; + } + if let Some(value) = preprocessor_config.get_extra::("n_fft") { + params.n_fft = value; + } + if let Some(value) = preprocessor_config.get_extra::("hop_length") { + params.hop_length = value; + } + if let Some(value) = preprocessor_config + .get_extra::("n_window") + .or_else(|| { + find_model_usize( + model_config, + &[ + &["thinker_config", "audio_config", "n_window"], + &["audio_config", "n_window"], + ], + ) + }) + { + params.n_window = value; + } + if let Some(value) = preprocessor_config.get_extra::("padding_value") { + params.padding_value = value; + } + // Qwen processors set `padding=true, truncation=false`: n_samples is + // the default padding target, not an input limit. Honor it only when a + // checkpoint or deployment explicitly enables truncation. A custom + // max_samples remains available as an operational hard limit. + params.max_samples = preprocessor_config + .get_extra::("max_samples") + .or_else(|| { + preprocessor_config + .get_extra::("truncation") + .filter(|enabled| *enabled) + .and_then(|_| { + preprocessor_config + .get_extra::("n_samples") + .or_else(|| { + preprocessor_config + .get_extra::("chunk_length") + .and_then(|seconds| seconds.checked_mul(params.sample_rate)) + }) + }) + }); + + params + } +} + +fn find_model_usize(config: &Value, paths: &[&[&str]]) -> Option { + paths.iter().find_map(|path| { + let mut value = config; + for key in *path { + value = value.get(*key)?; + } + value.as_u64().and_then(|value| usize::try_from(value).ok()) + }) +} + +#[derive(Debug, Clone)] +pub struct Qwen3AudioProcessor { + params: Qwen3AudioParams, +} + +impl Default for Qwen3AudioProcessor { + fn default() -> Self { + Self::new() + } +} + +impl Qwen3AudioProcessor { + pub fn new() -> Self { + Self { + params: Qwen3AudioParams::default(), + } + } + + pub fn with_params(params: Qwen3AudioParams) -> Self { + Self { params } + } + + pub fn from_configs(model_config: &Value, preprocessor_config: &PreProcessorConfig) -> Self { + Self::with_params(Qwen3AudioParams::from_configs( + model_config, + preprocessor_config, + )) + } + + pub fn params(&self) -> &Qwen3AudioParams { + &self.params + } + + pub fn preprocess_decoded_clips( + &self, + clips: Vec, + ) -> Result { + self.validate_params()?; + if clips.is_empty() { + return Err(TransformError::EmptyBatch); + } + + let mut waveforms = Vec::with_capacity(clips.len()); + for clip in clips { + if clip.sample_rate == 0 { + return Err(TransformError::ShapeError( + "decoded audio sample rate must be positive".to_string(), + )); + } + if clip.samples.is_empty() { + return Err(TransformError::ShapeError( + "decoded audio contains no samples".to_string(), + )); + } + if clip.samples.iter().any(|sample| !sample.is_finite()) { + return Err(TransformError::ShapeError( + "decoded audio contains a non-finite sample".to_string(), + )); + } + let samples = if clip.sample_rate == self.params.sample_rate { + clip.samples + } else { + bandlimited_resample(&clip.samples, clip.sample_rate, self.params.sample_rate)? + }; + let mut samples = samples; + if let Some(max_samples) = self.params.max_samples { + samples.truncate(max_samples); + } + if samples.is_empty() { + return Err(TransformError::ShapeError( + "decoded audio contains no samples after truncation".to_string(), + )); + } + waveforms.push(samples); + } + + let max_samples = waveforms.iter().map(Vec::len).max().unwrap_or(0); + // Whisper's centered STFT yields floor(samples / hop) + 1 frames and + // then drops the final frame. This matches padding=True with the batch + // padded to its longest waveform. + let max_frames = max_samples / self.params.hop_length; + if max_frames == 0 { + return Err(TransformError::ShapeError(format!( + "Qwen3 audio requires at least {} samples after resampling", + self.params.hop_length + ))); + } + + let batch_size = waveforms.len(); + let feature_values = batch_size + .checked_mul(self.params.n_mels) + .and_then(|value| value.checked_mul(max_frames)) + .ok_or_else(|| { + TransformError::ShapeError("Qwen3 audio feature size overflow".to_string()) + })?; + let mut all_features = Vec::with_capacity(feature_values); + let mut attention_mask = Vec::with_capacity(batch_size * max_frames); + let mut feature_lengths = Vec::with_capacity(batch_size); + let mut token_counts = Vec::with_capacity(batch_size); + let mut item_sizes = Vec::with_capacity(batch_size); + let mut planner = FftPlanner::::new(); + let fft = planner.plan_fft_forward(self.params.n_fft); + + for waveform in waveforms { + let original_samples = waveform.len(); + let feature_length = original_samples + .div_ceil(self.params.hop_length) + .min(max_frames); + let mut padded = waveform; + padded.resize(max_samples, self.params.padding_value); + let features = whisper_log_mel(&padded, max_frames, &self.params, fft.as_ref())?; + all_features.extend(features.into_raw_vec_and_offset().0); + + attention_mask.extend((0..max_frames).map(|frame| i64::from(frame < feature_length))); + feature_lengths.push(feature_length as i64); + token_counts.push(qwen3_audio_output_length( + feature_length, + self.params.n_window, + )); + item_sizes.push((self.params.n_mels as u32, feature_length as u32)); + } + + let encoder_input = + Array3::from_shape_vec((batch_size, self.params.n_mels, max_frames), all_features) + .map_err(|error| { + TransformError::ShapeError(format!( + "failed to create Qwen3 audio input [{batch_size}, {}, {max_frames}]: {error}", + self.params.n_mels + )) + })?; + + Ok( + PreprocessedEncoderInputs::new(encoder_input, token_counts, item_sizes) + .with_extra( + "feature_attention_mask", + ModelSpecificValue::int_2d(attention_mask, batch_size, max_frames), + ) + .with_extra( + "audio_feature_lengths", + ModelSpecificValue::int_1d(feature_lengths), + ), + ) + } + + pub fn preprocess_decoded(&self, decoded: DecodedAudio) -> Result, TransformError> { + let output = self.preprocess_decoded_clips(vec![decoded])?; + output + .encoder_input + .into_dimensionality::() + .map_err(|error| TransformError::ShapeError(error.to_string()))? + .index_axis_move(ndarray::Axis(0), 0) + .into_dimensionality::() + .map_err(|error| TransformError::ShapeError(error.to_string())) + } + + fn validate_params(&self) -> Result<(), TransformError> { + if self.params.sample_rate == 0 + || self.params.n_mels == 0 + || self.params.n_fft == 0 + || self.params.hop_length == 0 + || self.params.n_window == 0 + { + return Err(TransformError::ShapeError( + "Qwen3 audio sample rate, mel bins, FFT size, hop length, and window size must be positive" + .to_string(), + )); + } + if self.params.n_fft < self.params.hop_length { + return Err(TransformError::ShapeError(format!( + "Qwen3 audio n_fft ({}) must be at least hop_length ({})", + self.params.n_fft, self.params.hop_length + ))); + } + if !self.params.padding_value.is_finite() { + return Err(TransformError::ShapeError( + "Qwen3 audio padding_value must be finite".to_string(), + )); + } + if self.params.max_samples == Some(0) { + return Err(TransformError::ShapeError( + "Qwen3 audio max_samples must be positive".to_string(), + )); + } + Ok(()) + } +} + +impl AudioPreProcessor for Qwen3AudioProcessor { + fn preprocess( + &self, + clips: &[Arc], + ) -> Result { + self.preprocess_decoded_clips(clips.iter().map(|clip| clip.decoded().clone()).collect()) + } +} + +fn qwen_audio_cnn_output_length(mut input_length: usize) -> usize { + for _ in 0..3 { + input_length = input_length.div_ceil(2); + } + input_length +} + +/// Output tokens produced by Qwen's chunked audio encoder for a log-mel length. +fn qwen3_audio_output_length(input_length: usize, n_window: usize) -> usize { + debug_assert!(n_window > 0); + let chunk_size = 2 * n_window; + let full_windows = input_length / chunk_size; + let remainder = input_length % chunk_size; + full_windows * qwen_audio_cnn_output_length(chunk_size) + + qwen_audio_cnn_output_length(remainder) +} + +fn whisper_log_mel( + samples: &[f32], + frame_count: usize, + params: &Qwen3AudioParams, + fft: &dyn Fft, +) -> Result, TransformError> { + let center_pad = params.n_fft / 2; + let padded = reflect_pad(samples, center_pad); + let fft_bins = params.n_fft / 2 + 1; + let window = hann_window(params.n_fft); + let mel_filters = mel_basis(params.sample_rate, params.n_fft, params.n_mels); + let mut buffer = vec![Complex32::new(0.0, 0.0); params.n_fft]; + let mut output = vec![0.0_f32; params.n_mels * frame_count]; + + for frame in 0..frame_count { + let start = frame * params.hop_length; + let end = start + params.n_fft; + let frame_samples = padded.get(start..end).ok_or_else(|| { + TransformError::ShapeError(format!( + "Qwen3 audio STFT frame {frame} lies outside padded waveform" + )) + })?; + for index in 0..params.n_fft { + buffer[index] = Complex32::new(frame_samples[index] * window[index], 0.0); + } + fft.process(&mut buffer); + + for mel in 0..params.n_mels { + let filter = &mel_filters[mel * fft_bins..(mel + 1) * fft_bins]; + let mut value = 0.0_f32; + for bin in 0..fft_bins { + let fft_value = buffer[bin]; + let power = fft_value + .re + .mul_add(fft_value.re, fft_value.im * fft_value.im); + value = filter[bin].mul_add(power, value); + } + output[mel * frame_count + frame] = value.max(1e-10).log10(); + } + } + + let peak = output.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let floor = peak - 8.0; + for value in &mut output { + *value = (value.max(floor) + 4.0) / 4.0; + } + + Array2::from_shape_vec((params.n_mels, frame_count), output).map_err(|error| { + TransformError::ShapeError(format!( + "failed to create Qwen log-mel input [{}, {frame_count}]: {error}", + params.n_mels + )) + }) +} + +fn reflect_pad(samples: &[f32], padding: usize) -> Vec { + if padding == 0 { + return samples.to_vec(); + } + if samples.len() == 1 { + return vec![samples[0]; samples.len() + 2 * padding]; + } + + let mut padded = Vec::with_capacity(samples.len() + 2 * padding); + for position in -(padding as isize)..(samples.len() + padding) as isize { + padded.push(samples[reflect_index(position, samples.len())]); + } + padded +} + +fn reflect_index(mut index: isize, len: usize) -> usize { + let last = len as isize - 1; + while index < 0 || index > last { + if index < 0 { + index = -index; + } + if index > last { + index = 2 * last - index; + } + } + index as usize +} + +#[cfg(test)] +mod tests { + use super::*; + + fn decoded(samples: usize) -> DecodedAudio { + DecodedAudio { + samples: vec![0.0; samples], + sample_rate: 16_000, + } + } + + #[test] + fn silence_matches_whisper_normalization() { + let features = Qwen3AudioProcessor::new() + .preprocess_decoded(decoded(1600)) + .unwrap(); + assert_eq!(features.shape(), &[128, 10]); + assert!(features.iter().all(|value| (*value + 1.5).abs() < 1e-6)); + } + + #[test] + fn log_mel_matches_numpy_whisper_reference() { + let samples = (0..1000) + .map(|index| ((index % 23) as f32 - 11.0) / 32.0) + .collect(); + let features = Qwen3AudioProcessor::new() + .preprocess_decoded(DecodedAudio { + samples, + sample_rate: 16_000, + }) + .unwrap(); + assert_eq!(features.shape(), &[128, 6]); + + for ((mel, frame), expected) in [ + ((0, 0), 0.71467185), + ((10, 0), 0.735_666_9), + ((20, 1), 0.36943567), + ((32, 3), 0.79836977), + ((64, 5), -0.582_248_9), + ((100, 5), 0.60166824), + ((127, 5), 0.52557164), + ] { + assert!( + (features[[mel, frame]] - expected).abs() < 2e-4, + "log-mel mismatch at ({mel}, {frame}): {} vs {expected}", + features[[mel, frame]] + ); + } + let sum: f64 = features.iter().map(|&value| f64::from(value)).sum(); + assert!((sum - 125.08224487).abs() < 0.02, "feature sum {sum}"); + } + + #[test] + fn log_mel_boundary_matches_hf_whisper_reference() { + let samples = (0..1600) + .map(|index| ((index % 23) as f32 - 11.0) / 32.0) + .collect(); + let features = Qwen3AudioProcessor::new() + .preprocess_decoded(DecodedAudio { + samples, + sample_rate: 16_000, + }) + .unwrap(); + + assert_eq!(features.shape(), &[128, 10]); + for (mel, expected) in [(64, 0.017_234_564), (100, 0.601_944_7), (127, 0.525_99)] { + assert!( + (features[[mel, 9]] - expected).abs() < 2e-4, + "last-frame log-mel mismatch at ({mel}, 9): {} vs {expected}", + features[[mel, 9]] + ); + } + } + + #[test] + fn batches_variable_lengths_with_feature_mask() { + let output = Qwen3AudioProcessor::new() + .preprocess_decoded_clips(vec![decoded(1000), decoded(800)]) + .unwrap(); + + assert_eq!(output.encoder_input.shape(), &[2, 128, 6]); + assert_eq!(output.feature_token_counts, vec![1, 1]); + assert_eq!(output.item_sizes, vec![(128, 6), (128, 5)]); + assert!(matches!( + output.model_specific.get("audio_feature_lengths"), + Some(ModelSpecificValue::IntTensor { data, shape }) + if data == &vec![6, 5] && shape == &vec![2] + )); + assert!(matches!( + output.model_specific.get("feature_attention_mask"), + Some(ModelSpecificValue::IntTensor { data, shape }) + if data == &vec![1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0] + && shape == &vec![2, 6] + )); + } + + #[test] + fn reads_preprocessor_and_nested_model_config() { + let preprocessor = PreProcessorConfig::from_json( + r#"{"sampling_rate": 16000, "feature_size": 64, "n_fft": 320, "hop_length": 80, "n_samples": 4800000}"#, + ) + .unwrap(); + let params = Qwen3AudioParams::from_configs( + &serde_json::json!({ + "thinker_config": {"audio_config": {"num_mel_bins": 96, "n_window": 64}} + }), + &preprocessor, + ); + assert_eq!(params.sample_rate, 16_000); + assert_eq!(params.n_mels, 64); + assert_eq!(params.n_fft, 320); + assert_eq!(params.hop_length, 80); + assert_eq!(params.n_window, 64); + assert_eq!(params.max_samples, None); + + let truncating_preprocessor = PreProcessorConfig::from_json( + r#"{"sampling_rate": 16000, "n_samples": 4800000, "truncation": true}"#, + ) + .unwrap(); + let params = + Qwen3AudioParams::from_configs(&serde_json::json!({}), &truncating_preprocessor); + assert_eq!(params.max_samples, Some(4_800_000)); + + let params = Qwen3AudioParams::from_configs( + &serde_json::json!({ + "thinker_config": {"audio_config": {"num_mel_bins": 96}} + }), + &PreProcessorConfig::default(), + ); + assert_eq!(params.n_mels, 96); + assert_eq!(params.max_samples, None); + } + + #[test] + fn n_samples_is_padding_target_not_implicit_truncation() { + let preprocessor = PreProcessorConfig::from_json(r#"{"n_samples": 320}"#).unwrap(); + let processor = Qwen3AudioProcessor::from_configs(&serde_json::json!({}), &preprocessor); + let output = processor + .preprocess_decoded_clips(vec![decoded(800)]) + .unwrap(); + + assert_eq!(output.encoder_input.shape(), &[1, 128, 5]); + assert_eq!(output.item_sizes, vec![(128, 5)]); + } + + #[test] + fn truncates_to_explicit_audio_limit() { + let processor = Qwen3AudioProcessor::with_params(Qwen3AudioParams { + max_samples: Some(320), + ..Default::default() + }); + let output = processor + .preprocess_decoded_clips(vec![decoded(800)]) + .unwrap(); + + assert_eq!(output.encoder_input.shape(), &[1, 128, 2]); + assert_eq!(output.feature_token_counts, vec![1]); + assert_eq!(output.item_sizes, vec![(128, 2)]); + } + + #[test] + fn qwen_chunked_encoder_output_lengths_match_reference_formula() { + assert_eq!(qwen3_audio_output_length(0, 50), 0); + assert_eq!(qwen3_audio_output_length(1, 50), 1); + assert_eq!(qwen3_audio_output_length(8, 50), 1); + assert_eq!(qwen3_audio_output_length(9, 50), 2); + assert_eq!(qwen3_audio_output_length(99, 50), 13); + assert_eq!(qwen3_audio_output_length(100, 50), 13); + assert_eq!(qwen3_audio_output_length(101, 50), 14); + assert_eq!(qwen3_audio_output_length(3000, 50), 390); + assert_eq!(qwen3_audio_output_length(17, 4), 3); + } + + #[test] + fn reflect_padding_matches_numpy_convention() { + assert_eq!( + reflect_pad(&[1.0, 2.0, 3.0], 2), + vec![3.0, 2.0, 1.0, 2.0, 3.0, 2.0, 1.0] + ); + assert_eq!(reflect_pad(&[2.0], 2), vec![2.0; 5]); + } +} diff --git a/crates/multimodal/src/audio/transforms.rs b/crates/multimodal/src/audio/transforms.rs new file mode 100644 index 000000000..60f87e3dc --- /dev/null +++ b/crates/multimodal/src/audio/transforms.rs @@ -0,0 +1,159 @@ +//! Shared digital-signal-processing transforms for audio frontends. + +use std::{f32::consts::PI as PI_F32, f64::consts::PI}; + +use crate::error::TransformError; + +pub(super) fn hann_window(window_size: usize) -> Vec { + (0..window_size) + .map(|i| (0.5 - 0.5 * (2.0 * PI * i as f64 / window_size as f64).cos()) as f32) + .collect() +} + +fn hz_to_mel(frequency: f64) -> f64 { + let f_sp = 200.0 / 3.0; + let min_log_hz = 1000.0; + let min_log_mel = min_log_hz / f_sp; + let logstep = 6.4_f64.ln() / 27.0; + if frequency >= min_log_hz { + min_log_mel + (frequency / min_log_hz).ln() / logstep + } else { + frequency / f_sp + } +} + +fn mel_to_hz(mel: f64) -> f64 { + let f_sp = 200.0 / 3.0; + let min_log_hz = 1000.0; + let min_log_mel = min_log_hz / f_sp; + let logstep = 6.4_f64.ln() / 27.0; + if mel >= min_log_mel { + min_log_hz * (logstep * (mel - min_log_mel)).exp() + } else { + mel * f_sp + } +} + +/// Build a Slaney-normalized mel filter bank. +pub(super) fn mel_basis(sample_rate: usize, n_fft: usize, n_mels: usize) -> Vec { + let fft_bins = n_fft / 2 + 1; + let mut fft_freqs = Vec::with_capacity(fft_bins); + for bin in 0..fft_bins { + fft_freqs.push(bin as f64 * sample_rate as f64 / n_fft as f64); + } + + let mel_min = hz_to_mel(0.0); + let mel_max = hz_to_mel(sample_rate as f64 / 2.0); + let mut mel_edges = Vec::with_capacity(n_mels + 2); + for i in 0..n_mels + 2 { + let t = i as f64 / (n_mels + 1) as f64; + mel_edges.push(mel_to_hz(mel_min + (mel_max - mel_min) * t)); + } + + let mel_widths: Vec = mel_edges.windows(2).map(|w| w[1] - w[0]).collect(); + let mut weights = vec![0.0_f32; n_mels * fft_bins]; + for mel in 0..n_mels { + let enorm = 2.0 / (mel_edges[mel + 2] - mel_edges[mel]); + for (bin, &freq) in fft_freqs.iter().enumerate() { + let lower = (freq - mel_edges[mel]) / mel_widths[mel]; + let upper = (mel_edges[mel + 2] - freq) / mel_widths[mel + 1]; + weights[mel * fft_bins + bin] = lower.min(upper).max(0.0).mul_add(enorm, 0.0) as f32; + } + } + weights +} + +/// Match torchaudio's default `functional.resample`: band-limited sinc +/// interpolation with a Hann window, filter width 6, and rolloff 0.99. +pub(super) fn bandlimited_resample( + samples: &[f32], + src_sample_rate: usize, + dst_sample_rate: usize, +) -> Result, TransformError> { + const LOWPASS_FILTER_WIDTH: f32 = 6.0; + const ROLLOFF: f32 = 0.99; + + if src_sample_rate == 0 || dst_sample_rate == 0 { + return Err(TransformError::ShapeError( + "audio resampling rates must be positive".to_string(), + )); + } + if samples.is_empty() || src_sample_rate == dst_sample_rate { + return Ok(samples.to_vec()); + } + + let gcd = greatest_common_divisor(src_sample_rate, dst_sample_rate); + let orig_freq = src_sample_rate / gcd; + let new_freq = dst_sample_rate / gcd; + let base_freq = orig_freq.min(new_freq) as f32 * ROLLOFF; + let width = (LOWPASS_FILTER_WIDTH * orig_freq as f32 / base_freq).ceil() as usize; + let kernel_len = width + .checked_mul(2) + .and_then(|value| value.checked_add(orig_freq)) + .ok_or_else(|| TransformError::ShapeError("audio resample kernel is too large".into()))?; + let kernel_values = new_freq + .checked_mul(kernel_len) + .ok_or_else(|| TransformError::ShapeError("audio resample kernel size overflow".into()))?; + let mut kernels = Vec::new(); + kernels.try_reserve_exact(kernel_values).map_err(|error| { + TransformError::ShapeError(format!("failed to allocate audio resample kernel: {error}")) + })?; + + let orig_freq_f32 = orig_freq as f32; + let new_freq_f32 = new_freq as f32; + let scale = base_freq / orig_freq_f32; + for phase in 0..new_freq { + for kernel_index in 0..kernel_len { + let idx = (kernel_index as f32 - width as f32) / orig_freq_f32; + let mut t = (idx - phase as f32 / new_freq_f32) * base_freq; + t = t.clamp(-LOWPASS_FILTER_WIDTH, LOWPASS_FILTER_WIDTH); + let window = (t * PI_F32 / LOWPASS_FILTER_WIDTH / 2.0).cos().powi(2); + let radians = t * PI_F32; + let sinc = if radians == 0.0 { + 1.0 + } else { + radians.sin() / radians + }; + kernels.push(sinc * window * scale); + } + } + + let target_len_u128 = (samples.len() as u128 * new_freq as u128).div_ceil(orig_freq as u128); + let target_len = usize::try_from(target_len_u128).map_err(|_| { + TransformError::ShapeError("resampled audio length exceeds usize".to_string()) + })?; + let mut output = Vec::new(); + output.try_reserve_exact(target_len).map_err(|error| { + TransformError::ShapeError(format!("failed to allocate resampled audio: {error}")) + })?; + + for block in 0..samples.len().div_ceil(orig_freq) { + let input_start = block * orig_freq; + for phase in 0..new_freq { + if output.len() == target_len { + return Ok(output); + } + let kernel = &kernels[phase * kernel_len..(phase + 1) * kernel_len]; + let mut value = 0.0_f32; + for (kernel_index, &coefficient) in kernel.iter().enumerate() { + let padded_index = input_start + kernel_index; + if padded_index >= width { + let sample_index = padded_index - width; + if let Some(&sample) = samples.get(sample_index) { + value = sample.mul_add(coefficient, value); + } + } + } + output.push(value); + } + } + + Ok(output) +} + +fn greatest_common_divisor(mut lhs: usize, mut rhs: usize) -> usize { + while rhs != 0 { + (lhs, rhs) = (rhs, lhs % rhs); + } + lhs +} diff --git a/crates/multimodal/src/encoder_inputs.rs b/crates/multimodal/src/encoder_inputs.rs new file mode 100644 index 000000000..ed4600412 --- /dev/null +++ b/crates/multimodal/src/encoder_inputs.rs @@ -0,0 +1,328 @@ +//! Shared encoder-input types for all encoder-backed modalities. + +use std::{borrow::Cow, collections::HashMap}; + +use anyhow::{Context, Result as AnyhowResult}; +use ndarray::{Array, ArrayD, Dimension}; + +use crate::types::FieldLayout; + +/// Model-specific auxiliary output values. +#[derive(Debug, Clone)] +pub enum ModelSpecificValue { + /// A tensor with shape information (data as flat vec, shape as dims) + Tensor { data: Vec, shape: Vec }, + + /// A tensor of integers (e.g., aspect_ratio_ids) + IntTensor { data: Vec, shape: Vec }, + + /// A tensor of unsigned integers (e.g., image_grid_thw) + UintTensor { data: Vec, shape: Vec }, + + /// Simple integer value + Int(i64), + + /// Simple float value + Float(f64), + + /// List of integers + IntVec(Vec), + + /// List of unsigned integers + UintVec(Vec), + + /// List of floats + FloatVec(Vec), + + /// List of tuples (e.g., media item sizes) + TupleVec(Vec<(u32, u32)>), + + /// Boolean flag + Bool(bool), +} + +impl ModelSpecificValue { + /// Create a 1D uint tensor from a vector. + pub fn uint_1d(data: Vec) -> Self { + let len = data.len(); + Self::UintTensor { + data, + shape: vec![len], + } + } + + /// Create a 2D uint tensor. + pub fn uint_2d(data: Vec, rows: usize, cols: usize) -> Self { + Self::UintTensor { + data, + shape: vec![rows, cols], + } + } + + /// Create a 1D int tensor from a vector. + pub fn int_1d(data: Vec) -> Self { + let len = data.len(); + Self::IntTensor { + data, + shape: vec![len], + } + } + + /// Create a 2D int tensor. + pub fn int_2d(data: Vec, rows: usize, cols: usize) -> Self { + Self::IntTensor { + data, + shape: vec![rows, cols], + } + } + + /// Interpret this value as per-item flat sizes. + pub fn as_flat_sizes(&self) -> AnyhowResult> { + match self { + Self::IntTensor { data, .. } => data + .iter() + .map(|&v| usize::try_from(v).context("negative flat size")) + .collect(), + Self::UintTensor { data, .. } => Ok(data.iter().map(|&v| v as usize).collect()), + Self::IntVec(values) => values + .iter() + .map(|&v| usize::try_from(v).context("negative flat size")) + .collect(), + Self::UintVec(values) => Ok(values.iter().map(|&v| v as usize).collect()), + _ => Err(anyhow::anyhow!("unsupported flat sizes value type")), + } + } + + /// Slice item-batched metadata along the first dimension. + pub fn slice_first_dim(&self, start: usize, len: usize) -> AnyhowResult { + match self { + Self::Tensor { data, shape } => { + let (data, shape) = slice_tensor_first_dim(data, shape, start, len)?; + Ok(Self::Tensor { data, shape }) + } + Self::IntTensor { data, shape } => { + let (data, shape) = slice_tensor_first_dim(data, shape, start, len)?; + Ok(Self::IntTensor { data, shape }) + } + Self::UintTensor { data, shape } => { + let (data, shape) = slice_tensor_first_dim(data, shape, start, len)?; + Ok(Self::UintTensor { data, shape }) + } + Self::IntVec(values) => Ok(Self::IntVec(slice_1d(values, start, len)?.to_vec())), + Self::UintVec(values) => Ok(Self::UintVec(slice_1d(values, start, len)?.to_vec())), + Self::FloatVec(values) => Ok(Self::FloatVec(slice_1d(values, start, len)?.to_vec())), + Self::TupleVec(values) => Ok(Self::TupleVec(slice_1d(values, start, len)?.to_vec())), + _ => Ok(self.clone()), + } + } +} + +fn slice_tensor_first_dim( + data: &[T], + shape: &[usize], + start: usize, + len: usize, +) -> AnyhowResult<(Vec, Vec)> { + let first_dim = *shape + .first() + .ok_or_else(|| anyhow::anyhow!("cannot slice scalar tensor"))?; + let end = start + .checked_add(len) + .ok_or_else(|| anyhow::anyhow!("tensor slice range overflow"))?; + anyhow::ensure!( + end <= first_dim, + "tensor first-dimension slice {start}..{end} exceeds {first_dim}" + ); + let row_width = shape[1..] + .iter() + .try_fold(1usize, |acc, &dim| acc.checked_mul(dim)) + .ok_or_else(|| anyhow::anyhow!("tensor row width overflow"))?; + let data_start = start + .checked_mul(row_width) + .ok_or_else(|| anyhow::anyhow!("tensor data start overflow"))?; + let data_len = len + .checked_mul(row_width) + .ok_or_else(|| anyhow::anyhow!("tensor data length overflow"))?; + let data_end = data_start + .checked_add(data_len) + .ok_or_else(|| anyhow::anyhow!("tensor data end overflow"))?; + anyhow::ensure!( + data_end <= data.len(), + "tensor slice data range {data_start}..{data_end} exceeds {}", + data.len() + ); + let mut new_shape = shape.to_vec(); + new_shape[0] = len; + Ok((data[data_start..data_end].to_vec(), new_shape)) +} + +fn slice_1d(values: &[T], start: usize, len: usize) -> AnyhowResult<&[T]> { + let end = start + .checked_add(len) + .ok_or_else(|| anyhow::anyhow!("slice range overflow"))?; + values + .get(start..end) + .ok_or_else(|| anyhow::anyhow!("slice range {start}..{end} exceeds {}", values.len())) +} + +/// Preprocessed encoder inputs ready for model consumption. +#[derive(Debug, Clone)] +pub struct PreprocessedEncoderInputs { + /// Primary encoder input as a dynamic-dimensional float32 tensor. + pub encoder_input: ArrayD, + + /// Number of encoder feature tokens per media item in the batch. + pub feature_token_counts: Vec, + + /// Modality-specific item size metadata before preprocessing. + /// + /// The exact tuple order follows each processor/model contract. Auxiliary + /// shape tensors that need a fixed order should be emitted in + /// `model_specific`. + pub item_sizes: Vec<(u32, u32)>, + + /// Model-specific auxiliary outputs. + pub model_specific: HashMap, +} + +impl PreprocessedEncoderInputs { + /// Create encoder inputs backed by a tensor of any dimensionality. + pub fn new( + encoder_input: Array, + feature_token_counts: Vec, + item_sizes: Vec<(u32, u32)>, + ) -> Self { + Self { + encoder_input: encoder_input.into_dyn(), + feature_token_counts, + item_sizes, + model_specific: HashMap::new(), + } + } + + /// Add a model-specific value. + pub fn with_extra(mut self, key: impl Into, value: ModelSpecificValue) -> Self { + self.model_specific.insert(key.into(), value); + self + } + + /// Get the number of media items represented by this preprocessed batch. + pub fn batch_size(&self) -> usize { + self.item_sizes.len() + } + + /// Get the number of dimensions of encoder_input. + pub fn ndim(&self) -> usize { + self.encoder_input.ndim() + } + + /// Get total number of encoder feature tokens across all media items. + pub fn total_feature_tokens(&self) -> usize { + self.feature_token_counts.iter().sum() + } + + /// Get the primary encoder input as a flat f32 slice without copying if possible. + pub fn encoder_input_flat(&self) -> Cow<'_, [f32]> { + match self.encoder_input.as_slice() { + Some(slice) => Cow::Borrowed(slice), + None => Cow::Owned(self.encoder_input.iter().copied().collect()), + } + } + + /// Get the shape of the primary encoder input as a vector. + pub fn encoder_input_shape(&self) -> Vec { + self.encoder_input.shape().to_vec() + } + + /// Extract batched tensor keys from explicit field layout declarations. + pub fn batched_keys(layouts: &HashMap) -> Vec { + layouts + .iter() + .filter(|(_, layout)| matches!(layout, FieldLayout::Batched)) + .map(|(key, _)| key.clone()) + .collect() + } + + /// Extract flat-slicing tensor keys from explicit field layout declarations. + /// + /// Returns a map of tensor name to sizes tensor name. + pub fn flat_keys(layouts: &HashMap) -> HashMap { + layouts + .iter() + .filter_map(|(key, layout)| match layout { + FieldLayout::Flat { sizes_key } => Some((key.clone(), sizes_key.clone())), + FieldLayout::Batched => None, + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use ndarray::Array4; + + use super::*; + + #[test] + fn encoder_input_accessors_are_modality_neutral() { + let inputs = PreprocessedEncoderInputs::new( + Array4::::zeros((2, 3, 4, 5)), + vec![6, 7], + vec![(4, 5), (8, 9)], + ); + + assert_eq!(inputs.batch_size(), 2); + assert_eq!(inputs.ndim(), 4); + assert_eq!(inputs.total_feature_tokens(), 13); + assert_eq!(inputs.encoder_input_shape(), vec![2, 3, 4, 5]); + } + + #[test] + fn encoder_inputs_accept_model_specific_values() { + let inputs = PreprocessedEncoderInputs::new( + Array4::::zeros((1, 3, 224, 224)), + vec![196], + vec![(224, 224)], + ) + .with_extra( + "image_grid_thw", + ModelSpecificValue::uint_1d(vec![1, 16, 16]), + ) + .with_extra("aspect_ratio_id", ModelSpecificValue::Int(0)); + + assert!(inputs.model_specific.contains_key("image_grid_thw")); + assert!(inputs.model_specific.contains_key("aspect_ratio_id")); + } + + #[test] + fn model_specific_value_tensor_constructors_set_shapes() { + assert!(matches!( + ModelSpecificValue::uint_1d(vec![1, 2, 3]), + ModelSpecificValue::UintTensor { data, shape } + if data == vec![1, 2, 3] && shape == vec![3] + )); + assert!(matches!( + ModelSpecificValue::uint_2d(vec![1, 2, 3, 4], 2, 2), + ModelSpecificValue::UintTensor { data, shape } + if data == vec![1, 2, 3, 4] && shape == vec![2, 2] + )); + assert!(matches!( + ModelSpecificValue::int_1d(vec![1, 2, 3]), + ModelSpecificValue::IntTensor { data, shape } + if data == vec![1, 2, 3] && shape == vec![3] + )); + assert!(matches!( + ModelSpecificValue::int_2d(vec![1, 2, 3, 4], 2, 2), + ModelSpecificValue::IntTensor { data, shape } + if data == vec![1, 2, 3, 4] && shape == vec![2, 2] + )); + } + + #[test] + fn encoder_input_flat_preserves_values() { + let encoder_input = Array4::from_shape_vec((1, 1, 2, 2), vec![1.0, 2.0, 3.0, 4.0]).unwrap(); + let inputs = PreprocessedEncoderInputs::new(encoder_input, vec![4], vec![(2, 2)]); + + assert_eq!(inputs.encoder_input_flat(), vec![1.0, 2.0, 3.0, 4.0]); + } +} diff --git a/crates/multimodal/src/error.rs b/crates/multimodal/src/error.rs index 402b262c2..6e9d75c5c 100644 --- a/crates/multimodal/src/error.rs +++ b/crates/multimodal/src/error.rs @@ -4,6 +4,25 @@ use thiserror::Error; pub type MultiModalResult = Result; +/// Errors that can occur while transforming media into encoder inputs. +#[derive(Debug, Error)] +pub enum TransformError { + #[error("Invalid tensor shape: expected {expected}, got {actual:?}")] + InvalidShape { + expected: String, + actual: Vec, + }, + + #[error("Empty batch: cannot stack zero tensors")] + EmptyBatch, + + #[error("Inconsistent tensor shapes in batch")] + InconsistentShapes, + + #[error("Shape error: {0}")] + ShapeError(String), +} + #[derive(Debug, Error)] pub enum MediaConnectorError { #[error("unsupported media scheme: {0}")] @@ -22,10 +41,14 @@ pub enum MediaConnectorError { Base64Decode(#[from] base64::DecodeError), #[error("data URL parse error: {0}")] DataUrl(String), + #[error("{media} payload exceeds the maximum size of {limit} bytes")] + PayloadTooLarge { media: &'static str, limit: usize }, #[error("media decode task failed: {0}")] Blocking(#[from] tokio::task::JoinError), #[error("image decode error: {0}")] Image(#[from] image::ImageError), + #[error("audio decode error: {0}")] + AudioDecode(String), #[error("video decode error: {0}")] VideoDecode(String), #[error("media fetch timed out after {0:?}")] diff --git a/crates/multimodal/src/hasher.rs b/crates/multimodal/src/hasher.rs index 756e6d204..2019a3f68 100644 --- a/crates/multimodal/src/hasher.rs +++ b/crates/multimodal/src/hasher.rs @@ -12,6 +12,11 @@ pub fn hash_video(raw_bytes: &[u8]) -> String { blake3::hash(raw_bytes).to_hex().to_string() } +/// Compute a blake3 hex-digest hash for a single audio payload's raw bytes. +pub fn hash_audio(raw_bytes: &[u8]) -> String { + blake3::hash(raw_bytes).to_hex().to_string() +} + /// Compute per-image hashes keyed by modality. /// /// Returns a `BTreeMap` of per-modality hash lists, diff --git a/crates/multimodal/src/lib.rs b/crates/multimodal/src/lib.rs index 2815548fc..1aa91053f 100644 --- a/crates/multimodal/src/lib.rs +++ b/crates/multimodal/src/lib.rs @@ -1,3 +1,5 @@ +pub mod audio; +pub mod encoder_inputs; pub mod error; pub mod hasher; pub mod hub; @@ -10,19 +12,21 @@ pub mod tracker; pub mod types; pub mod vision; -pub use error::{MediaConnectorError, MultiModalError, MultiModalResult}; +pub use audio::{AudioPreProcessor, AudioProcessorRegistry}; +pub use encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}; +pub use error::{MediaConnectorError, MultiModalError, MultiModalResult, TransformError}; pub use media::{ ImageFetchConfig, MediaConnector, MediaConnectorConfig, MediaSource, VideoFetchConfig, }; pub use registry::{ModelMetadata, ModelProcessorSpec, ModelRegistry}; pub use tracker::{AsyncMultiModalTracker, TrackerOutput}; pub use types::{ - FieldLayout, ImageDetail, ImageFrame, ImageSize, ImageSource, MediaContentPart, Modality, - MultiModalData, MultiModalUUIDs, PlaceholderRange, PromptReplacement, RgbFrameRef, TokenId, - TrackedMedia, VideoClip, VideoSource, + AudioClip, AudioSource, EncoderFieldLayouts, FieldLayout, ImageDetail, ImageFrame, ImageSize, + ImageSource, MediaContentPart, Modality, MultiModalData, MultiModalUUIDs, PlaceholderRange, + PromptReplacement, RgbFrameRef, TokenId, TrackedMedia, VideoClip, VideoSource, }; // Re-export vision processing components pub use vision::{ - LlavaNextProcessor, LlavaProcessor, ModelSpecificValue, PreProcessorConfig, - PreprocessedEncoderInputs, TransformError, VisionPreProcessor, VisionProcessorRegistry, + LlavaNextProcessor, LlavaProcessor, PreProcessorConfig, VisionPreProcessor, + VisionProcessorRegistry, }; diff --git a/crates/multimodal/src/media.rs b/crates/multimodal/src/media.rs index c04ce7c15..b50675633 100644 --- a/crates/multimodal/src/media.rs +++ b/crates/multimodal/src/media.rs @@ -10,7 +10,7 @@ use std::{ }; use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; -use bytes::Bytes; +use bytes::{Bytes, BytesMut}; #[cfg(feature = "opencv-video")] use opencv::{ core::{Mat, Vector}, @@ -22,12 +22,16 @@ use tokio::{fs, process::Command, task, time}; use tracing::info; use url::Url; +use crate::audio::decode_audio_mono_f32; + const DEFAULT_VIDEO_PROCESS_TIMEOUT: Duration = Duration::from_secs(30); const DEFAULT_VIDEO_MAX_DECODED_BYTES: usize = 1024 * 1024 * 1024; +const DEFAULT_AUDIO_MAX_INPUT_BYTES: usize = 256 * 1024 * 1024; static VIDEO_DECODE_BACKEND: OnceLock> = OnceLock::new(); static LOG_VIDEO_DECODE_TIMING: OnceLock = OnceLock::new(); static VIDEO_PROCESS_TIMEOUT: OnceLock = OnceLock::new(); static VIDEO_MAX_DECODED_BYTES: OnceLock = OnceLock::new(); +static AUDIO_MAX_INPUT_BYTES: OnceLock = OnceLock::new(); #[cfg(feature = "opencv-video")] static ACTIVE_OPENCV_DECODES: AtomicUsize = AtomicUsize::new(0); #[cfg(feature = "opencv-video")] @@ -48,8 +52,8 @@ const OPENCV_HIGH_CONCURRENCY_CPU_BUDGET_DENOMINATOR: usize = 7; use super::{ error::MediaConnectorError, types::{ - DecodedRgbFrame, DecodedRgbVideo, ImageDetail, ImageFrame, ImageSource, VideoClip, - VideoSource, + AudioClip, AudioSource, DecodedRgbFrame, DecodedRgbVideo, ImageDetail, ImageFrame, + ImageSource, VideoClip, VideoSource, }, }; @@ -174,6 +178,21 @@ impl MediaConnector { } } + pub async fn fetch_audio( + &self, + source: MediaSource, + ) -> Result, MediaConnectorError> { + match source { + MediaSource::Url(url) => self.fetch_http_audio(url).await, + MediaSource::DataUrl(data_url) => self.fetch_audio_data_url(data_url).await, + MediaSource::InlineBytes(bytes) => { + self.decode_audio(bytes.into(), AudioSource::InlineBytes) + .await + } + MediaSource::File(path) => self.fetch_audio_file(path).await, + } + } + async fn fetch_http_image( &self, url: String, @@ -249,6 +268,26 @@ impl MediaConnector { .await } + async fn fetch_audio_data_url( + &self, + data_url: String, + ) -> Result, MediaConnectorError> { + let (metadata, data) = data_url + .split_once(',') + .ok_or_else(|| MediaConnectorError::DataUrl("missing comma in data url".into()))?; + + if !metadata.ends_with(";base64") { + return Err(MediaConnectorError::DataUrl( + "only base64 encoded data URLs are supported".into(), + )); + } + + let data = data.trim(); + let decoded = BASE64_STANDARD.decode(data)?; + self.decode_audio(decoded.into(), AudioSource::DataUrl) + .await + } + async fn fetch_file( &self, path: PathBuf, @@ -308,6 +347,34 @@ impl MediaConnector { .await } + async fn fetch_http_audio(&self, url: String) -> Result, MediaConnectorError> { + let parsed = Url::parse(&url).map_err(|_| MediaConnectorError::InvalidUrl(url.clone()))?; + self.ensure_domain_allowed(&parsed)?; + + let mut req = self.client.get(parsed.as_str()); + if self.fetch_timeout > Duration::ZERO { + req = req.timeout(self.fetch_timeout); + } + + let resp = req.send().await.map_err(|err| { + if err.is_timeout() { + MediaConnectorError::Timeout(self.fetch_timeout) + } else { + MediaConnectorError::Http(err) + } + })?; + + let resp = resp.error_for_status()?; + let bytes = collect_http_body_with_limit(resp, audio_max_input_bytes(), "audio").await?; + self.decode_audio( + bytes, + AudioSource::Url { + url: parsed.to_string(), + }, + ) + .await + } + async fn fetch_video_file( &self, path: PathBuf, @@ -330,6 +397,24 @@ impl MediaConnector { .await } + async fn fetch_audio_file(&self, path: PathBuf) -> Result, MediaConnectorError> { + let allowed_root = self + .allowed_local_media_path + .as_ref() + .ok_or_else(|| MediaConnectorError::DisallowedLocalPath(path.display().to_string()))?; + + let canonical = fs::canonicalize(&path).await?; + if !canonical.starts_with(allowed_root) { + return Err(MediaConnectorError::DisallowedLocalPath( + path.display().to_string(), + )); + } + + let bytes = fs::read(&canonical).await?; + self.decode_audio(bytes.into(), AudioSource::File { path: canonical }) + .await + } + fn ensure_domain_allowed(&self, url: &Url) -> Result<(), MediaConnectorError> { if let Some(allowed) = &self.allowed_domains { let host = url @@ -375,6 +460,18 @@ impl MediaConnector { ))) } + async fn decode_audio( + &self, + bytes: Bytes, + source: AudioSource, + ) -> Result, MediaConnectorError> { + let hash = crate::hasher::hash_audio(&bytes); + let decoded = decode_audio_mono_f32(&bytes) + .await + .map_err(|e| MediaConnectorError::AudioDecode(e.to_string()))?; + Ok(Arc::new(AudioClip::new(bytes, decoded, source, hash))) + } + async fn decode_video( &self, bytes: Bytes, @@ -406,18 +503,68 @@ impl MediaConnector { let decoded = decode_video_frames(bytes.clone(), cfg).await?; let clip = match decoded { - DecodedVideoFrames::Images(frames) => VideoClip::new(frames, bytes, source, hash), - DecodedVideoFrames::Rgb(rgb_video) => { - VideoClip::new_rgb(rgb_video, bytes, source, hash) + DecodedVideoFrames::Images { frames, sample_fps } => { + VideoClip::new_with_sample_fps(frames, bytes, source, hash, sample_fps) + } + DecodedVideoFrames::Rgb { video, sample_fps } => { + VideoClip::new_rgb_with_sample_fps(video, bytes, source, hash, sample_fps) } }; Ok(Arc::new(clip)) } } +async fn collect_http_body_with_limit( + mut response: reqwest::Response, + limit: usize, + media: &'static str, +) -> Result { + if response + .content_length() + .is_some_and(|length| length > limit as u64) + { + return Err(MediaConnectorError::PayloadTooLarge { media, limit }); + } + + let mut body = BytesMut::new(); + while let Some(chunk) = response.chunk().await? { + checked_payload_length(body.len(), chunk.len(), limit, media)?; + body.extend_from_slice(&chunk); + } + Ok(body.freeze()) +} + +fn checked_payload_length( + current: usize, + additional: usize, + limit: usize, + media: &'static str, +) -> Result { + current + .checked_add(additional) + .filter(|length| *length <= limit) + .ok_or(MediaConnectorError::PayloadTooLarge { media, limit }) +} + +fn audio_max_input_bytes() -> usize { + *AUDIO_MAX_INPUT_BYTES.get_or_init(|| { + std::env::var("SMG_AUDIO_MAX_INPUT_BYTES") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|bytes| *bytes > 0) + .unwrap_or(DEFAULT_AUDIO_MAX_INPUT_BYTES) + }) +} + enum DecodedVideoFrames { - Images(Vec), - Rgb(DecodedRgbVideo), + Images { + frames: Vec, + sample_fps: f32, + }, + Rgb { + video: DecodedRgbVideo, + sample_fps: f32, + }, } async fn decode_video_frames( @@ -903,7 +1050,14 @@ where MediaConnectorError::VideoDecode("OpenCV produced no RGB output".to_string()) })? .into_bytes(); - Ok(DecodedVideoFrames::Rgb(DecodedRgbVideo::new(data, frames))) + let sample_fps = effective_sample_fps( + (fps.is_finite() && fps > 0.0).then_some(total_frames as f64 / fps), + cfg, + ); + Ok(DecodedVideoFrames::Rgb { + video: DecodedRgbVideo::new(data, frames), + sample_fps, + }) } #[cfg(feature = "opencv-video")] @@ -1067,11 +1221,15 @@ async fn decode_video_with_ffmpeg( cfg: VideoFetchConfig, ) -> Result { if let Ok(metadata) = probe_video_metadata(input_path).await { + let sample_fps = effective_sample_fps(metadata.duration_seconds, cfg); let started = Instant::now(); match decode_video_with_ffmpeg_ppm(input_path, cfg, metadata).await { Ok(rgb_video) => { log_video_decode_backend_timing("ffmpeg_ppm_file", started, input_bytes, cfg, None); - return Ok(DecodedVideoFrames::Rgb(rgb_video)); + return Ok(DecodedVideoFrames::Rgb { + video: rgb_video, + sample_fps, + }); } Err(error) => { log_video_decode_backend_timing( @@ -1088,7 +1246,10 @@ async fn decode_video_with_ffmpeg( match decode_video_with_ffmpeg_raw(input_path, cfg, metadata).await { Ok(rgb_video) => { log_video_decode_backend_timing("ffmpeg_raw_file", started, input_bytes, cfg, None); - return Ok(DecodedVideoFrames::Rgb(rgb_video)); + return Ok(DecodedVideoFrames::Rgb { + video: rgb_video, + sample_fps, + }); } Err(error) => { log_video_decode_backend_timing( @@ -1104,9 +1265,9 @@ async fn decode_video_with_ffmpeg( let started = Instant::now(); match decode_video_with_ffmpeg_png(input_path, cfg).await { - Ok(frames) => { + Ok((frames, sample_fps)) => { log_video_decode_backend_timing("ffmpeg_png_file", started, input_bytes, cfg, None); - Ok(DecodedVideoFrames::Images(frames)) + Ok(DecodedVideoFrames::Images { frames, sample_fps }) } Err(error) => { log_video_decode_backend_timing( @@ -1358,8 +1519,8 @@ async fn decode_video_with_ffmpeg_raw( async fn decode_video_with_ffmpeg_png( input_path: &std::path::Path, cfg: VideoFetchConfig, -) -> Result, MediaConnectorError> { - let fps_filter = fps_filter_for_video(input_path, cfg).await; +) -> Result<(Vec, f32), MediaConnectorError> { + let (fps_filter, sample_fps) = sampling_filter_for_video(input_path, cfg).await; let max_frames = cfg.max_frames.to_string(); let output_limit = video_max_decoded_bytes().to_string(); let mut command = Command::new("ffmpeg"); @@ -1405,7 +1566,7 @@ async fn decode_video_with_ffmpeg_png( "ffmpeg produced no frames".to_string(), )); } - Ok(frames) + Ok((frames, sample_fps)) } #[derive(Debug, Clone, Copy)] @@ -1415,21 +1576,44 @@ struct VideoMetadata { duration_seconds: Option, } +#[derive(Debug, Clone, Copy)] +struct ProbedVideoInfo { + width: Option, + height: Option, + duration_seconds: Option, +} + async fn probe_video_metadata( input_path: &std::path::Path, ) -> Result { + let info = probe_video_info(input_path).await?; + let width = info.width.ok_or_else(|| { + MediaConnectorError::VideoDecode("ffprobe did not return video width".to_string()) + })?; + let height = info.height.ok_or_else(|| { + MediaConnectorError::VideoDecode("ffprobe did not return video height".to_string()) + })?; + Ok(VideoMetadata { + width, + height, + duration_seconds: info.duration_seconds, + }) +} + +async fn probe_video_info( + input_path: &std::path::Path, +) -> Result { let mut command = Command::new("ffprobe"); command .args([ "-v", "error", - "-nostdin", "-select_streams", "v:0", "-show_entries", - "stream=width,height:format=duration", + "stream=width,height,duration,duration_ts,time_base:format=duration", "-of", - "default=noprint_wrappers=1", + "json", ]) .arg(input_path); let output = run_video_command_output(command, "ffprobe").await?; @@ -1441,35 +1625,72 @@ async fn probe_video_metadata( ))); } - let stdout = String::from_utf8_lossy(&output.stdout); - let mut width = None; - let mut height = None; - let mut duration_seconds = None; - for line in stdout.lines() { - let Some((key, value)) = line.split_once('=') else { - continue; - }; - match key { - "width" => width = value.parse::().ok(), - "height" => height = value.parse::().ok(), - "duration" if value != "N/A" => duration_seconds = value.parse::().ok(), - _ => {} - } - } + parse_ffprobe_video_info(&output.stdout) +} - let width = width.ok_or_else(|| { - MediaConnectorError::VideoDecode("ffprobe did not return video width".to_string()) - })?; - let height = height.ok_or_else(|| { - MediaConnectorError::VideoDecode("ffprobe did not return video height".to_string()) +fn parse_ffprobe_video_info(stdout: &[u8]) -> Result { + let probe: serde_json::Value = serde_json::from_slice(stdout).map_err(|error| { + MediaConnectorError::VideoDecode(format!("failed to parse ffprobe output: {error}")) })?; - Ok(VideoMetadata { + let video_stream = probe + .get("streams") + .and_then(serde_json::Value::as_array) + .and_then(|streams| streams.first()); + + let width = video_stream + .and_then(|stream| stream.get("width")) + .and_then(json_u32); + let height = video_stream + .and_then(|stream| stream.get("height")) + .and_then(json_u32); + let stream_duration = video_stream + .and_then(|stream| stream.get("duration")) + .and_then(json_positive_f64); + let stream_time_base_duration = video_stream.and_then(|stream| { + let duration_ts = stream.get("duration_ts").and_then(json_positive_f64)?; + let time_base = stream + .get("time_base") + .and_then(serde_json::Value::as_str) + .and_then(parse_time_base)?; + let duration = duration_ts * time_base; + (duration.is_finite() && duration > 0.0).then_some(duration) + }); + let format_duration = probe + .get("format") + .and_then(|format| format.get("duration")) + .and_then(json_positive_f64); + + Ok(ProbedVideoInfo { width, height, - duration_seconds, + duration_seconds: stream_duration + .or(stream_time_base_duration) + .or(format_duration), }) } +fn json_u32(value: &serde_json::Value) -> Option { + value + .as_u64() + .and_then(|value| u32::try_from(value).ok()) + .or_else(|| value.as_str()?.parse::().ok()) +} + +fn json_positive_f64(value: &serde_json::Value) -> Option { + value + .as_f64() + .or_else(|| value.as_str()?.parse::().ok()) + .filter(|value| value.is_finite() && *value > 0.0) +} + +fn parse_time_base(value: &str) -> Option { + let (numerator, denominator) = value.split_once('/')?; + let numerator = numerator.parse::().ok()?; + let denominator = denominator.parse::().ok()?; + let time_base = numerator / denominator; + (time_base.is_finite() && time_base > 0.0).then_some(time_base) +} + fn fps_filter_for_metadata(metadata: VideoMetadata, cfg: VideoFetchConfig) -> String { if let Some(duration) = metadata.duration_seconds { if let Some(filter) = fps_filter_for_duration(duration, cfg) { @@ -1491,6 +1712,19 @@ fn expected_sampled_frame_count(metadata: VideoMetadata, cfg: VideoFetchConfig) cfg.max_frames } +fn effective_sample_fps(duration_seconds: Option, cfg: VideoFetchConfig) -> f32 { + duration_seconds + .filter(|duration| duration.is_finite() && *duration > 0.0) + .map(|duration| { + let target_frames = (duration * cfg.sample_fps as f64) + .round() + .clamp(cfg.min_frames as f64, cfg.max_frames as f64); + (target_frames / duration) as f32 + }) + .filter(|fps| fps.is_finite() && *fps > 0.0) + .unwrap_or(cfg.sample_fps) +} + fn fps_filter_for_duration(duration: f64, cfg: VideoFetchConfig) -> Option { if !duration.is_finite() || duration <= 0.0 { return None; @@ -1502,38 +1736,27 @@ fn fps_filter_for_duration(duration: f64, cfg: VideoFetchConfig) -> Option String { +async fn sampling_filter_for_video( + input_path: &std::path::Path, + cfg: VideoFetchConfig, +) -> (String, f32) { if let Ok(duration) = probe_video_duration_seconds(input_path).await { if let Some(filter) = fps_filter_for_duration(duration, cfg) { - return filter; + return (filter, effective_sample_fps(Some(duration), cfg)); } } - format!("fps={}", cfg.sample_fps) + (format!("fps={}", cfg.sample_fps), cfg.sample_fps) } async fn probe_video_duration_seconds( input_path: &std::path::Path, ) -> Result { - let mut command = Command::new("ffprobe"); - command - .args([ - "-v", - "error", - "-nostdin", - "-show_entries", - "format=duration", - "-of", - "default=noprint_wrappers=1:nokey=1", - ]) - .arg(input_path); - match run_video_command_output(command, "ffprobe").await { - Ok(output) if output.status.success() => { - let stdout = String::from_utf8_lossy(&output.stdout); - stdout.trim().parse::().map_err(|err| { - MediaConnectorError::VideoDecode(format!("failed to parse ffprobe duration: {err}")) - }) - } + match probe_video_info(input_path).await { + Ok(ProbedVideoInfo { + duration_seconds: Some(duration), + .. + }) => Ok(duration), Ok(_) | Err(_) => probe_video_duration_seconds_with_ffmpeg(input_path).await, } } @@ -1786,7 +2009,9 @@ fn skip_ppm_whitespace_and_comments(bytes: &[u8], pos: &mut usize) { #[cfg(test)] mod tests { use super::{ - parse_ffmpeg_duration_seconds, parse_ppm_stream, split_png_stream, video_temp_suffix, + checked_payload_length, effective_sample_fps, expected_sampled_frame_count, + fps_filter_for_metadata, parse_ffmpeg_duration_seconds, parse_ffprobe_video_info, + parse_ppm_stream, split_png_stream, video_temp_suffix, VideoFetchConfig, VideoMetadata, }; const TINY_PNG: &[u8] = &[ @@ -1816,6 +2041,53 @@ mod tests { assert_eq!(parse_ffmpeg_duration_seconds(stderr), Some(83.45)); } + #[test] + fn ffprobe_metadata_prefers_short_video_stream_over_long_container() { + let output = br#"{ + "streams": [{ + "width": 320, + "height": 240, + "duration": "1.000000", + "duration_ts": 30, + "time_base": "1/30" + }], + "format": {"duration": "120.000000"} + }"#; + let info = parse_ffprobe_video_info(output).expect("valid ffprobe output"); + assert_eq!(info.duration_seconds, Some(1.0)); + + let cfg = VideoFetchConfig { + min_frames: 4, + max_frames: 8, + sample_fps: 2.0, + }; + let metadata = VideoMetadata { + width: info.width.expect("video width"), + height: info.height.expect("video height"), + duration_seconds: info.duration_seconds, + }; + assert_eq!(expected_sampled_frame_count(metadata, cfg), 4); + assert_eq!(fps_filter_for_metadata(metadata, cfg), "fps=4.000000"); + } + + #[test] + fn ffprobe_metadata_uses_stream_time_base_before_container_duration() { + let output = br#"{ + "streams": [{ + "width": "640", + "height": "360", + "duration": "N/A", + "duration_ts": 45, + "time_base": "1/30" + }], + "format": {"duration": "90.000000"} + }"#; + let info = parse_ffprobe_video_info(output).expect("valid ffprobe output"); + assert_eq!(info.width, Some(640)); + assert_eq!(info.height, Some(360)); + assert_eq!(info.duration_seconds, Some(1.5)); + } + #[test] fn detects_video_temp_suffix_from_container_header() { let mut mp4 = vec![0; 12]; @@ -1843,6 +2115,20 @@ mod tests { assert_eq!(frames[1].height(), 2); } + #[test] + fn effective_sample_fps_tracks_min_and_max_frame_clamps() { + let cfg = VideoFetchConfig { + min_frames: 4, + max_frames: 8, + sample_fps: 2.0, + }; + + assert_eq!(effective_sample_fps(Some(1.0), cfg), 4.0); + assert!((effective_sample_fps(Some(10.0), cfg) - 0.8).abs() < 1e-6); + assert_eq!(effective_sample_fps(Some(3.0), cfg), 2.0); + assert_eq!(effective_sample_fps(None, cfg), 2.0); + } + #[test] fn rejects_truncated_ppm_stream() { assert!(parse_ppm_stream(b"P6\n2 1\n255\n\x01\x02").is_err()); @@ -1865,10 +2151,17 @@ mod tests { assert!(parse_ppm_stream(b"P6\n4294967295 4294967295\n255\n").is_err()); } + #[test] + fn enforces_http_media_payload_limit() { + assert_eq!(checked_payload_length(4, 4, 8, "audio").unwrap(), 8); + assert!(checked_payload_length(4, 5, 8, "audio").is_err()); + assert!(checked_payload_length(usize::MAX, 1, usize::MAX, "audio").is_err()); + } + #[cfg(feature = "opencv-video")] #[test] fn opencv_sampling_preserves_min_frames_for_short_clips() { - let cfg = super::VideoFetchConfig { + let cfg = VideoFetchConfig { min_frames: 4, max_frames: 8, sample_fps: 2.0, diff --git a/crates/multimodal/src/registry/kimi_k25.rs b/crates/multimodal/src/registry/kimi_k25.rs index 423bf0cf0..2abbff858 100644 --- a/crates/multimodal/src/registry/kimi_k25.rs +++ b/crates/multimodal/src/registry/kimi_k25.rs @@ -3,9 +3,9 @@ use std::collections::HashMap; use serde_json::{json, Value}; use crate::{ + encoder_inputs::PreprocessedEncoderInputs, registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, types::{FieldLayout, Modality, PromptReplacement, TokenId}, - vision::processor::PreprocessedEncoderInputs, }; pub(super) struct KimiK25VisionSpec; diff --git a/crates/multimodal/src/registry/llama4.rs b/crates/multimodal/src/registry/llama4.rs index 8389466c5..3e04058a1 100644 --- a/crates/multimodal/src/registry/llama4.rs +++ b/crates/multimodal/src/registry/llama4.rs @@ -3,9 +3,9 @@ use std::collections::HashMap; use serde_json::{json, Value}; use crate::{ + encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}, registry::{ModelMetadata, ModelProcessorSpec, RegistryResult}, types::{FieldLayout, Modality, PromptReplacement, TokenId}, - vision::processor::{ModelSpecificValue, PreprocessedEncoderInputs}, }; pub(super) struct Llama4Spec; diff --git a/crates/multimodal/src/registry/llava.rs b/crates/multimodal/src/registry/llava.rs index c8328136e..2bab71ab9 100644 --- a/crates/multimodal/src/registry/llava.rs +++ b/crates/multimodal/src/registry/llava.rs @@ -3,9 +3,9 @@ use std::collections::HashMap; use serde_json::{json, Value}; use crate::{ + encoder_inputs::PreprocessedEncoderInputs, registry::{ModelMetadata, ModelProcessorSpec, RegistryResult}, types::{FieldLayout, Modality, PromptReplacement, TokenId}, - vision::processor::PreprocessedEncoderInputs, }; pub(super) struct LlavaSpec; diff --git a/crates/multimodal/src/registry/mod.rs b/crates/multimodal/src/registry/mod.rs index e9a0e9108..6fb5a1902 100644 --- a/crates/multimodal/src/registry/mod.rs +++ b/crates/multimodal/src/registry/mod.rs @@ -2,6 +2,8 @@ mod kimi_k25; mod llama4; mod llava; mod phi3_v; +mod qwen3_asr; +mod qwen3_omni; mod qwen3_vl; mod qwen_vl; mod traits; @@ -11,6 +13,8 @@ use llama4::Llama4Spec; use llava::{LlavaNextSpec, LlavaSpec}; use once_cell::sync::Lazy; use phi3_v::Phi3VisionSpec; +use qwen3_asr::Qwen3AsrSpec; +use qwen3_omni::Qwen3OmniSpec; use qwen3_vl::Qwen3VLVisionSpec; use qwen_vl::QwenVLVisionSpec; // Re-export public API from traits. @@ -29,6 +33,8 @@ impl ModelRegistry { // LlavaNext must be registered before Llava so "llava_next" model_type matches first. LazySpec::new(|| Box::new(LlavaNextSpec)), LazySpec::new(|| Box::new(LlavaSpec)), + LazySpec::new(|| Box::new(Qwen3AsrSpec)), + LazySpec::new(|| Box::new(Qwen3OmniSpec)), // Qwen3-VL must be registered before QwenVL so "qwen3" matches first. LazySpec::new(|| Box::new(Qwen3VLVisionSpec)), LazySpec::new(|| Box::new(QwenVLVisionSpec)), @@ -78,8 +84,8 @@ pub(super) mod test_helpers { use once_cell::sync::Lazy; use crate::{ + encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}, types::ImageSize, - vision::processor::{ModelSpecificValue, PreprocessedEncoderInputs}, }; pub struct TestTokenizer { diff --git a/crates/multimodal/src/registry/phi3_v.rs b/crates/multimodal/src/registry/phi3_v.rs index a31009722..b5229d18d 100644 --- a/crates/multimodal/src/registry/phi3_v.rs +++ b/crates/multimodal/src/registry/phi3_v.rs @@ -3,9 +3,9 @@ use std::collections::HashMap; use serde_json::{json, Value}; use crate::{ + encoder_inputs::PreprocessedEncoderInputs, registry::{ModelMetadata, ModelProcessorSpec, RegistryResult}, types::{FieldLayout, Modality, PromptReplacement, TokenId}, - vision::processor::PreprocessedEncoderInputs, }; pub(super) struct Phi3VisionSpec; diff --git a/crates/multimodal/src/registry/qwen3_asr.rs b/crates/multimodal/src/registry/qwen3_asr.rs new file mode 100644 index 000000000..5d5ee6132 --- /dev/null +++ b/crates/multimodal/src/registry/qwen3_asr.rs @@ -0,0 +1,204 @@ +use std::collections::HashMap; + +use serde_json::{json, Value}; + +use crate::{ + encoder_inputs::PreprocessedEncoderInputs, + registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, + types::{EncoderFieldLayouts, FieldLayout, Modality, PromptReplacement, TokenId}, +}; + +const AUDIO_PAD_TOKEN: &str = "<|audio_pad|>"; + +pub(super) struct Qwen3AsrSpec; + +impl Qwen3AsrSpec { + fn audio_token_id(metadata: &ModelMetadata) -> RegistryResult { + metadata + .config_u32(&["thinker_config", "audio_token_id"]) + .or_else(|| metadata.config_u32(&["audio_token_id"])) + .map(|value| value as TokenId) + .map_or_else(|| metadata.token_id(AUDIO_PAD_TOKEN), Ok) + } +} + +impl ModelProcessorSpec for Qwen3AsrSpec { + fn name(&self) -> &'static str { + "qwen3_asr" + } + + fn matches(&self, metadata: &ModelMetadata) -> bool { + let model_id = metadata.model_id.to_ascii_lowercase(); + model_id.contains("qwen3-asr") + || model_id.contains("qwen3_asr") + || metadata + .config_model_type() + .is_some_and(|model_type| model_type == "qwen3_asr") + } + + fn placeholder_token(&self, metadata: &ModelMetadata) -> RegistryResult { + self.placeholder_token_for(metadata, Modality::Audio) + } + + fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { + self.placeholder_token_id_for(metadata, Modality::Audio) + } + + fn placeholder_token_for( + &self, + metadata: &ModelMetadata, + modality: Modality, + ) -> RegistryResult { + match modality { + Modality::Audio => { + let token_id = Self::audio_token_id(metadata)?; + match metadata.tokenizer.id_to_token(token_id as u32) { + Some(token) => Ok(token), + None => { + metadata.token_id(AUDIO_PAD_TOKEN)?; + Ok(AUDIO_PAD_TOKEN.to_string()) + } + } + } + Modality::Image | Modality::Video | Modality::ImageEmbeds => { + Err(ModelRegistryError::UnsupportedModality { + spec: self.name(), + modality, + }) + } + } + } + + fn placeholder_token_id_for( + &self, + metadata: &ModelMetadata, + modality: Modality, + ) -> RegistryResult { + match modality { + Modality::Audio => Self::audio_token_id(metadata), + Modality::Image | Modality::Video | Modality::ImageEmbeds => { + Err(ModelRegistryError::UnsupportedModality { + spec: self.name(), + modality, + }) + } + } + } + + fn modality_limits( + &self, + _metadata: &ModelMetadata, + ) -> RegistryResult> { + Ok(HashMap::from([(Modality::Audio, 10)])) + } + + fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok(json!({})) + } + + fn prompt_replacements( + &self, + metadata: &ModelMetadata, + preprocessed: &PreprocessedEncoderInputs, + ) -> RegistryResult> { + self.prompt_replacements_for(metadata, preprocessed, Modality::Audio) + } + + fn prompt_replacements_for( + &self, + metadata: &ModelMetadata, + preprocessed: &PreprocessedEncoderInputs, + modality: Modality, + ) -> RegistryResult> { + match modality { + Modality::Audio => { + let token_id = Self::audio_token_id(metadata)?; + let token = self.placeholder_token_for(metadata, Modality::Audio)?; + Ok(preprocessed + .feature_token_counts + .iter() + .map(|&count| { + PromptReplacement::repeated(Modality::Audio, &token, token_id, count) + }) + .collect()) + } + Modality::Image | Modality::Video | Modality::ImageEmbeds => { + Err(ModelRegistryError::UnsupportedModality { + spec: self.name(), + modality, + }) + } + } + } + + fn encoder_field_layouts_for(&self, modality: Modality) -> EncoderFieldLayouts { + match modality { + Modality::Audio => EncoderFieldLayouts::new( + FieldLayout::Batched, + HashMap::from([ + ("feature_attention_mask".to_string(), FieldLayout::Batched), + ("audio_feature_lengths".to_string(), FieldLayout::Batched), + ]), + ), + Modality::Image | Modality::Video | Modality::ImageEmbeds => { + EncoderFieldLayouts::default() + } + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::{ + registry::{test_helpers::*, ModelRegistry}, + types::ImageSize, + }; + + #[test] + fn asr_matches_and_expands_nested_audio_token() { + let tokenizer = TestTokenizer::new(&[(AUDIO_PAD_TOKEN, 151676)]); + let config = json!({ + "model_type": "qwen3_asr", + "thinker_config": {"audio_token_id": 151676} + }); + let metadata = ModelMetadata { + model_id: "Qwen/Qwen3-ASR-1.7B", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).unwrap(); + assert_eq!(spec.name(), "qwen3_asr"); + assert_eq!( + spec.placeholder_token(&metadata).unwrap(), + spec.placeholder_token_for(&metadata, Modality::Audio) + .unwrap() + ); + assert_eq!( + spec.placeholder_token_id(&metadata).unwrap(), + spec.placeholder_token_id_for(&metadata, Modality::Audio) + .unwrap() + ); + + let replacements = spec + .prompt_replacements_for( + &metadata, + &test_preprocessed_with_tokens(&[ImageSize::new(128, 100)], &[13]), + Modality::Audio, + ) + .unwrap(); + assert_eq!(replacements[0].tokens, vec![151676; 13]); + assert_eq!( + spec.encoder_field_layouts_for(Modality::Audio) + .encoder_input, + FieldLayout::Batched + ); + assert_eq!( + spec.modality_limits(&metadata).unwrap(), + HashMap::from([(Modality::Audio, 10)]) + ); + } +} diff --git a/crates/multimodal/src/registry/qwen3_omni.rs b/crates/multimodal/src/registry/qwen3_omni.rs new file mode 100644 index 000000000..2c4ecce68 --- /dev/null +++ b/crates/multimodal/src/registry/qwen3_omni.rs @@ -0,0 +1,309 @@ +use std::collections::HashMap; + +use serde_json::{json, Value}; + +use crate::{ + encoder_inputs::PreprocessedEncoderInputs, + registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, + types::{EncoderFieldLayouts, FieldLayout, Modality, PromptReplacement, TokenId}, +}; + +const IMAGE_PAD_TOKEN: &str = "<|image_pad|>"; +const VIDEO_PAD_TOKEN: &str = "<|video_pad|>"; +const AUDIO_PAD_TOKEN: &str = "<|audio_pad|>"; + +pub(super) struct Qwen3OmniSpec; + +impl Qwen3OmniSpec { + fn token_id(metadata: &ModelMetadata, field: &str) -> RegistryResult { + metadata + .config_u32(&["thinker_config", field]) + .or_else(|| metadata.config_u32(&[field])) + .map(|value| value as TokenId) + .ok_or_else(|| ModelRegistryError::MissingConfigField { + field: format!("thinker_config.{field}"), + }) + } + + fn token(metadata: &ModelMetadata, field: &str, fallback: &str) -> RegistryResult { + let token_id = Self::token_id(metadata, field)?; + if let Some(token) = metadata.tokenizer.id_to_token(token_id as u32) { + return Ok(token); + } + metadata.token_id(fallback)?; + Ok(fallback.to_string()) + } + + fn replacements( + metadata: &ModelMetadata, + preprocessed: &PreprocessedEncoderInputs, + modality: Modality, + field: &str, + fallback: &str, + ) -> RegistryResult> { + let token_id = Self::token_id(metadata, field)?; + let token = Self::token(metadata, field, fallback)?; + Ok(preprocessed + .feature_token_counts + .iter() + .map(|&count| PromptReplacement::repeated(modality, &token, token_id, count)) + .collect()) + } +} + +impl ModelProcessorSpec for Qwen3OmniSpec { + fn name(&self) -> &'static str { + "qwen3_omni" + } + + fn matches(&self, metadata: &ModelMetadata) -> bool { + let model_id = metadata.model_id.to_ascii_lowercase(); + model_id.contains("qwen3-omni") + || model_id.contains("qwen3_omni") + || metadata.config_model_type().is_some_and(|model_type| { + model_type == "qwen3_omni_moe" || model_type == "qwen3_omni_moe_thinker" + }) + } + + fn placeholder_token(&self, metadata: &ModelMetadata) -> RegistryResult { + self.placeholder_token_for(metadata, Modality::Image) + } + + fn placeholder_token_id(&self, metadata: &ModelMetadata) -> RegistryResult { + self.placeholder_token_id_for(metadata, Modality::Image) + } + + fn placeholder_token_for( + &self, + metadata: &ModelMetadata, + modality: Modality, + ) -> RegistryResult { + match modality { + Modality::Image => Self::token(metadata, "image_token_id", IMAGE_PAD_TOKEN), + Modality::Video => Self::token(metadata, "video_token_id", VIDEO_PAD_TOKEN), + Modality::Audio => Self::token(metadata, "audio_token_id", AUDIO_PAD_TOKEN), + Modality::ImageEmbeds => Err(ModelRegistryError::UnsupportedModality { + spec: self.name(), + modality, + }), + } + } + + fn placeholder_token_id_for( + &self, + metadata: &ModelMetadata, + modality: Modality, + ) -> RegistryResult { + match modality { + Modality::Image => Self::token_id(metadata, "image_token_id"), + Modality::Video => Self::token_id(metadata, "video_token_id"), + Modality::Audio => Self::token_id(metadata, "audio_token_id"), + Modality::ImageEmbeds => Err(ModelRegistryError::UnsupportedModality { + spec: self.name(), + modality, + }), + } + } + + fn modality_limits( + &self, + _metadata: &ModelMetadata, + ) -> RegistryResult> { + Ok(HashMap::from([ + (Modality::Image, 10), + (Modality::Video, 1), + (Modality::Audio, 10), + ])) + } + + fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok(json!({"use_audio_in_video": false})) + } + + fn prompt_replacements( + &self, + metadata: &ModelMetadata, + preprocessed: &PreprocessedEncoderInputs, + ) -> RegistryResult> { + self.prompt_replacements_for(metadata, preprocessed, Modality::Image) + } + + fn prompt_replacements_for( + &self, + metadata: &ModelMetadata, + preprocessed: &PreprocessedEncoderInputs, + modality: Modality, + ) -> RegistryResult> { + match modality { + Modality::Image => Self::replacements( + metadata, + preprocessed, + modality, + "image_token_id", + IMAGE_PAD_TOKEN, + ), + Modality::Video => Self::replacements( + metadata, + preprocessed, + modality, + "video_token_id", + VIDEO_PAD_TOKEN, + ), + Modality::Audio => Self::replacements( + metadata, + preprocessed, + modality, + "audio_token_id", + AUDIO_PAD_TOKEN, + ), + Modality::ImageEmbeds => Err(ModelRegistryError::UnsupportedModality { + spec: self.name(), + modality, + }), + } + } + + fn encoder_field_layouts_for(&self, modality: Modality) -> EncoderFieldLayouts { + match modality { + Modality::Image => EncoderFieldLayouts::new( + FieldLayout::flat("patches_per_image"), + HashMap::from([ + ("image_grid_thw".to_string(), FieldLayout::Batched), + ("patches_per_image".to_string(), FieldLayout::Batched), + ]), + ), + Modality::Video => EncoderFieldLayouts::new( + FieldLayout::flat("patches_per_video"), + HashMap::from([ + ("video_grid_thw".to_string(), FieldLayout::Batched), + ("patches_per_video".to_string(), FieldLayout::Batched), + ("video_second_per_grid".to_string(), FieldLayout::Batched), + ]), + ), + Modality::Audio => EncoderFieldLayouts::new( + FieldLayout::Batched, + HashMap::from([ + ("feature_attention_mask".to_string(), FieldLayout::Batched), + ("audio_feature_lengths".to_string(), FieldLayout::Batched), + ]), + ), + Modality::ImageEmbeds => EncoderFieldLayouts::default(), + } + } + + fn keep_on_cpu_keys(&self) -> Vec { + vec!["image_grid_thw".to_string(), "video_grid_thw".to_string()] + } + + fn keep_on_cpu_keys_for(&self, modality: Modality) -> Vec { + match modality { + Modality::Image => vec!["image_grid_thw".to_string()], + Modality::Video => vec!["video_grid_thw".to_string()], + Modality::Audio | Modality::ImageEmbeds => vec![], + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::{ + registry::{test_helpers::*, ModelRegistry}, + types::ImageSize, + }; + + fn omni_tokenizer() -> TestTokenizer { + TestTokenizer::new(&[ + (AUDIO_PAD_TOKEN, 151675), + (IMAGE_PAD_TOKEN, 151655), + (VIDEO_PAD_TOKEN, 151656), + ]) + } + + #[test] + fn omni_accepts_mixed_modalities_and_uses_nested_tokens() { + let tokenizer = omni_tokenizer(); + let config = json!({ + "model_type": "qwen3_omni_moe", + "thinker_config": { + "audio_token_id": 151675, + "image_token_id": 151655, + "video_token_id": 151656 + } + }); + let metadata = ModelMetadata { + model_id: "Qwen/Qwen3-Omni-30B-A3B-Thinking", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).unwrap(); + assert_eq!(spec.name(), "qwen3_omni"); + assert_eq!( + spec.placeholder_token(&metadata).unwrap(), + spec.placeholder_token_for(&metadata, Modality::Image) + .unwrap() + ); + assert_eq!( + spec.placeholder_token_id(&metadata).unwrap(), + spec.placeholder_token_id_for(&metadata, Modality::Image) + .unwrap() + ); + spec.validate_media_request( + &metadata, + &[ + (Modality::Image, 2), + (Modality::Video, 1), + (Modality::Audio, 2), + ], + ) + .unwrap(); + + for (modality, expected) in [ + (Modality::Image, 151655), + (Modality::Video, 151656), + (Modality::Audio, 151675), + ] { + let replacements = spec + .prompt_replacements_for( + &metadata, + &test_preprocessed_with_tokens(&[ImageSize::new(32, 32)], &[3]), + modality, + ) + .unwrap(); + assert_eq!(replacements[0].tokens, vec![expected; 3]); + } + } + + #[test] + fn omni_layouts_are_modality_specific() { + let image = Qwen3OmniSpec.encoder_field_layouts_for(Modality::Image); + assert_eq!(image.encoder_input, FieldLayout::flat("patches_per_image")); + assert!(image.model_specific.contains_key("image_grid_thw")); + + let video = Qwen3OmniSpec.encoder_field_layouts_for(Modality::Video); + assert_eq!(video.encoder_input, FieldLayout::flat("patches_per_video")); + assert!(video.model_specific.contains_key("video_grid_thw")); + + let audio = Qwen3OmniSpec.encoder_field_layouts_for(Modality::Audio); + assert_eq!(audio.encoder_input, FieldLayout::Batched); + assert!(audio.model_specific.contains_key("feature_attention_mask")); + } + + #[test] + fn omni_keep_on_cpu_keys_are_modality_specific() { + assert_eq!( + Qwen3OmniSpec.keep_on_cpu_keys_for(Modality::Image), + vec!["image_grid_thw"] + ); + assert_eq!( + Qwen3OmniSpec.keep_on_cpu_keys_for(Modality::Video), + vec!["video_grid_thw"] + ); + assert!(Qwen3OmniSpec + .keep_on_cpu_keys_for(Modality::Audio) + .is_empty()); + } +} diff --git a/crates/multimodal/src/registry/qwen3_vl.rs b/crates/multimodal/src/registry/qwen3_vl.rs index 4d832157b..3e2ff0342 100644 --- a/crates/multimodal/src/registry/qwen3_vl.rs +++ b/crates/multimodal/src/registry/qwen3_vl.rs @@ -4,9 +4,9 @@ use llm_tokenizer::Encoding; use serde_json::{json, Value}; use crate::{ + encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}, registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, types::{FieldLayout, Modality, PromptReplacement, TokenId}, - vision::processor::{ModelSpecificValue, PreprocessedEncoderInputs}, }; pub(super) struct Qwen3VLVisionSpec; @@ -293,6 +293,7 @@ impl ModelProcessorSpec for Qwen3VLVisionSpec { ("patches_per_image".to_string(), FieldLayout::Batched), ("video_grid_thw".to_string(), FieldLayout::Batched), ("patches_per_video".to_string(), FieldLayout::Batched), + ("video_second_per_grid".to_string(), FieldLayout::Batched), ]) } @@ -306,9 +307,9 @@ mod tests { use serde_json::json; use crate::{ + encoder_inputs::ModelSpecificValue, registry::{test_helpers::*, ModelMetadata, ModelRegistry}, types::ImageSize, - vision::processor::ModelSpecificValue, }; #[test] diff --git a/crates/multimodal/src/registry/qwen_vl.rs b/crates/multimodal/src/registry/qwen_vl.rs index ce5823459..e991a6a94 100644 --- a/crates/multimodal/src/registry/qwen_vl.rs +++ b/crates/multimodal/src/registry/qwen_vl.rs @@ -3,9 +3,9 @@ use std::collections::HashMap; use serde_json::{json, Value}; use crate::{ + encoder_inputs::PreprocessedEncoderInputs, registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, types::{FieldLayout, Modality, PromptReplacement, TokenId}, - vision::processor::PreprocessedEncoderInputs, }; pub(super) struct QwenVLVisionSpec; diff --git a/crates/multimodal/src/registry/traits.rs b/crates/multimodal/src/registry/traits.rs index 6f5c6dc47..9b2a33f19 100644 --- a/crates/multimodal/src/registry/traits.rs +++ b/crates/multimodal/src/registry/traits.rs @@ -5,11 +5,11 @@ use serde_json::Value; use thiserror::Error; use crate::{ - types::{FieldLayout, Modality, PromptReplacement, TokenId}, - vision::processor::PreprocessedEncoderInputs, + encoder_inputs::PreprocessedEncoderInputs, + types::{EncoderFieldLayouts, FieldLayout, Modality, PromptReplacement, TokenId}, }; -#[derive(Debug, Error)] +#[derive(Debug, Error, PartialEq, Eq)] pub enum ModelRegistryError { #[error("unsupported model: {0}")] UnsupportedModel(String), @@ -22,6 +22,18 @@ pub enum ModelRegistryError { spec: &'static str, modality: Modality, }, + #[error("model spec {spec} supports at most {limit} {modality} inputs; got {requested}")] + ModalityLimitExceeded { + spec: &'static str, + modality: Modality, + limit: usize, + requested: usize, + }, + #[error("modality {modality} appears more than once in the request for model spec {spec}")] + DuplicateModality { + spec: &'static str, + modality: Modality, + }, } pub type RegistryResult = Result; @@ -93,6 +105,51 @@ pub trait ModelProcessorSpec: Send + Sync { } fn modality_limits(&self, metadata: &ModelMetadata) -> RegistryResult>; + + /// Validate the active modalities and item counts in one media request. + /// + /// Any subset of the modalities declared by [`Self::modality_limits`] is + /// accepted. Each modality may appear once in `requested`; zero-count + /// entries are ignored. + fn validate_media_request( + &self, + metadata: &ModelMetadata, + requested: &[(Modality, usize)], + ) -> RegistryResult<()> { + let limits = self.modality_limits(metadata)?; + let mut active = Vec::with_capacity(requested.len()); + + for &(modality, count) in requested { + if count == 0 { + continue; + } + if active.contains(&modality) { + return Err(ModelRegistryError::DuplicateModality { + spec: self.name(), + modality, + }); + } + active.push(modality); + + let Some(&limit) = limits.get(&modality) else { + return Err(ModelRegistryError::UnsupportedModality { + spec: self.name(), + modality, + }); + }; + if count > limit { + return Err(ModelRegistryError::ModalityLimitExceeded { + spec: self.name(), + modality, + limit, + requested: count, + }); + } + } + + Ok(()) + } + fn processor_kwargs(&self, metadata: &ModelMetadata) -> RegistryResult; /// Compute per-media prompt replacement token sequences. /// @@ -129,6 +186,15 @@ pub trait ModelProcessorSpec: Send + Sync { HashMap::from([("pixel_values".to_string(), FieldLayout::Batched)]) } + /// Declare the neutral primary/side-tensor layout contract for one modality. + /// + /// The default converts the legacy HF/vLLM-shaped field map so existing + /// vision specs remain source-compatible. New multimodal specs should + /// override this method and keep backend-specific field names at adapters. + fn encoder_field_layouts_for(&self, _modality: Modality) -> EncoderFieldLayouts { + EncoderFieldLayouts::from_legacy_fields(self.field_layouts()) + } + /// Tensor keys that should remain on CPU (not transferred to GPU). /// /// In vLLM, certain model-specific tensors are marked `keep_on_cpu=True` @@ -138,4 +204,116 @@ pub trait ModelProcessorSpec: Send + Sync { fn keep_on_cpu_keys(&self) -> Vec { vec![] } + + /// Tensor keys that should remain on CPU for one modality. + /// + /// The default preserves the legacy model-wide declaration. + fn keep_on_cpu_keys_for(&self, _modality: Modality) -> Vec { + self.keep_on_cpu_keys() + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::registry::test_helpers::TestTokenizer; + + struct TestSpec; + + impl ModelProcessorSpec for TestSpec { + fn name(&self) -> &'static str { + "test" + } + + fn matches(&self, _metadata: &ModelMetadata) -> bool { + true + } + + fn placeholder_token(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok("".to_string()) + } + + fn placeholder_token_id(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok(1) + } + + fn modality_limits( + &self, + _metadata: &ModelMetadata, + ) -> RegistryResult> { + Ok(HashMap::from([(Modality::Image, 2), (Modality::Audio, 1)])) + } + + fn processor_kwargs(&self, _metadata: &ModelMetadata) -> RegistryResult { + Ok(json!({})) + } + + fn prompt_replacements( + &self, + _metadata: &ModelMetadata, + _preprocessed: &PreprocessedEncoderInputs, + ) -> RegistryResult> { + Ok(vec![]) + } + } + + fn validate( + spec: &dyn ModelProcessorSpec, + requested: &[(Modality, usize)], + ) -> RegistryResult<()> { + let tokenizer = TestTokenizer::new(&[]); + let config = json!({}); + let metadata = ModelMetadata { + model_id: "test-model", + tokenizer: &tokenizer, + config: &config, + }; + spec.validate_media_request(&metadata, requested) + } + + #[test] + fn validation_accepts_any_declared_modality_subset() { + assert_eq!(validate(&TestSpec, &[(Modality::Image, 2)]), Ok(())); + assert_eq!( + validate(&TestSpec, &[(Modality::Image, 1), (Modality::Audio, 1)]), + Ok(()) + ); + } + + #[test] + fn validation_rejects_undeclared_modality() { + assert_eq!( + validate(&TestSpec, &[(Modality::Video, 1)]), + Err(ModelRegistryError::UnsupportedModality { + spec: "test", + modality: Modality::Video, + }) + ); + } + + #[test] + fn validation_rejects_count_above_limit() { + assert_eq!( + validate(&TestSpec, &[(Modality::Image, 3)]), + Err(ModelRegistryError::ModalityLimitExceeded { + spec: "test", + modality: Modality::Image, + limit: 2, + requested: 3, + }) + ); + } + + #[test] + fn validation_rejects_duplicate_modality_counts() { + assert_eq!( + validate(&TestSpec, &[(Modality::Image, 1), (Modality::Image, 1)]), + Err(ModelRegistryError::DuplicateModality { + spec: "test", + modality: Modality::Image, + }) + ); + } } diff --git a/crates/multimodal/src/tracker.rs b/crates/multimodal/src/tracker.rs index dc23dfe97..10fd4386f 100644 --- a/crates/multimodal/src/tracker.rs +++ b/crates/multimodal/src/tracker.rs @@ -58,6 +58,20 @@ impl AsyncMultiModalTracker { MediaContentPart::ImageEmbeds { .. } => { return Err(MultiModalError::UnsupportedContent("image_embeds")); } + MediaContentPart::AudioUrl { url, uuid } => { + let source = match url::Url::parse(&url) { + Ok(parsed) if parsed.scheme() == "data" => MediaSource::DataUrl(url), + _ => MediaSource::Url(url), + }; + self.enqueue_audio(source, uuid); + } + MediaContentPart::AudioData { + data, + mime_type: _, + uuid, + } => { + self.enqueue_audio(MediaSource::InlineBytes(data), uuid); + } MediaContentPart::VideoUrl { url, uuid } => { let source = match url::Url::parse(&url) { Ok(parsed) if parsed.scheme() == "data" => MediaSource::DataUrl(url), @@ -130,4 +144,21 @@ impl AsyncMultiModalTracker { self.pending.entry(modality).or_default().push(handle); } + + fn enqueue_audio(&mut self, source: MediaSource, uuid: Option) { + let modality = Modality::Audio; + self.uuids.entry(modality).or_default().push(uuid); + + let connector = Arc::clone(&self.media_connector); + #[expect( + clippy::disallowed_methods, + reason = "spawn handle is stored in self.pending and awaited in finalize(); fire-and-forget is intentional for concurrent media fetching" + )] + let handle = tokio::spawn(async move { + let clip = connector.fetch_audio(source).await?; + Ok(TrackedMedia::Audio(clip)) + }); + + self.pending.entry(modality).or_default().push(handle); + } } diff --git a/crates/multimodal/src/types.rs b/crates/multimodal/src/types.rs index 47a0b4e1f..fe49170e8 100644 --- a/crates/multimodal/src/types.rs +++ b/crates/multimodal/src/types.rs @@ -4,6 +4,8 @@ use image::{DynamicImage, RgbImage}; use serde::{Deserialize, Serialize}; use serde_json::Value; +use crate::audio::DecodedAudio; + /// Supported multimodal modalities. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -63,6 +65,18 @@ pub enum MediaContentPart { #[serde(skip_serializing_if = "Option::is_none")] uuid: Option, }, + AudioUrl { + url: String, + #[serde(skip_serializing_if = "Option::is_none")] + uuid: Option, + }, + AudioData { + data: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + mime_type: Option, + #[serde(skip_serializing_if = "Option::is_none")] + uuid: Option, + }, VideoUrl { url: String, #[serde(skip_serializing_if = "Option::is_none")] @@ -87,6 +101,16 @@ pub enum ImageSource { File { path: PathBuf }, } +/// Audio source metadata (useful for hashing & tracing). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum AudioSource { + Url { url: String }, + DataUrl, + InlineBytes, + File { path: PathBuf }, +} + /// Video source metadata (useful for hashing & tracing). #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] @@ -108,11 +132,23 @@ pub struct ImageFrame { pub hash: String, } +/// Decoded audio payload captured by the media connector. +#[derive(Debug, Clone)] +pub struct AudioClip { + pub raw_bytes: bytes::Bytes, + pub decoded: DecodedAudio, + pub source: AudioSource, + /// Blake3 hex-digest of raw_bytes, computed at decode time. + pub hash: String, +} + /// Decoded video payload captured by the media connector. #[derive(Debug, Clone)] pub struct VideoClip { pub frames: Vec, pub rgb_video: Option, + /// Effective frame rate after connector-side sampling and frame-count clamps. + pub sample_fps: f32, pub raw_bytes: bytes::Bytes, pub source: VideoSource, /// Blake3 hex-digest of raw_bytes, computed at decode time. @@ -199,10 +235,21 @@ impl VideoClip { raw_bytes: bytes::Bytes, source: VideoSource, hash: String, + ) -> Self { + Self::new_with_sample_fps(frames, raw_bytes, source, hash, 2.0) + } + + pub fn new_with_sample_fps( + frames: Vec, + raw_bytes: bytes::Bytes, + source: VideoSource, + hash: String, + sample_fps: f32, ) -> Self { Self { frames, rgb_video: None, + sample_fps, raw_bytes, source, hash, @@ -214,10 +261,21 @@ impl VideoClip { raw_bytes: bytes::Bytes, source: VideoSource, hash: String, + ) -> Self { + Self::new_rgb_with_sample_fps(rgb_video, raw_bytes, source, hash, 2.0) + } + + pub fn new_rgb_with_sample_fps( + rgb_video: DecodedRgbVideo, + raw_bytes: bytes::Bytes, + source: VideoSource, + hash: String, + sample_fps: f32, ) -> Self { Self { frames: Vec::new(), rgb_video: Some(rgb_video), + sample_fps, raw_bytes, source, hash, @@ -232,6 +290,10 @@ impl VideoClip { self.rgb_video.as_ref() } + pub fn sample_fps(&self) -> f32 { + self.sample_fps + } + pub fn materialized_frames(&self) -> Result, String> { if !self.frames.is_empty() { return Ok(self.frames.clone()); @@ -251,6 +313,34 @@ impl VideoClip { } } +impl AudioClip { + pub fn new( + raw_bytes: bytes::Bytes, + decoded: DecodedAudio, + source: AudioSource, + hash: String, + ) -> Self { + Self { + raw_bytes, + decoded, + source, + hash, + } + } + + pub fn raw_bytes(&self) -> &[u8] { + &self.raw_bytes + } + + pub fn decoded(&self) -> &DecodedAudio { + &self.decoded + } + + pub fn source(&self) -> &AudioSource { + &self.source + } +} + impl ImageFrame { pub fn new( image: DynamicImage, @@ -289,9 +379,9 @@ impl ImageFrame { #[derive(Debug, Clone)] pub enum TrackedMedia { Image(Arc), + Audio(Arc), Video(Arc), /// Placeholder variants for future modalities. - Audio, Embeddings, } @@ -302,8 +392,8 @@ pub type TokenId = i32; /// Declares how a multimodal tensor's first dimension maps to media items. /// -/// Used by [`ModelProcessorSpec::field_layouts`] to tell the backend how to -/// split tensors for per-item scheduling (vLLM `MultiModalFieldConfig`). +/// Used by [`crate::registry::ModelProcessorSpec::encoder_field_layouts_for`] to tell the backend +/// how to split tensors for per-item scheduling (vLLM `MultiModalFieldConfig`). #[derive(Debug, Clone, PartialEq, Eq)] pub enum FieldLayout { /// First dimension equals number of media items (one slice per item). @@ -322,6 +412,43 @@ impl FieldLayout { } } +/// Layout contract for one modality's encoder inputs. +/// +/// The primary encoder input is transported independently from named, +/// model-specific side tensors. Keeping its layout typed avoids leaking a +/// vision-specific field name into audio and other modality processors. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EncoderFieldLayouts { + pub encoder_input: FieldLayout, + pub model_specific: HashMap, +} + +impl EncoderFieldLayouts { + pub fn new(encoder_input: FieldLayout, model_specific: HashMap) -> Self { + Self { + encoder_input, + model_specific, + } + } + + /// Convert the legacy HF/vLLM-shaped field map into the neutral contract. + /// + /// Existing vision specs use `pixel_values` for the primary encoder input. + /// New specs should construct [`Self`] directly instead. + pub fn from_legacy_fields(mut fields: HashMap) -> Self { + let encoder_input = fields + .remove("pixel_values") + .unwrap_or(FieldLayout::Batched); + Self::new(encoder_input, fields) + } +} + +impl Default for EncoderFieldLayouts { + fn default() -> Self { + Self::new(FieldLayout::Batched, HashMap::new()) + } +} + #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] pub struct ImageSize { pub width: u32, @@ -409,4 +536,24 @@ mod tests { let rep = PromptReplacement::repeated(Modality::Image, "", 100, 3); assert_eq!(rep.tokens, vec![100, 100, 100]); } + + #[test] + fn legacy_encoder_fields_are_split_into_typed_layouts() { + let layouts = EncoderFieldLayouts::from_legacy_fields(HashMap::from([ + ( + "pixel_values".to_string(), + FieldLayout::flat("patches_per_image"), + ), + ("image_grid_thw".to_string(), FieldLayout::Batched), + ])); + + assert_eq!( + layouts.encoder_input, + FieldLayout::flat("patches_per_image") + ); + assert_eq!( + layouts.model_specific, + HashMap::from([("image_grid_thw".to_string(), FieldLayout::Batched)]) + ); + } } diff --git a/crates/multimodal/src/vision/mod.rs b/crates/multimodal/src/vision/mod.rs index 83143818b..56cd062fd 100644 --- a/crates/multimodal/src/vision/mod.rs +++ b/crates/multimodal/src/vision/mod.rs @@ -9,9 +9,12 @@ //! //! - `transforms`: Core image transformations (resize, normalize, crop, etc.) //! - `preprocessor_config`: HuggingFace config parsing -//! - `processor`: Trait and output types for processors +//! - `processor`: Vision processor trait and registry //! - `processors`: Model-specific implementations (LLaVA, Qwen-VL, etc.) //! +//! Modality-neutral encoder outputs live in [`crate::encoder_inputs`], while +//! shared errors live in [`crate::error`]. +//! //! # Usage //! //! ```rust,ignore @@ -36,13 +39,15 @@ pub mod processors; pub(crate) mod scratch; pub mod transforms; -// Re-export commonly used types +// Re-export commonly used types, including compatibility paths for shared +// preprocessing outputs. pub use preprocessor_config::PreProcessorConfig; pub use processor::{ ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor, VisionProcessorRegistry, }; pub use processors::{ Llama4VisionProcessor, LlavaNextProcessor, LlavaProcessor, Phi3VisionProcessor, - Phi4VisionProcessor, PixtralProcessor, Qwen2VLProcessor, Qwen3VLProcessor, + Phi4VisionProcessor, PixtralProcessor, Qwen2VLProcessor, Qwen3OmniVisionProcessor, + Qwen3VLProcessor, }; pub use transforms::TransformError; diff --git a/crates/multimodal/src/vision/processor.rs b/crates/multimodal/src/vision/processor.rs index 000a1d505..de025dddb 100644 --- a/crates/multimodal/src/vision/processor.rs +++ b/crates/multimodal/src/vision/processor.rs @@ -1,16 +1,15 @@ -//! Vision processor trait and encoder-output types. +//! Vision processor trait and registry. //! -//! This module defines the interface for model-specific vision processors -//! and the common output format for preprocessed encoder inputs. +//! Shared encoder output types live in [`crate::encoder_inputs`] and are re-exported +//! here for compatibility. -use std::{borrow::Cow, collections::HashMap}; +use std::collections::HashMap; -use anyhow::{Context, Result as AnyhowResult}; use image::DynamicImage; -use ndarray::{Array4, ArrayD}; use super::{preprocessor_config::PreProcessorConfig, transforms::TransformError}; -use crate::types::{FieldLayout, RgbFrameRef}; +pub use crate::encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}; +use crate::types::RgbFrameRef; /// Helper to extract a dimension from encoder_input given an ndim-dependent axis index. /// Returns `Err` if the ndim is not 4 or 5. @@ -30,249 +29,7 @@ fn dim_for_ndim( } } -/// Model-specific output values that vary by architecture. -/// -/// Different vision models require different auxiliary outputs beyond encoder_input. -/// 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, shape: Vec }, - - /// A tensor of integers (e.g., aspect_ratio_ids) - IntTensor { data: Vec, shape: Vec }, - - /// A tensor of unsigned integers (e.g., image_grid_thw) - UintTensor { data: Vec, shape: Vec }, - - /// Simple integer value - Int(i64), - - /// Simple float value - Float(f64), - - /// List of integers - IntVec(Vec), - - /// List of unsigned integers - UintVec(Vec), - - /// List of floats - FloatVec(Vec), - - /// List of tuples (e.g., 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) -> Self { - let len = data.len(); - Self::UintTensor { - data, - shape: vec![len], - } - } - - /// Create a 2D uint tensor. - pub fn uint_2d(data: Vec, rows: usize, cols: usize) -> Self { - Self::UintTensor { - data, - shape: vec![rows, cols], - } - } - - /// Create a 1D int tensor from a vector. - pub fn int_1d(data: Vec) -> Self { - let len = data.len(); - Self::IntTensor { - data, - shape: vec![len], - } - } - - /// Create a 2D int tensor. - pub fn int_2d(data: Vec, rows: usize, cols: usize) -> Self { - Self::IntTensor { - data, - shape: vec![rows, cols], - } - } - - /// Interpret this value as per-item flat sizes. - pub fn as_flat_sizes(&self) -> AnyhowResult> { - match self { - Self::IntTensor { data, .. } => data - .iter() - .map(|&v| usize::try_from(v).context("negative flat size")) - .collect(), - Self::UintTensor { data, .. } => Ok(data.iter().map(|&v| v as usize).collect()), - Self::IntVec(values) => values - .iter() - .map(|&v| usize::try_from(v).context("negative flat size")) - .collect(), - Self::UintVec(values) => Ok(values.iter().map(|&v| v as usize).collect()), - _ => Err(anyhow::anyhow!("unsupported flat sizes value type")), - } - } - - /// Slice item-batched metadata along the first dimension. - pub fn slice_first_dim(&self, start: usize, len: usize) -> AnyhowResult { - match self { - Self::Tensor { data, shape } => { - let (data, shape) = slice_tensor_first_dim(data, shape, start, len)?; - Ok(Self::Tensor { data, shape }) - } - Self::IntTensor { data, shape } => { - let (data, shape) = slice_tensor_first_dim(data, shape, start, len)?; - Ok(Self::IntTensor { data, shape }) - } - Self::UintTensor { data, shape } => { - let (data, shape) = slice_tensor_first_dim(data, shape, start, len)?; - Ok(Self::UintTensor { data, shape }) - } - Self::IntVec(values) => Ok(Self::IntVec(slice_1d(values, start, len)?.to_vec())), - Self::UintVec(values) => Ok(Self::UintVec(slice_1d(values, start, len)?.to_vec())), - Self::FloatVec(values) => Ok(Self::FloatVec(slice_1d(values, start, len)?.to_vec())), - Self::TupleVec(values) => Ok(Self::TupleVec(slice_1d(values, start, len)?.to_vec())), - _ => Ok(self.clone()), - } - } -} - -fn slice_tensor_first_dim( - data: &[T], - shape: &[usize], - start: usize, - len: usize, -) -> AnyhowResult<(Vec, Vec)> { - let first_dim = *shape - .first() - .ok_or_else(|| anyhow::anyhow!("cannot slice scalar tensor"))?; - let end = start - .checked_add(len) - .ok_or_else(|| anyhow::anyhow!("tensor slice range overflow"))?; - anyhow::ensure!( - end <= first_dim, - "tensor first-dimension slice {start}..{end} exceeds {first_dim}" - ); - let row_width = shape[1..] - .iter() - .try_fold(1usize, |acc, &dim| acc.checked_mul(dim)) - .ok_or_else(|| anyhow::anyhow!("tensor row width overflow"))?; - let data_start = start - .checked_mul(row_width) - .ok_or_else(|| anyhow::anyhow!("tensor data start overflow"))?; - let data_len = len - .checked_mul(row_width) - .ok_or_else(|| anyhow::anyhow!("tensor data length overflow"))?; - let data_end = data_start - .checked_add(data_len) - .ok_or_else(|| anyhow::anyhow!("tensor data end overflow"))?; - anyhow::ensure!( - data_end <= data.len(), - "tensor slice data range {data_start}..{data_end} exceeds {}", - data.len() - ); - let mut new_shape = shape.to_vec(); - new_shape[0] = len; - Ok((data[data_start..data_end].to_vec(), new_shape)) -} - -fn slice_1d(values: &[T], start: usize, len: usize) -> AnyhowResult<&[T]> { - let end = start - .checked_add(len) - .ok_or_else(|| anyhow::anyhow!("slice range overflow"))?; - values - .get(start..end) - .ok_or_else(|| anyhow::anyhow!("slice range {start}..{end} exceeds {}", values.len())) -} - -/// Preprocessed encoder inputs ready for model consumption. -/// -/// This struct contains the processor outputs needed by serving backends to -/// construct `MultimodalInputs` for the model. Vision processors currently -/// produce this from images or sampled video frames; future modality processors -/// can reuse the same output contract for audio features or other encoder inputs. -#[derive(Debug, Clone)] -pub struct PreprocessedEncoderInputs { - /// Primary encoder input as a dynamic-dimensional float32 tensor. - /// - /// For vision models this is typically the preprocessed image/video tensor. - /// Shape varies by model and modality: - /// - Standard: [B, C, H, W] (4D) - /// - Phi3-Vision: [B, num_crops+1, C, H, W] (5D) - pub encoder_input: ArrayD, - - /// Number of encoder feature tokens per media item in the batch. - /// - /// Used to expand placeholder tokens in the text input. For vision this is - /// usually an image patch or video patch count; for audio this could be an - /// audio feature-frame count. - /// For example, LLaVA with 336x336 and patch_size=14 produces 576 tokens. - pub feature_token_counts: Vec, - - /// Modality-specific item size metadata before preprocessing. - /// - /// Vision processors use this for image/frame dimensions, but the exact - /// tuple order follows each processor/model contract. Model-specific shape - /// tensors that need a fixed order should also be emitted in `model_specific`. - pub item_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` auxiliary metadata - pub model_specific: HashMap, -} - impl PreprocessedEncoderInputs { - /// Create a new PreprocessedEncoderInputs with required fields (4D encoder input). - pub fn new( - encoder_input: Array4, - feature_token_counts: Vec, - item_sizes: Vec<(u32, u32)>, - ) -> Self { - Self { - encoder_input: encoder_input.into_dyn(), - feature_token_counts, - item_sizes, - model_specific: HashMap::new(), - } - } - - /// Create a new PreprocessedEncoderInputs with dynamic-dimensional encoder input. - /// - /// Use this for models like Phi3-Vision that have 5D tensors. - pub fn new_dynamic( - encoder_input: ArrayD, - feature_token_counts: Vec, - item_sizes: Vec<(u32, u32)>, - ) -> Self { - Self { - encoder_input, - feature_token_counts, - item_sizes, - model_specific: HashMap::new(), - } - } - - /// Add a model-specific value. - pub fn with_extra(mut self, key: impl Into, value: ModelSpecificValue) -> Self { - self.model_specific.insert(key.into(), value); - self - } - - /// Get the number of media items represented by this preprocessed batch. - pub fn batch_size(&self) -> usize { - self.item_sizes.len() - } - /// Get the number of channels. /// /// For 4D tensors [B, C, H, W], returns shape[1]. @@ -305,56 +62,6 @@ impl PreprocessedEncoderInputs { pub fn width(&self) -> Result { dim_for_ndim(self.encoder_input.ndim(), 3, 4, self.encoder_input.shape()) } - - /// Get the number of dimensions of encoder_input. - pub fn ndim(&self) -> usize { - self.encoder_input.ndim() - } - - /// Get total number of encoder feature tokens across all media items. - pub fn total_feature_tokens(&self) -> usize { - self.feature_token_counts.iter().sum() - } - - /// Get the primary encoder input as a flat f32 slice without copying if possible. - pub fn encoder_input_flat(&self) -> Cow<'_, [f32]> { - match self.encoder_input.as_slice() { - Some(slice) => Cow::Borrowed(slice), - None => Cow::Owned(self.encoder_input.iter().copied().collect()), - } - } - - /// Get the shape of the primary encoder input as a vector. - pub fn encoder_input_shape(&self) -> Vec { - self.encoder_input.shape().to_vec() - } - - /// Number of media items in this batch. - pub fn num_media_items(&self) -> usize { - self.item_sizes.len() - } - - /// Extract batched tensor keys from explicit field layout declarations. - pub fn batched_keys(layouts: &HashMap) -> Vec { - layouts - .iter() - .filter(|(_, l)| matches!(l, FieldLayout::Batched)) - .map(|(k, _)| k.clone()) - .collect() - } - - /// Extract flat-slicing tensor keys from explicit field layout declarations. - /// - /// Returns a map of tensor name → sizes tensor name. - pub fn flat_keys(layouts: &HashMap) -> HashMap { - layouts - .iter() - .filter_map(|(k, l)| match l { - FieldLayout::Flat { sizes_key } => Some((k.clone(), sizes_key.clone())), - FieldLayout::Batched => None, - }) - .collect() - } } /// Trait for model-specific vision preprocessors. @@ -536,6 +243,17 @@ impl VisionProcessorRegistry { Box::new(super::processors::Qwen3VLProcessor::new()), ); + // Qwen3-Omni uses the same patchification and normalization contract + // as Qwen3-VL for its image and video towers. + registry.register( + "qwen3-omni", + Box::new(super::processors::Qwen3OmniVisionProcessor::new()), + ); + registry.register( + "qwen3_omni", + Box::new(super::processors::Qwen3OmniVisionProcessor::new()), + ); + // Qwen3.5 family (and Qwen3.6: same arch) reuses Qwen3-VL preprocessing. registry.register( "qwen3.5", @@ -624,7 +342,7 @@ mod tests { use crate::vision::processors::LlavaProcessor; #[test] - fn test_preprocessed_encoder_inputs_accessors() { + fn test_preprocessed_encoder_inputs_geometry_accessors() { let encoder_input = Array4::::zeros((2, 3, 336, 336)); let inputs = PreprocessedEncoderInputs::new( encoder_input, @@ -632,78 +350,9 @@ mod tests { vec![(640, 480), (800, 600)], ); - assert_eq!(inputs.batch_size(), 2); assert_eq!(inputs.channels().unwrap(), 3); assert_eq!(inputs.height().unwrap(), 336); assert_eq!(inputs.width().unwrap(), 336); - assert_eq!(inputs.total_feature_tokens(), 1152); - } - - #[test] - fn test_preprocessed_encoder_inputs_with_extra() { - let encoder_input = Array4::::zeros((1, 3, 224, 224)); - let inputs = PreprocessedEncoderInputs::new(encoder_input, vec![196], vec![(224, 224)]) - .with_extra( - "image_grid_thw", - ModelSpecificValue::uint_1d(vec![1, 16, 16]), - ) - .with_extra("aspect_ratio_id", ModelSpecificValue::Int(0)); - - assert!(inputs.model_specific.contains_key("image_grid_thw")); - assert!(inputs.model_specific.contains_key("aspect_ratio_id")); - } - - #[test] - fn 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"), - } - - let int_1d = ModelSpecificValue::int_1d(vec![1, 2, 3]); - match int_1d { - ModelSpecificValue::IntTensor { data, shape } => { - assert_eq!(data, vec![1, 2, 3]); - assert_eq!(shape, vec![3]); - } - _ => panic!("Expected IntTensor"), - } - - let int_2d = ModelSpecificValue::int_2d(vec![1, 2, 3, 4], 2, 2); - match int_2d { - ModelSpecificValue::IntTensor { data, shape } => { - assert_eq!(data, vec![1, 2, 3, 4]); - assert_eq!(shape, vec![2, 2]); - } - _ => panic!("Expected IntTensor"), - } - } - - #[test] - fn test_encoder_input_flat() { - let mut encoder_input = Array4::::zeros((1, 1, 2, 2)); - encoder_input[[0, 0, 0, 0]] = 1.0; - encoder_input[[0, 0, 0, 1]] = 2.0; - encoder_input[[0, 0, 1, 0]] = 3.0; - encoder_input[[0, 0, 1, 1]] = 4.0; - - let inputs = PreprocessedEncoderInputs::new(encoder_input, vec![4], vec![(2, 2)]); - let flat = inputs.encoder_input_flat(); - - assert_eq!(flat, vec![1.0, 2.0, 3.0, 4.0]); } #[test] diff --git a/crates/multimodal/src/vision/processors/kimi_k25.rs b/crates/multimodal/src/vision/processors/kimi_k25.rs index 1d9adf2bd..bcdc58ad4 100644 --- a/crates/multimodal/src/vision/processors/kimi_k25.rs +++ b/crates/multimodal/src/vision/processors/kimi_k25.rs @@ -310,19 +310,16 @@ impl VisionPreProcessor for KimiK25Processor { )) })?; - let result = PreprocessedEncoderInputs::new_dynamic( - encoder_input.into_dyn(), - feature_token_counts, - item_sizes, - ) - .with_extra( - "grid_thws", - ModelSpecificValue::int_2d(grid_thw_data, images.len(), 3), - ) - .with_extra( - "patches_per_image", - ModelSpecificValue::int_1d(patches_per_image), - ); + let result = + PreprocessedEncoderInputs::new(encoder_input, feature_token_counts, item_sizes) + .with_extra( + "grid_thws", + ModelSpecificValue::int_2d(grid_thw_data, images.len(), 3), + ) + .with_extra( + "patches_per_image", + ModelSpecificValue::int_1d(patches_per_image), + ); Ok(result) } diff --git a/crates/multimodal/src/vision/processors/llava.rs b/crates/multimodal/src/vision/processors/llava.rs index fb955b4ea..1d0699135 100644 --- a/crates/multimodal/src/vision/processors/llava.rs +++ b/crates/multimodal/src/vision/processors/llava.rs @@ -594,11 +594,8 @@ impl VisionPreProcessor for LlavaNextProcessor { .flat_map(|&(w, h)| [h as i64, w as i64]) .collect(); - let mut result = PreprocessedEncoderInputs::new_dynamic( - encoder_input.into_dyn(), - feature_token_counts, - item_sizes, - ); + let mut result = + PreprocessedEncoderInputs::new(encoder_input, feature_token_counts, item_sizes); result.model_specific.insert( "image_sizes".to_string(), ModelSpecificValue::int_2d(image_sizes_flat, num_images, 2), diff --git a/crates/multimodal/src/vision/processors/mod.rs b/crates/multimodal/src/vision/processors/mod.rs index daad2f1f0..4c8c4135e 100644 --- a/crates/multimodal/src/vision/processors/mod.rs +++ b/crates/multimodal/src/vision/processors/mod.rs @@ -10,6 +10,8 @@ //! - **Qwen2-VL** (`qwen2_vl`): Dynamic resolution with smart resizing //! - **Qwen2.5-VL** (`qwen2_vl`): Same processor as Qwen2-VL (identical preprocessing) //! - **Qwen3-VL** (`qwen3_vl`): Similar to Qwen2-VL but with patch_size=16 and [0.5,0.5,0.5] normalization +//! - **Qwen3-Omni** (`qwen3_omni_vision`): Qwen3 vision preprocessing with Omni video limits and timing metadata +//! - **Kimi-K2.5** (`kimi_k25`): MoonViT resize and zero-padding to patch alignment //! - **Phi3-Vision** (`phi3_vision`): Dynamic HD transform with 336x336 tiles //! - **Phi4-Vision** (`phi4_vision`): Dynamic HD transform with 448x448 tiles and SiGLIP encoder //! - **LLaMA 4 Vision** (`llama4_vision`): Tile-based processing with 336x336 tiles and global tile @@ -22,6 +24,7 @@ pub mod phi3_vision; pub mod phi4_vision; pub mod pixtral; pub mod qwen2_vl; +pub mod qwen3_omni_vision; pub mod qwen3_vl; pub mod qwen_vl_base; @@ -32,4 +35,5 @@ pub use phi3_vision::Phi3VisionProcessor; pub use phi4_vision::Phi4VisionProcessor; pub use pixtral::PixtralProcessor; pub use qwen2_vl::Qwen2VLProcessor; +pub use qwen3_omni_vision::Qwen3OmniVisionProcessor; pub use qwen3_vl::Qwen3VLProcessor; diff --git a/crates/multimodal/src/vision/processors/qwen2_vl.rs b/crates/multimodal/src/vision/processors/qwen2_vl.rs index 6e2cfff80..90fb0234f 100644 --- a/crates/multimodal/src/vision/processors/qwen2_vl.rs +++ b/crates/multimodal/src/vision/processors/qwen2_vl.rs @@ -21,7 +21,7 @@ use std::ops::Deref; use image::DynamicImage; -use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase}; +use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase, QwenVideoResizeMode}; use crate::vision::{ preprocessor_config::PreProcessorConfig, processor::{PreprocessedEncoderInputs, VisionPreProcessor}, @@ -84,6 +84,9 @@ impl Qwen2VLProcessor { merge_size: DEFAULT_MERGE_SIZE, min_pixels: DEFAULT_MIN_PIXELS, max_pixels: DEFAULT_MAX_PIXELS, + video_min_pixels: DEFAULT_MIN_PIXELS, + video_max_pixels: DEFAULT_MAX_PIXELS, + video_resize_mode: QwenVideoResizeMode::TotalVolume, temporal_patch_size: DEFAULT_TEMPORAL_PATCH_SIZE, mean: CLIP_MEAN, std: CLIP_STD, @@ -106,6 +109,9 @@ impl Qwen2VLProcessor { merge_size, min_pixels, max_pixels, + video_min_pixels: min_pixels, + video_max_pixels: max_pixels, + video_resize_mode: QwenVideoResizeMode::TotalVolume, temporal_patch_size, mean: CLIP_MEAN, std: CLIP_STD, @@ -122,6 +128,9 @@ impl Qwen2VLProcessor { merge_size: config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), min_pixels: config.min_pixels.unwrap_or(DEFAULT_MIN_PIXELS), max_pixels: config.max_pixels.unwrap_or(DEFAULT_MAX_PIXELS), + video_min_pixels: config.min_pixels.unwrap_or(DEFAULT_MIN_PIXELS), + video_max_pixels: config.max_pixels.unwrap_or(DEFAULT_MAX_PIXELS), + video_resize_mode: QwenVideoResizeMode::TotalVolume, temporal_patch_size: config .temporal_patch_size .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), diff --git a/crates/multimodal/src/vision/processors/qwen3_omni_vision.rs b/crates/multimodal/src/vision/processors/qwen3_omni_vision.rs new file mode 100644 index 000000000..52458b396 --- /dev/null +++ b/crates/multimodal/src/vision/processors/qwen3_omni_vision.rs @@ -0,0 +1,314 @@ +//! Qwen3-Omni image and video preprocessing. +//! +//! This processor shares patchification and normalization through +//! `QwenVLProcessorBase`, but cannot use `Qwen3VLProcessor` directly. Omni +//! applies video pixel limits per frame rather than across the sampled clip, +//! gives its video preprocessor config precedence for video-specific limits, +//! and emits timing metadata required by mixed-modality M-RoPE. + +use image::DynamicImage; + +use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase, QwenVideoResizeMode}; +use crate::{ + types::RgbFrameRef, + vision::{ + preprocessor_config::PreProcessorConfig, + processor::{PreprocessedEncoderInputs, VisionPreProcessor}, + transforms::TransformError, + }, +}; + +pub const QWEN3_OMNI_MEAN: [f64; 3] = [0.5; 3]; +pub const QWEN3_OMNI_STD: [f64; 3] = [0.5; 3]; +pub const DEFAULT_IMAGE_MIN_PIXELS: usize = 3_136; +pub const DEFAULT_IMAGE_MAX_PIXELS: usize = 12_845_056; +pub const DEFAULT_VIDEO_MIN_PIXELS: usize = 128 * 32 * 32; +pub const DEFAULT_VIDEO_MAX_PIXELS: usize = 768 * 32 * 32; +pub const DEFAULT_PATCH_SIZE: usize = 16; +pub const DEFAULT_MERGE_SIZE: usize = 2; +pub const DEFAULT_TEMPORAL_PATCH_SIZE: usize = 2; + +#[derive(Debug, Clone)] +pub struct Qwen3OmniVisionProcessor { + inner: QwenVLProcessorBase, +} + +impl Default for Qwen3OmniVisionProcessor { + fn default() -> Self { + Self::new() + } +} + +impl Qwen3OmniVisionProcessor { + pub fn new() -> Self { + Self::with_limits( + DEFAULT_IMAGE_MIN_PIXELS, + DEFAULT_IMAGE_MAX_PIXELS, + DEFAULT_VIDEO_MIN_PIXELS, + DEFAULT_VIDEO_MAX_PIXELS, + DEFAULT_PATCH_SIZE, + DEFAULT_MERGE_SIZE, + DEFAULT_TEMPORAL_PATCH_SIZE, + ) + } + + fn with_limits( + image_min_pixels: usize, + image_max_pixels: usize, + video_min_pixels: usize, + video_max_pixels: usize, + patch_size: usize, + merge_size: usize, + temporal_patch_size: usize, + ) -> Self { + Self { + inner: QwenVLProcessorBase::new(QwenVLConfig { + patch_size, + merge_size, + min_pixels: image_min_pixels, + max_pixels: image_max_pixels, + video_min_pixels, + video_max_pixels, + video_resize_mode: QwenVideoResizeMode::PerFrame, + temporal_patch_size, + mean: QWEN3_OMNI_MEAN, + std: QWEN3_OMNI_STD, + model_name: "qwen3-omni", + }), + } + } + + fn from_preprocessor_config(config: &PreProcessorConfig) -> Self { + let configured_min = config.min_pixels.or_else(|| config.get_shortest_edge()); + let configured_max = config.max_pixels.or_else(|| config.get_longest_edge()); + Self::with_limits( + configured_min.unwrap_or(DEFAULT_IMAGE_MIN_PIXELS), + configured_max.unwrap_or(DEFAULT_IMAGE_MAX_PIXELS), + DEFAULT_VIDEO_MIN_PIXELS, + DEFAULT_VIDEO_MAX_PIXELS, + config.get_patch_size(DEFAULT_PATCH_SIZE), + config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), + config + .temporal_patch_size + .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), + ) + } + + fn from_video_preprocessor_config(config: &PreProcessorConfig) -> Self { + let configured_min = config.min_pixels.or_else(|| config.get_shortest_edge()); + let configured_max = config.max_pixels.or_else(|| config.get_longest_edge()); + Self::with_limits( + DEFAULT_IMAGE_MIN_PIXELS, + DEFAULT_IMAGE_MAX_PIXELS, + configured_min.unwrap_or(DEFAULT_VIDEO_MIN_PIXELS), + configured_max.unwrap_or(DEFAULT_VIDEO_MAX_PIXELS), + config.get_patch_size(DEFAULT_PATCH_SIZE), + config.merge_size.unwrap_or(DEFAULT_MERGE_SIZE), + config + .temporal_patch_size + .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), + ) + } + + fn contains_image_only_processor_config(config: &PreProcessorConfig) -> bool { + config + .image_processor_type + .as_deref() + .map(str::to_ascii_lowercase) + .is_some_and(|processor| { + processor.contains("imageprocessor") && !processor.contains("video") + }) + } + + fn has_structural_overrides(config: &PreProcessorConfig) -> bool { + config.patch_size.is_some() + || config.merge_size.is_some() + || config.min_pixels.is_some() + || config.max_pixels.is_some() + || config.temporal_patch_size.is_some() + || config.size.is_some() + } + + fn with_image_preprocessor_config(&self, config: &PreProcessorConfig) -> Self { + if Self::has_structural_overrides(config) { + Self::from_preprocessor_config(config) + } else { + self.clone() + } + } + + fn with_video_preprocessor_config(&self, config: &PreProcessorConfig) -> Self { + if Self::has_structural_overrides(config) { + if Self::contains_image_only_processor_config(config) { + // Qwen3-Omni's shared preprocessor_config.json carries image + // limits. The HF processor supplies separate video defaults at + // call time, so those image limits must not become a per-frame + // video budget here. + Self::from_preprocessor_config(config) + } else { + Self::from_video_preprocessor_config(config) + } + } else { + self.clone() + } + } +} + +impl VisionPreProcessor for Qwen3OmniVisionProcessor { + fn default_mean(&self) -> [f64; 3] { + self.inner.default_mean() + } + + fn default_std(&self) -> [f64; 3] { + self.inner.default_std() + } + + fn preprocess( + &self, + images: &[DynamicImage], + config: &PreProcessorConfig, + ) -> Result { + self.with_image_preprocessor_config(config) + .inner + .preprocess(images, config) + } + + fn preprocess_video( + &self, + frames: &[DynamicImage], + config: &PreProcessorConfig, + ) -> Result { + self.with_video_preprocessor_config(config) + .inner + .preprocess_video(frames, config) + } + + fn preprocess_video_rgb( + &self, + frames: &[RgbFrameRef<'_>], + config: &PreProcessorConfig, + ) -> Result { + self.with_video_preprocessor_config(config) + .inner + .preprocess_video_rgb(frames, config) + } + + fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { + self.with_image_preprocessor_config(config) + .inner + .calculate_num_tokens(width, height, config) + } + + fn model_name(&self) -> &'static str { + self.inner.model_name() + } + + fn get_processed_size(&self, config: &PreProcessorConfig) -> Option<(u32, u32)> { + self.inner.get_processed_size(config) + } +} + +#[cfg(test)] +mod tests { + use image::{DynamicImage, RgbImage}; + + use super::*; + use crate::vision::{processor::ModelSpecificValue, processors::Qwen3VLProcessor}; + + #[test] + fn omni_multiframe_resize_differs_from_qwen3_volume_budget() { + let omni = Qwen3OmniVisionProcessor::new(); + let qwen3_vl = Qwen3VLProcessor::with_config( + DEFAULT_PATCH_SIZE, + DEFAULT_MERGE_SIZE, + DEFAULT_VIDEO_MIN_PIXELS, + DEFAULT_VIDEO_MAX_PIXELS, + DEFAULT_TEMPORAL_PATCH_SIZE, + ); + + assert_eq!(omni.inner.min_pixels(), DEFAULT_IMAGE_MIN_PIXELS); + assert_eq!(omni.inner.max_pixels(), DEFAULT_IMAGE_MAX_PIXELS); + assert_eq!(omni.inner.video_min_pixels(), DEFAULT_VIDEO_MIN_PIXELS); + assert_eq!(omni.inner.video_max_pixels(), DEFAULT_VIDEO_MAX_PIXELS); + assert_eq!( + omni.inner.video_resize_mode(), + QwenVideoResizeMode::PerFrame + ); + + let omni_size = omni.inner.smart_resize_video(16, 720, 1280).unwrap(); + let volume_size = qwen3_vl.smart_resize_video(16, 720, 1280).unwrap(); + + assert_eq!(omni_size, (640, 1152)); + assert_eq!(volume_size, (160, 288)); + assert_ne!(omni_size, volume_size); + } + + #[test] + fn video_preprocessor_config_overrides_per_frame_limits() { + let config = PreProcessorConfig::from_json( + r#"{"size":{"shortest_edge":65536,"longest_edge":262144},"temporal_patch_size":2}"#, + ) + .unwrap(); + let processor = Qwen3OmniVisionProcessor::from_video_preprocessor_config(&config); + + assert_eq!(processor.inner.min_pixels(), DEFAULT_IMAGE_MIN_PIXELS); + assert_eq!(processor.inner.max_pixels(), DEFAULT_IMAGE_MAX_PIXELS); + assert_eq!(processor.inner.video_min_pixels(), 65_536); + assert_eq!(processor.inner.video_max_pixels(), 262_144); + assert_eq!( + processor.inner.smart_resize_video(32, 720, 1280).unwrap(), + (384, 672) + ); + } + + #[test] + fn shared_image_config_keeps_omni_video_defaults() { + let config = PreProcessorConfig::from_json( + r#"{"image_processor_type":"Qwen2VLImageProcessor","min_pixels":3136,"max_pixels":12845056,"patch_size":16,"merge_size":2,"temporal_patch_size":2}"#, + ) + .unwrap(); + let processor = Qwen3OmniVisionProcessor::new().with_video_preprocessor_config(&config); + + assert_eq!(processor.inner.video_min_pixels(), DEFAULT_VIDEO_MIN_PIXELS); + assert_eq!(processor.inner.video_max_pixels(), DEFAULT_VIDEO_MAX_PIXELS); + assert_eq!( + processor.inner.smart_resize_video(16, 720, 1280).unwrap(), + (640, 1152) + ); + } + + #[test] + fn empty_config_uses_omni_half_normalization() { + let image = DynamicImage::ImageRgb8(RgbImage::new(32, 32)); + let output = Qwen3OmniVisionProcessor::new() + .preprocess(&[image], &PreProcessorConfig::default()) + .unwrap(); + + assert!(output + .encoder_input + .iter() + .all(|value| (*value + 1.0).abs() < 1e-6)); + } + + #[test] + fn sampled_video_fps_controls_mrope_grid_timing() { + let config = PreProcessorConfig::from_json( + r#"{"size":{"shortest_edge":1024,"longest_edge":65536},"fps":4.0}"#, + ) + .unwrap(); + let frames = vec![ + DynamicImage::ImageRgb8(RgbImage::new(32, 32)), + DynamicImage::ImageRgb8(RgbImage::new(32, 32)), + ]; + + let output = Qwen3OmniVisionProcessor::new() + .preprocess_video(&frames, &config) + .unwrap(); + + assert!(matches!( + output.model_specific.get("video_second_per_grid"), + Some(ModelSpecificValue::Tensor { data, shape }) + if data == &vec![0.5] && shape == &vec![1] + )); + } +} diff --git a/crates/multimodal/src/vision/processors/qwen3_vl.rs b/crates/multimodal/src/vision/processors/qwen3_vl.rs index 05bda801a..669ea629c 100644 --- a/crates/multimodal/src/vision/processors/qwen3_vl.rs +++ b/crates/multimodal/src/vision/processors/qwen3_vl.rs @@ -20,7 +20,7 @@ use std::ops::Deref; use image::DynamicImage; -use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase}; +use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase, QwenVideoResizeMode}; use crate::{ types::RgbFrameRef, vision::{ @@ -88,6 +88,9 @@ impl Qwen3VLProcessor { merge_size: DEFAULT_MERGE_SIZE, min_pixels: DEFAULT_MIN_PIXELS, max_pixels: DEFAULT_MAX_PIXELS, + video_min_pixels: DEFAULT_MIN_PIXELS, + video_max_pixels: DEFAULT_MAX_PIXELS, + video_resize_mode: QwenVideoResizeMode::TotalVolume, temporal_patch_size: DEFAULT_TEMPORAL_PATCH_SIZE, mean: QWEN3_MEAN, std: QWEN3_STD, @@ -110,6 +113,9 @@ impl Qwen3VLProcessor { merge_size, min_pixels, max_pixels, + video_min_pixels: min_pixels, + video_max_pixels: max_pixels, + video_resize_mode: QwenVideoResizeMode::TotalVolume, temporal_patch_size, mean: QWEN3_MEAN, std: QWEN3_STD, @@ -132,6 +138,15 @@ impl Qwen3VLProcessor { .max_pixels .or_else(|| config.get_longest_edge()) .unwrap_or(DEFAULT_MAX_PIXELS), + video_min_pixels: config + .min_pixels + .or_else(|| config.get_shortest_edge()) + .unwrap_or(DEFAULT_MIN_PIXELS), + video_max_pixels: config + .max_pixels + .or_else(|| config.get_longest_edge()) + .unwrap_or(DEFAULT_MAX_PIXELS), + video_resize_mode: QwenVideoResizeMode::TotalVolume, temporal_patch_size: config .temporal_patch_size .unwrap_or(DEFAULT_TEMPORAL_PATCH_SIZE), @@ -548,6 +563,11 @@ mod tests { } else { panic!("Expected video_grid_thw to be IntTensor"); } + assert!(matches!( + result.model_specific.get("video_second_per_grid"), + Some(ModelSpecificValue::Tensor { data, shape }) + if data == &vec![1.0] && shape == &vec![1] + )); } #[test] diff --git a/crates/multimodal/src/vision/processors/qwen_vl_base.rs b/crates/multimodal/src/vision/processors/qwen_vl_base.rs index beb8ec42a..c76336af5 100644 --- a/crates/multimodal/src/vision/processors/qwen_vl_base.rs +++ b/crates/multimodal/src/vision/processors/qwen_vl_base.rs @@ -68,6 +68,12 @@ pub struct QwenVLConfig { pub min_pixels: usize, /// Maximum total pixels allowed pub max_pixels: usize, + /// Minimum video pixels, interpreted according to `video_resize_mode`. + pub video_min_pixels: usize, + /// Maximum video pixels, interpreted according to `video_resize_mode`. + pub video_max_pixels: usize, + /// Whether the video budget applies per frame or to the sampled volume. + pub video_resize_mode: QwenVideoResizeMode, /// Temporal patch size for video pub temporal_patch_size: usize, /// Normalization mean values @@ -78,6 +84,12 @@ pub struct QwenVLConfig { pub model_name: &'static str, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum QwenVideoResizeMode { + TotalVolume, + PerFrame, +} + #[derive(Clone)] struct VideoFrameRgb<'a> { width: usize, @@ -108,14 +120,29 @@ struct QwenVideoPlan { num_patches: usize, output_values: usize, tokens: usize, + second_per_grid: f32, filter: FilterType, do_resize: bool, lut: [[f32; 256]; 3], } -fn normalization_lut(config: &PreProcessorConfig) -> [[f32; 256]; 3] { - let mean = config.get_image_mean(); - let std = config.get_image_std(); +fn normalization_lut( + config: &PreProcessorConfig, + default_mean: [f64; 3], + default_std: [f64; 3], +) -> [[f32; 256]; 3] { + let mean = config + .image_mean + .as_ref() + .filter(|values| values.len() >= 3) + .map(|values| [values[0], values[1], values[2]]) + .unwrap_or(default_mean); + let std = config + .image_std + .as_ref() + .filter(|values| values.len() >= 3) + .map(|values| [values[0], values[1], values[2]]) + .unwrap_or(default_std); let do_normalize = config.do_normalize.unwrap_or(true); let scale: [f32; 3] = if do_normalize { std::array::from_fn(|channel| 1.0 / (255.0 * std[channel] as f32)) @@ -319,6 +346,18 @@ impl QwenVLProcessorBase { self.config.max_pixels } + pub fn video_min_pixels(&self) -> usize { + self.config.video_min_pixels + } + + pub fn video_max_pixels(&self) -> usize { + self.config.video_max_pixels + } + + pub fn video_resize_mode(&self) -> QwenVideoResizeMode { + self.config.video_resize_mode + } + /// Get the temporal patch size. pub fn temporal_patch_size(&self) -> usize { self.config.temporal_patch_size @@ -352,6 +391,14 @@ impl QwenVLProcessorBase { "Qwen video patch buffer size overflow: patches={num_patches}, features={patch_features}" )) })?; + // MediaConnector samples video at 2 fps by default. A checkpoint may + // override that value in video_preprocessor_config.json. + let sample_fps = config.get_extra::("fps").unwrap_or(2.0); + if !sample_fps.is_finite() || sample_fps <= 0.0 { + return Err(TransformError::ShapeError(format!( + "Qwen video fps must be finite and positive, got {sample_fps}" + ))); + } Ok(QwenVideoPlan { original_size: (width, height), @@ -364,9 +411,10 @@ impl QwenVLProcessorBase { num_patches, output_values, tokens: self.calculate_tokens_from_grid(grid_t, grid_h, grid_w), + second_per_grid: temporal_patch_size as f32 / sample_fps, filter: pil_to_filter(config.resampling.or(Some(3))), do_resize: config.do_resize.unwrap_or(true), - lut: normalization_lut(config), + lut: normalization_lut(config, self.config.mean, self.config.std), }) } @@ -384,8 +432,8 @@ impl QwenVLProcessorBase { }, )?; - Ok(PreprocessedEncoderInputs::new_dynamic( - encoder_input.into_dyn(), + Ok(PreprocessedEncoderInputs::new( + encoder_input, vec![plan.tokens], vec![plan.original_size], ) @@ -404,6 +452,13 @@ impl QwenVLProcessorBase { .with_extra( "patches_per_image", ModelSpecificValue::int_1d(vec![plan.num_patches as i64]), + ) + .with_extra( + "video_second_per_grid", + ModelSpecificValue::Tensor { + data: vec![plan.second_per_grid], + shape: vec![1], + }, )) } @@ -488,9 +543,9 @@ impl QwenVLProcessorBase { /// Smart resize for Qwen3-style video processors. /// - /// Unlike image resize, the pixel budget is applied to the full sampled - /// video volume (`T * H * W`), matching HuggingFace's Qwen3 video - /// processor. + /// `TotalVolume` applies the pixel budget to the padded sampled video + /// volume (`T * H * W`), while `PerFrame` applies it to each frame's + /// spatial area (`H * W`). pub fn smart_resize_video( &self, num_frames: usize, @@ -528,21 +583,26 @@ impl QwenVLProcessorBase { h_bar = h_bar.max(factor); w_bar = w_bar.max(factor); - let t_bar = - num_frames.div_ceil(self.config.temporal_patch_size) * self.config.temporal_patch_size; - let resized_pixels = t_bar as f64 * h_bar as f64 * w_bar as f64; - if resized_pixels > self.config.max_pixels as f64 { - // HF uses padded frames for the threshold but actual frames for beta. - let beta = (num_frames as f64 * height as f64 * width as f64 - / self.config.max_pixels as f64) - .sqrt(); + let (budget_scale, resized_pixels) = match self.config.video_resize_mode { + QwenVideoResizeMode::TotalVolume => { + let padded_frames = num_frames.div_ceil(self.config.temporal_patch_size) + * self.config.temporal_patch_size; + ( + num_frames as f64, + padded_frames as f64 * h_bar as f64 * w_bar as f64, + ) + } + QwenVideoResizeMode::PerFrame => (1.0, h_bar as f64 * w_bar as f64), + }; + let source_pixels = budget_scale * height as f64 * width as f64; + if resized_pixels > self.config.video_max_pixels as f64 { + let beta = (source_pixels / self.config.video_max_pixels as f64).sqrt(); h_bar = ((height as f64 / beta / factor as f64).floor() as usize) * factor; w_bar = ((width as f64 / beta / factor as f64).floor() as usize) * factor; h_bar = h_bar.max(factor); w_bar = w_bar.max(factor); - } else if resized_pixels < self.config.min_pixels as f64 { - let beta = - (self.config.min_pixels as f64 / (num_frames * height * width) as f64).sqrt(); + } else if resized_pixels < self.config.video_min_pixels as f64 { + let beta = (self.config.video_min_pixels as f64 / source_pixels).sqrt(); h_bar = ((height as f64 * beta / factor as f64).ceil() as usize) * factor; w_bar = ((width as f64 * beta / factor as f64).ceil() as usize) * factor; } @@ -1034,7 +1094,7 @@ impl VisionPreProcessor for QwenVLProcessorBase { let temporal_patch_size = self.config.temporal_patch_size; let patch_features = 3 * temporal_patch_size * patch_size * patch_size; let do_resize = config.do_resize.unwrap_or(true); - let lut = normalization_lut(config); + let lut = normalization_lut(config, self.config.mean, self.config.std); let mut image_plans = Vec::with_capacity(images.len()); let mut item_sizes = Vec::with_capacity(images.len()); @@ -1134,19 +1194,16 @@ impl VisionPreProcessor for QwenVLProcessorBase { )) })?; - let result = PreprocessedEncoderInputs::new_dynamic( - encoder_input.into_dyn(), - feature_token_counts, - item_sizes, - ) - .with_extra( - "image_grid_thw", - ModelSpecificValue::int_2d(grid_thw_data, images.len(), 3), - ) - .with_extra( - "patches_per_image", - ModelSpecificValue::int_1d(patches_per_image), - ); + let result = + PreprocessedEncoderInputs::new(encoder_input, feature_token_counts, item_sizes) + .with_extra( + "image_grid_thw", + ModelSpecificValue::int_2d(grid_thw_data, images.len(), 3), + ) + .with_extra( + "patches_per_image", + ModelSpecificValue::int_1d(patches_per_image), + ); Ok(result) } @@ -1335,6 +1392,9 @@ mod tests { merge_size: 2, min_pixels: 256 * 28 * 28, max_pixels: 1280 * 28 * 28, + video_min_pixels: 256 * 28 * 28, + video_max_pixels: 1280 * 28 * 28, + video_resize_mode: QwenVideoResizeMode::TotalVolume, temporal_patch_size: 2, mean: [0.5, 0.5, 0.5], std: [0.5, 0.5, 0.5], @@ -1348,6 +1408,9 @@ mod tests { merge_size: 1, min_pixels: 1, max_pixels: 1024 * 1024, + video_min_pixels: 1, + video_max_pixels: 1024 * 1024, + video_resize_mode: QwenVideoResizeMode::TotalVolume, temporal_patch_size: 2, mean: [0.5, 0.25, 0.75], std: [0.5, 0.25, 0.5], @@ -1671,6 +1734,9 @@ mod tests { merge_size: 2, min_pixels: 1, max_pixels: 16_777_216, + video_min_pixels: 1, + video_max_pixels: 16_777_216, + video_resize_mode: QwenVideoResizeMode::TotalVolume, temporal_patch_size: 2, mean: [0.5; 3], std: [0.5; 3], diff --git a/crates/multimodal/src/vision/transforms.rs b/crates/multimodal/src/vision/transforms.rs index ddb490f3c..ee0217497 100644 --- a/crates/multimodal/src/vision/transforms.rs +++ b/crates/multimodal/src/vision/transforms.rs @@ -11,34 +11,12 @@ use fast_image_resize::{ }; use image::{imageops::FilterType, DynamicImage, GenericImageView, Rgb, RgbImage}; use ndarray::{s, Array3, Array4}; -use thiserror::Error; use super::{ execution::{scope as parallel_scope, task_count}, scratch, }; - -/// 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, - }, - - #[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 use crate::error::TransformError; pub type Result = std::result::Result; diff --git a/crates/multimodal/tests/multimodal_tracker_test.rs b/crates/multimodal/tests/multimodal_tracker_test.rs index dea2b4f25..1efd82365 100644 --- a/crates/multimodal/tests/multimodal_tracker_test.rs +++ b/crates/multimodal/tests/multimodal_tracker_test.rs @@ -2,8 +2,8 @@ use std::{path::PathBuf, sync::Arc, time::Duration}; use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; use llm_multimodal::{ - AsyncMultiModalTracker, ImageFetchConfig, ImageSource, MediaConnector, MediaConnectorConfig, - MediaContentPart, MediaSource, Modality, + AsyncMultiModalTracker, AudioSource, ImageFetchConfig, ImageSource, MediaConnector, + MediaConnectorConfig, MediaContentPart, MediaSource, Modality, }; use reqwest::Client; use tempfile::tempdir; @@ -21,6 +21,27 @@ fn tiny_png_bytes() -> Vec { .expect("decode tiny png fixture") } +fn wav_i16_mono(sample_rate: u32, samples: &[i16]) -> Vec { + let data_bytes = samples.len() as u32 * 2; + let mut bytes = Vec::new(); + bytes.extend_from_slice(b"RIFF"); + bytes.extend_from_slice(&(36 + data_bytes).to_le_bytes()); + bytes.extend_from_slice(b"WAVEfmt "); + bytes.extend_from_slice(&16_u32.to_le_bytes()); + bytes.extend_from_slice(&1_u16.to_le_bytes()); + bytes.extend_from_slice(&1_u16.to_le_bytes()); + bytes.extend_from_slice(&sample_rate.to_le_bytes()); + bytes.extend_from_slice(&(sample_rate * 2).to_le_bytes()); + bytes.extend_from_slice(&2_u16.to_le_bytes()); + bytes.extend_from_slice(&16_u16.to_le_bytes()); + bytes.extend_from_slice(b"data"); + bytes.extend_from_slice(&data_bytes.to_le_bytes()); + for sample in samples { + bytes.extend_from_slice(&sample.to_le_bytes()); + } + bytes +} + #[expect( clippy::expect_used, reason = "test helper: panic on failure is intentional" @@ -98,6 +119,24 @@ async fn fetch_image_from_file() { } } +#[tokio::test] +async fn fetch_audio_from_inline_bytes_decodes_samples() { + let connector = test_connector(None); + let bytes = wav_i16_mono(16_000, &[0, 16_384, -16_384]); + let clip = connector + .fetch_audio(MediaSource::InlineBytes(bytes.clone())) + .await + .expect("inline audio"); + + assert_eq!(clip.raw_bytes(), bytes.as_slice()); + assert_eq!(clip.decoded().sample_rate, 16_000); + assert_eq!(clip.decoded().samples.len(), 3); + assert!(clip.decoded().samples[0].abs() < 1e-6); + assert!((clip.decoded().samples[1] - 0.5).abs() < 1e-4); + assert!((clip.decoded().samples[2] + 0.5).abs() < 1e-4); + assert!(matches!(clip.source(), AudioSource::InlineBytes)); +} + #[tokio::test] async fn tracker_fetches_images_and_records_uuids() { let connector = Arc::new(test_connector(None)); diff --git a/crates/protocols/src/chat.rs b/crates/protocols/src/chat.rs index 7d9b6148e..7c153c1cf 100644 --- a/crates/protocols/src/chat.rs +++ b/crates/protocols/src/chat.rs @@ -182,6 +182,9 @@ pub struct ChatCompletionRequest { /// Output types that you would like the model to generate for this request pub modalities: Option>, + /// Whether to return audio output. + pub return_audio: Option, + /// How many chat completion choices to generate for each input message #[validate(range(min = 1, max = 10))] pub n: Option, @@ -772,7 +775,21 @@ pub struct ChatStreamChoice { #[cfg(test)] mod tests { - use super::thinking_from_reasoning_effort; + use serde_json::{json, Value}; + + use super::{thinking_from_reasoning_effort, ChatCompletionRequest}; + + fn request_with_output_fields(fields: &[(&str, Value)]) -> ChatCompletionRequest { + let mut value = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}] + }); + let object = value.as_object_mut().expect("request must be an object"); + for (name, field_value) in fields { + object.insert((*name).to_string(), field_value.clone()); + } + serde_json::from_value(value).expect("request must deserialize") + } #[test] fn thinking_from_reasoning_effort_maps_disable_values() { @@ -787,4 +804,23 @@ mod tests { assert_eq!(thinking_from_reasoning_effort(None), None); assert_eq!(thinking_from_reasoning_effort(Some("bogus")), None); } + + #[test] + fn return_audio_preserves_explicit_values() { + for fields in [vec![], vec![("return_audio", Value::Null)]] { + let request = request_with_output_fields(&fields); + assert_eq!(request.return_audio, None); + assert!(!request.other.contains_key("return_audio")); + let serialized = serde_json::to_value(request).expect("request must serialize"); + assert!(serialized.get("return_audio").is_none()); + } + + for value in [false, true] { + let request = request_with_output_fields(&[("return_audio", json!(value))]); + assert_eq!(request.return_audio, Some(value)); + assert!(!request.other.contains_key("return_audio")); + let serialized = serde_json::to_value(request).expect("request must serialize"); + assert_eq!(serialized.get("return_audio"), Some(&Value::Bool(value))); + } + } } diff --git a/crates/protocols/src/common.rs b/crates/protocols/src/common.rs index 4446a652d..c9fe90c44 100644 --- a/crates/protocols/src/common.rs +++ b/crates/protocols/src/common.rs @@ -191,6 +191,10 @@ pub enum ContentPart { Text { text: String }, #[serde(rename = "image_url")] ImageUrl { image_url: ImageUrl }, + #[serde(rename = "audio_url")] + AudioUrl { audio_url: AudioUrl }, + #[serde(rename = "input_audio")] + InputAudio { input_audio: InputAudio }, #[serde(rename = "video_url")] VideoUrl { video_url: VideoUrl }, } @@ -202,6 +206,19 @@ pub struct ImageUrl { pub detail: Option, // "auto", "low", or "high" } +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, schemars::JsonSchema)] +pub struct AudioUrl { + pub url: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, schemars::JsonSchema)] +pub struct InputAudio { + /// Base64-encoded audio bytes. + pub data: String, + /// Encoded audio format. The OpenAI Chat API supports `wav` and `mp3`. + pub format: String, +} + #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, schemars::JsonSchema)] pub struct VideoUrl { pub url: String, @@ -815,6 +832,48 @@ mod tests { assert!(result.is_err()); } + #[test] + fn content_part_deserializes_audio_url() { + let value = json!({ + "type": "audio_url", + "audio_url": { + "url": "https://example.com/audio.wav" + } + }); + let part: ContentPart = serde_json::from_value(value).expect("audio_url content part"); + assert_eq!( + part, + ContentPart::AudioUrl { + audio_url: AudioUrl { + url: "https://example.com/audio.wav".to_string(), + }, + } + ); + } + + #[test] + fn content_part_round_trips_input_audio() { + let value = json!({ + "type": "input_audio", + "input_audio": { + "data": "UklGRg==", + "format": "wav" + } + }); + let part: ContentPart = + serde_json::from_value(value.clone()).expect("input_audio content part"); + assert_eq!( + part, + ContentPart::InputAudio { + input_audio: InputAudio { + data: "UklGRg==".to_string(), + format: "wav".to_string(), + }, + } + ); + assert_eq!(serde_json::to_value(part).unwrap(), value); + } + #[test] fn conversation_ref_deserializes_bare_string() { let v = json!("conv_abc"); diff --git a/crates/tokenizer/src/factory.rs b/crates/tokenizer/src/factory.rs index 975a75837..87d3d7fc9 100644 --- a/crates/tokenizer/src/factory.rs +++ b/crates/tokenizer/src/factory.rs @@ -63,7 +63,24 @@ pub fn create_tokenizer_with_chat_template( ); } - // Priority 2: tiktoken.model / *.tiktoken + let has_vocab_and_merges = + path.join("vocab.json").is_file() && path.join("merges.txt").is_file(); + + // Priority 2: Hugging Face Qwen2-style byte-level BPE files. Some + // official checkpoints (including Qwen3-ASR) intentionally omit + // tokenizer.json and ship only vocab.json + merges.txt. + if has_vocab_and_merges && has_qwen2_tokenizer_class(path) { + let final_chat_template = + resolve_and_log_chat_template(chat_template_path, path, file_path); + return Ok(Arc::new( + HuggingFaceTokenizer::from_vocab_and_merges_dir_with_chat_template( + path, + final_chat_template.as_deref(), + )?, + )); + } + + // Priority 3: tiktoken.model / *.tiktoken // Only forward the user's explicit chat_template_path — tiktoken handles // its own config/discovery (tokenizer_config.json → directory discovery). if has_tiktoken_file(path) { @@ -73,6 +90,19 @@ pub fn create_tokenizer_with_chat_template( )?)); } + // Preserve the specific validation error for vocab + merges layouts + // when there is no supported tokenizer to fall back to. + if has_vocab_and_merges { + let final_chat_template = + resolve_and_log_chat_template(chat_template_path, path, file_path); + return Ok(Arc::new( + HuggingFaceTokenizer::from_vocab_and_merges_dir_with_chat_template( + path, + final_chat_template.as_deref(), + )?, + )); + } + return Err(Error::msg(format!( "Directory '{file_path}' does not contain a valid tokenizer file (tokenizer.json, tiktoken.model, *.tiktoken, or vocab.json)" ))); @@ -115,6 +145,19 @@ pub fn create_tokenizer_with_chat_template( result } +fn has_qwen2_tokenizer_class(dir: &Path) -> bool { + std::fs::read_to_string(dir.join("tokenizer_config.json")) + .ok() + .and_then(|content| serde_json::from_str::(&content).ok()) + .and_then(|config| { + config + .get("tokenizer_class") + .and_then(serde_json::Value::as_str) + .map(|class| matches!(class, "Qwen2Tokenizer" | "Qwen2TokenizerFast")) + }) + .unwrap_or(false) +} + /// Auto-detect tokenizer type by examining file content fn auto_detect_tokenizer(file_path: &str) -> Result> { let mut file = File::open(file_path)?; diff --git a/crates/tokenizer/src/huggingface.rs b/crates/tokenizer/src/huggingface.rs index df3aaf182..a805b6893 100644 --- a/crates/tokenizer/src/huggingface.rs +++ b/crates/tokenizer/src/huggingface.rs @@ -1,9 +1,19 @@ -use std::collections::HashMap; +use std::{collections::HashMap, path::Path}; use anyhow::{Error, Result}; +use serde::Deserialize; use tokenizers::{ + models::bpe::BPE, + normalizers::unicode::NFC, + pre_tokenizers::{ + byte_level::ByteLevel, + sequence::Sequence, + split::{Split, SplitPattern}, + PreTokenizerWrapper, + }, processors::template::TemplateProcessing, - tokenizer::{step_decode_stream, Tokenizer as HfTokenizer}, + tokenizer::{step_decode_stream, SplitDelimiterBehavior, Tokenizer as HfTokenizer}, + AddedToken, }; use tracing::debug; @@ -36,11 +46,27 @@ pub struct HuggingFaceTokenizer { renderer: Renderer, } +const QWEN2_PRETOKENIZE_REGEX: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"; + +#[derive(Deserialize)] +struct AddedTokenConfig { + content: String, + #[serde(default)] + single_word: bool, + #[serde(default)] + lstrip: bool, + #[serde(default)] + rstrip: bool, + normalized: Option, + #[serde(default)] + special: bool, +} + impl HuggingFaceTokenizer { /// Create a tokenizer from a HuggingFace tokenizer JSON file pub fn from_file(file_path: &str) -> Result { // Try to auto-discover chat template if not explicitly provided - let path = std::path::Path::new(file_path); + let path = Path::new(file_path); let chat_template_path = path .parent() .and_then(crate::factory::discover_chat_template_in_dir); @@ -52,9 +78,150 @@ impl HuggingFaceTokenizer { file_path: &str, chat_template_path: Option<&str>, ) -> Result { - let mut tokenizer = HfTokenizer::from_file(file_path) + let tokenizer = HfTokenizer::from_file(file_path) .map_err(|e| Error::msg(format!("Failed to load tokenizer: {e}")))?; + Self::from_built_tokenizer(tokenizer, Path::new(file_path), chat_template_path) + } + /// Create a Qwen2-compatible byte-level BPE tokenizer from a Hugging Face + /// directory containing `vocab.json`, `merges.txt`, and + /// `tokenizer_config.json` but no `tokenizer.json`. + pub fn from_vocab_and_merges_dir(dir: &Path) -> Result { + let chat_template_path = crate::factory::discover_chat_template_in_dir(dir); + Self::from_vocab_and_merges_dir_with_chat_template(dir, chat_template_path.as_deref()) + } + + /// Create a Qwen2-compatible byte-level BPE tokenizer with an optional + /// explicit chat template. + pub fn from_vocab_and_merges_dir_with_chat_template( + dir: &Path, + chat_template_path: Option<&str>, + ) -> Result { + let tokenizer = Self::build_qwen2_bpe_tokenizer(dir)?; + // Shared initialization only needs this path to locate sibling config + // files. The file itself intentionally does not exist in this layout. + let logical_tokenizer_path = dir.join("tokenizer.json"); + Self::from_built_tokenizer(tokenizer, &logical_tokenizer_path, chat_template_path) + } + + fn build_qwen2_bpe_tokenizer(dir: &Path) -> Result { + let vocab_path = dir.join("vocab.json"); + let merges_path = dir.join("merges.txt"); + let config_path = dir.join("tokenizer_config.json"); + + let config_content = std::fs::read_to_string(&config_path).map_err(|error| { + Error::msg(format!("Failed to read {}: {error}", config_path.display())) + })?; + let config: serde_json::Value = serde_json::from_str(&config_content).map_err(|error| { + Error::msg(format!( + "Failed to parse {}: {error}", + config_path.display() + )) + })?; + let tokenizer_class = config + .get("tokenizer_class") + .and_then(serde_json::Value::as_str) + .ok_or_else(|| { + Error::msg(format!( + "{} is missing tokenizer_class; cannot infer vocab.json + merges.txt semantics", + config_path.display() + )) + })?; + if !matches!(tokenizer_class, "Qwen2Tokenizer" | "Qwen2TokenizerFast") { + return Err(Error::msg(format!( + "Unsupported vocab.json + merges.txt tokenizer_class '{tokenizer_class}' in {}", + config_path.display() + ))); + } + + let vocab_path_str = vocab_path.to_str().ok_or_else(|| { + Error::msg(format!("Tokenizer path is not valid UTF-8: {vocab_path:?}")) + })?; + let merges_path_str = merges_path.to_str().ok_or_else(|| { + Error::msg(format!( + "Tokenizer path is not valid UTF-8: {merges_path:?}" + )) + })?; + let bpe = BPE::builder() + .files(vocab_path_str.to_string(), merges_path_str.to_string()) + .build() + .map_err(|error| Error::msg(format!("Failed to build Qwen2 BPE model: {error}")))?; + let mut tokenizer = HfTokenizer::new(bpe); + + tokenizer + .with_normalizer(Some(NFC)) + .map_err(|error| Error::msg(format!("Failed to configure NFC normalizer: {error}")))?; + let split = Split::new( + SplitPattern::Regex(QWEN2_PRETOKENIZE_REGEX.to_string()), + SplitDelimiterBehavior::Isolated, + false, + ) + .map_err(|error| Error::msg(format!("Failed to build Qwen2 pre-tokenizer: {error}")))?; + let add_prefix_space = config + .get("add_prefix_space") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + let byte_level = ByteLevel::default() + .add_prefix_space(add_prefix_space) + .use_regex(false); + tokenizer.with_pre_tokenizer(Some(Sequence::new(vec![ + PreTokenizerWrapper::Split(split), + PreTokenizerWrapper::ByteLevel(byte_level), + ]))); + tokenizer.with_decoder(Some(ByteLevel::default())); + tokenizer.with_post_processor(Some(ByteLevel::default().trim_offsets(false))); + + let mut added_tokens = Vec::new(); + if let Some(decoder) = config + .get("added_tokens_decoder") + .and_then(serde_json::Value::as_object) + { + for (raw_id, raw_token) in decoder { + let id = raw_id.parse::().map_err(|error| { + Error::msg(format!( + "Invalid added token ID '{raw_id}' in {}: {error}", + config_path.display() + )) + })?; + let entry: AddedTokenConfig = + serde_json::from_value(raw_token.clone()).map_err(|error| { + Error::msg(format!( + "Invalid added token {raw_id} in {}: {error}", + config_path.display() + )) + })?; + let normalized = entry.normalized.unwrap_or(!entry.special); + let token = AddedToken::from(entry.content, entry.special) + .single_word(entry.single_word) + .lstrip(entry.lstrip) + .rstrip(entry.rstrip) + .normalized(normalized); + added_tokens.push((id, token)); + } + } + added_tokens.sort_unstable_by_key(|(id, _)| *id); + + tokenizer + .add_tokens(added_tokens.iter().map(|(_, token)| token.clone())) + .map_err(|error| Error::msg(format!("Failed to add configured tokens: {error}")))?; + for (expected_id, token) in &added_tokens { + let actual_id = tokenizer.token_to_id(&token.content); + if actual_id != Some(*expected_id) { + return Err(Error::msg(format!( + "Added token '{}' expected ID {expected_id}, got {actual_id:?}; non-contiguous explicit added-token IDs are unsupported", + token.content + ))); + } + } + + Ok(tokenizer) + } + + fn from_built_tokenizer( + mut tokenizer: HfTokenizer, + tokenizer_path: &Path, + chat_template_path: Option<&str>, + ) -> Result { // Build vocab mappings (include special tokens to get added_tokens like <|im_start|>) let vocab = tokenizer.get_vocab(true); // true = include special tokens and added_tokens let reverse_vocab: HashMap = vocab @@ -63,7 +230,7 @@ impl HuggingFaceTokenizer { .collect(); // Load tokenizer_config.json once for chat template, add_bos/eos, and special tokens - let config_result = Self::load_chat_template_and_config(file_path); + let config_result = Self::load_chat_template_and_config(&tokenizer_path.to_string_lossy()); let mut chat_template_str = config_result.chat_template; let add_bos_token = config_result.add_bos_token; let add_eos_token = config_result.add_eos_token; @@ -95,13 +262,13 @@ impl HuggingFaceTokenizer { } // Load merged EOS token IDs from config.json + generation_config.json - let eos_token_ids = std::path::Path::new(file_path) + let eos_token_ids = tokenizer_path .parent() .map(crate::eos::load_eos_token_ids) .unwrap_or_default(); // Detect a custom Python-encoder model from config.json::architectures. - let renderer = std::path::Path::new(file_path) + let renderer = tokenizer_path .parent() .map(detect_renderer_from_config) .unwrap_or(Renderer::Jinja); @@ -243,7 +410,7 @@ impl HuggingFaceTokenizer { /// Reads the file once and extracts everything needed by the tokenizer constructor. fn load_chat_template_and_config(tokenizer_path: &str) -> TokenizerConfigResult { (|| { - let path = std::path::Path::new(tokenizer_path); + let path = Path::new(tokenizer_path); let config_path = path.parent()?.join("tokenizer_config.json"); if !config_path.exists() { @@ -449,7 +616,7 @@ impl TokenizerTrait for HuggingFaceTokenizer { /// use. A missing or malformed file falls back to [`Renderer::Jinja`] without /// erroring (debug-logged), preserving backward compatibility for every model /// not in the architecture list. -fn detect_renderer_from_config(dir: &std::path::Path) -> Renderer { +fn detect_renderer_from_config(dir: &Path) -> Renderer { let path = dir.join("config.json"); if !path.exists() { return Renderer::Jinja; diff --git a/crates/tokenizer/tests/qwen2_vocab_merges.rs b/crates/tokenizer/tests/qwen2_vocab_merges.rs new file mode 100644 index 000000000..12ba3b947 --- /dev/null +++ b/crates/tokenizer/tests/qwen2_vocab_merges.rs @@ -0,0 +1,135 @@ +use std::{fs, path::Path}; + +use anyhow::Result; +use llm_tokenizer::{chat_template::ChatTemplateParams, create_tokenizer}; +use serde_json::json; +use tempfile::TempDir; + +const MIN_TIKTOKEN_MODEL: &str = "aGVsbG8= 0\n"; + +fn write_qwen2_files(dir: &Path, add_prefix_space: bool, first_added_id: u32) -> Result<()> { + let vocab = json!({ + "H": 0, + "e": 1, + "l": 2, + "o": 3, + "Ġ": 4, + "w": 5, + "r": 6, + "d": 7, + "!": 8 + }); + fs::write(dir.join("vocab.json"), serde_json::to_vec(&vocab)?)?; + fs::write(dir.join("merges.txt"), "#version: 0.2\n")?; + + let config = json!({ + "tokenizer_class": "Qwen2Tokenizer", + "add_prefix_space": add_prefix_space, + "add_bos_token": false, + "eos_token": "<|im_start|>", + "pad_token": "<|im_start|>", + "added_tokens_decoder": { + first_added_id.to_string(): { + "content": "<|im_start|>", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + (first_added_id + 1).to_string(): { + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": false + } + } + }); + fs::write( + dir.join("tokenizer_config.json"), + serde_json::to_vec(&config)?, + )?; + fs::write( + dir.join("chat_template.json"), + r#"{"chat_template":"{% for message in messages %}{{ message.role }}: {{ message.content }}\n{% endfor %}{% if add_generation_prompt %}assistant: {% endif %}"}"#, + )?; + Ok(()) +} + +#[test] +fn factory_loads_qwen2_vocab_and_merges_directory() -> Result<()> { + let dir = TempDir::new()?; + write_qwen2_files(dir.path(), false, 9)?; + + let tokenizer = create_tokenizer(dir.path().to_string_lossy().as_ref())?; + assert_eq!(tokenizer.vocab_size(), 9); + assert_eq!(tokenizer.token_to_id("<|im_start|>"), Some(9)); + assert_eq!(tokenizer.token_to_id(""), Some(10)); + + let text = "Hello world!"; + let encoded = tokenizer.encode(text, false)?; + assert_eq!(encoded.token_ids(), &[0, 1, 2, 2, 3, 4, 5, 3, 6, 2, 7, 8]); + assert_eq!(tokenizer.decode(encoded.token_ids(), false)?, text); + + let with_tokens = tokenizer.encode("<|im_start|>", false)?; + assert_eq!(with_tokens.token_ids(), &[9, 10]); + assert_eq!( + tokenizer.decode(with_tokens.token_ids(), true)?, + "" + ); + + let rendered = tokenizer.apply_chat_template( + &[json!({"role": "user", "content": "Hello"})], + ChatTemplateParams { + add_generation_prompt: true, + ..Default::default() + }, + )?; + assert_eq!(rendered, "user: Hello\nassistant: "); + Ok(()) +} + +#[test] +fn qwen2_loader_honors_add_prefix_space() -> Result<()> { + let dir = TempDir::new()?; + write_qwen2_files(dir.path(), true, 9)?; + + let tokenizer = create_tokenizer(dir.path().to_string_lossy().as_ref())?; + let encoded = tokenizer.encode("Hello", false)?; + assert_eq!(encoded.token_ids(), &[4, 0, 1, 2, 2, 3]); + assert_eq!(tokenizer.decode(encoded.token_ids(), false)?, " Hello"); + Ok(()) +} + +#[test] +fn qwen2_loader_rejects_unrepresentable_added_token_ids() -> Result<()> { + let dir = TempDir::new()?; + write_qwen2_files(dir.path(), false, 10)?; + + let error = match create_tokenizer(dir.path().to_string_lossy().as_ref()) { + Ok(_) => panic!("non-contiguous added token IDs must fail"), + Err(error) => error, + }; + assert!( + error.to_string().contains("expected ID 10, got Some(9)"), + "unexpected error: {error}" + ); + Ok(()) +} + +#[test] +fn mixed_non_qwen_directory_falls_back_to_tiktoken() -> Result<()> { + let dir = TempDir::new()?; + write_qwen2_files(dir.path(), false, 9)?; + fs::write( + dir.path().join("tokenizer_config.json"), + r#"{"tokenizer_class":"OtherTokenizer"}"#, + )?; + fs::write(dir.path().join("tiktoken.model"), MIN_TIKTOKEN_MODEL)?; + + let tokenizer = create_tokenizer(dir.path().to_string_lossy().as_ref())?; + assert_eq!(tokenizer.encode("hello", false)?.token_ids(), &[0]); + Ok(()) +} diff --git a/crates/tokenizer/tests/qwen3_asr_bpe_parity.rs b/crates/tokenizer/tests/qwen3_asr_bpe_parity.rs new file mode 100644 index 000000000..fe42c2c8f --- /dev/null +++ b/crates/tokenizer/tests/qwen3_asr_bpe_parity.rs @@ -0,0 +1,82 @@ +use anyhow::{anyhow, Result}; +use llm_tokenizer::{chat_template::ChatTemplateParams, create_tokenizer}; + +const TOKENIZER_DIR_ENV: &str = "QWEN3_ASR_TOKENIZER_DIR"; + +#[test] +#[ignore = "requires an official Qwen3-ASR-1.7B tokenizer snapshot"] +fn official_qwen3_asr_matches_transformers_qwen2_tokenizer() -> Result<()> { + let dir = std::env::var(TOKENIZER_DIR_ENV) + .map_err(|_| anyhow!("set {TOKENIZER_DIR_ENV} to the tokenizer snapshot"))?; + let tokenizer = create_tokenizer(&dir)?; + + let cases: &[(&str, &[u32])] = &[ + ("Hello, world!", &[9707, 11, 1879, 0]), + ( + "I'm testing Qwen2: 12345\nsecond line.", + &[ + 40, 2776, 7497, 1207, 16948, 17, 25, 220, 16, 17, 18, 19, 20, 198, 5569, 1555, 13, + ], + ), + ( + "Cafe\u{301} 中文🙂 — naïve", + &[34, 2577, 963, 72858, 16744, 145080, 1959, 94880, 586], + ), + ( + " leading and trailing ", + &[220, 6388, 220, 323, 27748, 256], + ), + ( + "<|im_start|>assistant\nHello<|im_end|>", + &[151644, 77091, 198, 9707, 151645], + ), + ( + "<|audio_start|><|audio_pad|><|audio_end|>", + &[151669, 151676, 151670], + ), + ( + "{\"name\":\"x\"}", + &[151657, 4913, 606, 3252, 87, 9207, 151658], + ), + ("", &[]), + ]; + + for (text, expected_ids) in cases { + let encoding = tokenizer.encode(text, false)?; + assert_eq!(encoding.token_ids(), *expected_ids, "text={text:?}"); + } + + assert_eq!(tokenizer.vocab_size(), 151643); + for (token, id) in [ + ("<|endoftext|>", 151643), + ("<|im_start|>", 151644), + ("<|im_end|>", 151645), + ("", 151657), + ("<|audio_start|>", 151669), + ("<|audio_end|>", 151670), + ("<|audio_pad|>", 151676), + ("", 151704), + ] { + assert_eq!(tokenizer.token_to_id(token), Some(id), "token={token}"); + } + + let special = tokenizer.encode("<|im_start|>assistant\nHello<|im_end|>", false)?; + assert_eq!( + tokenizer.decode(special.token_ids(), true)?, + "assistant\nHello" + ); + + let rendered = tokenizer.apply_chat_template( + &[serde_json::json!({ + "role": "user", + "content": [{"type": "audio", "audio": "fixture.wav"}] + })], + ChatTemplateParams { + add_generation_prompt: true, + ..Default::default() + }, + )?; + assert!(rendered.contains("<|audio_start|><|audio_pad|><|audio_end|>")); + assert!(rendered.ends_with("<|im_start|>assistant\n")); + Ok(()) +} diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py index 0122ee37c..39b8f35cf 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py @@ -1,9 +1,9 @@ """TokenSpeed EPD encode servicer. -Receives ``Encode`` RPCs from the gateway and forwards them to a vision-only +Receives ``Encode`` RPCs from the gateway and forwards them to an encoder-only encode worker (the engine's ``run_encode_loop``) over the same AsyncLLM -scheduler-input channel the LM uses. The encode worker runs the vision tower and -ships the resulting image embeddings to prefill workers over Mooncake; this +scheduler-input channel the LM uses. The encode worker runs the requested +multimodal tower and ships the resulting embeddings to prefill workers; this servicer only acks (the embeddings never flow back through the gateway). """ @@ -55,9 +55,9 @@ def __init__( health_servicer=None, ): self.async_llm = async_llm - # EPD_PIXEL_SHM: ship pixels to the scheduler process as POSIX-SHM + # EPD_PIXEL_SHM: ship encoder inputs to the scheduler process as POSIX-SHM # handles instead of pickling the raw tensor over ZMQ (the dominant - # per-image ingest cost). On by default (this servicer only runs in the + # per-item ingest cost). On by default (this servicer only runs in the # encode role, where SHM is always the right path); set EPD_PIXEL_SHM=0 to # fall back to the inline ZMQ pickle (e.g. a container with a tiny # /dev/shm). The decision is made once per item in _items_from_proto; the @@ -77,9 +77,6 @@ def __init__( self._bootstrap_host = get_local_ip_by_remote() self._bootstrap_port = server_args.disaggregation_bootstrap_port - # Spatial merge factor for post-merge token counts (Qwen vision default 2). - self._merge_size = self._resolve_merge_size() - self._rdma_pixel_puller = RdmaPixelPuller( agent_name=f"smg-encode-{self._bootstrap_host}-{self._bootstrap_port}", log_prefix="EPD RDMA", @@ -88,40 +85,50 @@ def __init__( self.async_llm.auto_create_handle_loop() logger.info("TokenSpeedEncoderServicer initialized") - def _resolve_merge_size(self) -> int: - hf_config = getattr(self.async_llm.model_config, "hf_config", None) - vision_config = getattr(hf_config, "vision_config", None) - return int(getattr(vision_config, "spatial_merge_size", 2) or 2) - def _items_from_proto(self, mm_inputs, bootstrap_room: int = 0): """Reconstruct the engine MultimodalDataItem(s) for the encode worker. Unlike the prefill leg, the encode worker NEEDS each item's encoder_input - (it runs the tower). It also needs each item's post-merge token count so the executor - can split the tower output; the gateway ships grid_thw but not - placeholders to encode, so derive the count from grid_thw and set it as - the item's single offset span (the offset positions are irrelevant to the - encode side, only the count matters). + (it runs the tower). Proto placeholders carry the output token spans the + executor uses to split the packed tower output back into items. """ Modality, MultimodalDataItem = _lazy_mm_item() model_dtype = getattr(self.async_llm.model_config, "dtype", None) - # mm_inputs is itemized (one MultimodalItem per image, each owning its + # mm_inputs is itemized (one MultimodalItem per media item, each owning its # encoder_input + model_specific_tensors). The gateway sends one item per # Encode RPC keyed by bootstrap_room, but iterate generally. items = [] for item_proto in mm_inputs.items: + if item_proto.modality == common_pb2.MODALITY_UNSPECIFIED: + item_modality = Modality.IMAGE + else: + item_modality = TokenSpeedSchedulerServicer._modality_from_proto( + item_proto.modality + ) + + if not item_proto.placeholders: + raise ValueError("encode MultimodalItem carried no placeholders") + if any(p.length <= 0 for p in item_proto.placeholders): + raise ValueError("encode MultimodalItem.placeholders.length must be > 0") + offsets = [(p.offset, p.offset + p.length - 1) for p in item_proto.placeholders] + + model_specific = { + name: TokenSpeedSchedulerServicer._tensor_from_proto(t) + for name, t in item_proto.model_specific_tensors.items() + } + # The feature's CROSS-PROCESS representation is decided here, once, for # both payload arms: a plain CPU tensor by default, or (EPD_PIXEL_SHM) a # POSIX-SHM handle so the ZMQ hop to the scheduler pickles ~KB instead of - # the 19-77MB pixels. The content hash is computed on the real bytes + # the full encoder tensor. The content hash is computed on the real bytes # before the swap and pre-set on the item. td = item_proto.encoder_input if td.WhichOneof("payload") == "remote": - # EPD RDMA: pull pixels from the gateway's exported NIXL memory. + # EPD RDMA: pull the encoder input from exported NIXL memory. # With EPD_PIXEL_SHM, the received slot is published directly to # scheduler SHM so the scheduler ingest path still avoids pickle - # copies of the full pixel tensor. + # copies of the full encoder tensor. feature, feat_hash = self._rdma_pixel_puller.feature_from_remote( td, explicit_room=bootstrap_room, @@ -139,42 +146,6 @@ def _items_from_proto(self, mm_inputs, bootstrap_room: int = 0): feat_hash = hash_feature(feature) feature = ShmTensorHandle.publish(feature) - model_specific = { - name: TokenSpeedSchedulerServicer._tensor_from_proto(t, cast_to=model_dtype) - for name, t in item_proto.model_specific_tensors.items() - } - - if item_proto.modality in ( - common_pb2.IMAGE, - common_pb2.MODALITY_UNSPECIFIED, - ): - item_modality = Modality.IMAGE - grid_key = "image_grid_thw" - elif item_proto.modality == common_pb2.VIDEO: - item_modality = Modality.VIDEO - grid_key = "video_grid_thw" - else: - raise ValueError(f"encode request modality={item_proto.modality} is not supported") - - grid = model_specific.get(grid_key) - if grid is None: - # Tolerate the legacy "grid_thws" key (older gateway builds emit it on - # the encode RPC); mirrors the engine kimi_k25 _grid() helper's tolerance. - grid = model_specific.get("grid_thws") - if grid is None: - raise ValueError( - f"encode request is missing {grid_key}/grid_thws; " - f"have keys={sorted(model_specific.keys())}" - ) - # grid is [num_media, 3] = (t, h, w) in patch units, per item. - merge = self._merge_size - offsets = [] - cursor = 0 - for row in grid.tolist(): - t, h, w = int(row[0]), int(row[1]), int(row[2]) - span = t * (h // merge) * (w // merge) - offsets.append((cursor, cursor + span - 1)) - cursor += span item = MultimodalDataItem( modality=item_modality, @@ -206,11 +177,11 @@ async def Encode(self, request, context): bootstrap_room = request.items[0].bootstrap_room if os.environ.get("EPD_INGEST_OFFLOOP", "1").lower() not in ("0", "false", "no"): - # Per-image ingest (proto->tensor + pickle) BLOCKS the lone asyncio + # Per-item ingest (proto->tensor + pickle) BLOCKS the lone asyncio # event loop, so grpc.aio cannot deliver the next Encode message until - # the previous one is fully ingested -- a per-worker serial pixel lane. + # the previous one is fully ingested -- a per-worker serial encoder-input lane. # Split it: parse + pickle on a worker thread (overlapping across - # images; the GIL is released in the tensor copy/cast), then the + # items; the GIL is released in the tensor copy/cast), then the # cheap zmq send back ON the loop -- send_to_scheduler is a # zmq.asyncio socket whose send() needs the running loop (and this # keeps it single-writer). @@ -233,7 +204,7 @@ async def Encode(self, request, context): return tokenspeed_encoder_pb2.EncodeResponse(accepted=True) def _build_encode_request(self, request, bootstrap_room): - """Proto -> engine EncodeRequest (the expensive per-image parse).""" + """Proto -> engine EncodeRequest (the expensive per-item parse).""" items = self._items_from_proto(request.mm_inputs, bootstrap_room) EncodeRequest = _lazy_encode_request() diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py index 39204afc2..c631acdf6 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py @@ -431,14 +431,11 @@ async def GetModelInfo( tokenizer_path = getattr(self.server_args, "tokenizer", None) or getattr( self.server_args, "tokenizer_path", "" ) - supports_vision = bool(getattr(model_config, "is_multimodal", False)) - image_modality = common_pb2.IMAGE - video_modality = common_pb2.VIDEO - supported_modalities = [] - if supports_vision: - supported_modalities.append(image_modality) - if hf_config is not None and getattr(hf_config, "video_token_id", None) is not None: - supported_modalities.append(video_modality) + supported_modalities = self._static_supported_modalities(model_config, hf_config) + supports_vision = any( + modality in (common_pb2.IMAGE, common_pb2.VIDEO) for modality in supported_modalities + ) + supports_multimodal = bool(supported_modalities) response_kwargs = dict( model_path=model_path, @@ -458,7 +455,7 @@ async def GetModelInfo( ) fields = tokenspeed_scheduler_pb2.GetModelInfoResponse.DESCRIPTOR.fields_by_name if "supports_multimodal" in fields: - response_kwargs["supports_multimodal"] = supports_vision + response_kwargs["supports_multimodal"] = supports_multimodal if "supported_modalities" in fields: response_kwargs["supported_modalities"] = supported_modalities dtype = self._torch_dtype_to_proto(getattr(model_config, "dtype", None)) @@ -468,6 +465,45 @@ async def GetModelInfo( response_kwargs["multimodal_encoder_dtype"] = dtype return tokenspeed_scheduler_pb2.GetModelInfoResponse(**response_kwargs) + @staticmethod + def _static_supported_modalities(model_config, hf_config) -> list[int]: + """Return modalities backed by towers present in the static model config.""" + + if getattr(model_config, "is_multimodal_active", None) is False: + return [] + + def config_value(config, name: str, default=None): + if isinstance(config, dict): + return config.get(name, default) + return getattr(config, name, default) + + image_modality = common_pb2.IMAGE + audio_modality = common_pb2.AUDIO + video_modality = common_pb2.VIDEO + model_type = config_value(hf_config, "model_type", "") if hf_config is not None else "" + + thinker_config = config_value(hf_config, "thinker_config") + if model_type == "qwen3_asr": + audio_config = config_value(thinker_config, "audio_config") + return [audio_modality] if audio_config is not None else [] + + if model_type.startswith("qwen3_omni"): + supported = [] + if config_value(thinker_config, "vision_config") is not None: + supported.append(image_modality) + if config_value(thinker_config, "audio_config") is not None: + supported.append(audio_modality) + if config_value(thinker_config, "video_token_id") is not None: + supported.append(video_modality) + return supported + + supported = [] + if bool(getattr(model_config, "is_multimodal", False)): + supported.append(image_modality) + if hf_config is not None and config_value(hf_config, "video_token_id") is not None: + supported.append(video_modality) + return supported + # ------------------------------------------------------------------ # GetServerInfo (unary) # ------------------------------------------------------------------ @@ -828,7 +864,7 @@ def _build_generate_req(self, request: tokenspeed_scheduler_pb2.GenerateRequest) ) # Decode the precomputed multimodal payload, if the request carries one. - # EPD requests carry mm metadata WITHOUT pixel_values (the image + # EPD requests carry mm metadata WITHOUT encoder inputs (the multimodal # embeddings arrive from encode workers over Mooncake), so build the mm # whenever mm_inputs is present, not only when pixel_values is. precomputed_mm = None @@ -841,14 +877,14 @@ def _build_generate_req(self, request: tokenspeed_scheduler_pb2.GenerateRequest) # EPD: per-item encode->prefill bootstrap info. Each entry tells the prefill # which Mooncake room to receive that item's embedding on, and the encode # worker's bootstrap endpoint to discover it. The gateway splits the mm - # payload one item per image, so ``item_index`` indexes ``mm_items`` 1:1; + # payload one RPC per item, so ``item_index`` indexes ``mm_items`` 1:1; # we attach each entry ONTO its item (``item.encode_handshake``) so it # rides with the item through the engine. Absent for non-EPD. if request.HasField("encode_bootstrap_info"): if precomputed_mm is None: raise ValueError( "GenerateRequest.encode_bootstrap_info present but mm_inputs " - "is missing; bootstrap info describes images that must be in mm_inputs" + "is missing; bootstrap info describes items that must be in mm_inputs" ) n_items = len(precomputed_mm.mm_items) for h in request.encode_bootstrap_info.items: @@ -1067,7 +1103,10 @@ def _mm_inputs_from_itemized_proto( ) model_started = time.perf_counter() if LOG_MM_TIMING else None model_specific_data = { - name: self._tensor_from_proto(tensor_data, cast_to=model_dtype) + # Side tensors are metadata, not encoder activations. Preserve + # their wire dtype: reducing video timing values to BF16 can + # change the integer M-RoPE positions at fractional frame rates. + name: self._tensor_from_proto(tensor_data) for name, tensor_data in item_proto.model_specific_tensors.items() } model_elapsed_ms = ( @@ -1152,7 +1191,7 @@ def _modality_from_proto(modality: int) -> Modality: if modality == common_pb2.VIDEO: return Modality.VIDEO if modality == common_pb2.AUDIO: - raise ValueError("TokenSpeed audio multimodal inputs are not supported yet") + return Modality.AUDIO raise ValueError(f"Unsupported multimodal item modality: {modality}") @staticmethod @@ -1179,6 +1218,8 @@ def _validate_item_tensor_consistency( raise ValueError("VIDEO MultimodalItem must not carry image_grid_thw") if modality == Modality.VIDEO and not has_video_grid: raise ValueError("VIDEO MultimodalItem must carry video_grid_thw") + if modality == Modality.AUDIO and (has_image_grid or has_video_grid): + raise ValueError("AUDIO MultimodalItem must not carry image/video grid tensors") @staticmethod def _tensor_from_proto( @@ -1201,7 +1242,7 @@ def _tensor_from_proto( f"TensorData byte length mismatch for bfloat16 shape={shape}: " f"expected {expected}, got {len(raw)}" ) - t = torch.from_numpy(np.frombuffer(raw, dtype=np.uint16).reshape(shape)).view( + t = torch.from_numpy(np.frombuffer(raw, dtype=np.uint16).copy().reshape(shape)).view( torch.bfloat16 ) else: @@ -1212,11 +1253,11 @@ def _tensor_from_proto( f"TensorData byte length mismatch for dtype={tensor_data.dtype}, " f"shape={shape}: expected {expected}, got {len(raw)}" ) - t = torch.from_numpy(np.frombuffer(raw, dtype=dtype).reshape(shape)) + t = torch.from_numpy(np.frombuffer(raw, dtype=dtype).copy().reshape(shape)) if cast_to is not None and t.dtype != cast_to and t.is_floating_point(): return t.to(cast_to) - return t.clone() + return t @staticmethod def _feature_from_proto( diff --git a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs index 699542b2b..f32a7f446 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -17,7 +17,7 @@ use crate::{ error, grpc::{ context::{EncodeWorkerAssignment, PreparationOutput, RequestContext, WorkerSelection}, - multimodal::{self, MultimodalIntermediate}, + multimodal, }, }, worker::{ @@ -442,9 +442,13 @@ impl WorkerSelectionStage { .iter() .map(|w| w.metadata().spec.runtime_type) .find(|runtime| { - all_decode - .iter() - .any(|w| w.metadata().spec.runtime_type == *runtime) + // The current EPD multimodal encoder adapter is TokenSpeed- + // specific. Do not select a shared SGLang/vLLM runtime only to + // reject it later during request building. + (!needs_encode || *runtime == RuntimeType::TokenSpeed) + && all_decode + .iter() + .any(|w| w.metadata().spec.runtime_type == *runtime) && (!needs_encode || all_encode .iter() @@ -563,10 +567,10 @@ fn encode_item_hashes(prep: &PreparationOutput) -> anyhow::Result>> } => processed_messages.multimodal_intermediate.as_ref(), _ => None, }; - let Some(MultimodalIntermediate::Precomputed(precomputed)) = intermediate else { + let Some(intermediate) = intermediate else { return Ok(Vec::new()); }; - multimodal::precomputed_encode_routing_hashes(precomputed) + multimodal::encode_routing_hashes(intermediate) } fn assign_encode_workers( diff --git a/model_gateway/src/routers/grpc/epd_encode.rs b/model_gateway/src/routers/grpc/epd_encode.rs index b37f1d708..c137f9c3a 100644 --- a/model_gateway/src/routers/grpc/epd_encode.rs +++ b/model_gateway/src/routers/grpc/epd_encode.rs @@ -15,7 +15,7 @@ use uuid::Uuid; use super::{ client::GrpcClient, context::{ClientSelection, WorkerSelection}, - multimodal::{assemble_tokenspeed, MultimodalIntermediate, PrecomputedMultimodalIntermediate}, + multimodal::{assemble_tokenspeed_for_encode, MultimodalIntermediate}, proto_wrapper::{ cleanup_mm_shm_handles, cleanup_tokenspeed_items_encoder_shm, collect_tokenspeed_multimodal_inputs_shm_handles, EncodeItemBootstrapInfo, @@ -159,20 +159,16 @@ pub(crate) fn build_plan_from_intermediate( clients: Option<&ClientSelection>, workers: Option<&WorkerSelection>, ) -> Result { - match intermediate { - MultimodalIntermediate::Precomputed(precomputed) => { - build_plan(precomputed, clients, workers) - } - } + build_plan(intermediate, clients, workers) } fn build_plan( - precomputed: &PrecomputedMultimodalIntermediate, + intermediate: &MultimodalIntermediate, clients: Option<&ClientSelection>, workers: Option<&WorkerSelection>, ) -> Result { let workers = workers.ok_or_else(|| anyhow!("Worker selection stage not completed"))?; - let items = prepare_items(precomputed, clients, Some(workers))?; + let items = prepare_items(intermediate, clients, Some(workers))?; if items.is_empty() { return Ok(EncodePlan { bootstrap_info: Vec::new(), @@ -233,7 +229,7 @@ fn build_plan( } pub(crate) fn prepare_items( - precomputed: &PrecomputedMultimodalIntermediate, + intermediate: &MultimodalIntermediate, clients: Option<&ClientSelection>, workers: Option<&WorkerSelection>, ) -> Result> { @@ -242,7 +238,7 @@ pub(crate) fn prepare_items( ClientSelection::Disaggregated { prefill: GrpcClient::TokenSpeed(_), .. - } => prepare_tokenspeed_items(precomputed, workers), + } => prepare_tokenspeed_items(intermediate, workers), ClientSelection::Disaggregated { prefill, .. } => Err(anyhow!( "EPD encode is not implemented for {} backend", backend_name(prefill) @@ -254,10 +250,10 @@ pub(crate) fn prepare_items( } fn prepare_tokenspeed_items( - precomputed: &PrecomputedMultimodalIntermediate, + intermediate: &MultimodalIntermediate, workers: Option<&WorkerSelection>, ) -> Result> { - let tokenspeed_mm = assemble_tokenspeed(precomputed, workers, false)?; + let tokenspeed_mm = assemble_tokenspeed_for_encode(intermediate, workers)?; let shm_enabled = tokenspeed_mm.shm_enabled; let shm_min_bytes = tokenspeed_mm.shm_min_bytes; Ok(tokenspeed_mm diff --git a/model_gateway/src/routers/grpc/harmony/builder.rs b/model_gateway/src/routers/grpc/harmony/builder.rs index f10efeba3..c902a06e2 100644 --- a/model_gateway/src/routers/grpc/harmony/builder.rs +++ b/model_gateway/src/routers/grpc/harmony/builder.rs @@ -30,6 +30,38 @@ use crate::routers::grpc::{proto_wrapper::ProtoOutputLogProbs, utils}; /// Global Harmony encoding (lazy-initialized) static HARMONY_ENCODING: OnceLock = OnceLock::new(); +fn reject_chat_audio(messages: &[ChatMessage]) -> Result<(), String> { + let has_audio = messages.iter().any(|message| match message { + ChatMessage::System { content, .. } + | ChatMessage::User { content, .. } + | ChatMessage::Tool { content, .. } + | ChatMessage::Developer { content, .. } => content_contains_audio(content), + ChatMessage::Assistant { content, .. } => { + content.as_ref().is_some_and(content_contains_audio) + } + ChatMessage::Function { .. } => false, + }); + + if has_audio { + Err( + "Harmony does not support audio content parts; use the regular multimodal path" + .to_string(), + ) + } else { + Ok(()) + } +} + +fn content_contains_audio(content: &MessageContent) -> bool { + matches!( + content, + MessageContent::Parts(parts) + if parts + .iter() + .any(|part| matches!(part, ContentPart::AudioUrl { .. } | ContentPart::InputAudio { .. })) + ) +} + /// Get or initialize the Harmony encoding /// /// Uses HarmonyGptOss encoding which supports the gpt-oss model family. @@ -262,6 +294,8 @@ impl HarmonyBuilder { &self, request: &ChatCompletionRequest, ) -> Result { + reject_chat_audio(&request.messages)?; + let mut all_messages = Vec::new(); let sys_msg = self.build_system_message_from_chat(request); @@ -1163,12 +1197,40 @@ mod tests { //! signature (`type image_generation = (_: …) => any;`) that //! gpt-oss is trained to emit calls against. - use openai_protocol::responses::{ - ImageGenerationTool, ResponseInput, ResponseTool, ResponsesRequest, + use openai_protocol::{ + common::{AudioUrl, InputAudio}, + responses::{ImageGenerationTool, ResponseInput, ResponseTool, ResponsesRequest}, }; use super::*; + #[test] + fn chat_audio_is_explicitly_rejected() { + let audio_parts = [ + ContentPart::AudioUrl { + audio_url: AudioUrl { + url: "https://example.com/audio.wav".to_string(), + }, + }, + ContentPart::InputAudio { + input_audio: InputAudio { + data: "UklGRg==".to_string(), + format: "wav".to_string(), + }, + }, + ]; + + for part in audio_parts { + let messages = vec![ChatMessage::User { + content: MessageContent::Parts(vec![part]), + name: None, + }]; + let error = reject_chat_audio(&messages).expect_err("Harmony must reject chat audio"); + assert!(error.contains("audio content parts")); + assert!(error.contains("regular multimodal path")); + } + } + /// Invariant: `image_generation` must never be advertised as a /// gpt-oss native builtin tool. If a future change re-adds it, /// gpt-oss's behavior becomes undefined (hallucinated tool call diff --git a/model_gateway/src/routers/grpc/mod.rs b/model_gateway/src/routers/grpc/mod.rs index b3d1c7d3a..3d1406b5b 100644 --- a/model_gateway/src/routers/grpc/mod.rs +++ b/model_gateway/src/routers/grpc/mod.rs @@ -1,6 +1,9 @@ //! gRPC router implementations -use openai_protocol::common::StringOrArray; +use axum::response::Response; +use openai_protocol::{chat::ChatCompletionRequest, common::StringOrArray}; + +use crate::routers::error; pub mod client; // Used by core/ pub(crate) mod common; @@ -19,6 +22,30 @@ pub mod utils; // Used by routers/http and bindings/golang // Re-export for convenience pub use proto_wrapper::{MultimodalData, TensorBytes}; +fn validate_text_only_output(request: &ChatCompletionRequest) -> Result<(), Box> { + if request.return_audio == Some(true) { + return Err(Box::new(error::bad_request( + "audio_output_not_supported", + "'return_audio' must be false because the gRPC backend only supports text output", + ))); + } + + if let Some(modality) = request.modalities.as_ref().and_then(|modalities| { + modalities + .iter() + .find(|modality| modality.as_str() != "text") + }) { + return Err(Box::new(error::bad_request( + "unsupported_output_modality", + format!( + "unsupported output modality '{modality}'; the gRPC backend only supports text output" + ), + ))); + } + + Ok(()) +} + /// Processed chat messages ready for gRPC generation #[derive(Debug)] pub struct ProcessedMessages { @@ -29,3 +56,45 @@ pub struct ProcessedMessages { pub(crate) multimodal_intermediate: Option, pub stop_sequences: Option, } + +#[cfg(test)] +mod tests { + use super::*; + use crate::routers::error::extract_error_code_from_response; + + #[test] + fn validates_grpc_text_only_output() { + for modalities in [ + None, + Some(vec![]), + Some(vec!["text".to_string()]), + Some(vec!["text".to_string(), "text".to_string()]), + ] { + let request = ChatCompletionRequest { + modalities, + ..Default::default() + }; + assert!(validate_text_only_output(&request).is_ok()); + } + + let audio_request = ChatCompletionRequest { + return_audio: Some(true), + ..Default::default() + }; + let response = validate_text_only_output(&audio_request).unwrap_err(); + assert_eq!( + extract_error_code_from_response(&response), + "audio_output_not_supported" + ); + + let modality_request = ChatCompletionRequest { + modalities: Some(vec!["text".to_string(), "audio".to_string()]), + ..Default::default() + }; + let response = validate_text_only_output(&modality_request).unwrap_err(); + assert_eq!( + extract_error_code_from_response(&response), + "unsupported_output_modality" + ); + } +} diff --git a/model_gateway/src/routers/grpc/multimodal/assemble.rs b/model_gateway/src/routers/grpc/multimodal/assemble.rs index 47bce6b1b..a18028392 100644 --- a/model_gateway/src/routers/grpc/multimodal/assemble.rs +++ b/model_gateway/src/routers/grpc/multimodal/assemble.rs @@ -4,11 +4,15 @@ //! For TokenSpeed this also splits the batched preprocessing output into //! per-item encoder inputs and per-item content hashes used for encode routing. -use std::{collections::HashMap, time::Instant}; +use std::{ + collections::{HashMap, HashSet}, + time::Instant, +}; use anyhow::{Context, Result}; use llm_multimodal::{ - FieldLayout, Modality, ModelSpecificValue, PlaceholderRange, PreprocessedEncoderInputs, + EncoderFieldLayouts, FieldLayout, Modality, ModelSpecificValue, PlaceholderRange, + PreprocessedEncoderInputs, }; use ndarray::ArrayViewD; use smg_grpc_client::common_proto as common; @@ -21,7 +25,7 @@ use super::{ serialize_encoder_input, serialize_model_specific, slice_array_axis0, }, transport::{mm_encoder_input_dtype, resolve_mm_shm_enabled, resolve_mm_shm_min_bytes}, - MultimodalIntermediate, PrecomputedMultimodalIntermediate, + MediaBatch, MultimodalIntermediate, PrecomputedMultimodalIntermediate, PromptBinding, }; use crate::routers::grpc::{ client::GrpcClient, @@ -55,45 +59,47 @@ pub(crate) async fn assemble_multimodal_data_after_encode( assemble_multimodal_data_impl(intermediate, client, workers, true).await } -#[expect( - clippy::unreachable, - reason = "MLX multimodal rejected by caller before reaching here" -)] async fn assemble_multimodal_data_impl( intermediate: MultimodalIntermediate, client: &GrpcClient, workers: Option<&WorkerSelection>, omit_prefill_pixels: bool, ) -> Result { - match intermediate { - MultimodalIntermediate::Precomputed(precomputed) => match client { - GrpcClient::Sglang(_) => { - ensure_image_only(&precomputed, "SGLang")?; - Ok(MultimodalData::Sglang(assemble_sglang(precomputed))) - } - GrpcClient::Vllm(_) => { - ensure_image_or_video(&precomputed, "vLLM")?; - Ok(MultimodalData::Vllm(assemble_vllm(precomputed, workers))) - } - GrpcClient::Trtllm(_) => { - ensure_image_only(&precomputed, "TRT-LLM")?; - Ok(MultimodalData::Trtllm(assemble_trtllm(precomputed))) - } - GrpcClient::TokenSpeed(_) => { - let options = - tokenspeed_assembly_options(precomputed.modality, workers, omit_prefill_pixels); - let pending = tokio::task::spawn_blocking(move || { - assemble_tokenspeed_with_options(&precomputed, options) - .map(PendingTokenSpeedAssembly::new) + validate_intermediate(&intermediate)?; + match client { + GrpcClient::Sglang(_) => { + let batch = into_single_image_batch(intermediate, "SGLang")?; + Ok(MultimodalData::Sglang(assemble_sglang(batch)?)) + } + GrpcClient::Vllm(_) => { + let batch = into_single_vision_batch(intermediate, "vLLM")?; + Ok(MultimodalData::Vllm(assemble_vllm(batch, workers)?)) + } + GrpcClient::Trtllm(_) => { + let batch = into_single_image_batch(intermediate, "TRT-LLM")?; + Ok(MultimodalData::Trtllm(assemble_trtllm(batch)?)) + } + GrpcClient::TokenSpeed(_) => { + let options = intermediate + .batches() + .iter() + .map(|batch| { + tokenspeed_assembly_options( + batch.media.modality(), + workers, + omit_prefill_pixels, + ) }) - .await - .context("TokenSpeed multimodal assembly task failed")??; - Ok(MultimodalData::TokenSpeed(pending.into_inner()?)) - } - GrpcClient::Mlx(_) => unreachable!( - "caller rejects multimodal for MLX in build_chat_request/build_messages_request" - ), - }, + .collect::>(); + let batches = intermediate.into_batches(); + let pending = tokio::task::spawn_blocking(move || { + assemble_tokenspeed_batches(&batches, options).map(PendingTokenSpeedAssembly::new) + }) + .await + .context("TokenSpeed multimodal assembly task failed")??; + Ok(MultimodalData::TokenSpeed(pending.into_inner()?)) + } + GrpcClient::Mlx(_) => anyhow::bail!("MLX does not support multimodal inputs"), } } @@ -126,103 +132,92 @@ impl Drop for PendingTokenSpeedAssembly { } } -fn ensure_image_only( - intermediate: &PrecomputedMultimodalIntermediate, +fn into_single_image_batch( + intermediate: MultimodalIntermediate, backend: &str, -) -> Result<()> { - if intermediate.modality != Modality::Image { - return Err(anyhow::anyhow!( - "{backend} multimodal path currently supports image inputs only; got {}", - intermediate.modality - )); - } - Ok(()) +) -> Result { + anyhow::ensure!( + intermediate.batches().len() == 1, + "{backend} multimodal path requires exactly one image batch; got {} batches", + intermediate.batches().len() + ); + let mut batches = intermediate.into_batches(); + let batch = batches + .pop() + .context("multimodal intermediate is missing its sole batch")?; + anyhow::ensure!( + matches!(&batch.media, MediaBatch::Images(_)), + "{backend} multimodal path currently supports image inputs only; got {}", + batch.media.modality() + ); + Ok(batch) } -/// Backends that accept both image and video (single-modality per request; -/// mixed image+video is already rejected upstream in `process`). Audio is not -/// supported on these paths. -fn ensure_image_or_video( - intermediate: &PrecomputedMultimodalIntermediate, +fn into_single_vision_batch( + intermediate: MultimodalIntermediate, backend: &str, -) -> Result<()> { - match intermediate.modality { - // Adds video to the previous image-only gate; ImageEmbeds/Audio stay - // rejected (unchanged from ensure_image_only). - Modality::Image | Modality::Video => Ok(()), - Modality::ImageEmbeds | Modality::Audio => Err(anyhow::anyhow!( - "{backend} multimodal path supports image and video inputs; got {}", - intermediate.modality - )), - } +) -> Result { + anyhow::ensure!( + intermediate.batches().len() == 1, + "{backend} multimodal path requires exactly one vision batch; got {} batches", + intermediate.batches().len() + ); + let mut batches = intermediate.into_batches(); + let batch = batches + .pop() + .context("multimodal intermediate is missing its sole batch")?; + anyhow::ensure!( + matches!(&batch.media, MediaBatch::Images(_) | MediaBatch::Videos(_)), + "{backend} multimodal path currently supports image and video inputs only; got {}", + batch.media.modality() + ); + Ok(batch) } -fn assemble_sglang(intermediate: PrecomputedMultimodalIntermediate) -> SglangMultimodalData { +fn assemble_sglang( + intermediate: PrecomputedMultimodalIntermediate, +) -> Result { let (pixel_values, pixel_values_shape) = serialize_encoder_input(&intermediate.preprocessed); let model_specific_tensors = serialize_model_specific(intermediate.preprocessed.model_specific); - let image_data = intermediate - .images - .iter() - .map(|f| f.raw_bytes.to_vec()) - .collect(); - // Use patch-only offsets when available and non-empty; fall back to full structural ranges. - let mm_placeholders = intermediate - .patch_offsets - .filter(|offsets| !offsets.is_empty()) - .unwrap_or_else(|| { - intermediate - .placeholders - .iter() - .map(|p| (p.offset as u32, p.length as u32)) - .collect() - }); + let MediaBatch::Images(images) = &intermediate.media else { + anyhow::bail!("SGLang assembly requires an image batch"); + }; + let image_data = images.iter().map(|f| f.raw_bytes.to_vec()).collect(); + let mm_placeholders = placeholders_for_bindings(&intermediate.bindings, true)?; - SglangMultimodalData { + Ok(SglangMultimodalData { image_data, pixel_values, pixel_values_shape, model_specific_tensors, im_token_id: intermediate.placeholder_token_id, mm_placeholders, - } + }) } fn assemble_vllm( intermediate: PrecomputedMultimodalIntermediate, workers: Option<&WorkerSelection>, -) -> VllmMultimodalData { +) -> Result { let (pixel_values, pixel_values_shape) = serialize_encoder_input(&intermediate.preprocessed); let model_specific_tensors = serialize_model_specific(intermediate.preprocessed.model_specific); - let modality = match intermediate.modality { - Modality::Video => common::Modality::Video, - Modality::Audio => common::Modality::Audio, - Modality::Image | Modality::ImageEmbeds => common::Modality::Image, - }; - let is_video = modality == common::Modality::Video; - // Hashes track the per-item media for encoder-output caching: videos for the - // video path, images otherwise. - let mm_hashes = if is_video { - intermediate - .videos - .iter() - .map(|video| video.hash.clone()) - .collect() - } else { - intermediate - .images - .iter() - .map(|frame| frame.hash.clone()) - .collect() + let (modality, mm_hashes) = match &intermediate.media { + MediaBatch::Images(images) => ( + common::Modality::Image, + images.iter().map(|frame| frame.hash.clone()).collect(), + ), + MediaBatch::Videos(videos) => ( + common::Modality::Video, + videos.iter().map(|video| video.hash.clone()).collect(), + ), + MediaBatch::Audios(_) => { + anyhow::bail!("vLLM assembly requires an image or video batch") + } }; - let mm_placeholders = intermediate - .placeholders - .iter() - .map(|p| (p.offset as u32, p.length as u32)) - .collect(); - let batched_keys = PreprocessedEncoderInputs::batched_keys(&intermediate.field_layouts); - let flat_keys = PreprocessedEncoderInputs::flat_keys(&intermediate.field_layouts); + let mm_placeholders = placeholders_for_bindings(&intermediate.bindings, false)?; + let (batched_keys, flat_keys) = vllm_field_layout_keys(&intermediate.field_layouts); - VllmMultimodalData { + Ok(VllmMultimodalData { pixel_values, pixel_values_shape, model_specific_tensors, @@ -235,27 +230,57 @@ fn assemble_vllm( modality, shm_enabled: resolve_mm_shm_enabled(workers, false), shm_min_bytes: resolve_mm_shm_min_bytes(workers), + }) +} + +/// Translate the neutral layout contract to vLLM's legacy HF field names. +fn vllm_field_layout_keys(layouts: &EncoderFieldLayouts) -> (Vec, HashMap) { + let mut batched_keys = PreprocessedEncoderInputs::batched_keys(&layouts.model_specific); + let mut flat_keys = PreprocessedEncoderInputs::flat_keys(&layouts.model_specific); + match &layouts.encoder_input { + FieldLayout::Batched => batched_keys.push("pixel_values".to_string()), + FieldLayout::Flat { sizes_key } => { + flat_keys.insert("pixel_values".to_string(), sizes_key.clone()); + } } + (batched_keys, flat_keys) } -fn assemble_trtllm(intermediate: PrecomputedMultimodalIntermediate) -> TrtllmMultimodalData { - let image_data = intermediate - .images - .iter() - .map(|f| f.raw_bytes.to_vec()) - .collect(); - TrtllmMultimodalData { image_data } +fn assemble_trtllm( + intermediate: PrecomputedMultimodalIntermediate, +) -> Result { + let MediaBatch::Images(images) = &intermediate.media else { + anyhow::bail!("TRT-LLM assembly requires an image batch"); + }; + let image_data = images.iter().map(|f| f.raw_bytes.to_vec()).collect(); + Ok(TrtllmMultimodalData { image_data }) } -pub(crate) fn assemble_tokenspeed( +#[cfg(test)] +fn assemble_tokenspeed( intermediate: &PrecomputedMultimodalIntermediate, workers: Option<&WorkerSelection>, skip_pixel_values: bool, ) -> Result { - let options = tokenspeed_assembly_options(intermediate.modality, workers, skip_pixel_values); + validate_precomputed_batch(intermediate)?; + let options = + tokenspeed_assembly_options(intermediate.media.modality(), workers, skip_pixel_values); assemble_tokenspeed_with_options(intermediate, options) } +pub(crate) fn assemble_tokenspeed_for_encode( + intermediate: &MultimodalIntermediate, + workers: Option<&WorkerSelection>, +) -> Result { + validate_intermediate(intermediate)?; + let options = intermediate + .batches() + .iter() + .map(|batch| tokenspeed_assembly_options(batch.media.modality(), workers, false)) + .collect(); + assemble_tokenspeed_batches(intermediate.batches(), options) +} + struct TokenSpeedAssemblyOptions { shm_enabled: bool, shm_min_bytes: usize, @@ -288,28 +313,38 @@ fn assemble_tokenspeed_with_options( encoder_input_dtype, skip_pixel_values, } = options; - // Use patch-only offsets when available and non-empty; fall back to full structural ranges. - let patch_offsets = intermediate - .patch_offsets - .clone() - .filter(|offsets| !offsets.is_empty()) - .unwrap_or_default(); - - let modality = match intermediate.modality { + let modality = match intermediate.media.modality() { Modality::Image => TokenSpeedModality::Image, Modality::Video => TokenSpeedModality::Video, Modality::Audio => TokenSpeedModality::Audio, Modality::ImageEmbeds => TokenSpeedModality::Image, }; - let item_count = precomputed_multimodal_item_count(intermediate)?; + let item_count = intermediate.media.len(); // Build items imperatively so that if any step fails partway we can unlink // the /dev/shm segments already created for prior items' encoder inputs // (and this item's, once created). `?`/`collect` would drop those // `TokenSpeedTensor::Shm` handles without ever reaching the send-path // cleanup, leaking files until the next sweep. + let mut ordered_bindings = intermediate.bindings.iter().collect::>(); + ordered_bindings.sort_by_key(|binding| binding.prompt_ordinal); let mut items: Vec = Vec::with_capacity(item_count); - for item_index in 0..item_count { + for binding in ordered_bindings { + let item_index = binding.item_index; + let mm_placeholders = match placeholders_for_binding(binding, true) { + Ok(value) => value, + Err(error) => { + cleanup_tokenspeed_items_encoder_shm(&items, None); + return Err(error); + } + }; + let content_hash = match content_hash_for_item(intermediate, item_index) { + Ok(value) => value, + Err(error) => { + cleanup_tokenspeed_items_encoder_shm(&items, None); + return Err(error); + } + }; let encoder_input_started = Instant::now(); // EPD prefill: the embedding arrives over Mooncake and this item's // encoder_input is stripped downstream (clear_mm_pixel_values), so skip @@ -319,7 +354,7 @@ fn assemble_tokenspeed_with_options( } else { let item_encoder_input = match encoder_input_for_item( &intermediate.preprocessed, - &intermediate.field_layouts, + &intermediate.field_layouts.encoder_input, item_index, ) { Ok(value) => value, @@ -339,7 +374,7 @@ fn assemble_tokenspeed_with_options( let model_specific_started = Instant::now(); let model_specific_tensors = match serialize_model_specific_for_item( &intermediate.preprocessed.model_specific, - &intermediate.field_layouts, + &intermediate.field_layouts.model_specific, item_index, ) { Ok(value) => value, @@ -351,10 +386,6 @@ fn assemble_tokenspeed_with_options( } }; let model_specific_serialize_ms = model_specific_started.elapsed().as_secs_f64() * 1000.0; - let mm_placeholders = - placeholders_for_item(item_index, &intermediate.placeholders, &patch_offsets); - let content_hash = content_hash_for_item(intermediate.modality, intermediate, item_index); - if log_timing { info!( modality = ?modality, @@ -395,61 +426,175 @@ fn assemble_tokenspeed_with_options( }) } -fn precomputed_multimodal_item_count( - intermediate: &PrecomputedMultimodalIntermediate, -) -> Result { - let media_count = match intermediate.modality { - Modality::Image | Modality::ImageEmbeds => intermediate.images.len(), - Modality::Video => intermediate.videos.len(), - Modality::Audio => 0, - }; - let token_count = intermediate.preprocessed.feature_token_counts.len(); - let placeholder_count = intermediate.placeholders.len(); - let item_count = token_count.max(media_count).max(placeholder_count); +fn assemble_tokenspeed_batches( + batches: &[PrecomputedMultimodalIntermediate], + options: Vec, +) -> Result { anyhow::ensure!( - item_count > 0, - "precomputed multimodal assembly requires at least one item" + batches.len() == options.len(), + "multimodal batch/assembly option count mismatch" ); - if media_count > 0 { - anyhow::ensure!( - media_count == item_count, - "precomputed multimodal assembly media count mismatch: modality={}, media_count={media_count}, item_count={item_count}", - intermediate.modality - ); + + let shm_enabled = options + .first() + .map(|opts| opts.shm_enabled) + .unwrap_or(false); + let shm_min_bytes = options.first().map(|opts| opts.shm_min_bytes).unwrap_or(0); + let mut ordered_items: Vec<(usize, TokenSpeedMultimodalItem)> = Vec::new(); + + for (batch, options) in batches.iter().zip(options) { + match assemble_tokenspeed_with_options(batch, options) { + Ok(data) => { + let mut ordinals = batch + .bindings + .iter() + .map(|binding| binding.prompt_ordinal) + .collect::>(); + ordinals.sort_unstable(); + if ordinals.len() != data.items.len() { + cleanup_tokenspeed_items_encoder_shm(&data.items, None); + cleanup_ordered_tokenspeed_items(&ordered_items); + return Err(anyhow::anyhow!( + "TokenSpeed binding/item count mismatch for {}", + batch.media.modality() + )); + } + ordered_items.extend(ordinals.into_iter().zip(data.items)); + } + Err(error) => { + cleanup_ordered_tokenspeed_items(&ordered_items); + return Err(error); + } + } } + let items = into_prompt_order(ordered_items); + + Ok(TokenSpeedMultimodalData { + items, + shm_enabled, + shm_min_bytes, + }) +} + +fn cleanup_ordered_tokenspeed_items(items: &[(usize, TokenSpeedMultimodalItem)]) { + for (_, item) in items { + cleanup_tokenspeed_items_encoder_shm(std::slice::from_ref(item), None); + } +} + +fn into_prompt_order(mut entries: Vec<(usize, T)>) -> Vec { + entries.sort_unstable_by_key(|(prompt_ordinal, _)| *prompt_ordinal); + entries.into_iter().map(|(_, value)| value).collect() +} + +fn validate_intermediate(intermediate: &MultimodalIntermediate) -> Result<()> { anyhow::ensure!( - token_count == item_count, - "precomputed multimodal assembly token count mismatch: modality={}, token_count={token_count}, item_count={item_count}", - intermediate.modality + !intermediate.batches().is_empty(), + "multimodal intermediate requires at least one batch" + ); + let total_bindings = intermediate + .batches() + .iter() + .map(|batch| batch.bindings.len()) + .sum::(); + let mut prompt_ordinals = HashSet::with_capacity(total_bindings); + for batch in intermediate.batches() { + validate_precomputed_batch(batch)?; + for binding in &batch.bindings { + anyhow::ensure!( + binding.prompt_ordinal < total_bindings, + "multimodal prompt ordinal {} is outside 0..{total_bindings}", + binding.prompt_ordinal + ); + anyhow::ensure!( + prompt_ordinals.insert(binding.prompt_ordinal), + "duplicate multimodal prompt ordinal {}", + binding.prompt_ordinal + ); + } + } + Ok(()) +} + +fn validate_precomputed_batch(intermediate: &PrecomputedMultimodalIntermediate) -> Result<()> { + let modality = intermediate.media.modality(); + let media_count = intermediate.media.len(); + let token_count = intermediate.preprocessed.feature_token_counts.len(); + let binding_count = intermediate.bindings.len(); + anyhow::ensure!( + media_count > 0, + "precomputed {modality} batch requires at least one media item" ); anyhow::ensure!( - placeholder_count == item_count, - "precomputed multimodal assembly placeholder count mismatch: modality={}, placeholder_count={placeholder_count}, item_count={item_count}", - intermediate.modality + token_count == media_count, + "precomputed multimodal token count mismatch: modality={modality}, token_count={token_count}, media_count={media_count}" ); - Ok(item_count) + anyhow::ensure!( + binding_count == media_count, + "precomputed multimodal binding count mismatch: modality={modality}, binding_count={binding_count}, media_count={media_count}" + ); + + let mut item_indices = HashSet::with_capacity(binding_count); + for binding in &intermediate.bindings { + anyhow::ensure!( + binding.item_index < media_count, + "precomputed {modality} binding item index {} is outside 0..{media_count}", + binding.item_index + ); + anyhow::ensure!( + item_indices.insert(binding.item_index), + "duplicate precomputed {modality} binding for item {}", + binding.item_index + ); + anyhow::ensure!( + binding.structural.length > 0, + "precomputed {modality} binding for item {} has an empty structural range", + binding.item_index + ); + let structural_end = binding + .structural + .offset + .checked_add(binding.structural.length) + .context("structural prompt range overflow")?; + for patch in &binding.patches { + let patch_end = patch + .offset + .checked_add(patch.length) + .context("patch prompt range overflow")?; + anyhow::ensure!( + patch.length > 0 + && patch.offset >= binding.structural.offset + && patch_end <= structural_end, + "precomputed {modality} patch range ({}, {}) lies outside structural range ({}, {})", + patch.offset, + patch.length, + binding.structural.offset, + binding.structural.length + ); + } + } + Ok(()) } -pub(crate) fn precomputed_encode_routing_hashes( - intermediate: &PrecomputedMultimodalIntermediate, -) -> Result>> { - let item_count = precomputed_multimodal_item_count(intermediate)?; - Ok((0..item_count) - .map(|item_index| content_hash_for_item(intermediate.modality, intermediate, item_index)) - .collect()) +pub(crate) fn encode_routing_hashes(intermediate: &MultimodalIntermediate) -> Result>> { + validate_intermediate(intermediate)?; + let mut hashes = Vec::new(); + for batch in intermediate.batches() { + for binding in &batch.bindings { + hashes.push(( + binding.prompt_ordinal, + content_hash_for_item(batch, binding.item_index)?, + )); + } + } + Ok(into_prompt_order(hashes)) } fn encoder_input_for_item<'a>( preprocessed: &'a PreprocessedEncoderInputs, - field_layouts: &HashMap, + layout: &FieldLayout, item_index: usize, ) -> Result> { - // The field layout key remains "pixel_values" because it mirrors the - // HuggingFace/vLLM vision kwargs contract. Internally this tensor is the - // modality encoder input we pass to TokenSpeed. - let layout = field_layouts - .get("pixel_values") - .unwrap_or(&FieldLayout::Batched); match layout { FieldLayout::Batched => slice_array_axis0(&preprocessed.encoder_input, item_index, 1), FieldLayout::Flat { sizes_key } => { @@ -489,46 +634,61 @@ fn serialize_model_specific_for_item( Ok(serialized) } -fn placeholders_for_item( - item_index: usize, - placeholders: &[PlaceholderRange], - patch_offsets: &[(u32, u32)], -) -> Vec<(u32, u32)> { - let Some(placeholder) = placeholders.get(item_index) else { - return Vec::new(); - }; - let start = placeholder.offset as u32; - let end = start + placeholder.length as u32; - let item_patch_offsets = patch_offsets - .iter() - .copied() - .filter(|(offset, length)| *offset >= start && offset.saturating_add(*length) <= end) - .collect::>(); - if item_patch_offsets.is_empty() { - vec![(start, end - start)] +fn placeholders_for_bindings( + bindings: &[PromptBinding], + prefer_patches: bool, +) -> Result> { + let mut ordered = bindings.iter().collect::>(); + ordered.sort_by_key(|binding| binding.prompt_ordinal); + ordered + .into_iter() + .flat_map(|binding| { + let ranges = if prefer_patches && !binding.patches.is_empty() { + binding.patches.as_slice() + } else { + std::slice::from_ref(&binding.structural) + }; + ranges.iter() + }) + .map(placeholder_range_to_u32) + .collect() +} + +fn placeholders_for_binding( + binding: &PromptBinding, + prefer_patches: bool, +) -> Result> { + let ranges = if prefer_patches && !binding.patches.is_empty() { + binding.patches.as_slice() } else { - item_patch_offsets - } + std::slice::from_ref(&binding.structural) + }; + ranges.iter().map(placeholder_range_to_u32).collect() +} + +fn placeholder_range_to_u32(range: &PlaceholderRange) -> Result<(u32, u32)> { + Ok(( + u32::try_from(range.offset).context("multimodal placeholder offset exceeds u32")?, + u32::try_from(range.length).context("multimodal placeholder length exceeds u32")?, + )) } fn content_hash_for_item( - modality: Modality, intermediate: &PrecomputedMultimodalIntermediate, item_index: usize, -) -> Vec { - match modality { - Modality::Image | Modality::ImageEmbeds => intermediate - .images - .get(item_index) - .map(|image| hash_hex_strings(std::iter::once(image.hash.as_str()))) - .unwrap_or_default(), - Modality::Video => intermediate - .videos - .get(item_index) - .map(|video| hash_hex_strings(std::iter::once(video.hash.as_str()))) - .unwrap_or_default(), - Modality::Audio => Vec::new(), +) -> Result> { + let hash = match &intermediate.media { + MediaBatch::Images(items) => items.get(item_index).map(|item| item.hash.as_str()), + MediaBatch::Videos(items) => items.get(item_index).map(|item| item.hash.as_str()), + MediaBatch::Audios(items) => items.get(item_index).map(|item| item.hash.as_str()), } + .ok_or_else(|| { + anyhow::anyhow!( + "missing {} media item {item_index} for content hash", + intermediate.media.modality() + ) + })?; + Ok(hash_hex_strings(std::iter::once(hash))) } fn tensor_sizes_from_model_specific( @@ -566,7 +726,9 @@ fn hash_hex_strings<'a>(hashes: impl Iterator) -> Vec { mod tests { use std::{mem::size_of, path::Path, sync::Arc}; - use llm_multimodal::{ImageDetail, ImageFrame, VideoClip}; + use llm_multimodal::{ + audio::DecodedAudio, AudioClip, AudioSource, ImageDetail, ImageFrame, VideoClip, + }; use ndarray::{ArrayD, IxDyn}; use super::*; @@ -665,30 +827,42 @@ mod tests { ]; let intermediate = PrecomputedMultimodalIntermediate { - modality: Modality::Image, preprocessed, - images, - videos: vec![], - placeholders: vec![ - PlaceholderRange { - offset: 10, - length: 2, + media: MediaBatch::Images(images), + bindings: vec![ + PromptBinding { + item_index: 0, + prompt_ordinal: 0, + structural: PlaceholderRange { + offset: 10, + length: 2, + }, + patches: vec![PlaceholderRange { + offset: 10, + length: 2, + }], }, - PlaceholderRange { - offset: 20, - length: 2, + PromptBinding { + item_index: 1, + prompt_ordinal: 1, + structural: PlaceholderRange { + offset: 20, + length: 2, + }, + patches: vec![PlaceholderRange { + offset: 20, + length: 2, + }], }, ], - patch_offsets: Some(vec![(10, 2), (20, 2)]), placeholder_token_id: Some(151655), - field_layouts: HashMap::from([ - ( - "pixel_values".to_string(), - FieldLayout::flat("patches_per_image"), - ), - ("patches_per_image".to_string(), FieldLayout::Batched), - ("image_grid_thw".to_string(), FieldLayout::Batched), - ]), + field_layouts: EncoderFieldLayouts::new( + FieldLayout::flat("patches_per_image"), + HashMap::from([ + ("patches_per_image".to_string(), FieldLayout::Batched), + ("image_grid_thw".to_string(), FieldLayout::Batched), + ]), + ), keep_on_cpu_keys: vec![], }; @@ -772,30 +946,42 @@ mod tests { ]; let intermediate = PrecomputedMultimodalIntermediate { - modality: Modality::Video, preprocessed, - images: vec![], - videos, - placeholders: vec![ - PlaceholderRange { - offset: 30, - length: 2, + media: MediaBatch::Videos(videos), + bindings: vec![ + PromptBinding { + item_index: 0, + prompt_ordinal: 0, + structural: PlaceholderRange { + offset: 30, + length: 2, + }, + patches: vec![PlaceholderRange { + offset: 30, + length: 2, + }], }, - PlaceholderRange { - offset: 40, - length: 2, + PromptBinding { + item_index: 1, + prompt_ordinal: 1, + structural: PlaceholderRange { + offset: 40, + length: 2, + }, + patches: vec![PlaceholderRange { + offset: 40, + length: 2, + }], }, ], - patch_offsets: Some(vec![(30, 2), (40, 2)]), placeholder_token_id: Some(151656), - field_layouts: HashMap::from([ - ( - "pixel_values".to_string(), - FieldLayout::flat("patches_per_video"), - ), - ("patches_per_video".to_string(), FieldLayout::Batched), - ("video_grid_thw".to_string(), FieldLayout::Batched), - ]), + field_layouts: EncoderFieldLayouts::new( + FieldLayout::flat("patches_per_video"), + HashMap::from([ + ("patches_per_video".to_string(), FieldLayout::Batched), + ("video_grid_thw".to_string(), FieldLayout::Batched), + ]), + ), keep_on_cpu_keys: vec![], }; @@ -833,4 +1019,272 @@ mod tests { vec![1, 3] ); } + + #[test] + fn assemble_tokenspeed_splits_audio_items_as_bfloat16_by_default() { + let mut model_specific = HashMap::new(); + model_specific.insert( + "row_lengths".to_string(), + ModelSpecificValue::IntTensor { + data: vec![2, 2], + shape: vec![2], + }, + ); + + let preprocessed = PreprocessedEncoderInputs { + encoder_input: ArrayD::from_shape_vec( + IxDyn(&[4, 2]), + vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], + ) + .unwrap(), + feature_token_counts: vec![2, 2], + item_sizes: vec![(2, 2), (2, 2)], + model_specific, + }; + + let audios = vec![ + Arc::new(AudioClip::new( + bytes::Bytes::from_static(b"a"), + DecodedAudio { + samples: Vec::new(), + sample_rate: 16_000, + }, + AudioSource::InlineBytes, + "audio-hash-a".to_string(), + )), + Arc::new(AudioClip::new( + bytes::Bytes::from_static(b"b"), + DecodedAudio { + samples: Vec::new(), + sample_rate: 16_000, + }, + AudioSource::InlineBytes, + "audio-hash-b".to_string(), + )), + ]; + + let intermediate = PrecomputedMultimodalIntermediate { + preprocessed, + media: MediaBatch::Audios(audios), + bindings: vec![ + PromptBinding { + item_index: 0, + prompt_ordinal: 0, + structural: PlaceholderRange { + offset: 30, + length: 2, + }, + patches: vec![PlaceholderRange { + offset: 30, + length: 2, + }], + }, + PromptBinding { + item_index: 1, + prompt_ordinal: 1, + structural: PlaceholderRange { + offset: 40, + length: 2, + }, + patches: vec![PlaceholderRange { + offset: 40, + length: 2, + }], + }, + ], + placeholder_token_id: Some(42), + field_layouts: EncoderFieldLayouts::new( + FieldLayout::flat("row_lengths"), + HashMap::from([("row_lengths".to_string(), FieldLayout::Batched)]), + ), + keep_on_cpu_keys: vec![], + }; + + let assembled = assemble_tokenspeed(&intermediate, None, false).unwrap(); + assert_eq!(assembled.items.len(), 2); + + let first = &assembled.items[0]; + assert_eq!(first.modality, TokenSpeedModality::Audio); + assert_eq!(first.encoder_input.dtype, "bfloat16"); + assert_eq!(first.encoder_input.shape, vec![2, 2]); + assert_eq!(first.encoder_input.nbytes(), 4 * size_of::()); + assert_eq!(first.mm_placeholders, vec![(30, 2)]); + assert_eq!( + first.content_hash, + hash_hex_strings(std::iter::once("audio-hash-a")) + ); + assert_eq!(first.model_specific_tensors["row_lengths"].shape, vec![1]); + + let second = &assembled.items[1]; + assert_eq!(second.encoder_input.dtype, "bfloat16"); + assert_eq!(second.encoder_input.shape, vec![2, 2]); + assert_eq!(second.mm_placeholders, vec![(40, 2)]); + assert_eq!( + second.content_hash, + hash_hex_strings(std::iter::once("audio-hash-b")) + ); + } + + #[test] + fn tokenspeed_epd_items_and_hashes_follow_prompt_binding_order() { + let one_item_inputs = || PreprocessedEncoderInputs { + encoder_input: ArrayD::from_shape_vec(IxDyn(&[1, 1]), vec![1.0]).unwrap(), + feature_token_counts: vec![1], + item_sizes: vec![(1, 1)], + model_specific: HashMap::new(), + }; + let image_batch = PrecomputedMultimodalIntermediate { + preprocessed: one_item_inputs(), + media: MediaBatch::Images(vec![Arc::new(ImageFrame::new( + image::DynamicImage::new_rgb8(1, 1), + bytes::Bytes::from_static(b"image"), + ImageDetail::Auto, + llm_multimodal::ImageSource::InlineBytes, + "image-hash".to_string(), + ))]), + // Deliberately earlier by offset but later by explicit ordinal. + bindings: vec![PromptBinding { + item_index: 0, + prompt_ordinal: 2, + structural: PlaceholderRange { + offset: 1, + length: 1, + }, + patches: vec![], + }], + placeholder_token_id: Some(10), + field_layouts: EncoderFieldLayouts::default(), + keep_on_cpu_keys: vec![], + }; + let video_batch = PrecomputedMultimodalIntermediate { + preprocessed: one_item_inputs(), + media: MediaBatch::Videos(vec![Arc::new(VideoClip::new( + vec![image::DynamicImage::new_rgb8(1, 1)], + bytes::Bytes::from_static(b"video"), + llm_multimodal::VideoSource::InlineBytes, + "video-hash".to_string(), + ))]), + bindings: vec![PromptBinding { + item_index: 0, + prompt_ordinal: 1, + structural: PlaceholderRange { + offset: 50, + length: 1, + }, + patches: vec![], + }], + placeholder_token_id: Some(30), + field_layouts: EncoderFieldLayouts::default(), + keep_on_cpu_keys: vec![], + }; + let audio_batch = PrecomputedMultimodalIntermediate { + preprocessed: one_item_inputs(), + media: MediaBatch::Audios(vec![Arc::new(AudioClip::new( + bytes::Bytes::from_static(b"audio"), + DecodedAudio { + samples: Vec::new(), + sample_rate: 16_000, + }, + AudioSource::InlineBytes, + "audio-hash".to_string(), + ))]), + bindings: vec![PromptBinding { + item_index: 0, + prompt_ordinal: 0, + structural: PlaceholderRange { + offset: 100, + length: 1, + }, + patches: vec![], + }], + placeholder_token_id: Some(20), + field_layouts: EncoderFieldLayouts::default(), + keep_on_cpu_keys: vec![], + }; + let intermediate = + MultimodalIntermediate::try_new(vec![image_batch, video_batch, audio_batch]).unwrap(); + let routing_hashes = encode_routing_hashes(&intermediate).unwrap(); + let assembled = assemble_tokenspeed_for_encode(&intermediate, None).unwrap(); + + assert_eq!(assembled.items.len(), 3); + assert_eq!(assembled.items[0].modality, TokenSpeedModality::Audio); + assert_eq!(assembled.items[0].mm_placeholders, vec![(100, 1)]); + assert_eq!(assembled.items[1].modality, TokenSpeedModality::Video); + assert_eq!(assembled.items[1].mm_placeholders, vec![(50, 1)]); + assert_eq!(assembled.items[2].modality, TokenSpeedModality::Image); + assert_eq!(assembled.items[2].mm_placeholders, vec![(1, 1)]); + assert_eq!( + routing_hashes, + assembled + .items + .iter() + .map(|item| item.content_hash.clone()) + .collect::>() + ); + } + + #[test] + fn qwen_audio_batched_layout_slices_items_not_mel_bins() { + let mut model_specific = HashMap::new(); + model_specific.insert( + "feature_attention_mask".to_string(), + ModelSpecificValue::int_2d(vec![1, 1, 1, 1, 1, 1, 0, 0], 2, 4), + ); + model_specific.insert( + "audio_feature_lengths".to_string(), + ModelSpecificValue::int_1d(vec![4, 2]), + ); + let preprocessed = PreprocessedEncoderInputs { + encoder_input: ArrayD::from_shape_vec( + IxDyn(&[2, 3, 4]), + (0..24).map(|value| value as f32).collect(), + ) + .unwrap(), + feature_token_counts: vec![1, 1], + item_sizes: vec![(3, 4), (3, 2)], + model_specific, + }; + let layouts = EncoderFieldLayouts::new( + FieldLayout::Batched, + HashMap::from([ + ("feature_attention_mask".to_string(), FieldLayout::Batched), + ("audio_feature_lengths".to_string(), FieldLayout::Batched), + ]), + ); + + let second = encoder_input_for_item(&preprocessed, &layouts.encoder_input, 1).unwrap(); + assert_eq!(second.shape(), &[1, 3, 4]); + assert_eq!( + second.iter().copied().collect::>(), + (12..24).map(|v| v as f32).collect::>() + ); + + let extras = serialize_model_specific_for_item( + &preprocessed.model_specific, + &layouts.model_specific, + 1, + ) + .unwrap(); + assert_eq!(extras["feature_attention_mask"].shape, vec![1, 4]); + assert_eq!(extras["audio_feature_lengths"].shape, vec![1]); + } + + #[test] + fn vllm_layout_adapter_restores_legacy_primary_field_name() { + let flat = EncoderFieldLayouts::new( + FieldLayout::flat("patches_per_image"), + HashMap::from([("image_grid_thw".to_string(), FieldLayout::Batched)]), + ); + let (batched_keys, flat_keys) = vllm_field_layout_keys(&flat); + assert_eq!(batched_keys, vec!["image_grid_thw"]); + assert_eq!( + flat_keys.get("pixel_values").map(String::as_str), + Some("patches_per_image") + ); + + let batched = EncoderFieldLayouts::default(); + let (batched_keys, flat_keys) = vllm_field_layout_keys(&batched); + assert_eq!(batched_keys, vec!["pixel_values"]); + assert!(flat_keys.is_empty()); + } } diff --git a/model_gateway/src/routers/grpc/multimodal/config.rs b/model_gateway/src/routers/grpc/multimodal/config.rs index d87c3e563..ffdd7f0f2 100644 --- a/model_gateway/src/routers/grpc/multimodal/config.rs +++ b/model_gateway/src/routers/grpc/multimodal/config.rs @@ -6,8 +6,8 @@ use std::{path::Path, sync::Arc}; use anyhow::{Context, Result}; use dashmap::DashMap; use llm_multimodal::{ - MediaConnector, MediaConnectorConfig, ModelRegistry, PreProcessorConfig, - VisionProcessorRegistry, + AudioProcessorRegistry, MediaConnector, MediaConnectorConfig, ModelRegistry, + PreProcessorConfig, VisionProcessorRegistry, }; use tracing::{debug, warn}; @@ -201,6 +201,7 @@ pub(crate) fn load_video_preprocessor_config(base_dir: &Path) -> Option, pub vision_processor_registry: Arc, + pub audio_processor_registry: Arc, pub model_registry: Arc, /// Shared reference to the app-level multimodal config cache. pub config_registry: Arc, @@ -222,6 +223,7 @@ impl MultimodalComponents { Ok(Self { media_connector: Arc::new(media_connector), vision_processor_registry: Arc::new(VisionProcessorRegistry::with_defaults()), + audio_processor_registry: Arc::new(AudioProcessorRegistry::with_defaults()), model_registry: Arc::new(ModelRegistry::default()), config_registry, pixel_cache: pixel_cache_from_env(), diff --git a/model_gateway/src/routers/grpc/multimodal/detect.rs b/model_gateway/src/routers/grpc/multimodal/detect.rs index 44734aafb..8cd30a2f4 100644 --- a/model_gateway/src/routers/grpc/multimodal/detect.rs +++ b/model_gateway/src/routers/grpc/multimodal/detect.rs @@ -4,54 +4,18 @@ //! pipeline (`InputMessage`) funnel into the shared processing core; only the //! detection and extraction differ, because the input message types differ. -use llm_multimodal::{ImageDetail, MediaContentPart, Modality}; +use llm_multimodal::{ImageDetail, MediaContentPart}; use openai_protocol::{ chat::{ChatMessage, MessageContent}, common::ContentPart, messages::{ImageSource, InputContent, InputContentBlock, InputMessage, Role}, }; -/// Return the multimodal modalities present in OpenAI chat messages. -pub(crate) fn chat_modalities(messages: &[ChatMessage]) -> Vec { - let mut modalities = Vec::new(); - let mut push_unique = |modality| { - if !modalities.contains(&modality) { - modalities.push(modality); - } - }; - - for msg in messages { - let content = match msg { - ChatMessage::User { content, .. } => Some(content), - ChatMessage::System { content, .. } => Some(content), - ChatMessage::Developer { content, .. } => Some(content), - ChatMessage::Tool { content, .. } => Some(content), - _ => None, - }; +use super::plan::MediaPlan; - if let Some(MessageContent::Parts(parts)) = content { - for part in parts { - match part { - ContentPart::ImageUrl { .. } => push_unique(Modality::Image), - ContentPart::VideoUrl { .. } => push_unique(Modality::Video), - ContentPart::Text { .. } => {} - } - } - } - } - - modalities -} - -/// Check if any messages in the request contain multimodal content. -#[cfg(test)] -pub(crate) fn has_multimodal_content(messages: &[ChatMessage]) -> bool { - !chat_modalities(messages).is_empty() -} - -/// Extract multimodal content parts from OpenAI chat messages, +/// Extract media parts from OpenAI chat messages, /// converting protocol `ContentPart` to multimodal crate `MediaContentPart`. -pub(super) fn extract_content_parts(messages: &[ChatMessage]) -> Vec { +fn extract_media_parts(messages: &[ChatMessage]) -> Vec { let mut parts = Vec::new(); for msg in messages { @@ -74,8 +38,21 @@ pub(super) fn extract_content_parts(messages: &[ChatMessage]) -> Vec { - parts.push(MediaContentPart::Text { text: text.clone() }); + ContentPart::Text { .. } => {} + ContentPart::AudioUrl { audio_url } => { + parts.push(MediaContentPart::AudioUrl { + url: audio_url.url.clone(), + uuid: None, + }); + } + ContentPart::InputAudio { input_audio } => { + parts.push(MediaContentPart::AudioUrl { + url: format!( + "data:audio/{};base64,{}", + input_audio.format, input_audio.data + ), + uuid: None, + }); } ContentPart::VideoUrl { video_url } => { parts.push(MediaContentPart::VideoUrl { @@ -91,6 +68,11 @@ pub(super) fn extract_content_parts(messages: &[ChatMessage]) -> Vec MediaPlan { + MediaPlan::new(extract_media_parts(messages)) +} + /// Parse OpenAI detail string to multimodal ImageDetail enum. fn parse_detail(detail: &str) -> Option { match detail.to_ascii_lowercase().as_str() { @@ -105,24 +87,9 @@ fn parse_detail(detail: &str) -> Option { // Messages API multimodal detection and extraction // --------------------------------------------------------------------------- -/// Check if any messages in a Messages API request contain multimodal content. -pub(crate) fn has_multimodal_content_messages(messages: &[InputMessage]) -> bool { - messages.iter().any(|msg| { - if msg.role != Role::User { - return false; - } - match &msg.content { - InputContent::Blocks(blocks) => blocks - .iter() - .any(|block| matches!(block, InputContentBlock::Image(_))), - InputContent::String(_) => false, - } - }) -} - -/// Extract multimodal content parts from Messages API input messages, +/// Extract media parts from Messages API input messages, /// converting `InputContentBlock::Image` to multimodal crate `MediaContentPart`. -pub(super) fn extract_content_parts_messages(messages: &[InputMessage]) -> Vec { +fn extract_media_parts_messages(messages: &[InputMessage]) -> Vec { let mut parts = Vec::new(); for msg in messages { @@ -154,11 +121,7 @@ pub(super) fn extract_content_parts_messages(messages: &[InputMessage]) -> Vec { - parts.push(MediaContentPart::Text { - text: text_block.text.clone(), - }); - } + InputContentBlock::Text(_) => {} _ => {} } } @@ -167,14 +130,20 @@ pub(super) fn extract_content_parts_messages(messages: &[InputMessage]) -> Vec MediaPlan { + MediaPlan::new(extract_media_parts_messages(messages)) +} + #[cfg(test)] mod tests { - use openai_protocol::common::{ImageUrl, VideoUrl}; + use llm_multimodal::Modality; + use openai_protocol::common::{AudioUrl, ImageUrl, InputAudio, VideoUrl}; use super::*; #[test] - fn test_has_multimodal_content_with_images() { + fn media_plan_detects_image() { let messages = vec![ChatMessage::User { content: MessageContent::Parts(vec![ ContentPart::Text { @@ -190,11 +159,11 @@ mod tests { name: None, }]; - assert!(has_multimodal_content(&messages)); + assert_eq!(media_plan_chat(&messages).modalities(), &[Modality::Image]); } #[test] - fn test_has_multimodal_content_with_video() { + fn media_plan_detects_video() { let messages = vec![ChatMessage::User { content: MessageContent::Parts(vec![ContentPart::VideoUrl { video_url: VideoUrl { @@ -204,22 +173,35 @@ mod tests { name: None, }]; - assert!(has_multimodal_content(&messages)); - assert_eq!(chat_modalities(&messages), vec![Modality::Video]); + assert_eq!(media_plan_chat(&messages).modalities(), &[Modality::Video]); } #[test] - fn test_has_multimodal_content_text_only() { + fn media_plan_detects_audio() { + let messages = vec![ChatMessage::User { + content: MessageContent::Parts(vec![ContentPart::AudioUrl { + audio_url: AudioUrl { + url: "https://example.com/clip.wav".to_string(), + }, + }]), + name: None, + }]; + + assert_eq!(media_plan_chat(&messages).modalities(), &[Modality::Audio]); + } + + #[test] + fn media_plan_is_empty_for_string_text() { let messages = vec![ChatMessage::User { content: MessageContent::Text("Hello".to_string()), name: None, }]; - assert!(!has_multimodal_content(&messages)); + assert!(media_plan_chat(&messages).is_empty()); } #[test] - fn test_has_multimodal_content_parts_text_only() { + fn media_plan_is_empty_for_text_parts() { let messages = vec![ChatMessage::User { content: MessageContent::Parts(vec![ContentPart::Text { text: "Just text".to_string(), @@ -227,11 +209,11 @@ mod tests { name: None, }]; - assert!(!has_multimodal_content(&messages)); + assert!(media_plan_chat(&messages).is_empty()); } #[test] - fn test_extract_content_parts() { + fn extracts_image_media_part() { let messages = vec![ ChatMessage::System { content: MessageContent::Text("You are helpful".to_string()), @@ -253,15 +235,10 @@ mod tests { }, ]; - let parts = extract_content_parts(&messages); - assert_eq!(parts.len(), 2); + let parts = extract_media_parts(&messages); + assert_eq!(parts.len(), 1); match &parts[0] { - MediaContentPart::Text { text } => assert_eq!(text, "Describe this:"), - _ => panic!("Expected Text part"), - } - - match &parts[1] { MediaContentPart::ImageUrl { url, detail, .. } => { assert_eq!(url, "https://example.com/image.jpg"); assert_eq!(*detail, Some(ImageDetail::High)); @@ -271,7 +248,7 @@ mod tests { } #[test] - fn test_extract_video_content_parts() { + fn extracts_video_media_part() { let messages = vec![ChatMessage::User { content: MessageContent::Parts(vec![ContentPart::VideoUrl { video_url: VideoUrl { @@ -281,7 +258,7 @@ mod tests { name: None, }]; - let parts = extract_content_parts(&messages); + let parts = extract_media_parts(&messages); assert_eq!(parts.len(), 1); match &parts[0] { MediaContentPart::VideoUrl { url, .. } => { @@ -291,6 +268,51 @@ mod tests { } } + #[test] + fn extracts_audio_url_media_part() { + let messages = vec![ChatMessage::User { + content: MessageContent::Parts(vec![ContentPart::AudioUrl { + audio_url: AudioUrl { + url: "https://example.com/audio.wav".to_string(), + }, + }]), + name: None, + }]; + + let parts = extract_media_parts(&messages); + assert_eq!(parts.len(), 1); + match &parts[0] { + MediaContentPart::AudioUrl { url, .. } => { + assert_eq!(url, "https://example.com/audio.wav"); + } + _ => panic!("Expected AudioUrl part"), + } + } + + #[test] + fn extracts_inline_audio_as_data_url() { + let messages = vec![ChatMessage::User { + content: MessageContent::Parts(vec![ContentPart::InputAudio { + input_audio: InputAudio { + data: "UklGRg==".to_string(), + format: "wav".to_string(), + }, + }]), + name: None, + }]; + + assert_eq!(media_plan_chat(&messages).modalities(), &[Modality::Audio]); + + let parts = extract_media_parts(&messages); + assert_eq!(parts.len(), 1); + match &parts[0] { + MediaContentPart::AudioUrl { url, .. } => { + assert_eq!(url, "data:audio/wav;base64,UklGRg=="); + } + _ => panic!("Expected AudioUrl part"), + } + } + #[test] fn test_parse_detail() { assert_eq!(parse_detail("auto"), Some(ImageDetail::Auto)); diff --git a/model_gateway/src/routers/grpc/multimodal/mod.rs b/model_gateway/src/routers/grpc/multimodal/mod.rs index b35fd0e9c..0dafd82b1 100644 --- a/model_gateway/src/routers/grpc/multimodal/mod.rs +++ b/model_gateway/src/routers/grpc/multimodal/mod.rs @@ -14,34 +14,37 @@ //! namespace verification. use std::{ - collections::HashMap, + collections::HashSet, sync::{Arc, OnceLock}, }; use llm_multimodal::{ - FieldLayout, ImageFrame, Modality, PlaceholderRange, PreprocessedEncoderInputs, VideoClip, + AudioClip, EncoderFieldLayouts, ImageFrame, Modality, PlaceholderRange, + PreprocessedEncoderInputs, VideoClip, }; mod assemble; mod config; mod detect; mod pixel_cache; +mod plan; mod process; mod serialize; mod transport; pub(crate) use assemble::{ - assemble_multimodal_data, assemble_multimodal_data_after_encode, assemble_tokenspeed, - precomputed_encode_routing_hashes, + assemble_multimodal_data, assemble_multimodal_data_after_encode, + assemble_tokenspeed_for_encode, encode_routing_hashes, }; pub(crate) use config::{ load_preprocessor_config_file, load_video_preprocessor_config, MultimodalComponents, MultimodalConfigRegistry, MultimodalModelConfig, }; -pub(crate) use detect::{chat_modalities, has_multimodal_content_messages}; -pub(crate) use process::{ - process_multimodal, process_multimodal_messages, resolve_placeholder_token, +pub(crate) use detect::{media_plan_chat, media_plan_messages}; +pub(crate) use plan::{ + prepare_placeholder_tokens, validate_rendered_media_anchors, PlaceholderTokens, }; +pub(crate) use process::process_multimodal_plan; pub(crate) use transport::init_mm_transport_defaults; #[cfg(feature = "mm-rdma")] pub(crate) use transport::mm_default_transport_is_rdma; @@ -73,28 +76,97 @@ pub(crate) struct MultimodalOutput { /// The assembly stage converts this into a backend-specific `MultimodalData` /// variant once the target backend is known (after worker selection). #[derive(Debug)] -pub(crate) enum MultimodalIntermediate { - Precomputed(PrecomputedMultimodalIntermediate), +pub(crate) struct MultimodalIntermediate { + /// Independently preprocessed modality batches sharing one expanded prompt. + /// A single-modality request is represented by a one-element vector. + batches: Vec, +} + +impl MultimodalIntermediate { + pub(crate) fn try_new(batches: Vec) -> anyhow::Result { + anyhow::ensure!( + !batches.is_empty(), + "multimodal intermediate requires at least one batch" + ); + let mut modalities = HashSet::with_capacity(batches.len()); + for batch in &batches { + let modality = batch.media.modality(); + anyhow::ensure!( + modalities.insert(modality), + "multimodal intermediate contains duplicate {modality} batches" + ); + anyhow::ensure!( + batch.media.len() > 0, + "multimodal intermediate contains an empty {modality} batch" + ); + } + Ok(Self { batches }) + } + + pub(crate) fn batches(&self) -> &[PrecomputedMultimodalIntermediate] { + &self.batches + } + + pub(crate) fn into_batches(self) -> Vec { + self.batches + } +} + +/// Raw media for one preprocessed batch. +/// +/// Encoding the modality in the enum prevents contradictory states such as an +/// audio batch carrying images or an image batch carrying both images and +/// videos. +#[derive(Debug, Clone)] +pub(crate) enum MediaBatch { + Images(Vec>), + Audios(Vec>), + Videos(Vec>), +} + +impl MediaBatch { + pub(crate) fn modality(&self) -> Modality { + match self { + Self::Images(_) => Modality::Image, + Self::Audios(_) => Modality::Audio, + Self::Videos(_) => Modality::Video, + } + } + + pub(crate) fn len(&self) -> usize { + match self { + Self::Images(items) => items.len(), + Self::Audios(items) => items.len(), + Self::Videos(items) => items.len(), + } + } +} + +/// Explicit association between one media item and its expanded prompt span. +#[derive(Debug, Clone)] +pub(crate) struct PromptBinding { + /// Index of the media/preprocessed item within its modality batch. + pub item_index: usize, + /// Position of this media item among all modalities in the rendered prompt. + pub prompt_ordinal: usize, + /// Full replacement span, including structural tokens. + pub structural: PlaceholderRange, + /// Patch-only spans within `structural`. + pub patches: Vec, } #[derive(Debug)] pub(crate) struct PrecomputedMultimodalIntermediate { - /// Active modality for this preprocessed payload. - pub modality: Modality, /// Preprocessed encoder input and model-specific tensors (not yet serialized). pub preprocessed: PreprocessedEncoderInputs, - /// Raw image frames (bytes + blake3 hashes). - pub images: Vec>, - /// Raw video clips (bytes + blake3 hashes + sampled frames). - pub videos: Vec>, - /// Full structural placeholder ranges (offset, length). - pub placeholders: Vec, - /// Patch-only placeholder offsets for sglang. - pub patch_offsets: Option>, + /// Raw media whose variant determines this batch's modality. + pub media: MediaBatch, + /// Exact media-to-prompt associations for this batch. + pub bindings: Vec, /// Placeholder token ID from model config for the active modality. pub placeholder_token_id: Option, - /// Per-tensor field layout classification from the model spec. - pub field_layouts: HashMap, + /// Primary encoder input and model-specific side-tensor layouts. + pub field_layouts: EncoderFieldLayouts, /// Tensor keys that should remain on CPU (vLLM `keep_on_cpu` hint). pub keep_on_cpu_keys: Vec, } diff --git a/model_gateway/src/routers/grpc/multimodal/plan.rs b/model_gateway/src/routers/grpc/multimodal/plan.rs new file mode 100644 index 000000000..88965808d --- /dev/null +++ b/model_gateway/src/routers/grpc/multimodal/plan.rs @@ -0,0 +1,253 @@ +//! Canonical multimodal request planning shared by protocol adapters, rendering, +//! and preprocessing. + +use std::collections::HashMap; + +use anyhow::{Context, Result}; +use llm_multimodal::{MediaContentPart, Modality, ModelMetadata}; +use llm_tokenizer::TokenizerTrait; + +use super::config::MultimodalComponents; + +/// Ordered media extracted from an API request. +/// +/// This is the single hand-off between protocol-specific parsing and the shared +/// multimodal pipeline. Text remains in the message representation used by the +/// chat template; media is kept here in authored order for fetching and count +/// validation. +#[derive(Debug, Clone, Default)] +pub(crate) struct MediaPlan { + parts: Vec, + modalities: Vec, + counts: HashMap, +} + +impl MediaPlan { + pub(crate) fn new(parts: impl IntoIterator) -> Self { + let mut plan = Self::default(); + for part in parts { + let modality = match &part { + MediaContentPart::ImageUrl { .. } | MediaContentPart::ImageData { .. } => { + Some(Modality::Image) + } + MediaContentPart::ImageEmbeds { .. } => Some(Modality::ImageEmbeds), + MediaContentPart::AudioUrl { .. } | MediaContentPart::AudioData { .. } => { + Some(Modality::Audio) + } + MediaContentPart::VideoUrl { .. } | MediaContentPart::VideoData { .. } => { + Some(Modality::Video) + } + MediaContentPart::Text { .. } => None, + }; + + let Some(modality) = modality else { + continue; + }; + if !plan.modalities.contains(&modality) { + plan.modalities.push(modality); + } + *plan.counts.entry(modality).or_default() += 1; + plan.parts.push(part); + } + plan + } + + pub(crate) fn is_empty(&self) -> bool { + self.parts.is_empty() + } + + pub(crate) fn modalities(&self) -> &[Modality] { + &self.modalities + } + + pub(crate) fn count(&self, modality: Modality) -> usize { + self.counts.get(&modality).copied().unwrap_or_default() + } + + pub(crate) fn into_parts(self) -> Vec { + self.parts + } +} + +/// Model-specific structural anchor strings keyed by modality. +/// +/// String-format templates need the actual anchor string, while OpenAI-format +/// templates receive canonical `image` / `audio` / `video` parts. Keeping the +/// mapping typed prevents the former `image_placeholder` argument from being +/// accidentally reused for every modality. +#[derive(Debug, Clone, Default)] +pub(crate) struct PlaceholderTokens { + tokens: HashMap, +} + +impl PlaceholderTokens { + pub(crate) fn insert(&mut self, modality: Modality, token: String) { + self.tokens.insert(modality, token); + } + + pub(crate) fn get(&self, modality: Modality) -> Option<&str> { + self.tokens.get(&modality).map(String::as_str) + } +} + +/// Validate a multimodal request against the model spec and resolve the +/// structural anchors for its active modalities in one config/spec lookup. +pub(crate) async fn prepare_placeholder_tokens( + plan: &MediaPlan, + model_id: &str, + tokenizer: &dyn TokenizerTrait, + components: &MultimodalComponents, + tokenizer_id: &str, + tokenizer_source: &str, +) -> Result { + anyhow::ensure!(!plan.is_empty(), "multimodal media plan is empty"); + let model_config = components + .config_registry + .get_or_load(tokenizer_id, tokenizer_source) + .await?; + let metadata = ModelMetadata { + model_id, + tokenizer, + config: &model_config.config, + }; + let spec = components + .model_registry + .lookup(&metadata) + .with_context(|| format!("multimodal not supported for model: {model_id}"))?; + let requested = plan + .modalities() + .iter() + .map(|&modality| (modality, plan.count(modality))) + .collect::>(); + spec.validate_media_request(&metadata, &requested) + .map_err(|error| { + anyhow::anyhow!("invalid media request for model {}: {error}", spec.name()) + })?; + let mut placeholders = PlaceholderTokens::default(); + for &modality in plan.modalities() { + let token = spec + .placeholder_token_for(&metadata, modality) + .map_err(|error| { + anyhow::anyhow!( + "model {} supports {modality} but its placeholder token could not be resolved: {error}", + spec.name() + ) + })?; + anyhow::ensure!( + tokenizer.token_to_id(&token).is_some(), + "{modality} placeholder token '{token}' is missing from the tokenizer vocabulary" + ); + placeholders.insert(modality, token); + } + + Ok(placeholders) +} + +/// Verify a rendered/tokenized multimodal prompt contains exactly one +/// structural anchor for every planned media item before any media is fetched +/// or preprocessed. +/// +/// This catches stale/custom templates, adapter omissions, and literal internal +/// anchors in user text at the cheapest point in the pipeline. +pub(crate) fn validate_rendered_media_anchors( + plan: &MediaPlan, + placeholders: &PlaceholderTokens, + tokenizer: &dyn TokenizerTrait, + token_ids: &[u32], +) -> Result<()> { + for &modality in plan.modalities() { + let token = placeholders + .get(modality) + .ok_or_else(|| anyhow::anyhow!("missing resolved {modality} placeholder token"))?; + let token_id = tokenizer.token_to_id(token).ok_or_else(|| { + anyhow::anyhow!( + "{modality} placeholder token '{token}' is missing from the tokenizer vocabulary" + ) + })?; + let expected = plan.count(modality); + let actual = token_ids + .iter() + .filter(|&&candidate| candidate == token_id) + .count(); + anyhow::ensure!( + actual == expected, + "rendered {modality} anchor count mismatch: expected {expected}, found {actual}; verify the chat template contract and escape literal media anchors in user text" + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use llm_tokenizer::mock::MockTokenizer; + + use super::*; + + #[test] + fn media_plan_preserves_media_order_and_counts() { + let plan = MediaPlan::new([ + MediaContentPart::Text { + text: "ignored".to_string(), + }, + MediaContentPart::AudioUrl { + url: "audio".to_string(), + uuid: None, + }, + MediaContentPart::ImageUrl { + url: "image".to_string(), + detail: None, + uuid: None, + }, + MediaContentPart::AudioUrl { + url: "audio-2".to_string(), + uuid: None, + }, + ]); + + assert_eq!(plan.modalities(), &[Modality::Audio, Modality::Image]); + assert_eq!(plan.count(Modality::Audio), 2); + assert_eq!(plan.count(Modality::Image), 1); + + let parts = plan.into_parts(); + assert_eq!(parts.len(), 3); + assert!(matches!( + &parts[0], + MediaContentPart::AudioUrl { url, .. } if url == "audio" + )); + assert!(matches!( + &parts[1], + MediaContentPart::ImageUrl { url, .. } if url == "image" + )); + assert!(matches!( + &parts[2], + MediaContentPart::AudioUrl { url, .. } if url == "audio-2" + )); + } + + #[test] + fn rendered_anchor_validation_is_exact_per_modality() { + let plan = MediaPlan::new([ + MediaContentPart::ImageUrl { + url: "image".to_string(), + detail: None, + uuid: None, + }, + MediaContentPart::AudioUrl { + url: "audio".to_string(), + uuid: None, + }, + ]); + let mut placeholders = PlaceholderTokens::default(); + placeholders.insert(Modality::Image, "<|im_start|>".to_string()); + placeholders.insert(Modality::Audio, "<|im_end|>".to_string()); + let tokenizer = MockTokenizer::new(); + + validate_rendered_media_anchors(&plan, &placeholders, &tokenizer, &[1001, 7, 1002]) + .unwrap(); + + let error = + validate_rendered_media_anchors(&plan, &placeholders, &tokenizer, &[1001, 1001, 1002]) + .unwrap_err(); + assert!(error.to_string().contains("image anchor count mismatch")); + } +} diff --git a/model_gateway/src/routers/grpc/multimodal/process.rs b/model_gateway/src/routers/grpc/multimodal/process.rs index c45d4756b..05fa80fdc 100644 --- a/model_gateway/src/routers/grpc/multimodal/process.rs +++ b/model_gateway/src/routers/grpc/multimodal/process.rs @@ -1,116 +1,43 @@ //! Multimodal processing core: fetch media → preprocess pixels → expand //! placeholder tokens → build the lightweight [`MultimodalIntermediate`]. //! -//! The chat and Messages API pipelines share `process_multimodal_parts`; only -//! the content extraction differs (see [`super::detect`]). +//! Protocol adapters first normalize media into a [`MediaPlan`], so this module +//! has one request-independent processing entry point. -use std::{sync::Arc, time::Instant}; +use std::{collections::HashMap, sync::Arc, time::Instant}; use anyhow::Result; +use futures::future::try_join_all; use llm_multimodal::{ - AsyncMultiModalTracker, ImageFrame, Modality, ModelMetadata, PlaceholderRange, - PreProcessorConfig, PreprocessedEncoderInputs, PromptReplacement, TrackedMedia, TrackerOutput, - VideoClip, VisionProcessorRegistry, + AsyncMultiModalTracker, AudioClip, EncoderFieldLayouts, ImageFrame, Modality, ModelMetadata, + PlaceholderRange, PreProcessorConfig, PreprocessedEncoderInputs, PromptReplacement, + TrackedMedia, TrackerOutput, VideoClip, VisionProcessorRegistry, }; use llm_tokenizer::TokenizerTrait; -use openai_protocol::{chat::ChatMessage, messages::InputMessage}; use tracing::{debug, info, warn}; use super::{ - config::MultimodalComponents, - detect::{extract_content_parts, extract_content_parts_messages}, + config::{MultimodalComponents, MultimodalModelConfig}, log_mm_timing_enabled, pixel_cache::{config_fingerprint, CachedPreprocessedItem, PixelCache, PixelCacheKey}, - MultimodalIntermediate, MultimodalOutput, PrecomputedMultimodalIntermediate, + plan::MediaPlan, + MediaBatch, MultimodalIntermediate, MultimodalOutput, PrecomputedMultimodalIntermediate, + PromptBinding, }; -/// Resolve the placeholder token string for a multimodal model. -/// -/// Loads the model config (via the shared registry, keyed by `tokenizer_id`) -/// and looks up the model spec to get the placeholder token (e.g. -/// `"<|image|>"` for Phi-3-vision). Returns `None` if the model is not -/// recognized as multimodal. -pub(crate) async fn resolve_placeholder_token( - model_id: &str, - tokenizer: &dyn TokenizerTrait, - components: &MultimodalComponents, - tokenizer_id: &str, - tokenizer_source: &str, - modality: Modality, -) -> Result> { - let model_config = components - .config_registry - .get_or_load(tokenizer_id, tokenizer_source) - .await?; - let metadata = ModelMetadata { - model_id, - tokenizer, - config: &model_config.config, - }; - let spec = match components.model_registry.lookup(&metadata) { - Some(s) => s, - None => return Ok(None), - }; - Ok(Some( - spec.placeholder_token_for(&metadata, modality) - .map_err(|e| anyhow::anyhow!("Failed to get placeholder token: {e}"))?, - )) -} - -/// Process multimodal content from Messages API input messages. -pub(crate) async fn process_multimodal_messages( - messages: &[InputMessage], - model_id: &str, - tokenizer: &dyn TokenizerTrait, - token_ids: Vec, - components: &MultimodalComponents, - tokenizer_id: &str, - tokenizer_source: &str, -) -> Result { - let content_parts = extract_content_parts_messages(messages); - process_multimodal_parts( - content_parts, - model_id, - tokenizer, - token_ids, - components, - tokenizer_id, - tokenizer_source, - ) - .await -} - -/// Process multimodal content: fetch images, preprocess pixels, expand tokens, collect hashes. -/// -/// Single entry point called from preparation.rs. Handles the full pipeline: -pub(crate) async fn process_multimodal( - messages: &[ChatMessage], - model_id: &str, - tokenizer: &dyn TokenizerTrait, - token_ids: Vec, - components: &MultimodalComponents, - tokenizer_id: &str, - tokenizer_source: &str, -) -> Result { - let content_parts = extract_content_parts(messages); - process_multimodal_parts( - content_parts, - model_id, - tokenizer, - token_ids, - components, - tokenizer_id, - tokenizer_source, - ) - .await +struct PreparedMultimodalPart { + preprocessed: PreprocessedEncoderInputs, + media: MediaBatch, + prompt_replacements: Vec, + search_token_id: Option, + placeholder_token_id: Option, + field_layouts: EncoderFieldLayouts, + keep_on_cpu_keys: Vec, } -/// Shared multimodal processing core. -/// -/// Takes pre-extracted `MediaContentPart`s (from either chat or messages pipeline) -/// and runs the full processing chain: fetch → preprocess → expand → build intermediate. -async fn process_multimodal_parts( - content_parts: Vec, +/// Process a protocol-independent, ordered media plan. +pub(crate) async fn process_multimodal_plan( + plan: MediaPlan, model_id: &str, tokenizer: &dyn TokenizerTrait, token_ids: Vec, @@ -123,7 +50,7 @@ async fn process_multimodal_parts( let media_started = Instant::now(); let mut tracker = AsyncMultiModalTracker::new(components.media_connector.clone()); - for part in content_parts { + for part in plan.into_parts() { tracker .push_part(part) .map_err(|e| anyhow::anyhow!("Failed to push content part: {e}"))?; @@ -162,44 +89,76 @@ async fn process_multimodal_parts( }) .unwrap_or_default(); + let audios: Vec> = tracker_output + .data + .get(&Modality::Audio) + .map(|media_vec| { + media_vec + .iter() + .filter_map(|m| match m { + TrackedMedia::Audio(clip) => Some(clip.clone()), + _ => None, + }) + .collect() + }) + .unwrap_or_default(); + let media_elapsed_ms = media_started.elapsed().as_secs_f64() * 1000.0; - let modality = match (images.is_empty(), videos.is_empty()) { - (false, true) => Modality::Image, - (true, false) => Modality::Video, - (false, false) => { - return Err(anyhow::anyhow!( - "Mixed image and video multimodal requests are not supported yet" - )); - } - (true, true) => { - return Err(anyhow::anyhow!( - "No media was successfully fetched for multimodal request" - )); + let image_count = images.len(); + let audio_count = audios.len(); + let video_count = videos.len(); + let video_frame_count = videos.first().map_or(0, |video| { + if video.frames().is_empty() { + video + .rgb_video() + .map_or(0, |rgb_video| rgb_video.frames.len()) + } else { + video.frames().len() } - }; - - if modality == Modality::Video && videos.len() != 1 { + }); + let mut media_batches = Vec::with_capacity(3); + if !images.is_empty() { + media_batches.push(MediaBatch::Images(images)); + } + if !videos.is_empty() { + media_batches.push(MediaBatch::Videos(videos)); + } + if !audios.is_empty() { + media_batches.push(MediaBatch::Audios(audios)); + } + if media_batches.is_empty() { return Err(anyhow::anyhow!( - "Exactly one video is supported per request for the initial video path" + "No media was successfully fetched for multimodal request" )); } + let present_modalities = media_batches + .iter() + .map(MediaBatch::modality) + .collect::>(); - match modality { - Modality::Image => { - debug!( - image_count = images.len(), - item_sizes = ?images.iter().map(|f| (f.image.width(), f.image.height())).collect::>(), - "Fetched images for multimodal processing" - ); - } - Modality::Video => { - debug!( - video_count = videos.len(), - frame_count = videos.first().map_or(0, |v| v.frames.len()), - "Fetched video for multimodal processing" - ); + for batch in &media_batches { + match batch { + MediaBatch::Images(images) => { + debug!( + image_count = images.len(), + item_sizes = ?images.iter().map(|f| (f.image.width(), f.image.height())).collect::>(), + "Fetched images for multimodal processing" + ); + } + MediaBatch::Videos(videos) => { + debug!( + video_count = videos.len(), + frame_count = video_frame_count, + "Fetched video for multimodal processing" + ); + } + MediaBatch::Audios(audios) => { + debug!( + audio_count = audios.len(), + "Fetched audios for multimodal processing" + ); + } } - _ => {} } // Step 2: Resolve model spec and preprocess media. @@ -223,187 +182,134 @@ async fn process_multimodal_parts( .ok_or_else(|| anyhow::anyhow!("Multimodal not supported for model: {model_id}"))?; let config_elapsed_ms = config_started.elapsed().as_secs_f64() * 1000.0; - // Run CPU-intensive vision preprocessing on a blocking thread pool so it - // doesn't block the tokio async runtime under concurrent load. - // TODO: consider making the thread pool size configurable. - let pp_config = match modality { - Modality::Video => model_config - .video_preprocessor_config - .clone() - .unwrap_or_else(|| model_config.preprocessor_config.clone()), - _ => model_config.preprocessor_config.clone(), - }; let preprocess_started = Instant::now(); - - let preprocessed: PreprocessedEncoderInputs = if let (Some(cache), Modality::Image, 1) = - (components.pixel_cache.clone(), modality, images.len()) - { - preprocess_image_cached( - cache, - &images[0], - components.vision_processor_registry.clone(), - model_id.to_string(), - model_type.map(String::from), - pp_config, - config_fingerprint(tokenizer_id, &model_config.config), + let mut prepared_parts = Vec::with_capacity(media_batches.len()); + // Every modality batch is independent until prompt expansion. Poll all + // preprocessors concurrently and preserve the batch order in the returned + // vector so the media/preprocessed zip below remains exact for any model- + // validated modality combination. + let preprocessed_parts = try_join_all(media_batches.iter().map(|media| { + preprocess_modality( + media, + components, + model_id, + model_type, + spec.name(), + tokenizer_id, + &model_config, ) - .await? - } else { - let registry = components.vision_processor_registry.clone(); - let model_id_owned = model_id.to_string(); - let model_type_owned = model_type.map(String::from); - let images_for_preprocess = images.clone(); // cheap Arc refcount bumps - let videos_for_preprocess = videos.clone(); // cheap Arc refcount bumps - let preprocessed: PreprocessedEncoderInputs = tokio::task::spawn_blocking(move || { - let processor = registry - .find(&model_id_owned, model_type_owned.as_deref()) - .ok_or_else(|| { - anyhow::anyhow!("No vision processor found for model: {model_id_owned}") - })?; + })) + .await?; - match modality { - Modality::Image => { - // Extract DynamicImages inside the blocking closure so the expensive - // clone happens off the tokio async runtime. - let raw_images: Vec = images_for_preprocess - .iter() - .map(|f| f.image.clone()) - .collect(); - processor - .preprocess(&raw_images, &pp_config) - .map_err(|e| anyhow::anyhow!("Image preprocessing failed: {e}")) - } - Modality::Video => { - let video = videos_for_preprocess - .first() - .ok_or_else(|| anyhow::anyhow!("No video available for preprocessing"))?; - - if !video.frames().is_empty() { - return processor - .preprocess_video(video.frames(), &pp_config) - .map_err(|e| anyhow::anyhow!("Video preprocessing failed: {e}")); - } + for (media, preprocessed) in media_batches.into_iter().zip(preprocessed_parts) { + let modality = media.modality(); + debug!( + ?modality, + item_count = preprocessed.feature_token_counts.len(), + total_tokens = preprocessed.feature_token_counts.iter().sum::(), + "Multimodal preprocessing complete" + ); - if let Some(rgb_video) = video.rgb_video() { - match rgb_video.frame_refs() { - Ok(frame_refs) => { - match processor.preprocess_video_rgb(&frame_refs, &pp_config) { - Ok(preprocessed) => return Ok(preprocessed), - Err(error) => { - warn!( - error = %error, - "RGB video preprocessing fast path failed; falling back to materialized frames" - ); - } - } - } - Err(error) => { - warn!( - error = %error, - "RGB video frame refs are invalid; falling back to materialized frames" - ); - } - } - } + let prompt_replacements = spec + .prompt_replacements_for(&metadata, &preprocessed, modality) + .map_err(|e| anyhow::anyhow!("Failed to compute prompt replacements: {e}"))?; - let frames = video - .materialized_frames() - .map_err(|e| anyhow::anyhow!("Video frame materialization failed: {e}"))?; - processor - .preprocess_video(&frames, &pp_config) - .map_err(|e| anyhow::anyhow!("Video preprocessing failed: {e}")) - } - _ => Err(anyhow::anyhow!( - "Unsupported modality for preprocessing: {modality}" - )), + let media_count = media.len(); + anyhow::ensure!( + preprocessed.feature_token_counts.len() == media_count, + "Preprocessing item count mismatch for {modality}: {} media items, {} feature-token counts", + media_count, + preprocessed.feature_token_counts.len() + ); + anyhow::ensure!( + prompt_replacements.len() == media_count, + "Prompt replacement count mismatch for {modality}: {} media items, {} replacements", + media_count, + prompt_replacements.len() + ); + + // Two token IDs may differ for the same placeholder: + // - search_token_id: what the tokenizer actually emits (e.g. 200090 for "<|image|>") + // - placeholder_token_id: what the model config declares (e.g. image_token_id/video_token_id) + let placeholder_token = spec + .placeholder_token_for(&metadata, modality) + .map_err(|e| anyhow::anyhow!("Failed to get placeholder token: {e}"))?; + let search_token_id = tokenizer.token_to_id(&placeholder_token); + let placeholder_token_id: Option = match spec + .placeholder_token_id_for(&metadata, modality) + { + Ok(id) => Some(u32::try_from(id).map_err(|_| { + anyhow::anyhow!( + "Invalid negative placeholder token ID {id} for modality {modality}" + ) + })?), + Err(e) => { + warn!( + error = %e, + ?search_token_id, + "Failed to resolve placeholder_token_id from config, falling back to tokenizer lookup" + ); + search_token_id } - }) - .await - .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))??; - preprocessed - }; - let preprocess_elapsed_ms = preprocess_started.elapsed().as_secs_f64() * 1000.0; + }; - debug!( - ?modality, - item_count = preprocessed.feature_token_counts.len(), - total_tokens = preprocessed.feature_token_counts.iter().sum::(), - "Multimodal preprocessing complete" - ); + prepared_parts.push(PreparedMultimodalPart { + preprocessed, + media, + prompt_replacements, + search_token_id, + placeholder_token_id, + field_layouts: spec.encoder_field_layouts_for(modality), + keep_on_cpu_keys: spec.keep_on_cpu_keys_for(modality), + }); + } + let preprocess_elapsed_ms = preprocess_started.elapsed().as_secs_f64() * 1000.0; // Step 3: Compute prompt replacements and expand tokens. let expansion_started = Instant::now(); - let prompt_replacements = spec - .prompt_replacements_for(&metadata, &preprocessed, modality) - .map_err(|e| anyhow::anyhow!("Failed to compute prompt replacements: {e}"))?; - - // Two token IDs may differ for the same placeholder: - // - search_token_id: what the tokenizer actually emits (e.g. 200090 for "<|image|>") - // - placeholder_token_id: what the model config declares (e.g. image_token_id/video_token_id) - let placeholder_token = spec - .placeholder_token_for(&metadata, modality) - .map_err(|e| anyhow::anyhow!("Failed to get placeholder token: {e}"))?; - let search_token_id = tokenizer.token_to_id(&placeholder_token); - let placeholder_token_id: Option = match spec.placeholder_token_id_for(&metadata, modality) - { - Ok(id) => Some(id as u32), - Err(e) => { - warn!( - error = %e, - ?search_token_id, - "Failed to resolve placeholder_token_id from config, falling back to tokenizer lookup" - ); - search_token_id - } - }; - - let expanded = expand_tokens( - &token_ids, - search_token_id, - placeholder_token_id, - &prompt_replacements, - ); + let expansions = prepared_parts + .iter() + .map(|part| ModalityExpansion { + modality: part.media.modality(), + search_token_id: part.search_token_id, + placeholder_token_id: part.placeholder_token_id, + replacements: &part.prompt_replacements, + }) + .collect::>(); + let expanded = expand_tokens_for_modalities(&token_ids, &expansions)?; + let placeholder_count = expanded.bindings.iter().map(Vec::len).sum::(); debug!( original_len = token_ids.len(), expanded_len = expanded.token_ids.len(), - placeholder_count = expanded.placeholders.len(), - ?search_token_id, - ?placeholder_token_id, + placeholder_count, + modality_count = prepared_parts.len(), "Token expansion complete" ); let expansion_elapsed_ms = expansion_started.elapsed().as_secs_f64() * 1000.0; - let image_count = images.len(); - let video_count = videos.len(); - let video_frame_count = videos.first().map_or(0, |video| { - if video.frames().is_empty() { - video - .rgb_video() - .map_or(0, |rgb_video| rgb_video.frames.len()) - } else { - video.frames().len() - } - }); let original_tokens = token_ids.len(); let expanded_tokens = expanded.token_ids.len(); // Step 4: Build lightweight intermediate (defers tensor serialization to assembly) - let intermediate = MultimodalIntermediate::Precomputed(PrecomputedMultimodalIntermediate { - modality, - preprocessed, - images, - videos, - placeholders: expanded.placeholders, - patch_offsets: expanded.patch_offsets, - placeholder_token_id, - field_layouts: spec.field_layouts(), - keep_on_cpu_keys: spec.keep_on_cpu_keys(), - }); + let batches = prepared_parts + .into_iter() + .zip(expanded.bindings) + .map(|(part, bindings)| PrecomputedMultimodalIntermediate { + preprocessed: part.preprocessed, + media: part.media, + bindings, + placeholder_token_id: part.placeholder_token_id, + field_layouts: part.field_layouts, + keep_on_cpu_keys: part.keep_on_cpu_keys, + }) + .collect::>(); + let intermediate = MultimodalIntermediate::try_new(batches)?; if log_timing { info!( - modality = ?modality, + modalities = ?present_modalities, image_count, + audio_count, video_count, video_frame_count, media_fetch_decode_ms = media_elapsed_ms, @@ -413,7 +319,7 @@ async fn process_multimodal_parts( total_ms = total_started.elapsed().as_secs_f64() * 1000.0, original_tokens, expanded_tokens, - "smg_mm_timing process_multimodal_parts" + "smg_mm_timing process_multimodal_plan" ); } @@ -423,6 +329,151 @@ async fn process_multimodal_parts( }) } +async fn preprocess_modality( + media: &MediaBatch, + components: &MultimodalComponents, + model_id: &str, + model_type: Option<&str>, + model_spec: &str, + tokenizer_id: &str, + model_config: &MultimodalModelConfig, +) -> Result { + // Run CPU-intensive preprocessing on a blocking thread pool so it doesn't + // block the tokio async runtime under concurrent load. + // TODO: consider making the thread pool size configurable. + let modality = media.modality(); + let pp_config = match modality { + Modality::Video => model_config + .video_preprocessor_config + .clone() + .unwrap_or_else(|| model_config.preprocessor_config.clone()), + _ => model_config.preprocessor_config.clone(), + }; + + if let MediaBatch::Images(images) = media { + if let (Some(cache), [image]) = (components.pixel_cache.clone(), images.as_slice()) { + return preprocess_image_cached( + cache, + image, + components.vision_processor_registry.clone(), + model_id.to_string(), + model_type.map(String::from), + pp_config, + config_fingerprint(tokenizer_id, &model_config.config), + ) + .await; + } + } + + let registry = components.vision_processor_registry.clone(); + let model_id_owned = model_id.to_string(); + let model_type_owned = model_type.map(String::from); + let media_for_preprocess = media.clone(); // cheap Arc refcount bumps + let audio_processor = if modality == Modality::Audio { + Some( + components + .audio_processor_registry + .create( + model_spec, + &model_config.config, + &model_config.preprocessor_config, + ) + .ok_or_else(|| { + anyhow::anyhow!("No audio processor registered for model spec: {model_spec}") + })?, + ) + } else { + None + }; + + tokio::task::spawn_blocking(move || match media_for_preprocess { + MediaBatch::Images(images) => { + let processor = registry + .find(&model_id_owned, model_type_owned.as_deref()) + .ok_or_else(|| { + anyhow::anyhow!("No vision processor found for model: {model_id_owned}") + })?; + // Extract DynamicImages inside the blocking closure so the expensive + // clone happens off the tokio async runtime. + let raw_images: Vec = + images.iter().map(|frame| frame.image.clone()).collect(); + processor + .preprocess(&raw_images, &pp_config) + .map_err(|e| anyhow::anyhow!("Image preprocessing failed: {e}")) + } + MediaBatch::Videos(videos) => { + // VisionPreProcessor currently models one decoded clip per call. + // Video-capable model specs therefore declare a per-request limit + // of one; a future batched-video processor can lift this without + // adding any modality-combination policy here. + let [video] = videos.as_slice() else { + anyhow::bail!( + "Video preprocessing currently requires exactly one clip per modality batch; got {}", + videos.len() + ); + }; + let processor = registry + .find(&model_id_owned, model_type_owned.as_deref()) + .ok_or_else(|| { + anyhow::anyhow!("No vision processor found for model: {model_id_owned}") + })?; + let video_pp_config = with_video_sample_fps(pp_config.clone(), video); + + if !video.frames().is_empty() { + return processor + .preprocess_video(video.frames(), &video_pp_config) + .map_err(|e| anyhow::anyhow!("Video preprocessing failed: {e}")); + } + + if let Some(rgb_video) = video.rgb_video() { + match rgb_video.frame_refs() { + Ok(frame_refs) => match processor + .preprocess_video_rgb(&frame_refs, &video_pp_config) + { + Ok(preprocessed) => return Ok(preprocessed), + Err(error) => { + warn!( + error = %error, + "RGB video preprocessing fast path failed; falling back to materialized frames" + ); + } + }, + Err(error) => { + warn!( + error = %error, + "RGB video frame refs are invalid; falling back to materialized frames" + ); + } + } + } + + let frames = video + .materialized_frames() + .map_err(|e| anyhow::anyhow!("Video frame materialization failed: {e}"))?; + processor + .preprocess_video(&frames, &video_pp_config) + .map_err(|e| anyhow::anyhow!("Video preprocessing failed: {e}")) + } + MediaBatch::Audios(audios) => { + let processor = audio_processor.ok_or_else(|| { + anyhow::anyhow!("Model did not provide an audio processor") + })?; + processor + .preprocess(&audios) + .map_err(|e| anyhow::anyhow!("Audio preprocessing failed: {e}")) + } + }) + .await + .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))? +} + +fn with_video_sample_fps(mut config: PreProcessorConfig, video: &VideoClip) -> PreProcessorConfig { + config + .extra + .insert("fps".to_string(), serde_json::json!(video.sample_fps())); + config +} + /// Pixel-cache image preprocessing for single-image requests. async fn preprocess_image_cached( cache: Arc, @@ -478,120 +529,180 @@ async fn preprocess_image_batch( .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))? } -/// Output of token expansion, containing both full structural and patch-only ranges. -struct ExpandedTokens { - /// The expanded token ID sequence. +struct ModalityExpansion<'a> { + modality: Modality, + search_token_id: Option, + placeholder_token_id: Option, + replacements: &'a [PromptReplacement], +} + +#[derive(Debug)] +struct ExpandedMultimodalTokens { token_ids: Vec, - /// Full structural placeholder ranges (offset, length) covering the entire - /// replacement including structural tokens. Used by vLLM (which filters via is_embed). - placeholders: Vec, - /// Patch-only placeholder ranges: contiguous runs of `im_token_id` within each - /// expansion. Used by sglang (which expects offsets aligned 1:1 with vision - /// encoder output). `None` when `im_token_id` is not set. - patch_offsets: Option>, + bindings: Vec>, } -/// Expand placeholder tokens in the token ID sequence. -/// -/// For each placeholder token found, replace it with the expanded token sequence -/// from the corresponding `PromptReplacement`. Also track both the full structural -/// placeholder ranges and patch-only offsets (contiguous runs of `im_token_id`) -/// in a single pass — no extra iteration needed. -fn expand_tokens( +fn expand_tokens_for_modalities( token_ids: &[u32], - placeholder_token_id: Option, - im_token_id: Option, - replacements: &[PromptReplacement], -) -> ExpandedTokens { - let Some(placeholder_id) = placeholder_token_id else { - // If we can't resolve the placeholder token, return unchanged - warn!("Could not resolve placeholder token ID; skipping token expansion"); - return ExpandedTokens { - token_ids: token_ids.to_vec(), - placeholders: vec![], - patch_offsets: None, + expansions: &[ModalityExpansion<'_>], +) -> Result { + let mut anchor_to_expansion = HashMap::with_capacity(expansions.len()); + for (idx, expansion) in expansions.iter().enumerate() { + let Some(anchor_id) = expansion.search_token_id else { + anyhow::ensure!( + expansion.replacements.is_empty(), + "Could not resolve prompt anchor token ID for {} ({} replacements)", + expansion.modality, + expansion.replacements.len() + ); + continue; }; - }; + if let Some(previous_idx) = anchor_to_expansion.insert(anchor_id, idx) { + return Err(anyhow::anyhow!( + "Prompt anchor token ID {anchor_id} is shared by {} and {}; anchors must be unique", + expansions[previous_idx].modality, + expansion.modality + )); + } + for (item_index, replacement) in expansion.replacements.iter().enumerate() { + anyhow::ensure!( + replacement.modality == expansion.modality, + "Prompt replacement {item_index} has modality {}, expected {}", + replacement.modality, + expansion.modality + ); + anyhow::ensure!( + !replacement.tokens.is_empty(), + "Prompt replacement {item_index} for {} is empty", + expansion.modality + ); + } + } let mut expanded = Vec::with_capacity(token_ids.len()); - let mut placeholders = Vec::new(); - let mut patch_offsets: Option> = im_token_id.map(|_| Vec::new()); - let mut replacement_idx = 0; - let mut extra_placeholders = 0usize; - - for &token in token_ids { - if token == placeholder_id && replacement_idx < replacements.len() { - let repl = &replacements[replacement_idx]; - let offset = expanded.len(); - - // Track patch-only runs while extending - if let (Some(im_id), Some(ref mut offsets)) = (im_token_id, &mut patch_offsets) { - let mut run_start: Option = None; - for (i, &t) in repl.tokens.iter().enumerate() { - let pos = (offset + i) as u32; - if t as u32 == im_id { - if run_start.is_none() { - run_start = Some(pos); - } - } else if let Some(s) = run_start { - offsets.push((s, pos - s)); - run_start = None; - } - } - if let Some(s) = run_start { - offsets.push((s, (offset + repl.tokens.len()) as u32 - s)); - } - } + let mut bindings = vec![Vec::new(); expansions.len()]; + let mut replacement_indices = vec![0usize; expansions.len()]; + let mut prompt_ordinal = 0usize; - // PromptReplacement uses TokenId = i32, convert to u32 - expanded.extend(repl.tokens.iter().map(|&t| t as u32)); - // Fold any template-emitted structural prefix (already in `expanded`, - // e.g. Qwen's leading <|vision_start|>) into the reported range so a - // backend that scans the range for structural markers — vLLM's video - // mrope walks each frame from <|vision_start|> — starts on the marker. - // `offset` (used by the sglang patch_offsets pass above) is untouched. - let prefix = repl.structural_prefix.min(offset); - placeholders.push(PlaceholderRange { - offset: offset - prefix, - length: repl.tokens.len() + prefix, + for (prompt_offset, &token) in token_ids.iter().enumerate() { + if let Some(&idx) = anchor_to_expansion.get(&token) { + let expansion = &expansions[idx]; + let item_index = replacement_indices[idx]; + let replacement = expansion.replacements.get(item_index).ok_or_else(|| { + anyhow::anyhow!( + "Extra prompt anchor for {} at input token offset {prompt_offset}: expected {} anchors", + expansion.modality, + expansion.replacements.len() + ) + })?; + let replacement_tokens = replacement + .tokens + .iter() + .enumerate() + .map(|(replacement_offset, &token)| { + u32::try_from(token).map_err(|_| { + anyhow::anyhow!( + "Invalid negative token ID {token} in {} replacement {item_index} at offset {replacement_offset}", + expansion.modality + ) + }) + }) + .collect::>>()?; + let offset = expanded.len(); + let length = replacement_tokens.len(); + let patches = patch_ranges(offset, &replacement_tokens, expansion.placeholder_token_id); + expanded.extend(replacement_tokens); + let prefix = replacement.structural_prefix.min(offset); + bindings[idx].push(PromptBinding { + item_index, + prompt_ordinal, + structural: PlaceholderRange { + offset: offset - prefix, + length: length + prefix, + }, + patches, }); - replacement_idx += 1; + prompt_ordinal = prompt_ordinal + .checked_add(1) + .ok_or_else(|| anyhow::anyhow!("Prompt binding ordinal overflow"))?; + replacement_indices[idx] += 1; } else { - // A placeholder token seen after all replacements are consumed is - // left in place (unchanged behavior) but counted so we can warn. - if token == placeholder_id { - extra_placeholders += 1; - } expanded.push(token); } } - if replacement_idx < replacements.len() { - warn!( - expected = replacements.len(), - found = replacement_idx, - "Fewer placeholder tokens found in sequence than expected" - ); - } - if extra_placeholders > 0 { - warn!( - extra_placeholders, - replacements = replacements.len(), - "More placeholder tokens than replacements; extra placeholders left unexpanded" + for (idx, expansion) in expansions.iter().enumerate() { + anyhow::ensure!( + replacement_indices[idx] == expansion.replacements.len(), + "Missing prompt anchors for {}: expected {}, found {}", + expansion.modality, + expansion.replacements.len(), + replacement_indices[idx] ); } - ExpandedTokens { + Ok(ExpandedMultimodalTokens { token_ids: expanded, - placeholders, - patch_offsets, + bindings, + }) +} + +fn patch_ranges( + offset: usize, + replacement_tokens: &[u32], + placeholder_token_id: Option, +) -> Vec { + let Some(placeholder_id) = placeholder_token_id else { + return Vec::new(); + }; + + let mut ranges = Vec::new(); + let mut run_start: Option = None; + for (i, &token) in replacement_tokens.iter().enumerate() { + let pos = offset + i; + if token == placeholder_id { + if run_start.is_none() { + run_start = Some(pos); + } + } else if let Some(start) = run_start.take() { + ranges.push(PlaceholderRange { + offset: start, + length: pos - start, + }); + } + } + if let Some(start) = run_start { + let end = offset + replacement_tokens.len(); + ranges.push(PlaceholderRange { + offset: start, + length: end - start, + }); } + ranges } #[cfg(test)] mod tests { + use bytes::Bytes; + use llm_multimodal::VideoSource; + use super::*; + #[test] + fn decoded_video_sample_fps_overrides_processor_default() { + let video = VideoClip::new_with_sample_fps( + Vec::new(), + Bytes::new(), + VideoSource::InlineBytes, + "video-hash".to_string(), + 0.8, + ); + + let config = with_video_sample_fps(PreProcessorConfig::default(), &video); + + assert!((config.get_extra::("fps").unwrap() - 0.8).abs() < 1e-6); + } + #[test] fn test_expand_tokens_basic() { let token_ids = vec![1, 2, 100, 3, 4]; // 100 is the placeholder @@ -602,13 +713,19 @@ mod tests { structural_prefix: 0, }]; - let result = expand_tokens(&token_ids, Some(100), None, &replacements); + let expansion = ModalityExpansion { + modality: Modality::Image, + search_token_id: Some(100), + placeholder_token_id: None, + replacements: &replacements, + }; + let result = expand_tokens_for_modalities(&token_ids, &[expansion]).unwrap(); assert_eq!(result.token_ids, vec![1, 2, 50, 50, 50, 50, 3, 4]); - assert_eq!(result.placeholders.len(), 1); - assert_eq!(result.placeholders[0].offset, 2); - assert_eq!(result.placeholders[0].length, 4); - assert!(result.patch_offsets.is_none()); + assert_eq!(result.bindings[0].len(), 1); + assert_eq!(result.bindings[0][0].structural.offset, 2); + assert_eq!(result.bindings[0][0].structural.length, 4); + assert!(result.bindings[0][0].patches.is_empty()); } #[test] @@ -626,23 +743,35 @@ mod tests { structural_prefix: 1, }]; - let result = expand_tokens(&token_ids, Some(100), None, &replacements); + let expansion = ModalityExpansion { + modality: Modality::Video, + search_token_id: Some(100), + placeholder_token_id: None, + replacements: &replacements, + }; + let result = expand_tokens_for_modalities(&token_ids, &[expansion]).unwrap(); assert_eq!(result.token_ids, vec![1, 2, 777, 50, 50, 50, 778, 3]); - assert_eq!(result.placeholders.len(), 1); + assert_eq!(result.bindings[0].len(), 1); // Range starts on (index 2) and covers it + the 3 video tokens. - assert_eq!(result.placeholders[0].offset, 2); - assert_eq!(result.placeholders[0].length, 4); + assert_eq!(result.bindings[0][0].structural.offset, 2); + assert_eq!(result.bindings[0][0].structural.length, 4); } #[test] fn test_expand_tokens_no_placeholder() { let token_ids = vec![1, 2, 3]; - let result = expand_tokens(&token_ids, None, None, &[]); + let replacements = Vec::new(); + let expansion = ModalityExpansion { + modality: Modality::Image, + search_token_id: None, + placeholder_token_id: None, + replacements: &replacements, + }; + let result = expand_tokens_for_modalities(&token_ids, &[expansion]).unwrap(); assert_eq!(result.token_ids, vec![1, 2, 3]); - assert!(result.placeholders.is_empty()); - assert!(result.patch_offsets.is_none()); + assert!(result.bindings[0].is_empty()); } #[test] @@ -663,14 +792,20 @@ mod tests { }, ]; - let result = expand_tokens(&token_ids, Some(100), None, &replacements); + let expansion = ModalityExpansion { + modality: Modality::Image, + search_token_id: Some(100), + placeholder_token_id: None, + replacements: &replacements, + }; + let result = expand_tokens_for_modalities(&token_ids, &[expansion]).unwrap(); assert_eq!(result.token_ids, vec![1, 50, 50, 2, 60, 60, 60, 3]); - assert_eq!(result.placeholders.len(), 2); - assert_eq!(result.placeholders[0].offset, 1); - assert_eq!(result.placeholders[0].length, 2); - assert_eq!(result.placeholders[1].offset, 4); - assert_eq!(result.placeholders[1].length, 3); + assert_eq!(result.bindings[0].len(), 2); + assert_eq!(result.bindings[0][0].structural.offset, 1); + assert_eq!(result.bindings[0][0].structural.length, 2); + assert_eq!(result.bindings[0][1].structural.offset, 4); + assert_eq!(result.bindings[0][1].structural.length, 3); } #[test] @@ -685,24 +820,224 @@ mod tests { structural_prefix: 0, }]; - let result = expand_tokens(&token_ids, Some(100), Some(92), &replacements); + let expansion = ModalityExpansion { + modality: Modality::Image, + search_token_id: Some(100), + placeholder_token_id: Some(92), + replacements: &replacements, + }; + let result = expand_tokens_for_modalities(&token_ids, &[expansion]).unwrap(); // Full structural range - assert_eq!(result.placeholders.len(), 1); - assert_eq!(result.placeholders[0].offset, 1); - assert_eq!(result.placeholders[0].length, 9); + assert_eq!(result.bindings[0].len(), 1); + assert_eq!(result.bindings[0][0].structural.offset, 1); + assert_eq!(result.bindings[0][0].structural.length, 9); // Patch-only offsets: two runs of token 92 - let patch = result.patch_offsets.unwrap(); + let patch = &result.bindings[0][0].patches; assert_eq!(patch.len(), 2); - assert_eq!(patch[0], (2, 3)); // offset=2, length=3 - assert_eq!(patch[1], (6, 3)); // offset=6, length=3 + assert_eq!((patch[0].offset, patch[0].length), (2, 3)); + assert_eq!((patch[1].offset, patch[1].length), (6, 3)); + } + + #[test] + fn test_expand_tokens_preserves_template_owned_audio_end() { + // The template owns the audio end marker. Expansion preserves the + // anchor, adds feature tokens, and must not inject another end marker. + let audio_anchor: u32 = 100; + let audio_placeholder: u32 = 101; + let audio_end: u32 = 102; + let message_end: u32 = 103; + let token_ids = vec![1, audio_anchor, audio_end, message_end]; + let replacements = vec![PromptReplacement { + modality: Modality::Audio, + placeholder_token: "