rustytorch/crates/models/rtx-csm/src/audio_io.rs

//! Audio I/O: load arbitrary-format audio → 24 kHz mono f32.
//!
//! Mimi consumes 24 kHz mono. We read via symphonia (broad codec coverage) and
//! resample with rubato (high-quality sinc interpolation). For pure WAV writes
//! we use hound directly.

use crate::error::{CsmError, Result};
use rubato::{
    Resampler, SincFixedIn, SincInterpolationParameters, SincInterpolationType, WindowFunction,
};
use std::fs::File;
use std::path::Path;
use symphonia::core::audio::{AudioBufferRef, Signal};
use symphonia::core::codecs::DecoderOptions;
use symphonia::core::formats::FormatOptions;
use symphonia::core::io::MediaSourceStream;
use symphonia::core::meta::MetadataOptions;
use symphonia::core::probe::Hint;

pub const TARGET_SAMPLE_RATE: u32 = 24_000;

/// Load any symphonia-supported audio file as mono f32 at 24 kHz.
pub fn load_mono_24k<P: AsRef<Path>>(path: P) -> Result<Vec<f32>> {
    let file = File::open(path.as_ref())?;
    let mss = MediaSourceStream::new(Box::new(file), Default::default());

    let hint = Hint::new();
    let probed = symphonia::default::get_probe().format(
        &hint,
        mss,
        &FormatOptions::default(),
        &MetadataOptions::default(),
    )?;
    let mut format = probed.format;

    let track = format
        .default_track()
        .ok_or_else(|| CsmError::Config("no default audio track".into()))?;
    let track_id = track.id;
    let codec_params = track.codec_params.clone();

    let mut decoder =
        symphonia::default::get_codecs().make(&codec_params, &DecoderOptions::default())?;

    let src_rate = codec_params
        .sample_rate
        .ok_or_else(|| CsmError::Config("missing sample rate".into()))?;
    let channels = codec_params
        .channels
        .ok_or_else(|| CsmError::Config("missing channel layout".into()))?
        .count();

    let mut mono: Vec<f32> = Vec::new();

    loop {
        let packet = match format.next_packet() {
            Ok(p) => p,
            Err(symphonia::core::errors::Error::IoError(e))
                if e.kind() == std::io::ErrorKind::UnexpectedEof =>
            {
                break;
            }
            Err(e) => return Err(e.into()),
        };
        if packet.track_id() != track_id {
            continue;
        }

        let decoded = decoder.decode(&packet)?;
        append_mono_f32(&decoded, channels, &mut mono);
    }

    if src_rate == TARGET_SAMPLE_RATE {
        Ok(mono)
    } else {
        resample_to_24k(&mono, src_rate)
    }
}

fn append_mono_f32(buf: &AudioBufferRef<'_>, channels: usize, out: &mut Vec<f32>) {
    macro_rules! mix {
        ($buf:expr, $convert:expr) => {{
            let frames = $buf.frames();
            for f in 0..frames {
                let mut acc = 0.0f32;
                for c in 0..channels {
                    acc += $convert($buf.chan(c)[f]);
                }
                out.push(acc / channels as f32);
            }
        }};
    }
    match buf {
        AudioBufferRef::F32(b) => mix!(b, |x: f32| x),
        AudioBufferRef::F64(b) => mix!(b, |x: f64| x as f32),
        AudioBufferRef::S16(b) => mix!(b, |x: i16| x as f32 / i16::MAX as f32),
        AudioBufferRef::S32(b) => mix!(b, |x: i32| x as f32 / i32::MAX as f32),
        AudioBufferRef::U8(b) => mix!(b, |x: u8| (x as f32 - 128.0) / 128.0),
        AudioBufferRef::U16(b) => mix!(b, |x: u16| (x as f32 - 32768.0) / 32768.0),
        AudioBufferRef::U32(b) => mix!(b, |x: u32| (x as f32 - 2_147_483_648.0) / 2_147_483_648.0),
        AudioBufferRef::S8(b) => mix!(b, |x: i8| x as f32 / i8::MAX as f32),
        AudioBufferRef::S24(b) => mix!(b, |x: symphonia::core::sample::i24| x.0 as f32
            / 8_388_607.0),
        AudioBufferRef::U24(b) => mix!(b, |x: symphonia::core::sample::u24| (x.0 as f32
            - 8_388_608.0)
            / 8_388_608.0),
    }
}

