merge main

src/models/feature_extractor/config.rs

@@ -0,0 +1,21 @@
+use serde::Deserialize;
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct FeatureExtractor {
+    pub chunk_length: usize,
+    pub dither: f64,
+    pub feature_size: usize,
+    pub hop_length: usize,
+    pub n_fft: usize,
+    pub n_samples: usize,
+    pub nb_max_frames: usize,
+    pub padding_side: String,
+    pub padding_value: f32,
+    pub return_attention_mask: bool,
+    #[serde(default = "default_sampling_rate")]
+    pub sampling_rate: usize,
+}
+
+fn default_sampling_rate() -> usize {
+    16000
+}

src/models/feature_extractor/feature_extraction_whisper.rs

@@ -0,0 +1,113 @@
+use anyhow::Result;
+use candle_core::{D, Device, Tensor};
+
+use crate::utils::{
+    audio_utils::{create_hann_window, mel_filter_bank, torch_stft},
+    tensor_utils::{log10, pad_reflect_last_dim},
+};
+
+pub struct WhisperFeatureExtractor {
+    feature_size: usize,
+    hop_length: usize,
+    chunk_length: usize,
+    n_samples: usize,
+    n_fft: usize,
+    dither: f64,
+    padding_value: f32,
+    sampling_rate: usize,
+    mel_filters: Tensor,
+    window: Tensor,
+}
+
+impl WhisperFeatureExtractor {
+    pub fn new(
+        feature_size: usize,
+        hop_length: usize,
+        chunk_length: usize,
+        n_fft: usize,
+        dither: f64,
+        padding_value: f32,
+        sampling_rate: usize,
+        device: &Device,
+    ) -> Result<Self> {
+        let window = create_hann_window(n_fft, candle_core::DType::F32, device)?;
+        let window = window.unsqueeze(0)?.unsqueeze(0)?;
+        let mel_filters = mel_filter_bank(
+            1 + n_fft / 2,
+            feature_size,
+            0.0,
+            8000.0,
+            sampling_rate as f32,
+            Some("slaney"),
+            crate::utils::audio_utils::MelScale::Slaney,
+            false,
+            device,
+        )?
+        .t()?;
+        let n_samples = chunk_length * sampling_rate;
+        Ok(Self {
+            feature_size,
+            hop_length,
+            chunk_length,
+            n_samples,
+            n_fft,
+            dither,
+            padding_value,
+            sampling_rate,
+            mel_filters,
+            window,
+        })
+    }
+
+    pub fn call(
+        &self,
+        raw_speech: &Tensor,
+        sampling_rate: usize,
+        // do_normalize: bool,
+        return_attention_mask: bool,
+    ) -> Result<(Tensor, Option<Tensor>)> {
+        // raw_speech: 重采样后的音频，shape: (bs, raw_len)
+        // sampling_rate: 音频采样率，验证是否与模型的预处理采样率一致
+        if sampling_rate != self.sampling_rate {
+            return Err(anyhow::anyhow!(
+                "The model feature extractor was trained sampling rate {} not equal to audio sample rate {}",
+                self.sampling_rate,
+                sampling_rate
+            ));
+        }
+
+        let input_features = self.extract_fbank_features(raw_speech)?;
+        let mask_len = input_features.dim(2)?;
+        let mask = if return_attention_mask {
+            let mask = Tensor::new(1u32, input_features.device())?.broadcast_as((1, mask_len))?;
+            Some(mask)
+        } else {
+            None
+        };
+        Ok((input_features, mask))
+    }
+
+    pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result<Tensor> {
+        let mut waveform = waveform.clone();
+        if self.dither != 0.0 {
+            waveform = waveform.add(&waveform.randn_like(0.0, 1.0)?.affine(self.dither, 0.0)?)?;
+        }
+        let pad = self.n_fft / 2;
+        let waveform = pad_reflect_last_dim(&waveform, (pad, pad))?;
+        let (_, samples) = waveform.dims2()?;
+
+        let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)?
+            .transpose(D::Minus1, D::Minus2)?;
+        let n_frames = (samples - self.n_fft) / self.hop_length + 1;
+        let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?;
+        let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?;
+        let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?;
+        // let ln_spec = mel_spec.log()?;
+        // let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?;
+        let log10_spec = log10(&mel_spec)?;
+        let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?;
+        let log10_spec = log10_spec.broadcast_maximum(&max_val)?;
+        let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?;
+        Ok(log_spec)
+    }
+}

