User specified

Author: EricLBuehler

mistralrs-core/src/engine/add_request.rs

@@ -122,10 +122,21 @@ impl Engine {
                 ref images,
                 messages: _,
                 enable_thinking: _,
+                audios: _,
             } => Some(images.clone()),
             _ => None,
         };
 
+        let audios = match request.messages {
+            RequestMessage::VisionChat {
+                images: _,
+                messages: _,
+                enable_thinking: _,
+                ref audios,
+            } => Some(audios.clone()),
+            _ => None,
+        };
+
         let matcher = Arc::new(handle_seq_error!(
             ToolCallingMatcher::new(request.tool_choice.unwrap_or(ToolChoice::Auto),),
             request.response
@@ -157,6 +168,7 @@ impl Engine {
             }
             | RequestMessage::VisionChat {
                 images: _,
+                audios: _,
                 messages,
                 enable_thinking,
             } => {
@@ -497,6 +509,7 @@ impl Engine {
                     None
                 },
                 images.clone(),
+                audios.clone(),
                 block_size,
                 Some(matcher.clone()),
                 image_generation_format,

mistralrs-core/src/lib.rs

@@ -97,9 +97,10 @@ pub use pipeline::{
     UQFF_MULTI_FILE_DELIMITER,
 };
 pub use request::{
-    ApproximateUserLocation, Constraint, DetokenizationRequest, ImageGenerationResponseFormat,
-    LlguidanceGrammar, MessageContent, NormalRequest, Request, RequestMessage, SearchContextSize,
-    TokenizationRequest, WebSearchOptions, WebSearchUserLocation,
+    ApproximateUserLocation, AudioInput, Constraint, DetokenizationRequest,
+    ImageGenerationResponseFormat, LlguidanceGrammar, MessageContent, NormalRequest, Request,
+    RequestMessage, SearchContextSize, TokenizationRequest, WebSearchOptions,
+    WebSearchUserLocation,
 };
 pub use response::*;
 pub use sampler::{

mistralrs-core/src/pipeline/amoe.rs

@@ -584,6 +584,7 @@ fn new_dummy_seq(
         None,
         None,
         images,
+        None,
         None, // TODO incorrect for PagedAttention
         None,
         None,

mistralrs-core/src/request.rs

@@ -50,8 +50,10 @@ pub enum RequestMessage {
     },
     CompletionTokens(Vec<u32>),
     VisionChat {
-        #[serde(skip)] // TODO!!!!
+        #[serde(skip)] // TODO
         images: Vec<image::DynamicImage>,
+        #[serde(skip)] // TODO
+        audios: Vec<AudioInput>,
         messages: Vec<IndexMap<String, MessageContent>>,
         enable_thinking: Option<bool>,
     },
@@ -116,6 +118,42 @@ pub struct WebSearchOptions {
     pub extract_description: Option<String>,
 }
 
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
+/// Raw audio input consisting of PCM samples and a sample rate.
+pub struct AudioInput {
+    pub samples: Vec<f32>,
+    pub sample_rate: u32,
+}
+
+impl AudioInput {
+    pub fn read_wav(wav_path: &str) -> anyhow::Result<Self> {
+        let mut reader = hound::WavReader::open(wav_path)
+            .map_err(|e| anyhow::Error::msg(format!("Failed to load audio: {}", e)))?;
+        let spec = reader.spec();
+
+        let samples: Vec<f32> = match spec.sample_format {
+            hound::SampleFormat::Float => reader
+                .samples::<f32>()
+                .map(|s| s.map_err(|e| anyhow::Error::msg(e.to_string())))
+                .collect::<std::result::Result<_, _>>()?,
+
+            hound::SampleFormat::Int => reader
+                .samples::<i16>() // read as integers
+                .map(|s| {
+                    s.map(|v| v as f32 / i16::MAX as f32) // scale to –1.0…1.0
+                        .map_err(|e| candle_core::Error::Msg(e.to_string()))
+                })
+                .collect::<std::result::Result<_, _>>()?,
+        };
+
+        Ok(Self {
+            samples,
+            sample_rate: spec.sample_rate,
+        })
+    }
+}
+
 #[derive(Clone, Serialize, Deserialize)]
 /// A normal request request to the `MistralRs`.
 /// - `messages`: Messages for the request

mistralrs-core/src/sequence.rs

@@ -4,7 +4,7 @@ use crate::{
     pipeline::{text_models_inputs_processor::PagedAttentionMeta, LayerCaches},
     response::{ChatCompletionChunkResponse, Choice, ChunkChoice, Response, SYSTEM_FINGERPRINT},
     sampler::{Logprobs, Sampler},
-    ChatCompletionResponse, Usage,
+    AudioInput, ChatCompletionResponse, Usage,
 };
 use crate::{
     paged_attention::{BlockEngineSequence, LogicalTokenBlock},
@@ -171,6 +171,53 @@ pub struct SequenceImages {
     hashes: Vec<u64>,
 }
 
+#[derive(Clone)]
+pub struct SequenceAudios {
+    audios: Vec<AudioInput>,
+    hashes: Vec<u64>,
+}
+
+impl SequenceAudios {
+    fn new(input_audios: Vec<AudioInput>) -> Self {
+        let hashes = input_audios.iter().map(|a| {
+            let mut hasher = DefaultHasher::new();
+            for s in &a.samples {
+                s.to_bits().hash(&mut hasher);
+            }
+            a.sample_rate.hash(&mut hasher);
+            hasher.finish()
+        });
+        Self {
+            hashes: hashes.collect(),
+            audios: input_audios,
+        }
+    }
+
+    fn clone_audios(&self) -> Vec<AudioInput> {
+        self.audios.clone()
+    }
+
+    fn audios(&self) -> &[AudioInput] {
+        &self.audios
+    }
+
+    fn audios_mut(&mut self) -> &mut Vec<AudioInput> {
+        &mut self.audios
+    }
+
+    fn hashes(&self) -> &[u64] {
+        &self.hashes
+    }
+
+    fn keep_num_audios(&mut self, audios_to_keep: usize) {
+        if self.audios.len() > audios_to_keep {
+            let start = self.audios.len() - audios_to_keep;
+            self.audios = self.audios[start..].to_vec();
+            self.hashes = self.hashes[start..].to_vec();
+        }
+    }
+}
+
 impl SequenceImages {
     fn new(input_images: Vec<image::DynamicImage>) -> Self {
         let hashes = input_images.iter().map(|x| {
@@ -211,6 +258,7 @@ impl SequenceImages {
 // Holds all multimodal (vision/diffusion) data for a Sequence.
 pub struct MultimodalData {
     pub input_images: Option<SequenceImages>,
+    pub input_audios: Option<SequenceAudios>,
     pub cached_pixel_values: Option<Tensor>,
     pub cached_img_thw: Option<Tensor>,
     pub cached_vid_thw: Option<Tensor>,
@@ -222,11 +270,13 @@ pub struct MultimodalData {
 impl MultimodalData {
     pub fn new(
         input_images: Option<Vec<image::DynamicImage>>,
+        input_audios: Option<Vec<AudioInput>>,
         image_gen_response_format: Option<ImageGenerationResponseFormat>,
         diffusion_params: Option<DiffusionGenerationParams>,
     ) -> Self {
         MultimodalData {
             input_images: input_images.map(SequenceImages::new),
+            input_audios: input_audios.map(SequenceAudios::new),
             cached_pixel_values: None,
             cached_img_thw: None,
             cached_vid_thw: None,
@@ -268,6 +318,40 @@ impl MultimodalData {
             .is_some_and(|imgs| !imgs.images().is_empty())
     }
 
+    pub fn take_audios(&mut self) -> Option<Vec<AudioInput>> {
+        if let Some(input_audios) = self.input_audios.as_mut() {
+            let mut audios = Vec::new();
+            std::mem::swap(&mut audios, input_audios.audios_mut());
+            Some(audios)
+        } else {
+            None
+        }
+    }
+
+    pub fn clone_audios(&self) -> Option<Vec<AudioInput>> {
+        self.input_audios.as_ref().map(|a| a.clone_audios())
+    }
+
+    pub fn audios(&self) -> Option<&[AudioInput]> {
+        self.input_audios.as_ref().map(|a| a.audios())
+    }
+
+    pub fn audio_hashes(&self) -> Option<&[u64]> {
+        self.input_audios.as_ref().map(|a| a.hashes())
+    }
+
+    pub fn has_audios(&self) -> bool {
+        self.input_audios
+            .as_ref()
+            .is_some_and(|a| !a.audios().is_empty())
+    }
+
+    pub fn keep_num_audios(&mut self, audios_to_keep: usize) {
+        if let Some(auds) = self.input_audios.as_mut() {
+            auds.keep_num_audios(audios_to_keep)
+        }
+    }
+
     pub fn keep_num_images(&mut self, images_to_keep: usize) {
         if let Some(imgs) = self.input_images.as_mut() {
             imgs.keep_num_images(images_to_keep)
@@ -422,6 +506,7 @@ impl Sequence {
         suffix: Option<String>,
         prefix: Option<String>,
         input_images: Option<Vec<image::DynamicImage>>,
+        input_audios: Option<Vec<AudioInput>>,
         // Paged attention
         block_size: Option<usize>,
         //
@@ -492,6 +577,7 @@ impl Sequence {
             // Multimodal data
             multimodal: MultimodalData::new(
                 input_images,
+                input_audios,
                 image_gen_response_format,
                 diffusion_params,
             ),
@@ -967,6 +1053,30 @@ impl Sequence {
         self.multimodal.has_images()
     }
 
+    pub fn take_audios(&mut self) -> Option<Vec<AudioInput>> {
+        self.multimodal.take_audios()
+    }
+
+    pub fn clone_audios(&self) -> Option<Vec<AudioInput>> {
+        self.multimodal.clone_audios()
+    }
+
+    pub fn audios(&self) -> Option<&[AudioInput]> {
+        self.multimodal.audios()
+    }
+
+    pub fn audio_hashes(&self) -> Option<&[u64]> {
+        self.multimodal.audio_hashes()
+    }
+
+    pub fn has_audios(&self) -> bool {
+        self.multimodal.has_audios()
+    }
+
+    pub fn keep_num_audios(&mut self, audios_to_keep: usize) {
+        self.multimodal.keep_num_audios(audios_to_keep)
+    }
+
     /// Keep these last n images
     pub fn keep_num_images(&mut self, images_to_keep: usize) {
         self.multimodal.keep_num_images(images_to_keep)

mistralrs-core/src/vision_models/phi4/inputs_processor.rs

@@ -720,7 +720,7 @@ impl Phi4MMInputsProcessor {
 
     fn process_audio_for_sequences(
         &self,
-        input_seqs: &[&mut Sequence],
+        input_seqs: &mut [&mut Sequence],
         device: &Device,
     ) -> AudioProcessingResult {
         // Check if any sequence has audio tokens
@@ -737,12 +737,19 @@ impl Phi4MMInputsProcessor {
         let mut audio_frames_list = Vec::new();
 
         // Process audio for each sequence that needs it
-        for seq in input_seqs.iter() {
+        for seq in input_seqs.iter_mut() {
             let has_audio = seq.get_toks().contains(&(AUDIO_SPECIAL_TOKEN_ID as u32));
 
             if has_audio {
-                // Load dummy audio (TODO: make this per-sequence)
-                let (audio_data, sample_rate) = self.load_dummy_audio()?;
+                let (audio_data, sample_rate) = if let Some(mut audios) = seq.take_audios() {
+                    if let Some(audio) = audios.pop() {
+                        (audio.samples, audio.sample_rate)
+                    } else {
+                        self.load_dummy_audio()?
+                    }
+                } else {
+                    self.load_dummy_audio()?
+                };
 
                 // Extract features
                 let features = self.extract_audio_features(&audio_data, sample_rate)?;

mistralrs-pyo3/src/lib.rs

@@ -1069,6 +1069,7 @@ impl Runner {
                         RequestMessage::VisionChat {
                             messages: messages_vec,
                             images,
+                            audios: Vec::new(),
                             enable_thinking: request.enable_thinking,
                         }
                     } else {

mistralrs-server-core/src/chat_completion.rs

@@ -442,6 +442,7 @@ pub async fn parse_request(
                 RequestMessage::VisionChat {
                     messages,
                     images,
+                    audios: Vec::new(),
                     enable_thinking: oairequest.enable_thinking,
                 }
             } else {