feat: stabilize voice activity and audio routing

This commit is contained in:
Edison Jwa
2026-05-25 01:19:09 +09:00
parent eb9014cd81
commit 5515ff6643
34 changed files with 3054 additions and 1751 deletions
-74
View File
@@ -7,7 +7,6 @@
pub mod resampler;
pub mod silero_onnx;
pub mod ten_onnx;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{OnceLock, RwLock};
@@ -17,7 +16,6 @@ use crate::AudioError;
use resampler::{Downsampler48to16, INPUT_FRAME_10MS};
pub use silero_onnx::SileroOnnxVad;
pub use ten_onnx::{TenOnnxVad, TenOnnxVadWorker};
/// Voice activity detector output for one 10 ms frame.
#[derive(Debug, Clone, Copy)]
@@ -108,17 +106,11 @@ pub fn process_i16_10ms(detector: &mut dyn VoiceActivityDetector, samples: &[i16
static SILERO_MODEL_PATH_OVERRIDE: OnceLock<RwLock<Option<String>>> = OnceLock::new();
static SILERO_MODEL_EPOCH: AtomicU64 = AtomicU64::new(0);
static TEN_MODEL_PATH_OVERRIDE: OnceLock<RwLock<Option<String>>> = OnceLock::new();
static TEN_MODEL_EPOCH: AtomicU64 = AtomicU64::new(0);
fn silero_model_path_override() -> &'static RwLock<Option<String>> {
SILERO_MODEL_PATH_OVERRIDE.get_or_init(|| RwLock::new(None))
}
fn ten_model_path_override() -> &'static RwLock<Option<String>> {
TEN_MODEL_PATH_OVERRIDE.get_or_init(|| RwLock::new(None))
}
/// Configure the preferred Silero ONNX model path.
///
/// The path is validated eagerly. A successful call increments the
@@ -149,32 +141,6 @@ pub fn silero_model_epoch() -> u64 {
SILERO_MODEL_EPOCH.load(Ordering::Relaxed)
}
/// Configure the preferred TEN VAD ONNX model path.
pub fn set_ten_model_path(path: &str) -> Result<(), AudioError> {
let path = path.trim();
if path.is_empty() {
return Err(AudioError::InvalidAudioProcessingConfig(
"ten vad model path must not be empty".to_string(),
));
}
if !std::path::Path::new(path).is_file() {
return Err(AudioError::InvalidAudioProcessingConfig(format!(
"ten vad model path does not exist or is not a file: {path}"
)));
}
let mut guard = ten_model_path_override()
.write()
.map_err(|_| AudioError::Backend("ten vad model path lock poisoned".to_string()))?;
*guard = Some(path.to_string());
TEN_MODEL_EPOCH.fetch_add(1, Ordering::Relaxed);
Ok(())
}
/// Monotonic counter incremented whenever the configured TEN model path changes.
pub fn ten_model_epoch() -> u64 {
TEN_MODEL_EPOCH.load(Ordering::Relaxed)
}
/// Return the expected path of the Silero VAD v6 ONNX model.
/// The model is shipped as a Flutter asset and copied to the app's
/// data directory by the Dart-side asset loader.
@@ -231,46 +197,6 @@ pub fn silero_model_bundle_path() -> String {
}
}
/// Return the expected path of the TEN VAD ONNX model copied by Flutter.
pub fn ten_model_bundle_path() -> String {
if let Ok(guard) = ten_model_path_override().read() {
if let Some(path) = guard.as_ref() {
return path.clone();
}
}
#[cfg(any(target_os = "ios", target_os = "macos"))]
{
if let Ok(home) = std::env::var("HOME") {
let docs = format!("{home}/Documents/ten_vad.onnx");
if std::path::Path::new(&docs).exists() {
return docs;
}
let bundle = format!("{home}/../Library/ten_vad.onnx");
if std::path::Path::new(&bundle).exists() {
return bundle;
}
}
"ten_vad.onnx".to_string()
}
#[cfg(target_os = "android")]
{
"ten_vad.onnx".to_string()
}
#[cfg(not(any(target_os = "ios", target_os = "macos", target_os = "android")))]
{
if let Ok(cwd) = std::env::current_dir() {
let local = cwd.join("ten_vad.onnx");
if local.exists() {
return local.to_string_lossy().to_string();
}
}
"ten_vad.onnx".to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
-452
View File
@@ -1,452 +0,0 @@
//! TEN VAD ONNX backend.
//!
//! TEN's ONNX graph does not accept raw PCM. It expects the same feature
//! stack produced by TEN's `AUP_Aed_aivad_proc`: three context frames of
//! 40 log-mel powers plus one pitch feature, followed by four recurrent
//! state tensors. This module ports that preprocessing path to Rust and
//! keeps ONNX Runtime off the realtime callback where possible.
use crate::frame::f32_to_i16;
use rustfft::{num_complex::Complex32, FftPlanner};
use super::resampler::{Downsampler48to16, INPUT_FRAME_10MS};
use super::{VadOutput, VoiceActivityDetector};
const SAMPLE_RATE_16K: f32 = 16_000.0;
const HOP_16K: usize = 256;
const WINDOW_16K: usize = 768;
const FFT_SIZE: usize = 1024;
const N_BINS: usize = FFT_SIZE / 2 + 1;
const MEL_BANDS: usize = 40;
const FEATURE_LEN: usize = 41;
const CONTEXT: usize = 3;
const HIDDEN: usize = 64;
const POWER_NORMALIZER: f32 = 32768.0 * 32768.0;
const EPS: f32 = 1.0e-20;
const FEATURE_MEANS: [f32; FEATURE_LEN] = [
-8.198236, -6.2657166, -5.4838185, -4.7586913, -4.417089, -4.142893, -3.9128504, -3.845928,
-3.6570904, -3.7234187, -3.8761342, -3.843891, -3.6904051, -3.7560658, -3.6986961, -3.650463,
-3.7004688, -3.5673213, -3.4989002, -3.477807, -3.458816, -3.4449239, -3.4013286, -3.3062613,
-3.2785568, -3.2332509, -3.198616, -3.2045264, -3.2087986, -3.257838, -3.3813767, -3.5340214,
-3.640868, -3.7268589, -3.773731, -3.8046672, -3.832901, -3.8711205, -3.990593, -4.4802895,
92.3569,
];
const FEATURE_STDS: [f32; FEATURE_LEN] = [
5.166064, 4.9772096, 4.698896, 4.6306214, 4.634348, 4.641156, 4.6406765, 4.666367, 4.6505346,
4.640021, 4.6374, 4.620099, 4.5963163, 4.562655, 4.5543604, 4.5669107, 4.56249, 4.5624127,
4.5852995, 4.6001797, 4.592846, 4.5859227, 4.5834966, 4.626093, 4.626958, 4.6262894, 4.637006,
4.683016, 4.726814, 4.7342896, 4.753227, 4.849723, 4.869435, 4.884483, 4.921327, 4.9592123,
4.996619, 5.0448236, 5.072217, 5.0964394, 115.21369,
];
/// TEN VAD using ONNX Runtime and Rust-ported TEN feature preprocessing.
pub struct TenOnnxVad {
session: ort::session::Session,
downsampler: Downsampler48to16,
hop_accum: Vec<f32>,
sample_fifo: Vec<f32>,
feature_stack: [[f32; FEATURE_LEN]; CONTEXT],
states: [[f32; HIDDEN]; 4],
mel_filters: Vec<[f32; N_BINS]>,
fft: std::sync::Arc<dyn rustfft::Fft<f32>>,
fft_buffer: Vec<Complex32>,
last_probability: f32,
last_speech: bool,
}
unsafe impl Send for TenOnnxVad {}
impl TenOnnxVad {
/// Load TEN VAD ONNX model.
pub fn try_new(model_path: &str) -> Option<Self> {
if !std::path::Path::new(model_path).exists() {
tracing::warn!(target: "chanora_audio", path = model_path, "TEN VAD ONNX model not found");
return None;
}
let session = match std::panic::catch_unwind(|| {
ort::session::Session::builder().and_then(|mut b| b.commit_from_file(model_path))
}) {
Ok(Ok(session)) => session,
Ok(Err(error)) => {
tracing::warn!(target: "chanora_audio", %error, path = model_path, "TEN VAD ONNX model load failed");
return None;
}
Err(_) => {
tracing::warn!(target: "chanora_audio", path = model_path, "TEN VAD ONNX Runtime panicked during load");
return None;
}
};
let mut fft_planner = FftPlanner::<f32>::new();
let fft = fft_planner.plan_fft_forward(FFT_SIZE);
tracing::info!(target: "chanora_audio", path = model_path, "TEN VAD ONNX model loaded");
Some(Self {
session,
downsampler: Downsampler48to16::default(),
hop_accum: Vec::with_capacity(HOP_16K + super::resampler::OUTPUT_FRAME_10MS),
sample_fifo: Vec::with_capacity(WINDOW_16K + HOP_16K),
feature_stack: [[0.0; FEATURE_LEN]; CONTEXT],
states: [[0.0; HIDDEN]; 4],
mel_filters: build_mel_filters(),
fft,
fft_buffer: vec![Complex32::ZERO; FFT_SIZE],
last_probability: 0.0,
last_speech: false,
})
}
fn process_hop(&mut self, hop: &[f32]) {
self.sample_fifo.extend_from_slice(hop);
let frame = if self.sample_fifo.len() >= WINDOW_16K {
let start = self.sample_fifo.len() - WINDOW_16K;
self.sample_fifo[start..].to_vec()
} else {
let mut padded = vec![0.0; WINDOW_16K - self.sample_fifo.len()];
padded.extend_from_slice(&self.sample_fifo);
padded
};
if self.sample_fifo.len() > WINDOW_16K {
let excess = self.sample_fifo.len() - WINDOW_16K;
self.sample_fifo.drain(..excess);
}
let feature = compute_feature(
&self.mel_filters,
self.fft.as_ref(),
&mut self.fft_buffer,
&frame,
);
self.feature_stack.copy_within(1..CONTEXT, 0);
self.feature_stack[CONTEXT - 1] = feature;
self.run_onnx();
}
fn run_onnx(&mut self) {
use ndarray::{Array, IxDyn};
use ort::value::Value;
let input: Vec<f32> = self.feature_stack.iter().flatten().copied().collect();
let input_arr = match Array::from_shape_vec(IxDyn(&[1, CONTEXT, FEATURE_LEN]), input) {
Ok(v) => v,
Err(_) => return,
};
let state_arrs = [0, 1, 2, 3]
.map(|idx| Array::from_shape_vec(IxDyn(&[1, HIDDEN]), self.states[idx].to_vec()));
let input_val = match Value::from_array(input_arr) {
Ok(v) => v,
Err(error) => {
tracing::warn!(target: "chanora_audio", %error, "TEN VAD input tensor error");
return;
}
};
let state_vals = match state_arrs {
[Ok(a), Ok(b), Ok(c), Ok(d)] => [a, b, c, d],
_ => return,
};
let state_vals = match state_vals.map(Value::from_array) {
[Ok(a), Ok(b), Ok(c), Ok(d)] => [a, b, c, d],
_ => return,
};
let outputs = match self.session.run([
(&input_val).into(),
(&state_vals[0]).into(),
(&state_vals[1]).into(),
(&state_vals[2]).into(),
(&state_vals[3]).into(),
]) {
Ok(outputs) => outputs,
Err(error) => {
tracing::warn!(target: "chanora_audio", %error, "TEN VAD ONNX inference failed");
return;
}
};
if let Ok((_, prob)) = outputs["output_1"].try_extract_tensor::<f32>() {
if let Some(&p) = prob.first() {
self.last_probability = p.clamp(0.0, 1.0);
self.last_speech = self.last_probability >= 0.5;
}
}
for (idx, name) in ["output_2", "output_3", "output_6", "output_7"]
.iter()
.enumerate()
{
if let Ok((_, state)) = outputs[*name].try_extract_tensor::<f32>() {
let copy_len = state.len().min(HIDDEN);
self.states[idx][..copy_len].copy_from_slice(&state[..copy_len]);
}
}
}
}
fn compute_feature(
mel_filters: &[[f32; N_BINS]],
fft: &dyn rustfft::Fft<f32>,
fft_buffer: &mut [Complex32],
frame: &[f32],
) -> [f32; FEATURE_LEN] {
let power = power_spectrum(fft, fft_buffer, frame);
let mut feature = [0.0; FEATURE_LEN];
for band in 0..MEL_BANDS {
let energy = mel_filters[band]
.iter()
.zip(power.iter())
.map(|(w, p)| w * p)
.sum::<f32>()
/ POWER_NORMALIZER;
let log_energy = (energy + EPS).ln();
feature[band] = (log_energy - FEATURE_MEANS[band]) / (FEATURE_STDS[band] + EPS);
}
let pitch_hz = estimate_pitch_hz(frame);
feature[MEL_BANDS] = (pitch_hz - FEATURE_MEANS[MEL_BANDS]) / (FEATURE_STDS[MEL_BANDS] + EPS);
feature
}
impl VoiceActivityDetector for TenOnnxVad {
fn process_10ms(&mut self, samples: &[f32]) -> VadOutput {
debug_assert_eq!(samples.len(), INPUT_FRAME_10MS);
let mut input = [0.0_f32; INPUT_FRAME_10MS];
input.copy_from_slice(samples);
let downsampled = self.downsampler.process_frame_10ms(&input);
self.hop_accum.extend_from_slice(&downsampled);
while self.hop_accum.len() >= HOP_16K {
let hop: Vec<f32> = self.hop_accum[..HOP_16K].to_vec();
self.hop_accum.drain(..HOP_16K);
self.process_hop(&hop);
}
VadOutput {
probability: self.last_probability,
speech: self.last_speech,
}
}
}
fn hz_to_mel(hz: f32) -> f32 {
2595.0 * (1.0 + hz / 700.0).log10()
}
fn mel_to_hz(mel: f32) -> f32 {
700.0 * (10.0_f32.powf(mel / 2595.0) - 1.0)
}
fn build_mel_filters() -> Vec<[f32; N_BINS]> {
let low_mel = hz_to_mel(0.0);
let high_mel = hz_to_mel(8000.0);
let mut bins = [0_usize; MEL_BANDS + 2];
for idx in 0..bins.len() {
let mel = idx as f32 * (high_mel - low_mel) / (MEL_BANDS as f32 + 1.0) + low_mel;
let hz = mel_to_hz(mel);
let mut bin = ((FFT_SIZE as f32 + 1.0) * hz / SAMPLE_RATE_16K).floor() as usize;
bin = bin.min(N_BINS - 1);
if idx > 0 && bin == bins[idx - 1] {
bin = (bin + 1).min(N_BINS - 1);
}
bins[idx] = bin;
}
let mut filters = vec![[0.0_f32; N_BINS]; MEL_BANDS];
for band in 0..MEL_BANDS {
let left = bins[band];
let center = bins[band + 1].max(left + 1);
let right = bins[band + 2].max(center + 1).min(N_BINS - 1);
for (i, weight) in filters[band]
.iter_mut()
.enumerate()
.take(center.min(N_BINS))
.skip(left)
{
*weight = (i - left) as f32 / (center - left) as f32;
}
for (i, weight) in filters[band]
.iter_mut()
.enumerate()
.take(right + 1)
.skip(center)
{
*weight = (right - i) as f32 / (right - center).max(1) as f32;
}
}
filters
}
fn power_spectrum(
fft: &dyn rustfft::Fft<f32>,
fft_buffer: &mut [Complex32],
frame: &[f32],
) -> [f32; N_BINS] {
debug_assert_eq!(fft_buffer.len(), FFT_SIZE);
fft_buffer.fill(Complex32::ZERO);
for (idx, sample) in frame.iter().take(WINDOW_16K).enumerate() {
let hann = 0.5 - 0.5 * (2.0 * std::f32::consts::PI * idx as f32 / WINDOW_16K as f32).cos();
fft_buffer[idx].re = f32_to_i16(*sample) as f32 * hann;
}
fft.process(fft_buffer);
let mut out = [0.0_f32; N_BINS];
for (dst, bin) in out.iter_mut().zip(fft_buffer.iter()) {
*dst = bin.norm_sqr();
}
out
}
fn estimate_pitch_hz(frame: &[f32]) -> f32 {
let min_lag = (SAMPLE_RATE_16K / 400.0) as usize;
let max_lag = (SAMPLE_RATE_16K / 60.0) as usize;
let mut best_lag = 0_usize;
let mut best_corr = 0.0_f32;
for lag in min_lag..=max_lag.min(frame.len().saturating_sub(1)) {
let mut corr = 0.0_f32;
let mut energy = 0.0_f32;
for i in lag..frame.len() {
corr += frame[i] * frame[i - lag];
energy += frame[i - lag] * frame[i - lag];
}
let norm = if energy > 1.0e-8 {
corr / energy.sqrt()
} else {
0.0
};
if norm > best_corr {
best_corr = norm;
best_lag = lag;
}
}
if best_lag == 0 || best_corr < 0.01 {
0.0
} else {
SAMPLE_RATE_16K / best_lag as f32
}
}
// ---------------------------------------------------------------------------
// Background worker — same pattern as SileroOnnxVadWorker so the realtime
// callback never blocks on STFT / pitch / ONNX inference.
// ---------------------------------------------------------------------------
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64};
use std::sync::Arc;
use std::thread::JoinHandle;
/// Maximum number of 10 ms frames the worker may lag before the callback
/// treats its output as stale and uses WebRTC fallback instead.
const TEN_MAX_STALE_FRAMES: u64 = 8;
struct TenFrameMessage {
seq: u64,
frame: [f32; INPUT_FRAME_10MS],
}
/// Background TEN VAD worker. The realtime callback only enqueues 10 ms
/// frames and reads the latest probability atomically.
pub struct TenOnnxVadWorker {
tx: Option<std::sync::mpsc::SyncSender<TenFrameMessage>>,
latest_probability: Arc<AtomicU32>,
latest_processed_seq: Arc<AtomicU64>,
alive: Arc<AtomicBool>,
handle: Option<JoinHandle<()>>,
}
impl TenOnnxVadWorker {
/// Start a background TEN worker if the model loads.
pub fn try_new(model_path: &str) -> Option<Self> {
let vad = TenOnnxVad::try_new(model_path)?;
let latest_probability = Arc::new(AtomicU32::new(0.0_f32.to_bits()));
let latest_processed_seq = Arc::new(AtomicU64::new(u64::MAX));
let alive = Arc::new(AtomicBool::new(true));
let (tx, rx) = std::sync::mpsc::sync_channel::<TenFrameMessage>(128);
let prob_arc = latest_probability.clone();
let seq_arc = latest_processed_seq.clone();
let alive_arc = alive.clone();
let handle = std::thread::Builder::new()
.name("chanora-ten-vad".to_string())
.spawn(move || {
let mut vad = vad;
while alive_arc.load(std::sync::atomic::Ordering::Relaxed) {
let msg = match rx.recv() {
Ok(m) => m,
Err(_) => break,
};
let mut frame_f32 = [0.0_f32; INPUT_FRAME_10MS];
frame_f32.copy_from_slice(&msg.frame);
let out = VoiceActivityDetector::process_10ms(&mut vad, &frame_f32);
prob_arc.store(
out.probability.clamp(0.0, 1.0).to_bits(),
std::sync::atomic::Ordering::Relaxed,
);
seq_arc.store(msg.seq, std::sync::atomic::Ordering::Relaxed);
}
})
.ok()?;
Some(Self {
tx: Some(tx),
latest_probability,
latest_processed_seq,
alive,
handle: Some(handle),
})
}
/// Best-effort enqueue of a 10 ms frame for background inference.
pub fn try_send(&self, seq: u64, frame: &[f32; INPUT_FRAME_10MS]) -> bool {
let Some(tx) = &self.tx else {
return false;
};
tx.try_send(TenFrameMessage { seq, frame: *frame }).is_ok()
}
/// Latest probability published by the background worker.
pub fn latest_probability(&self) -> f32 {
f32::from_bits(
self.latest_probability
.load(std::sync::atomic::Ordering::Relaxed),
)
}
/// True when the worker is too far behind to trust its output.
pub fn is_stale(&self, capture_seq: u64) -> bool {
let latest = self
.latest_processed_seq
.load(std::sync::atomic::Ordering::Relaxed);
latest == u64::MAX || capture_seq.saturating_sub(latest) > TEN_MAX_STALE_FRAMES
}
}
impl Drop for TenOnnxVadWorker {
fn drop(&mut self) {
self.alive
.store(false, std::sync::atomic::Ordering::Relaxed);
drop(self.tx.take());
if let Some(h) = self.handle.take() {
let _ = h.join();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn mel_filter_bank_has_expected_shape() {
let filters = build_mel_filters();
assert_eq!(filters.len(), MEL_BANDS);
assert!(filters.iter().all(|f| f.iter().any(|&v| v > 0.0)));
}
#[test]
fn preprocessing_produces_finite_features() {
let filters = build_mel_filters();
let mut planner = FftPlanner::<f32>::new();
let fft = planner.plan_fft_forward(FFT_SIZE);
let mut fft_buffer = vec![Complex32::ZERO; FFT_SIZE];
let frame = vec![0.0_f32; WINDOW_16K];
let feature = compute_feature(&filters, fft.as_ref(), &mut fft_buffer, &frame);
assert!(feature.iter().all(|v| v.is_finite()));
}
}