src/models/feature_extractor/mod.rs

@@ -1 +1,3 @@
-pub mod seamless_m4t_feature_extractor;
+pub mod seamless_m4t_feature_extractor;
+pub mod feature_extraction_whisper;
+pub mod config;

src/models/feature_extractor/seamless_m4t_feature_extractor.rs

@@ -7,7 +7,7 @@ use crate::utils::{
 };
 
 pub struct SeamlessM4TFeatureExtractor {
-    feature_size: usize,
+    // feature_size: usize,
     num_mel_bins: usize,
     padding_side: PaddingSide,
     padding_value: f32,
@@ -19,7 +19,7 @@ pub struct SeamlessM4TFeatureExtractor {
 
 impl SeamlessM4TFeatureExtractor {
     pub fn new(
-        feature_size: usize,
+        // feature_size: usize,
         num_mel_bins: usize,
         padding_side: PaddingSide,
         padding_value: f32,
@@ -40,7 +40,7 @@ impl SeamlessM4TFeatureExtractor {
         )?;
         let window = create_povey_window(400, candle_core::DType::F32, device)?;
         Ok(Self {
-            feature_size,
+            // feature_size,
             num_mel_bins,
             padding_side,
             padding_value,

src/models/glm_asr_nano/config.rs

@@ -1,6 +1,8 @@
 use candle_nn::Activation;
 use serde::Deserialize;
 
+use crate::models::feature_extractor::config::FeatureExtractor;
+
 #[derive(Debug, Clone, PartialEq, Deserialize)]
 pub struct GlmAsrNanoProcessorConfig {
     pub audio_token: String,
@@ -9,20 +11,7 @@ pub struct GlmAsrNanoProcessorConfig {
     pub max_audio_len: usize,
 }
 
-#[derive(Debug, Clone, PartialEq, Deserialize)]
-pub struct FeatureExtractor {
-    pub chunk_length: usize,
-    pub dither: f32,
-    pub feature_size: usize,
-    pub hop_length: usize,
-    pub n_fft: usize,
-    pub n_samples: usize,
-    pub nb_max_frames: usize,
-    pub padding_side: String,
-    pub padding_value: f32,
-    pub return_attention_mask: bool,
-    pub sampling_rate: usize,
-}
+
 
 #[derive(Debug, Clone, PartialEq, Deserialize)]
 pub struct GlmAsrNanoConfig {

src/models/glm_asr_nano/processor.rs

@@ -5,7 +5,10 @@ use anyhow::Result;
 use candle_core::{D, DType, Device, IndexOp, Tensor};
 
 use crate::{
-    models::glm_asr_nano::config::GlmAsrNanoProcessorConfig,
+    models::{
+        feature_extractor::feature_extraction_whisper::WhisperFeatureExtractor,
+        glm_asr_nano::config::GlmAsrNanoProcessorConfig,
+    },
     utils::{
         audio_utils::{
             apply_stft, create_hann_window, extract_audios, extract_frames, mel_filter_bank,
@@ -28,6 +31,7 @@ pub struct GlmAsrNanoProcessor {
     max_audio_len: usize,
     // default_transcription_prompt: String,
     device: Device,
+    whisper_feature_extrator: WhisperFeatureExtractor,
 }
 
 impl GlmAsrNanoProcessor {
@@ -67,6 +71,16 @@ impl GlmAsrNanoProcessor {
             device,
         )?
         .t()?;
+        let whisper_feature_extrator = WhisperFeatureExtractor::new(
+            processor_cfg.feature_extractor.feature_size,
+            processor_cfg.feature_extractor.hop_length,
+            processor_cfg.feature_extractor.chunk_length,
+            processor_cfg.feature_extractor.n_fft,
+            processor_cfg.feature_extractor.dither,
+            processor_cfg.feature_extractor.padding_value,
+            processor_cfg.feature_extractor.sampling_rate,
+            device,
+        )?;
         Ok(Self {
             sampling_rate,
             chunk_length,
@@ -80,34 +94,35 @@ impl GlmAsrNanoProcessor {
             max_audio_len,
             // default_transcription_prompt,
             device: device.clone(),
+            whisper_feature_extrator,
         })
     }
 
-    pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result<Tensor> {
-        let pad = self.n_fft / 2;
-        let waveform = pad_reflect_last_dim(waveform, (pad, pad))?;
-        let (_, samples) = waveform.dims2()?;
+    // pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result<Tensor> {
+    //     let pad = self.n_fft / 2;
+    //     let waveform = pad_reflect_last_dim(waveform, (pad, pad))?;
+    //     let (_, samples) = waveform.dims2()?;
 
-        // // (bs, n_frames, n_fft)
-        // let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?;
-        // // 应用汉明窗口
-        // let result = frames.broadcast_mul(&self.window)?;
-        // // 傅立叶变换
-        // let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?;
-        let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)?
-            .transpose(D::Minus1, D::Minus2)?;
-        let n_frames = (samples - self.n_fft) / self.hop_length + 1;
-        let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?;
-        let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?;
-        let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?;
-        // let ln_spec = mel_spec.log()?;
-        // let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?;
-        let log10_spec = log10(&mel_spec)?;
-        let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?;
-        let log10_spec = log10_spec.broadcast_maximum(&max_val)?;
-        let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?;
-        Ok(log_spec)
-    }
+    //     // // (bs, n_frames, n_fft)
+    //     // let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?;
+    //     // // 应用汉明窗口
+    //     // let result = frames.broadcast_mul(&self.window)?;
+    //     // // 傅立叶变换
+    //     // let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?;
+    //     let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)?
+    //         .transpose(D::Minus1, D::Minus2)?;
+    //     let n_frames = (samples - self.n_fft) / self.hop_length + 1;
+    //     let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?;
+    //     let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?;
+    //     let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?;
+    //     // let ln_spec = mel_spec.log()?;
+    //     // let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?;
+    //     let log10_spec = log10(&mel_spec)?;
+    //     let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?;
+    //     let log10_spec = log10_spec.broadcast_maximum(&max_val)?;
+    //     let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?;
+    //     Ok(log_spec)
+    // }
 
     pub fn feature_extractor(&self, raw_speech: Vec<Tensor>) -> Result<(Tensor, Tensor)> {
         let mut pad_audio = vec![];
@@ -132,7 +147,10 @@ impl GlmAsrNanoProcessor {
         }
         let input_features = Tensor::cat(&pad_audio, 0)?;
         let input_features_mask = Tensor::new(input_features_mask, input_features.device())?;
-        let input_features = self.extract_fbank_features(&input_features)?;
+        // let input_features = self.extract_fbank_features(&input_features)?;
+        let (input_features, _) =
+            self.whisper_feature_extrator
+                .call(&input_features, self.sampling_rate, false)?;
         let (_, audio_len) = input_features_mask.dims2()?;
         let mask_idx: Vec<u32> = (0..audio_len)
             .step_by(self.hop_length)

src/models/index_tts2/processor.rs

@@ -45,7 +45,7 @@ impl IndexTTS2Processor {
         dtype: DType,
     ) -> Result<Self> {
         let feature_extractor = SeamlessM4TFeatureExtractor::new(
-            80,
+            // 80,
             80,
             crate::utils::tensor_utils::PaddingSide::Right,
             1.0,

src/models/mod.rs

@@ -11,6 +11,7 @@ pub mod minicpm4;
 pub mod paddleocr_vl;
 pub mod qwen2_5vl;
 pub mod qwen3;
+pub mod qwen3_asr;
 pub mod qwen3vl;
 pub mod rmbg2_0;
 pub mod voxcpm;

src/models/qwen3_asr/model.rs

@@ -11,7 +11,7 @@ use crate::{
         qwen3::{config::Qwen3Config, model::Qwen3Model},
     },
     position_embed::sinusoidal_pe::SinusoidalPositionEncoderCat,
-    utils::tensor_utils::{get_equal_mask, mask_filled, masked_scatter_dim0},
+    utils::tensor_utils::{get_equal_mask, attn_masked_fill, masked_scatter_dim0},
 };
 
 pub struct MultiHeadedAttentionSANM {
@@ -124,9 +124,9 @@ impl MultiHeadedAttentionSANM {
             };
             // mask: rank = 2
             let mask = get_equal_mask(&mask, 0)?;
-            let scores = mask_filled(scores, &mask, f32::NEG_INFINITY)?;
+            let scores = attn_masked_fill(scores, &mask, f32::NEG_INFINITY)?;
             let attn = softmax_last_dim(&scores)?;
-            mask_filled(&attn, &mask, 0.0)?
+            attn_masked_fill(&attn, &mask, 0.0)?
         } else {
             softmax_last_dim(scores)?
         };