diff --git a/crates/multimodal/src/audio/mod.rs b/crates/multimodal/src/audio/mod.rs index 665a7d920..dff8c6d76 100644 --- a/crates/multimodal/src/audio/mod.rs +++ b/crates/multimodal/src/audio/mod.rs @@ -6,5 +6,5 @@ pub mod processors; pub(crate) mod transforms; pub use decode::{decode_audio_mono_f32, DecodedAudio}; -pub use processor::{AudioPreProcessor, AudioProcessorFactory, AudioProcessorRegistry}; +pub use processor::AudioPreProcessor; pub use processors::{Qwen3AudioParams, Qwen3AudioProcessor}; diff --git a/crates/multimodal/src/audio/processor.rs b/crates/multimodal/src/audio/processor.rs index fca17c33e..8e953ca5c 100644 --- a/crates/multimodal/src/audio/processor.rs +++ b/crates/multimodal/src/audio/processor.rs @@ -1,112 +1,16 @@ -use std::{collections::HashMap, sync::Arc}; +use std::sync::Arc; -use serde_json::Value; +use crate::{encoder_inputs::PreprocessedEncoderInputs, error::TransformError, types::AudioClip}; -use super::Qwen3AudioProcessor; -use crate::{ - encoder_inputs::PreprocessedEncoderInputs, error::TransformError, types::AudioClip, - vision::PreProcessorConfig, -}; - -/// Audio preprocessing contract selected by [`AudioProcessorRegistry`]. +/// Audio preprocessing contract for a model family. +/// +/// The concrete processor for a model is selected by its +/// [`ModelProcessorSpec::audio_processor`](crate::registry::ModelProcessorSpec::audio_processor), +/// which owns audio-processor selection alongside the model's prompt/placeholder +/// logic. 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/lib.rs b/crates/multimodal/src/lib.rs index 1aa91053f..ff6ce601c 100644 --- a/crates/multimodal/src/lib.rs +++ b/crates/multimodal/src/lib.rs @@ -12,7 +12,7 @@ pub mod tracker; pub mod types; pub mod vision; -pub use audio::{AudioPreProcessor, AudioProcessorRegistry}; +pub use audio::AudioPreProcessor; pub use encoder_inputs::{ModelSpecificValue, PreprocessedEncoderInputs}; pub use error::{MediaConnectorError, MultiModalError, MultiModalResult, TransformError}; pub use media::{ diff --git a/crates/multimodal/src/registry/llava.rs b/crates/multimodal/src/registry/llava.rs index 2bab71ab9..117091c00 100644 --- a/crates/multimodal/src/registry/llava.rs +++ b/crates/multimodal/src/registry/llava.rs @@ -199,4 +199,22 @@ mod tests { let spec = registry.lookup(&metadata).expect("llava alias"); assert_eq!(spec.name(), "llava"); } + + #[test] + fn llava_spec_has_no_audio_processor() { + use crate::vision::PreProcessorConfig; + + let tokenizer = TestTokenizer::new(&[("", 32000)]); + let config = json!({"model_type": "llava", "image_token_index": 32000}); + let metadata = ModelMetadata { + model_id: "llava-v1.5", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).expect("llava spec"); + assert!(spec + .audio_processor(&config, &PreProcessorConfig::default()) + .is_none()); + } } diff --git a/crates/multimodal/src/registry/qwen3_asr.rs b/crates/multimodal/src/registry/qwen3_asr.rs index 5d5ee6132..bdcf3c2f4 100644 --- a/crates/multimodal/src/registry/qwen3_asr.rs +++ b/crates/multimodal/src/registry/qwen3_asr.rs @@ -3,9 +3,11 @@ use std::collections::HashMap; use serde_json::{json, Value}; use crate::{ + audio::{AudioPreProcessor, Qwen3AudioProcessor}, encoder_inputs::PreprocessedEncoderInputs, registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, types::{EncoderFieldLayouts, FieldLayout, Modality, PromptReplacement, TokenId}, + vision::PreProcessorConfig, }; const AUDIO_PAD_TOKEN: &str = "<|audio_pad|>"; @@ -96,6 +98,17 @@ impl ModelProcessorSpec for Qwen3AsrSpec { Ok(json!({})) } + fn audio_processor( + &self, + model_config: &Value, + preprocessor_config: &PreProcessorConfig, + ) -> Option> { + Some(Box::new(Qwen3AudioProcessor::from_configs( + model_config, + preprocessor_config, + ))) + } + fn prompt_replacements( &self, metadata: &ModelMetadata, @@ -201,4 +214,46 @@ mod tests { HashMap::from([(Modality::Audio, 10)]) ); } + + #[test] + fn asr_spec_builds_qwen_audio_processor() { + use std::sync::Arc; + + use bytes::Bytes; + + use crate::{ + audio::DecodedAudio, + types::{AudioClip, AudioSource}, + }; + + let tokenizer = TestTokenizer::new(&[(AUDIO_PAD_TOKEN, 151676)]); + let config = json!({"model_type": "qwen3_asr"}); + let metadata = ModelMetadata { + model_id: "Qwen/Qwen3-ASR-1.7B", + tokenizer: &tokenizer, + config: &config, + }; + let registry = ModelRegistry::new(); + let spec = registry.lookup(&metadata).unwrap(); + + let preprocessor_config = PreProcessorConfig::from_json( + r#"{"feature_size": 16, "sampling_rate": 16000, "n_fft": 400, "hop_length": 160}"#, + ) + .unwrap(); + let processor = spec + .audio_processor(&config, &preprocessor_config) + .expect("qwen3_asr spec must provide an audio processor"); + + let clip = 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(), + )); + let result = processor.preprocess(&[clip]).unwrap(); + assert_eq!(result.encoder_input.shape(), &[1, 16, 5]); + } } diff --git a/crates/multimodal/src/registry/qwen3_omni.rs b/crates/multimodal/src/registry/qwen3_omni.rs index 2c4ecce68..1ec927c43 100644 --- a/crates/multimodal/src/registry/qwen3_omni.rs +++ b/crates/multimodal/src/registry/qwen3_omni.rs @@ -3,9 +3,11 @@ use std::collections::HashMap; use serde_json::{json, Value}; use crate::{ + audio::{AudioPreProcessor, Qwen3AudioProcessor}, encoder_inputs::PreprocessedEncoderInputs, registry::{ModelMetadata, ModelProcessorSpec, ModelRegistryError, RegistryResult}, types::{EncoderFieldLayouts, FieldLayout, Modality, PromptReplacement, TokenId}, + vision::PreProcessorConfig, }; const IMAGE_PAD_TOKEN: &str = "<|image_pad|>"; @@ -120,6 +122,17 @@ impl ModelProcessorSpec for Qwen3OmniSpec { Ok(json!({"use_audio_in_video": false})) } + fn audio_processor( + &self, + model_config: &Value, + preprocessor_config: &PreProcessorConfig, + ) -> Option> { + Some(Box::new(Qwen3AudioProcessor::from_configs( + model_config, + preprocessor_config, + ))) + } + fn prompt_replacements( &self, metadata: &ModelMetadata, @@ -306,4 +319,47 @@ mod tests { .keep_on_cpu_keys_for(Modality::Audio) .is_empty()); } + + #[test] + fn omni_spec_builds_qwen_audio_processor() { + use std::sync::Arc; + + use bytes::Bytes; + + use crate::{ + audio::DecodedAudio, + types::{AudioClip, AudioSource}, + vision::PreProcessorConfig, + }; + + let tokenizer = omni_tokenizer(); + let config = json!({"model_type": "qwen3_omni_moe"}); + 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(); + + let preprocessor_config = PreProcessorConfig::from_json( + r#"{"feature_size": 16, "sampling_rate": 16000, "n_fft": 400, "hop_length": 160}"#, + ) + .unwrap(); + let processor = spec + .audio_processor(&config, &preprocessor_config) + .expect("qwen3_omni spec must provide an audio processor"); + + let clip = 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(), + )); + let result = processor.preprocess(&[clip]).unwrap(); + assert_eq!(result.encoder_input.shape(), &[1, 16, 5]); + } } diff --git a/crates/multimodal/src/registry/traits.rs b/crates/multimodal/src/registry/traits.rs index 9b2a33f19..96c9059e2 100644 --- a/crates/multimodal/src/registry/traits.rs +++ b/crates/multimodal/src/registry/traits.rs @@ -5,8 +5,10 @@ use serde_json::Value; use thiserror::Error; use crate::{ + audio::AudioPreProcessor, encoder_inputs::PreprocessedEncoderInputs, types::{EncoderFieldLayouts, FieldLayout, Modality, PromptReplacement, TokenId}, + vision::PreProcessorConfig, }; #[derive(Debug, Error, PartialEq, Eq)] @@ -151,6 +153,24 @@ pub trait ModelProcessorSpec: Send + Sync { } fn processor_kwargs(&self, metadata: &ModelMetadata) -> RegistryResult; + + /// Build the audio preprocessor for this model, if it supports audio. + /// + /// This is the single source of truth for audio-processor selection: the + /// same spec that owns a model's prompt/placeholder logic also owns its + /// audio preprocessor factory, so there is no separate string-keyed + /// registry to keep in sync. Audio-less specs use the default (`None`). + /// + /// The processor is built from the current model config because its feature + /// shapes and quantization parameters can be checkpoint-specific. + fn audio_processor( + &self, + _model_config: &Value, + _preprocessor_config: &PreProcessorConfig, + ) -> Option> { + None + } + /// Compute per-media prompt replacement token sequences. /// /// Receives the full preprocessed output so each model can extract whatever diff --git a/model_gateway/src/routers/grpc/multimodal/config.rs b/model_gateway/src/routers/grpc/multimodal/config.rs index ffdd7f0f2..d87c3e563 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::{ - AudioProcessorRegistry, MediaConnector, MediaConnectorConfig, ModelRegistry, - PreProcessorConfig, VisionProcessorRegistry, + MediaConnector, MediaConnectorConfig, ModelRegistry, PreProcessorConfig, + VisionProcessorRegistry, }; use tracing::{debug, warn}; @@ -201,7 +201,6 @@ 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, @@ -223,7 +222,6 @@ 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/process.rs b/model_gateway/src/routers/grpc/multimodal/process.rs index 05fa80fdc..f41b7c490 100644 --- a/model_gateway/src/routers/grpc/multimodal/process.rs +++ b/model_gateway/src/routers/grpc/multimodal/process.rs @@ -10,8 +10,8 @@ use anyhow::Result; use futures::future::try_join_all; use llm_multimodal::{ AsyncMultiModalTracker, AudioClip, EncoderFieldLayouts, ImageFrame, Modality, ModelMetadata, - PlaceholderRange, PreProcessorConfig, PreprocessedEncoderInputs, PromptReplacement, - TrackedMedia, TrackerOutput, VideoClip, VisionProcessorRegistry, + ModelProcessorSpec, PlaceholderRange, PreProcessorConfig, PreprocessedEncoderInputs, + PromptReplacement, TrackedMedia, TrackerOutput, VideoClip, VisionProcessorRegistry, }; use llm_tokenizer::TokenizerTrait; use tracing::{debug, info, warn}; @@ -194,7 +194,7 @@ pub(crate) async fn process_multimodal_plan( components, model_id, model_type, - spec.name(), + spec, tokenizer_id, &model_config, ) @@ -334,7 +334,7 @@ async fn preprocess_modality( components: &MultimodalComponents, model_id: &str, model_type: Option<&str>, - model_spec: &str, + spec: &dyn ModelProcessorSpec, tokenizer_id: &str, model_config: &MultimodalModelConfig, ) -> Result { @@ -371,15 +371,12 @@ async fn preprocess_modality( 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, - ) + spec.audio_processor(&model_config.config, &model_config.preprocessor_config) .ok_or_else(|| { - anyhow::anyhow!("No audio processor registered for model spec: {model_spec}") + anyhow::anyhow!( + "No audio processor registered for model spec: {}", + spec.name() + ) })?, ) } else {