Add audio neural layers and model architectures for ClawSample integration
New nn layers: - ConvTranspose1d with stride, padding, output_padding (9 tests) - LSTM/BiLSTM with multi-layer support and hidden state (10 tests) Audio source separation: - Demucs ONNX inference with segmented overlap-add processing - Native HtDemucs architecture (encoder/decoder with BiLSTM bottleneck) - StemType enum: vocals, drums, bass, other, piano, guitar Audio generation: - Stable Audio Open ONNX inference scaffold - GenerationParams (prompt, duration, steps, cfg_scale, seed) ONNX export scripts: - export_demucs_onnx.py — Demucs v4 to ONNX with segment chunking - export_stable_audio_onnx.py — Stable Audio Open components - export_mert_onnx.py — MERT music understanding transformer Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -0,0 +1,442 @@
|
||||
//! ONNX-backed Demucs source separation model.
|
||||
//!
|
||||
//! Loads a pre-exported Demucs ONNX model (htdemucs 4-stem or htdemucs_6s 6-stem)
|
||||
//! and performs segmented inference with overlap-add for any-length audio input.
|
||||
//!
|
||||
//! # Usage
|
||||
//!
|
||||
//! ```rust,ignore
|
||||
//! let config = DemucsConfig::four_stem("path/to/htdemucs.onnx");
|
||||
//! let mut model = DemucsModel::load(config)?;
|
||||
//! let stems = model.separate(&mix_tensor)?;
|
||||
//! // stems: [StemOutput { stem_type: Drums, waveform }, ...]
|
||||
//! ```
|
||||
|
||||
use rtx_onnx::session::{OnnxSession, OnnxSessionConfig};
|
||||
use rtx_tensor::{Device, Tensor};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use tracing::{debug, info};
|
||||
|
||||
// StemType and StemOutput are defined in the parent module (mod.rs)
|
||||
use super::{StemOutput, StemType};
|
||||
|
||||
/// Configuration for the Demucs ONNX model.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DemucsConfig {
|
||||
/// Path to the ONNX model file.
|
||||
pub model_path: PathBuf,
|
||||
/// Expected sample rate of input audio (Demucs expects 44100 Hz).
|
||||
pub sample_rate: u32,
|
||||
/// Number of stems the model outputs (4 for htdemucs, 6 for htdemucs_6s).
|
||||
pub num_stems: usize,
|
||||
/// Number of audio channels (2 for stereo).
|
||||
pub channels: usize,
|
||||
/// Segment length in samples for chunked inference.
|
||||
/// Default: 441000 (10 seconds at 44.1 kHz).
|
||||
pub segment_length: usize,
|
||||
/// Overlap ratio between consecutive segments (0.0–1.0).
|
||||
/// Default: 0.25 (25% overlap).
|
||||
pub overlap: f32,
|
||||
/// ONNX session configuration (execution provider, threads, etc.).
|
||||
#[serde(skip)]
|
||||
pub onnx_config: Option<OnnxSessionConfig>,
|
||||
}
|
||||
|
||||
impl DemucsConfig {
|
||||
/// Configuration for 4-stem htdemucs model.
|
||||
pub fn four_stem(model_path: impl Into<PathBuf>) -> Self {
|
||||
Self {
|
||||
model_path: model_path.into(),
|
||||
sample_rate: 44100,
|
||||
num_stems: 4,
|
||||
channels: 2,
|
||||
segment_length: 44100 * 10, // 10 seconds
|
||||
overlap: 0.25,
|
||||
onnx_config: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for 6-stem htdemucs model (adds piano + guitar).
|
||||
pub fn six_stem(model_path: impl Into<PathBuf>) -> Self {
|
||||
Self {
|
||||
model_path: model_path.into(),
|
||||
sample_rate: 44100,
|
||||
num_stems: 6,
|
||||
channels: 2,
|
||||
segment_length: 44100 * 10,
|
||||
overlap: 0.25,
|
||||
onnx_config: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set a custom ONNX session config (execution provider, etc.).
|
||||
pub fn with_onnx_config(mut self, config: OnnxSessionConfig) -> Self {
|
||||
self.onnx_config = Some(config);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// ONNX-backed Demucs source separation model.
|
||||
///
|
||||
/// Wraps an `OnnxSession` and handles the segmented inference pipeline:
|
||||
/// normalization → chunking → inference → overlap-add → denormalization.
|
||||
pub struct DemucsModel {
|
||||
session: OnnxSession,
|
||||
config: DemucsConfig,
|
||||
}
|
||||
|
||||
impl DemucsModel {
|
||||
/// Load a Demucs ONNX model from disk.
|
||||
pub fn load(config: DemucsConfig) -> Result<Self, DemucsError> {
|
||||
let onnx_config = config.onnx_config.clone().unwrap_or_default();
|
||||
|
||||
info!(
|
||||
model = %config.model_path.display(),
|
||||
stems = config.num_stems,
|
||||
segment = config.segment_length,
|
||||
"Loading Demucs ONNX model"
|
||||
);
|
||||
|
||||
let session = OnnxSession::from_file(&config.model_path, onnx_config)
|
||||
.map_err(|e| DemucsError::ModelLoad(e.to_string()))?;
|
||||
|
||||
Ok(Self { session, config })
|
||||
}
|
||||
|
||||
/// Separate a stereo audio mix into individual stems.
|
||||
///
|
||||
/// Input: interleaved stereo samples at 44.1 kHz (f32).
|
||||
/// Output: one `StemOutput` per stem (drums, bass, vocals, other, +piano/guitar for 6-stem).
|
||||
///
|
||||
/// For audio longer than `segment_length`, the input is split into overlapping
|
||||
/// segments, each processed independently, then recombined via overlap-add with
|
||||
/// a triangular cross-fade window.
|
||||
pub fn separate(&mut self, waveform: &[f32], channels: usize) -> Result<Vec<StemOutput>, DemucsError> {
|
||||
if waveform.is_empty() {
|
||||
return Err(DemucsError::EmptyInput);
|
||||
}
|
||||
|
||||
let total_samples = waveform.len();
|
||||
let total_frames = total_samples / channels;
|
||||
|
||||
debug!(
|
||||
frames = total_frames,
|
||||
channels = channels,
|
||||
"Starting source separation"
|
||||
);
|
||||
|
||||
// 1. Normalize to unit variance
|
||||
let (normalized, scale) = normalize(waveform);
|
||||
|
||||
// 2. De-interleave to channel-first layout: [channels, frames]
|
||||
let channel_first = deinterleave(&normalized, channels);
|
||||
|
||||
// 3. Segment into overlapping chunks
|
||||
let segments = self.segment(&channel_first, channels, total_frames);
|
||||
|
||||
// 4. Run inference on each segment
|
||||
let stem_types = StemType::stems_for(self.config.num_stems);
|
||||
let mut stem_accumulators: Vec<Vec<f32>> = vec![vec![0.0; total_frames * channels]; self.config.num_stems];
|
||||
let mut weight_accumulator: Vec<f32> = vec![0.0; total_frames];
|
||||
|
||||
for (start_frame, chunk) in &segments {
|
||||
let chunk_frames = chunk.len() / channels;
|
||||
|
||||
// Pad chunk to segment_length if needed
|
||||
let padded = self.pad_to_segment(chunk, channels);
|
||||
|
||||
// Run ONNX inference: input [1, channels, segment_length] → output [1, num_stems, channels, segment_length]
|
||||
let output = self.run_inference(&padded, channels)?;
|
||||
|
||||
// Build triangular window for overlap-add
|
||||
let window = triangular_window(chunk_frames);
|
||||
|
||||
// Accumulate each stem
|
||||
for stem_idx in 0..self.config.num_stems {
|
||||
let stem_offset = stem_idx * channels * self.config.segment_length;
|
||||
for frame in 0..chunk_frames {
|
||||
let w = window[frame];
|
||||
for ch in 0..channels {
|
||||
let src_idx = stem_offset + ch * self.config.segment_length + frame;
|
||||
let dst_idx = (start_frame + frame) * channels + ch;
|
||||
if src_idx < output.len() && dst_idx < stem_accumulators[stem_idx].len() {
|
||||
stem_accumulators[stem_idx][dst_idx] += output[src_idx] * w;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Accumulate weights
|
||||
for frame in 0..chunk_frames {
|
||||
let dst = start_frame + frame;
|
||||
if dst < weight_accumulator.len() {
|
||||
weight_accumulator[dst] += window[frame];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Normalize by accumulated weights and denormalize
|
||||
let stems: Vec<StemOutput> = stem_types
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(idx, stem_type)| {
|
||||
let mut samples = stem_accumulators[idx].clone();
|
||||
for frame in 0..total_frames {
|
||||
let w = weight_accumulator[frame].max(1e-8);
|
||||
for ch in 0..channels {
|
||||
let i = frame * channels + ch;
|
||||
samples[i] = samples[i] / w * scale;
|
||||
}
|
||||
}
|
||||
StemOutput {
|
||||
stem_type,
|
||||
samples,
|
||||
channels,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
info!(
|
||||
stems = stems.len(),
|
||||
frames = total_frames,
|
||||
"Source separation complete"
|
||||
);
|
||||
|
||||
Ok(stems)
|
||||
}
|
||||
|
||||
/// Segment audio into overlapping chunks.
|
||||
fn segment(&self, channel_first: &[f32], channels: usize, total_frames: usize) -> Vec<(usize, Vec<f32>)> {
|
||||
let seg_len = self.config.segment_length;
|
||||
let hop = ((1.0 - self.config.overlap) * seg_len as f32) as usize;
|
||||
let hop = hop.max(1);
|
||||
|
||||
let mut segments = Vec::new();
|
||||
let mut start = 0;
|
||||
|
||||
while start < total_frames {
|
||||
let end = (start + seg_len).min(total_frames);
|
||||
let chunk_frames = end - start;
|
||||
|
||||
// Interleave back for this chunk
|
||||
let mut chunk = vec![0.0f32; chunk_frames * channels];
|
||||
for frame in 0..chunk_frames {
|
||||
for ch in 0..channels {
|
||||
chunk[frame * channels + ch] = channel_first[ch * total_frames + start + frame];
|
||||
}
|
||||
}
|
||||
|
||||
segments.push((start, chunk));
|
||||
start += hop;
|
||||
}
|
||||
|
||||
segments
|
||||
}
|
||||
|
||||
/// Pad a chunk to the model's expected segment length.
|
||||
fn pad_to_segment(&self, chunk: &[f32], channels: usize) -> Vec<f32> {
|
||||
let seg_samples = self.config.segment_length * channels;
|
||||
if chunk.len() >= seg_samples {
|
||||
return chunk[..seg_samples].to_vec();
|
||||
}
|
||||
let mut padded = chunk.to_vec();
|
||||
padded.resize(seg_samples, 0.0);
|
||||
padded
|
||||
}
|
||||
|
||||
/// Run a single segment through the ONNX model.
|
||||
///
|
||||
/// Input shape: [1, channels, segment_length]
|
||||
/// Output shape: [1, num_stems, channels, segment_length]
|
||||
fn run_inference(&mut self, segment: &[f32], channels: usize) -> Result<Vec<f32>, DemucsError> {
|
||||
let seg_len = self.config.segment_length;
|
||||
|
||||
// Convert interleaved to channel-first [1, channels, seg_len]
|
||||
let mut input_data = vec![0.0f32; channels * seg_len];
|
||||
for frame in 0..seg_len {
|
||||
for ch in 0..channels {
|
||||
let src = frame * channels + ch;
|
||||
let dst = ch * seg_len + frame;
|
||||
if src < segment.len() {
|
||||
input_data[dst] = segment[src];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let input_shape = vec![1, channels, seg_len];
|
||||
let input_tensor = Tensor::from_vec(input_data, &input_shape, &Device::Cpu)
|
||||
.map_err(|e| DemucsError::Inference(e.to_string()))?;
|
||||
|
||||
let mut inputs = HashMap::new();
|
||||
inputs.insert("mix".to_string(), &input_tensor);
|
||||
|
||||
let outputs = self.session.run(inputs)
|
||||
.map_err(|e| DemucsError::Inference(e.to_string()))?;
|
||||
|
||||
// Extract the output tensor (first output, whatever its name)
|
||||
let output_tensor = outputs.into_values().next()
|
||||
.ok_or_else(|| DemucsError::Inference("no output tensor from ONNX model".into()))?;
|
||||
|
||||
let output_data = output_tensor.to_vec_f32()
|
||||
.map_err(|e| DemucsError::Inference(e.to_string()))?;
|
||||
|
||||
Ok(output_data)
|
||||
}
|
||||
}
|
||||
|
||||
/// Normalize audio to unit variance, returning (normalized, scale_factor).
|
||||
fn normalize(samples: &[f32]) -> (Vec<f32>, f32) {
|
||||
if samples.is_empty() {
|
||||
return (vec![], 1.0);
|
||||
}
|
||||
|
||||
let mean_sq: f64 = samples.iter().map(|&s| (s as f64) * (s as f64)).sum::<f64>() / samples.len() as f64;
|
||||
let rms = mean_sq.sqrt() as f32;
|
||||
let scale = rms.max(1e-8);
|
||||
|
||||
let normalized: Vec<f32> = samples.iter().map(|&s| s / scale).collect();
|
||||
(normalized, scale)
|
||||
}
|
||||
|
||||
/// De-interleave audio from [frame0_L, frame0_R, frame1_L, ...] to channel-first [L0, L1, ..., R0, R1, ...].
|
||||
fn deinterleave(samples: &[f32], channels: usize) -> Vec<f32> {
|
||||
let frames = samples.len() / channels;
|
||||
let mut out = vec![0.0f32; samples.len()];
|
||||
for frame in 0..frames {
|
||||
for ch in 0..channels {
|
||||
out[ch * frames + frame] = samples[frame * channels + ch];
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Triangular window for overlap-add (linearly ramps up then down).
|
||||
fn triangular_window(length: usize) -> Vec<f32> {
|
||||
if length <= 1 {
|
||||
return vec![1.0; length];
|
||||
}
|
||||
let half = length as f32 / 2.0;
|
||||
(0..length)
|
||||
.map(|i| {
|
||||
let t = i as f32;
|
||||
if t < half {
|
||||
t / half
|
||||
} else {
|
||||
(length as f32 - t) / half
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Errors from Demucs source separation.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum DemucsError {
|
||||
#[error("failed to load Demucs model: {0}")]
|
||||
ModelLoad(String),
|
||||
#[error("inference error: {0}")]
|
||||
Inference(String),
|
||||
#[error("empty input audio")]
|
||||
EmptyInput,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn stem_type_names() {
|
||||
assert_eq!(StemType::Drums.name(), "drums");
|
||||
assert_eq!(StemType::Vocals.name(), "vocals");
|
||||
assert_eq!(StemType::Piano.name(), "piano");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn four_stem_types() {
|
||||
let stems = StemType::stems_for(4);
|
||||
assert_eq!(stems.len(), 4);
|
||||
assert!(stems.contains(&StemType::Drums));
|
||||
assert!(stems.contains(&StemType::Bass));
|
||||
assert!(stems.contains(&StemType::Vocals));
|
||||
assert!(stems.contains(&StemType::Other));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn six_stem_types() {
|
||||
let stems = StemType::stems_for(6);
|
||||
assert_eq!(stems.len(), 6);
|
||||
assert!(stems.contains(&StemType::Piano));
|
||||
assert!(stems.contains(&StemType::Guitar));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_unit_variance() {
|
||||
let input = vec![0.5, -0.5, 0.3, -0.3];
|
||||
let (normed, scale) = normalize(&input);
|
||||
assert!(scale > 0.0);
|
||||
// After normalization, RMS should be ~1.0
|
||||
let rms: f32 = (normed.iter().map(|&s| s * s).sum::<f32>() / normed.len() as f32).sqrt();
|
||||
assert!((rms - 1.0).abs() < 0.01, "rms={rms}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_silence() {
|
||||
let input = vec![0.0; 100];
|
||||
let (normed, scale) = normalize(&input);
|
||||
// Scale clamped to 1e-8, all values stay ~0
|
||||
assert!(scale > 0.0);
|
||||
assert!(normed.iter().all(|&s| s.abs() < 0.01));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deinterleave_stereo() {
|
||||
// Interleaved: [L0, R0, L1, R1, L2, R2]
|
||||
let interleaved = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
|
||||
let channel_first = deinterleave(&interleaved, 2);
|
||||
// Channel-first: [L0, L1, L2, R0, R1, R2]
|
||||
assert_eq!(channel_first, vec![1.0, 3.0, 5.0, 2.0, 4.0, 6.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn triangular_window_shape() {
|
||||
let w = triangular_window(100);
|
||||
assert_eq!(w.len(), 100);
|
||||
// Starts near 0, peaks in middle, ends near 0
|
||||
assert!(w[0] < 0.02);
|
||||
assert!(w[50] > 0.9);
|
||||
assert!(w[99] < 0.02);
|
||||
// Symmetric
|
||||
for i in 0..50 {
|
||||
assert!((w[i] - w[99 - i]).abs() < 0.02, "asymmetry at {i}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn triangular_window_single() {
|
||||
let w = triangular_window(1);
|
||||
assert_eq!(w, vec![1.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_four_stem() {
|
||||
let config = DemucsConfig::four_stem("/tmp/model.onnx");
|
||||
assert_eq!(config.num_stems, 4);
|
||||
assert_eq!(config.sample_rate, 44100);
|
||||
assert_eq!(config.segment_length, 441000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_six_stem() {
|
||||
let config = DemucsConfig::six_stem("/tmp/model.onnx");
|
||||
assert_eq!(config.num_stems, 6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stem_type_serializes() {
|
||||
let json = serde_json::to_string(&StemType::Vocals).unwrap();
|
||||
assert_eq!(json, "\"vocals\"");
|
||||
let back: StemType = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(back, StemType::Vocals);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user