fn resample_to_24k(input: &[f32], src_rate: u32) -> Result<Vec<f32>> {
    let params = SincInterpolationParameters {
        sinc_len: 256,
        f_cutoff: 0.95,
        interpolation: SincInterpolationType::Linear,
        oversampling_factor: 256,
        window: WindowFunction::BlackmanHarris2,
    };
    let chunk = 1024usize;
    let mut resampler = SincFixedIn::<f32>::new(
        TARGET_SAMPLE_RATE as f64 / src_rate as f64,
        2.0,
        params,
        chunk,
        1,
    )
    .map_err(|e| CsmError::Rubato(e.to_string()))?;

    let mut out = Vec::with_capacity(
        ((input.len() as f64 * TARGET_SAMPLE_RATE as f64 / src_rate as f64).ceil()) as usize
            + chunk,
    );

    let mut pos = 0usize;
    while pos + chunk <= input.len() {
        let frame_in = vec![input[pos..pos + chunk].to_vec()];
        let frame_out = resampler
            .process(&frame_in, None)
            .map_err(|e| CsmError::Rubato(e.to_string()))?;
        out.extend_from_slice(&frame_out[0]);
        pos += chunk;
    }
    if pos < input.len() {
        let mut tail = input[pos..].to_vec();
        tail.resize(chunk, 0.0);
        let frame_out = resampler
            .process(&[tail], None)
            .map_err(|e| CsmError::Rubato(e.to_string()))?;
        let kept = ((input.len() - pos) as f64 * TARGET_SAMPLE_RATE as f64 / src_rate as f64)
            .round() as usize;
        out.extend_from_slice(&frame_out[0][..kept.min(frame_out[0].len())]);
    }
    Ok(out)
}

/// Load any symphonia-supported audio file as mono f32 at an arbitrary
/// `target_rate`. Identical pipeline to `load_mono_24k` but accepts any
/// sample rate (e.g. 16 kHz for AudioSeal).
pub fn load_mono_at_rate<P: AsRef<Path>>(path: P, target_rate: u32) -> Result<Vec<f32>> {
    if target_rate == TARGET_SAMPLE_RATE {
        return load_mono_24k(path);
    }
    let raw = load_mono_24k(path)?;
    resample(&raw, TARGET_SAMPLE_RATE, target_rate)
}

/// Generic sinc resample (rubato). Public so callers can resample buffers
/// they already have in memory without hitting disk.
pub fn resample(input: &[f32], src_rate: u32, dst_rate: u32) -> Result<Vec<f32>> {
    if src_rate == dst_rate {
        return Ok(input.to_vec());
    }
    let params = SincInterpolationParameters {
        sinc_len: 256,
        f_cutoff: 0.95,
        interpolation: SincInterpolationType::Linear,
        oversampling_factor: 256,
        window: WindowFunction::BlackmanHarris2,
    };
    let chunk = 1024usize;
    let mut resampler =
        SincFixedIn::<f32>::new(dst_rate as f64 / src_rate as f64, 2.0, params, chunk, 1)
            .map_err(|e| CsmError::Rubato(e.to_string()))?;

    let mut out = Vec::with_capacity(
        ((input.len() as f64 * dst_rate as f64 / src_rate as f64).ceil()) as usize + chunk,
    );
    let mut pos = 0usize;
    while pos + chunk <= input.len() {
        let frame_in = vec![input[pos..pos + chunk].to_vec()];
        let frame_out = resampler
            .process(&frame_in, None)
            .map_err(|e| CsmError::Rubato(e.to_string()))?;
        out.extend_from_slice(&frame_out[0]);
        pos += chunk;
    }
    if pos < input.len() {
        let mut tail = input[pos..].to_vec();
        tail.resize(chunk, 0.0);
        let frame_out = resampler
            .process(&[tail], None)
            .map_err(|e| CsmError::Rubato(e.to_string()))?;
        let kept =
            ((input.len() - pos) as f64 * dst_rate as f64 / src_rate as f64).round() as usize;
        out.extend_from_slice(&frame_out[0][..kept.min(frame_out[0].len())]);
    }
    Ok(out)
}

/// Write a mono f32 slice as a 16-bit WAV at arbitrary sample rate.
pub fn write_wav_mono<P: AsRef<Path>>(path: P, samples: &[f32], sample_rate: u32) -> Result<()> {
    let spec = hound::WavSpec {
        channels: 1,
        sample_rate,
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };
    let mut writer = hound::WavWriter::create(path, spec)?;
    for &s in samples {
        let clipped = s.clamp(-1.0, 1.0);
        let v = (clipped * i16::MAX as f32) as i16;
        writer.write_sample(v)?;
    }
    writer.finalize()?;
    Ok(())
}

/// Write a mono f32 slice as a 24 kHz 16-bit WAV.
pub fn write_wav_24k_mono<P: AsRef<Path>>(path: P, samples: &[f32]) -> Result<()> {
    let spec = hound::WavSpec {
        channels: 1,
        sample_rate: TARGET_SAMPLE_RATE,
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };
    let mut writer = hound::WavWriter::create(path, spec)?;
    for &s in samples {
        let clipped = s.clamp(-1.0, 1.0);
        let v = (clipped * i16::MAX as f32) as i16;
        writer.write_sample(v)?;
    }
    writer.finalize()?;
    Ok(())
}