From fba96fc9a9d19b6cbd8cfc927097ef59e542eef9 Mon Sep 17 00:00:00 2001 From: vedanth Date: Mon, 26 Oct 2026 07:09:01 +0000 Subject: [PATCH] perf(gpui): accelerate transcription audio preparation --- apps/desktop-gpui/src/transcription.rs | 75 +++++++++++++---------- crates/audio/src/lib.rs | 2 + crates/audio/src/normalization.rs | 82 ++++++++++++++++++++++++++ 3 files changed, 127 insertions(+), 32 deletions(-) create mode 100644 crates/audio/src/normalization.rs diff --git a/apps/desktop-gpui/src/transcription.rs b/apps/desktop-gpui/src/transcription.rs index 92fa942a9cc..56d2747e88e 100644 --- a/apps/desktop-gpui/src/transcription.rs +++ b/apps/desktop-gpui/src/transcription.rs @@ -1112,35 +1112,7 @@ fn append_resampled_frame( /// `normalize_audio_for_transcription` (`captions.rs:885-915`). fn normalize_audio_for_transcription(samples: &mut [f32]) -> f32 { - if samples.is_empty() { - return 1.0; - } - - let peak = samples - .iter() - .fold(0.0_f32, |max, sample| max.max(sample.abs())); - if peak <= f32::EPSILON { - return 1.0; - } - - let rms = - (samples.iter().map(|sample| sample * sample).sum::() / samples.len() as f32).sqrt(); - if rms <= f32::EPSILON { - return 1.0; - } - - let target_rms = 0.08_f32; - let desired_gain = (target_rms / rms).clamp(1.0, 8.0); - let peak_limited_gain = 0.98 / peak; - let gain = desired_gain.min(peak_limited_gain); - - if (gain - 1.0).abs() > 0.01 { - for sample in samples { - *sample = (*sample * gain).clamp(-0.98, 0.98); - } - } - - gain + cap_audio::normalize_for_transcription(samples) } /// The recording-directory arm of `extract_audio_from_video` @@ -1262,6 +1234,14 @@ fn extract_project_audio(project_path: &Path, output_path: &Path) -> Result<(), /// The WAV encode half of the recording-directory arm /// (`captions.rs:372-515`): 48kHz f32 mono in, 16kHz s16le mono WAV out. fn write_wav_16k_mono(mixed_samples: &[f32], output_path: &Path) -> Result<(), String> { + write_wav_16k_mono_chunked(mixed_samples, output_path, 16_384) +} + +fn write_wav_16k_mono_chunked( + mixed_samples: &[f32], + output_path: &Path, + default_pcm_chunk: usize, +) -> Result<(), String> { let mut output = avformat::output(&output_path).map_err(|e| format!("Failed to create output file: {e}"))?; @@ -1302,7 +1282,11 @@ fn write_wav_16k_mono(mixed_samples: &[f32], output_path: &Path) -> Result<(), S .map_err(|e| format!("Failed to create resampler: {e}"))?; let frame_size = encoder.frame_size() as usize; - let frame_size = if frame_size == 0 { 1024 } else { frame_size }; + let frame_size = if frame_size == 0 { + default_pcm_chunk + } else { + frame_size + }; for chunk in mixed_samples.chunks(frame_size) { let mut input_frame = ffmpeg::frame::Audio::new( @@ -1857,8 +1841,8 @@ fn extract_audio_from_video(video_path: &Path, output_path: &Path) -> Result<(), // --------------------------------------------------------------------------- use cap_project::caption_words::{ - caption_text_from_words, caption_word_chunks, normalize_caption_words, - MAX_CAPTION_WORD_DURATION, + MAX_CAPTION_WORD_DURATION, caption_text_from_words, caption_word_chunks, + normalize_caption_words, }; // --------------------------------------------------------------------------- @@ -2770,4 +2754,31 @@ mod tests { ); assert_eq!(caption_generation_error_message("boom"), "boom"); } + #[test] + fn transcription_pcm_chunking_preserves_legacy_wav_bytes() { + ffmpeg::init().unwrap(); + let directory = tempfile::tempdir().unwrap(); + for len in [ + 0, 1, 2, 3, 15, 16, 17, 1023, 1024, 1025, 16383, 16384, 16385, 32769, + ] { + let mut state = 0x1947_cdef_u32; + let samples: Vec = (0..len) + .map(|_| { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + state as i32 as f32 / i32::MAX as f32 * 0.8 + }) + .collect(); + let reference = directory.path().join(format!("reference-{len}.wav")); + let actual = directory.path().join(format!("actual-{len}.wav")); + write_wav_16k_mono_chunked(&samples, &reference, 1024).unwrap(); + write_wav_16k_mono(&samples, &actual).unwrap(); + assert_eq!( + std::fs::read(&actual).unwrap(), + std::fs::read(&reference).unwrap(), + "{len} samples" + ); + } + } } diff --git a/crates/audio/src/lib.rs b/crates/audio/src/lib.rs index ef0e71c45e3..c0ab04f5b29 100644 --- a/crates/audio/src/lib.rs +++ b/crates/audio/src/lib.rs @@ -1,6 +1,7 @@ mod audio_data; mod calibration_store; mod latency; +mod normalization; mod progressive; mod renderer; mod streaming; @@ -12,6 +13,7 @@ mod voice_level; pub use audio_data::*; pub use calibration_store::*; pub use latency::*; +pub use normalization::normalize_for_transcription; pub use progressive::*; pub use renderer::*; pub use streaming::*; diff --git a/crates/audio/src/normalization.rs b/crates/audio/src/normalization.rs new file mode 100644 index 00000000000..cb1cf35ea7f --- /dev/null +++ b/crates/audio/src/normalization.rs @@ -0,0 +1,82 @@ +pub fn normalize_for_transcription(samples: &mut [f32]) -> f32 { + if samples.is_empty() { + return 1.0; + } + + let peak = samples + .iter() + .fold(0.0_f32, |max, sample| max.max(sample.abs())); + if peak <= f32::EPSILON { + return 1.0; + } + + let rms = + (samples.iter().map(|sample| sample * sample).sum::() / samples.len() as f32).sqrt(); + if rms <= f32::EPSILON { + return 1.0; + } + + let target_rms = 0.08_f32; + let desired_gain = (target_rms / rms).clamp(1.0, 8.0); + let peak_limited_gain = 0.98 / peak; + let gain = desired_gain.min(peak_limited_gain); + + if (gain - 1.0).abs() > 0.01 { + for sample in samples { + *sample = (*sample * gain).clamp(-0.98, 0.98); + } + } + + gain +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn silence_and_nominal_audio_remain_bit_identical() { + for mut samples in [vec![], vec![0.0_f32, -0.0], vec![0.1, -0.1]] { + let before: Vec<_> = samples.iter().map(|sample| sample.to_bits()).collect(); + assert_eq!(normalize_for_transcription(&mut samples), 1.0); + assert_eq!( + samples + .iter() + .map(|sample| sample.to_bits()) + .collect::>(), + before + ); + } + } + + #[test] + fn quiet_speech_gain_stops_at_eight() { + let mut samples = [0.01, -0.01]; + assert_eq!(normalize_for_transcription(&mut samples), 8.0); + assert_eq!(samples, [0.08, -0.08]); + } + + #[test] + fn peaks_are_limited_even_when_rms_gain_would_be_one() { + let mut samples = [2.0, -2.0]; + assert_eq!(normalize_for_transcription(&mut samples), 0.49); + assert_eq!(samples, [0.98, -0.98]); + } + + #[test] + fn gain_within_one_percent_does_not_mutate_samples() { + let mut samples = [0.985, -0.985]; + let before = samples.map(f32::to_bits); + let gain = normalize_for_transcription(&mut samples); + assert_eq!(gain.to_bits(), (0.98_f32 / 0.985).to_bits()); + assert_eq!(samples.map(f32::to_bits), before); + } + + #[test] + fn silent_nan_and_subnormal_inputs_keep_their_bits() { + let mut samples = [f32::from_bits(0x7fc0_1234), f32::from_bits(1), -0.0]; + let before = samples.map(f32::to_bits); + assert_eq!(normalize_for_transcription(&mut samples), 1.0); + assert_eq!(samples.map(f32::to_bits), before); + } +}