diff --git a/src/models/common/modules.rs b/src/models/common/modules.rs index bc0fb8b..1f335d0 100644 --- a/src/models/common/modules.rs +++ b/src/models/common/modules.rs @@ -16,8 +16,7 @@ use crate::{ pub struct VadFrameResult { pub is_speech: bool, // pub is_speech_start: bool, - pub is_i16: bool, - pub orig_audio: Option, + pub orig_audio: Option, // nedd f32 data type // pub kaldi_audio: Option, pub model_name: String, pub mode: String, diff --git a/src/models/fire_red_vad/processor.rs b/src/models/fire_red_vad/processor.rs index 145b39d..4b72452 100644 --- a/src/models/fire_red_vad/processor.rs +++ b/src/models/fire_red_vad/processor.rs @@ -133,8 +133,9 @@ impl AudioFeat { } pub fn extract_file(&self, audio_path: &str, device: &Device) -> Result<(Tensor, f32)> { - let wave_tensor = - load_audio_with_resample(audio_path, device, Some(16000), true)?.squeeze(0)?; + let wave_tensor = load_audio_with_resample(audio_path, device, Some(16000))?.squeeze(0)?; + // fire_red_vad need i16 type data + let wave_tensor = wave_tensor.affine(32768.0, 0.0)?; let dur = wave_tensor.dim(0)? as f32 / 16000.0; let fbank = self.extract(&wave_tensor)?; Ok((fbank, dur)) diff --git a/src/models/fire_red_vad/vad.rs b/src/models/fire_red_vad/vad.rs index 6631df8..90acd18 100644 --- a/src/models/fire_red_vad/vad.rs +++ b/src/models/fire_red_vad/vad.rs @@ -36,6 +36,11 @@ pub struct FireRedVad { caches: Option>, frame_length_sample: usize, speech_cache: Vec, + pred_cache: Vec, + min_speach_frames: usize, + look_back_frames: usize, + min_speach_ratio: f32, + end_silence_ratio: f32, } impl FireRedVad { @@ -78,6 +83,11 @@ impl FireRedVad { caches: None, frame_length_sample: 400, speech_cache: vec![], + pred_cache: vec![], + min_speach_frames: 30, // 约 250ms + look_back_frames: 15, // 约 80ms + min_speach_ratio: 0.1, + end_silence_ratio: 0.8, }) } @@ -89,7 +99,8 @@ impl FireRedVad { audio_frame.dim(0)? )); } - let feats = self.audio_feat.extract(audio_frame)?; + let wave_tensor = audio_frame.affine(32768.0, 0.0)?; + let feats = self.audio_feat.extract(&wave_tensor)?; let (probs, caches) = self .vad_model .forward(&feats.unsqueeze(0)?, self.caches.as_ref())?; @@ -101,42 +112,45 @@ impl FireRedVad { .to_dtype(DType::U32)?; let preds_sum = binary_preds.sum_all()?.to_scalar::()?; let probs_len = probs.dim(0)?; - // 输入数据中 is_speech > 0.1, 认为这帧数据可用 - let final_data = if preds_sum as f32 > probs_len as f32 * 0.1 { - // 通过最后10个数据，判断说话是否结束，如果数据长度小于10就取整个长度 - let select_len = if probs_len > 10 { 10 } else { probs_len }; - let last_10_preds_sum = binary_preds - .narrow(0, probs_len - select_len, select_len)? - .sum_all()? - .to_scalar::()?; - // 选中数据中，至少0.8个是 speech, 认为说话没有结束，缓存数据，等待下一帧 - if last_10_preds_sum >= (select_len as f32 * 0.8).ceil() as u32 { - self.speech_cache.push(audio_frame.clone()); + // 输入数据中 is_speech > 0.1, 认为这帧数据有人声 + let final_data = if preds_sum as f32 > probs_len as f32 * self.min_speach_ratio { + self.speech_cache.push(audio_frame.clone()); + let preds = binary_preds.to_vec1::()?; + self.pred_cache.extend_from_slice(&preds); + + // 人声缓存数据过少，等待下一帧 + if self.pred_cache.len() < self.min_speach_frames { None } else { - // 否则认为此次说话结束，结合缓存数据，一起返回 - let data = if self.speech_cache.is_empty() { - // 缓存数据为空，直接返回 - audio_frame.clone() + // 判断是否停止说话 + let start = self.pred_cache.len() - self.look_back_frames; + let look_back = self.pred_cache[start..].iter().sum::(); + // 判断结尾是否静音 + let silence_ratio = 1.0 - (look_back as f32 / self.look_back_frames as f32); + if silence_ratio >= self.end_silence_ratio { + // 静音返回缓存数据并清空缓存 + let speech = Tensor::cat(&self.speech_cache, 0)?; + self.speech_cache.clear(); + self.pred_cache.clear(); + Some(speech) } else { - // 缓存数据不为空，cat数据 - self.speech_cache.push(audio_frame.clone()); - let audio_frame = Tensor::cat(&self.speech_cache, 0)?; - self.speech_cache = vec![]; // 清空缓存 - audio_frame - }; - Some(data) + // 不是静音此次返回None + None + } } } else { - // 认为这帧数据不可用，是否有缓存数据 - if self.speech_cache.is_empty() { - // 没有返回None - None + // 认为这帧数据没有人声 + // 有缓存数据, 且数据够长 + if self.pred_cache.len() >= self.min_speach_frames { + let data = Tensor::cat(&self.speech_cache, 0)?; + self.speech_cache.clear(); + self.pred_cache.clear(); + Some(data) } else { - // 有缓存，返回缓存数据 - let data = Some(Tensor::cat(&self.speech_cache, 0)?); - self.speech_cache.clear(); // 清空缓存 - data + // 否则直接清空缓存 + self.speech_cache.clear(); + self.pred_cache.clear(); + None } }; @@ -145,7 +159,6 @@ impl FireRedVad { } else { Ok(Some(VadFrameResult { is_speech: true, - is_i16: true, orig_audio: final_data, model_name: self.model_name.clone(), mode: "speech".to_string(), @@ -168,7 +181,6 @@ impl FireRedVad { channels, orig_sr, Some(16000), - true, )? .squeeze(0)?; self.detect_frame(&audio_frame) @@ -179,7 +191,7 @@ impl FireRedVad { return Err(anyhow!("only stream model support detect_frame")); } let audio_frame = - resample_audio_from_bytes(audio_bytes, &self.device, Some(16000), true)?.squeeze(0)?; + resample_audio_from_bytes(audio_bytes, &self.device, Some(16000))?.squeeze(0)?; self.detect_frame(&audio_frame) } diff --git a/src/models/qwen3_asr/generate.rs b/src/models/qwen3_asr/generate.rs index 7a66d9a..5af3bd4 100644 --- a/src/models/qwen3_asr/generate.rs +++ b/src/models/qwen3_asr/generate.rs @@ -90,19 +90,16 @@ impl<'a> Qwen3AsrGenerateModel<'a> { return Ok(AsrResult::init_empty()); } if let Some(audio) = vad_res.orig_audio { - self.asr_audio(&audio, vad_res.is_i16) + self.asr_audio(&audio) } else { Ok(AsrResult::init_empty()) } } - pub fn asr_audio(&mut self, audio: &Tensor, is_i16: bool) -> Result { - let audio_data = self.processor.process_audio_tensor( - &self.default_template, - audio, - is_i16, - &self.tokenizer, - )?; + pub fn asr_audio(&mut self, audio: &Tensor) -> Result { + let audio_data = + self.processor + .process_audio_tensor(&self.default_template, audio, &self.tokenizer)?; let input_ids = audio_data.input_ids.clone(); let input_features = Some(audio_data.input_features.clone().to_dtype(self.dtype)?); let mut ctx = GenerationContext::new( diff --git a/src/models/qwen3_asr/processor.rs b/src/models/qwen3_asr/processor.rs index 24b225b..859022c 100644 --- a/src/models/qwen3_asr/processor.rs +++ b/src/models/qwen3_asr/processor.rs @@ -1,6 +1,5 @@ use crate::{ - models::common::modules::{VadFrameResult, float_range_normalize}, - params::chat::ChatCompletionParameters, + models::common::modules::float_range_normalize, params::chat::ChatCompletionParameters, }; use anyhow::{Result, anyhow}; use candle_core::{Device, Tensor}; @@ -101,7 +100,6 @@ impl Qwen3AsrProcessor { &self, render: &str, audio: &Tensor, - is_i16: bool, tokenizer: &TokenizerModel, ) -> Result { let audio_len = audio.dim(0)? as f32; @@ -109,9 +107,6 @@ impl Qwen3AsrProcessor { return Err(anyhow!("vad_res orig_audio is too long!")); } let mut audio = audio.unsqueeze(0)?; - if is_i16 { - audio = audio.affine(1.0 / 32768.0, 0.0)?; - } audio = float_range_normalize(&audio)?; let (input_features, _) = self.whisper_feature_extracor @@ -128,19 +123,6 @@ impl Qwen3AsrProcessor { Ok(audio_data) } - pub fn process_vad_res( - &self, - render: &str, - vad_res: VadFrameResult, - tokenizer: &TokenizerModel, - ) -> Result { - if let Some(audio) = &vad_res.orig_audio { - self.process_audio_tensor(render, audio, vad_res.is_i16, tokenizer) - } else { - Err(anyhow!("vad_res orig_audio is none!")) - } - } - pub fn process_info( &self, mes: &ChatCompletionParameters, diff --git a/src/models/voxcpm/model.rs b/src/models/voxcpm/model.rs index 4f6e3ba..16265b0 100644 --- a/src/models/voxcpm/model.rs +++ b/src/models/voxcpm/model.rs @@ -534,8 +534,7 @@ impl VoxCPMModel { let audio_start = Tensor::new(vec![self.audio_start_token], &self.device)?; let text_token = Tensor::cat(&[text_token, audio_start], D::Minus1)?; let text_length = text_token.dim(0)?; - let mut audio = - load_audio_with_resample(&path, &self.device, Some(self.sample_rate), false)?; + let mut audio = load_audio_with_resample(&path, &self.device, Some(self.sample_rate))?; let patch_len = self.patch_size * self.chunk_size; if audio.dim(1)? % patch_len != 0 { audio = @@ -575,8 +574,7 @@ impl VoxCPMModel { let audio_start = Tensor::new(vec![self.audio_start_token], &self.device)?; let text_token = Tensor::cat(&[text_token, audio_start], D::Minus1)?; let text_length = text_token.dim(0)?; - let mut audio = - load_audio_with_resample(&path, &self.device, Some(self.sample_rate), false)?; + let mut audio = load_audio_with_resample(&path, &self.device, Some(self.sample_rate))?; let patch_len = self.patch_size * self.chunk_size; if audio.dim(1)? % patch_len != 0 { audio = @@ -843,12 +841,8 @@ impl VoxCPMModel { ) -> Result> { let text_token = self.tokenizer.encode(prompt_text)?; let text_token = Tensor::from_slice(&text_token, text_token.len(), &self.device)?; - let mut audio = load_audio_with_resample( - &prompt_wav_path, - &self.device, - Some(self.sample_rate), - false, - )?; + let mut audio = + load_audio_with_resample(&prompt_wav_path, &self.device, Some(self.sample_rate))?; let patch_len = self.patch_size * self.chunk_size; if audio.dim(1)? % patch_len != 0 { audio = audio.pad_with_zeros(D::Minus1, 0, patch_len - audio.dim(1)? % patch_len)?; diff --git a/src/utils/audio_utils.rs b/src/utils/audio_utils.rs index 6ab5eef..d399849 100644 --- a/src/utils/audio_utils.rs +++ b/src/utils/audio_utils.rs @@ -462,11 +462,7 @@ pub fn get_audio_format_from_bytes(bytes: &[u8]) -> Result { } } -pub fn load_audio_use_symphonia( - audio_vec: Vec, - is_i16: bool, - device: &Device, -) -> Result<(Tensor, usize)> { +pub fn load_audio_use_symphonia(audio_vec: Vec, device: &Device) -> Result<(Tensor, usize)> { let extension = get_audio_format_from_bytes(&audio_vec)?; let content = Cursor::new(audio_vec); let mss = MediaSourceStream::new(Box::new(content), Default::default()); @@ -509,14 +505,7 @@ pub fn load_audio_use_symphonia( all_samples.push(Vec::new()); } let channel_data = buf.chan(channel); - if is_i16 { - // 将[-1.0, 1.0] => [-i16_max, i16_max] - let i16_data: Vec = - channel_data.iter().map(|&s| s * 32768.0).collect(); - all_samples[channel].extend_from_slice(&i16_data); - } else { - all_samples[channel].extend_from_slice(channel_data); - } + all_samples[channel].extend_from_slice(channel_data); } } AudioBufferRef::S16(buf) => { @@ -527,17 +516,10 @@ pub fn load_audio_use_symphonia( all_samples.push(Vec::new()); } let channel_data = buf.chan(channel); - let float_samples: Vec = if is_i16 { - channel_data - .iter() - .map(|&s| s as f32) // 转换为f32类型 - .collect() - } else { - channel_data - .iter() - .map(|&s| s as f32 / 32768.0) // 转换为[-1, 1] - .collect() - }; + let float_samples: Vec = channel_data + .iter() + .map(|&s| s as f32 / 32768.0) // 转换为[-1, 1] + .collect(); all_samples[channel].extend(float_samples); } } @@ -549,17 +531,10 @@ pub fn load_audio_use_symphonia( all_samples.push(Vec::new()); } let channel_data = buf.chan(channel); - let float_samples: Vec = if is_i16 { - channel_data - .iter() - .map(|&s| s.inner() as f32 / 8388608.0 * 32768.0) // 转换为[-i16_max, i16_max] - .collect() - } else { - channel_data - .iter() - .map(|&s| s.inner() as f32 / 8388608.0) // 转换为[-1, 1] - .collect() - }; + let float_samples: Vec = channel_data + .iter() + .map(|&s| s.inner() as f32 / 8388608.0) // 转换为[-1, 1] + .collect(); all_samples[channel].extend(float_samples); } } @@ -584,7 +559,7 @@ pub fn load_audio_use_symphonia( pub fn load_audio(path: &str, device: &Device) -> Result<(Tensor, usize)> { let audio_vec = get_audio_bytes_vec(path)?; - load_audio_use_symphonia(audio_vec, false, device) + load_audio_use_symphonia(audio_vec, device) } pub fn resample_audio_from_vec_f32( @@ -593,11 +568,11 @@ pub fn resample_audio_from_vec_f32( channels: usize, orig_sr: Option, target_sample_rate: Option, - is_i16: bool, ) -> Result { let frame_len = audio_vec.len() / channels; let audio = Tensor::new(&audio_vec[0..frame_len * channels], device)?; let mut audio = if channels > 1 { + // 交错格式（立体声时 LRLR...） audio .reshape((frame_len, channels))? .mean_keepdim(1)? @@ -613,9 +588,6 @@ pub fn resample_audio_from_vec_f32( { audio = resample_simple(&audio, sr as i64, target_sample_rate as i64)?; } - if is_i16 { - audio = audio.affine(32768.0, 0.0)?; - } Ok(audio) } @@ -623,9 +595,8 @@ pub fn resample_audio_from_bytes( audio_vec: Vec, device: &Device, target_sample_rate: Option, - is_i16: bool, ) -> Result { - let (mut audio, sr) = load_audio_use_symphonia(audio_vec, is_i16, device)?; + let (mut audio, sr) = load_audio_use_symphonia(audio_vec, device)?; if let Some(target_sample_rate) = target_sample_rate && target_sample_rate != sr { @@ -638,13 +609,12 @@ pub fn load_audio_with_resample( path: &str, device: &Device, target_sample_rate: Option, - is_i16: bool, ) -> Result { // hound 只支持wav文件 // let audio_path = get_audio_path(path)?; // let (mut audio, sr) = load_audio_use_hound(audio_path, device)?; let audio_vec = get_audio_bytes_vec(path)?; - resample_audio_from_bytes(audio_vec, device, target_sample_rate, is_i16) + resample_audio_from_bytes(audio_vec, device, target_sample_rate) } pub fn save_wav(audio: &Tensor, save_path: &str, sample_rate: u32) -> Result<()> { @@ -719,7 +689,7 @@ pub fn extract_audios( // 并行加载音频 audio_url_vec .par_iter() - .map(|url| load_audio_with_resample(url, device, target_sample_rate, false)) + .map(|url| load_audio_with_resample(url, device, target_sample_rate)) .collect() // #[cfg(not(feature = "ffmpeg"))] // { diff --git a/tests/weight_test.rs b/tests/weight_test.rs index 9e10bf3..90b4938 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -410,3 +410,18 @@ fn fire_red_vad_weight() -> Result<()> { Ok(()) } + +#[test] +fn silero_vad_weight() -> Result<()> { + // cargo test -F cuda --test weight_test silero_vad_weight -r -- --nocapture + let save_dir = + aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?; + let model_path = format!("{}/silero_vad/silero_vad_16k.safetensors", save_dir); + let device = get_device(None); + let weights = safetensors::load(model_path, &device)?; + for (key, tensor) in weights.iter() { + println!("=== {} === {:?}", key, tensor); + } + + Ok(()) +}