forked from Keelrleek/SoulStream
dj sets + lyrics: a mix that names its own tracks, and words pulled off the separated voice
This commit is contained in:
1 parent
4efe85b6bf
commit
1dd8573808
41 files changed
+9288
-647
No files matched your search
@@ -0,0 +1,499 @@
|
||||
//! Learned vocal separation for the lyric aligner — the "honest fix" the
|
||||
//! vocals.rs reducer comment promises: a real separator lifting the voice
|
||||
//! out of the band before CTC alignment, instead of band-pass pre-emphasis.
|
||||
//!
|
||||
//! Model: UVR-MDX-NET Voc_FT (KUIELab MDX-Net / TFC-TDF v2, the Ultimate
|
||||
//! Vocal Remover project's vocal fine-tune; MIT with credit). Chosen
|
||||
//! 2026-08-19 from the verified table in
|
||||
//! research/lyrics-vocal-separation/candidates.md: 63.7 MiB, ONNX opset 13 —
|
||||
//! under the inbox-WinML ceiling this machine PROVED by running the MMS
|
||||
//! aligner (opset 14) — static `input`/`output` [1,4,3072,256] f32, vocals
|
||||
//! SDR 9.64 on the MVSEP multisong benchmark (+1.5 dB over htdemucs at a
|
||||
//! fifth of the size; the RoFormer/Demucs ONNX exports are opset 17–23 and
|
||||
//! cannot load here as published).
|
||||
//!
|
||||
//! Posture is lyric_align.rs's exactly: WinML (no onnxruntime, no child
|
||||
//! processes), download only from the user's explicit Settings click, the
|
||||
//! session cached in a static — and NOTHING depends on this module. When
|
||||
//! the model is absent or anything fails, the aligner keeps its band-passed
|
||||
//! mixture input (the ctc-v4 behavior).
|
||||
//!
|
||||
//! The inference recipe is UVR's own non-overlap framing, kept bit-faithful
|
||||
//! so the model sees what it was trained on: stereo 44.1 kHz → per-chunk
|
||||
//! centered STFT (n_fft 7680, hop 1024, periodic Hann, reflect pad) → crop
|
||||
//! to 3072 bins → model (channels [L_re, L_im, R_re, R_im]) → vocal
|
||||
//! spectrogram → zero-pad the cropped bins back → WOLA iSTFT → trim
|
||||
//! n_fft/2 from each end → butt-join. Each chunk contributes 253,440 fresh
|
||||
//! samples (~5.75 s), so a four-minute track is ~42 model evaluations.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Instant;
|
||||
|
||||
use realfft::num_complex::Complex;
|
||||
use realfft::{ComplexToReal, RealFftPlanner, RealToComplex};
|
||||
use tauri::Emitter;
|
||||
use windows::{
|
||||
core::{Interface, HSTRING},
|
||||
AI::MachineLearning::{
|
||||
LearningModel, LearningModelBinding, LearningModelDevice, LearningModelDeviceKind,
|
||||
LearningModelSession, TensorFloat,
|
||||
},
|
||||
Win32::System::WinRT::{RoInitialize, RO_INIT_MULTITHREADED},
|
||||
};
|
||||
|
||||
const MODEL_URL: &str =
|
||||
"https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/UVR-MDX-NET-Voc_FT.onnx";
|
||||
/// Exact pinned size, shown in Settings and enforced after download — if the
|
||||
/// asset is ever republished with different bytes, refuse rather than trust.
|
||||
/// (Verified against the served bytes AND the file's own protobuf tail:
|
||||
/// opset 13, `input`/`output` [batch,4,3072,256]. Kim_Vocal_2 is 66,759,214
|
||||
/// — a first pin mistakenly used that sibling's size, and this guard caught
|
||||
/// it exactly as designed.)
|
||||
const MODEL_BYTES: u64 = 66_762_490;
|
||||
|
||||
// UVR MDX family constants (their model_data.json for this checkpoint:
|
||||
// dim_f 3072, dim_t 2^8, n_fft 7680; hop is the architecture's fixed 1024).
|
||||
const N_FFT: usize = 7_680;
|
||||
const HOP: usize = 1_024;
|
||||
const BINS: usize = N_FFT / 2 + 1; // 3841
|
||||
const DIM_F: usize = 3_072;
|
||||
const DIM_T: usize = 256;
|
||||
/// Samples one model window covers: 256 centered STFT frames at hop 1024.
|
||||
const CHUNK: usize = HOP * (DIM_T - 1); // 261,120
|
||||
/// STFT edge effects live inside half a window; UVR trims this from each
|
||||
/// chunk end and butt-joins the rest — no cross-fade, nothing to tune, and
|
||||
/// the kept region sits in full window-overlap coverage by construction.
|
||||
const TRIM: usize = N_FFT / 2; // 3840
|
||||
/// Fresh samples each chunk contributes (~5.75 s at 44.1 kHz).
|
||||
const GEN: usize = CHUNK - 2 * TRIM; // 253,440
|
||||
/// Flat f32 count of the model's input and output: [1, 4, 3072, 256].
|
||||
const TENSOR_LEN: usize = 4 * DIM_F * DIM_T;
|
||||
|
||||
struct Engine {
|
||||
session: LearningModelSession,
|
||||
/// Which device kind actually carried the smoke eval — for the log line.
|
||||
device: &'static str,
|
||||
}
|
||||
|
||||
static ENGINE: Mutex<Option<Engine>> = Mutex::new(None);
|
||||
|
||||
fn model_file(app: &tauri::AppHandle) -> Result<PathBuf, String> {
|
||||
Ok(crate::lyric_align::models_dir(app)?.join("uvr-mdx-voc-ft.onnx"))
|
||||
}
|
||||
|
||||
pub fn model_present(app: &tauri::AppHandle) -> bool {
|
||||
matches!(model_file(app), Ok(m) if m.is_file())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn vocal_sep_status(app: tauri::AppHandle) -> serde_json::Value {
|
||||
let size = model_file(&app)
|
||||
.ok()
|
||||
.and_then(|p| std::fs::metadata(p).ok())
|
||||
.map(|m| m.len())
|
||||
.unwrap_or(0);
|
||||
serde_json::json!({
|
||||
"present": model_present(&app),
|
||||
"bytes": size,
|
||||
"downloadBytes": MODEL_BYTES,
|
||||
})
|
||||
}
|
||||
|
||||
/// Download the separation model. Runs ONLY from the user's explicit
|
||||
/// Settings click — the consent-before-download charter rule. Streams to a
|
||||
/// temp file, verifies the exact pinned byte count, renames atomically,
|
||||
/// emits `lyrics:sep-model-progress` along the way.
|
||||
#[tauri::command]
|
||||
pub async fn vocal_sep_download(app: tauri::AppHandle) -> Result<(), String> {
|
||||
let client = reqwest::Client::builder()
|
||||
.user_agent("SoulStream/0.1 (vocal separation)")
|
||||
.build()
|
||||
.map_err(|e| e.to_string())?;
|
||||
let dest = model_file(&app)?;
|
||||
let tmp = dest.with_extension("onnx.downloading");
|
||||
let mut response = client
|
||||
.get(MODEL_URL)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| e.to_string())?
|
||||
.error_for_status()
|
||||
.map_err(|e| e.to_string())?;
|
||||
let total = response.content_length().unwrap_or(MODEL_BYTES);
|
||||
let mut file = std::fs::File::create(&tmp).map_err(|e| e.to_string())?;
|
||||
let mut received: u64 = 0;
|
||||
let mut last_emit: u64 = 0;
|
||||
use std::io::Write;
|
||||
while let Some(chunk) = response.chunk().await.map_err(|e| e.to_string())? {
|
||||
file.write_all(&chunk).map_err(|e| e.to_string())?;
|
||||
received += chunk.len() as u64;
|
||||
if received - last_emit > 2 * 1024 * 1024 {
|
||||
last_emit = received;
|
||||
let _ = app.emit(
|
||||
"lyrics:sep-model-progress",
|
||||
serde_json::json!({ "received": received, "total": total }),
|
||||
);
|
||||
}
|
||||
}
|
||||
drop(file);
|
||||
if received != MODEL_BYTES {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
return Err(format!(
|
||||
"download size mismatch: {received} bytes, expected {MODEL_BYTES}"
|
||||
));
|
||||
}
|
||||
std::fs::rename(&tmp, &dest).map_err(|e| e.to_string())?;
|
||||
let _ = app.emit(
|
||||
"lyrics:sep-model-progress",
|
||||
serde_json::json!({ "received": received, "total": total, "done": true }),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_engine(app: &tauri::AppHandle) -> Result<(), String> {
|
||||
let mut guard = ENGINE.lock().unwrap();
|
||||
if guard.is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
let model_path = model_file(app)?;
|
||||
if !model_path.is_file() {
|
||||
return Err("separation model not downloaded".into());
|
||||
}
|
||||
// Same WinRT posture as the aligner: initialize, tolerate an existing
|
||||
// apartment, never uninitialize someone else's.
|
||||
match unsafe { RoInitialize(RO_INIT_MULTITHREADED) } {
|
||||
Ok(()) => {}
|
||||
Err(e) if e.code().0 as u32 == 0x8001_0106 => {}
|
||||
Err(e) => return Err(format!("WinRT init failed: {e}")),
|
||||
}
|
||||
let model =
|
||||
LearningModel::LoadFromFilePath(&HSTRING::from(model_path.to_string_lossy().as_ref()))
|
||||
.map_err(|e| format!("could not load separation model: {e}"))?;
|
||||
// DirectX first (DirectML — the discrete GPU when there is one): a track
|
||||
// is ~42 conv-net evaluations and the GPU turns minutes into seconds.
|
||||
// The zeros smoke eval proves the device can actually RUN the graph —
|
||||
// "loads but cannot execute" is a real DML failure mode — then Cpu is
|
||||
// the quiet fallback, so this works everywhere the aligner does.
|
||||
for (kind, name) in [
|
||||
(LearningModelDeviceKind::DirectX, "DirectX"),
|
||||
(LearningModelDeviceKind::Cpu, "Cpu"),
|
||||
] {
|
||||
let Ok(device) = LearningModelDevice::Create(kind) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(session) = LearningModelSession::CreateFromModelOnDevice(&model, &device) else {
|
||||
continue;
|
||||
};
|
||||
let zeros = vec![0.0f32; TENSOR_LEN];
|
||||
if eval_chunk(&session, &zeros).is_ok() {
|
||||
*guard = Some(Engine {
|
||||
session,
|
||||
device: name,
|
||||
});
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
Err("no WinML device could run the separation model".into())
|
||||
}
|
||||
|
||||
fn eval_chunk(session: &LearningModelSession, flat: &[f32]) -> Result<Vec<f32>, String> {
|
||||
let input = TensorFloat::CreateFromShapeArrayAndDataArray(
|
||||
&[1, 4, DIM_F as i64, DIM_T as i64],
|
||||
flat,
|
||||
)
|
||||
.map_err(|e| format!("could not create input tensor: {e}"))?;
|
||||
let binding = LearningModelBinding::CreateFromSession(session)
|
||||
.map_err(|e| format!("could not create binding: {e}"))?;
|
||||
binding
|
||||
.Bind(&HSTRING::from("input"), &input)
|
||||
.map_err(|e| format!("could not bind input: {e}"))?;
|
||||
let result = session
|
||||
.Evaluate(&binding, &HSTRING::from("vocal-sep"))
|
||||
.map_err(|e| format!("inference failed: {e}"))?;
|
||||
if !result.Succeeded().unwrap_or(false) {
|
||||
return Err("inference did not succeed".into());
|
||||
}
|
||||
let out: TensorFloat = result
|
||||
.Outputs()
|
||||
.and_then(|o| o.Lookup(&HSTRING::from("output")))
|
||||
.map_err(|e| format!("output missing: {e}"))?
|
||||
.cast()
|
||||
.map_err(|e| format!("output wrong type: {e}"))?;
|
||||
let values: Vec<f32> = out
|
||||
.GetAsVectorView()
|
||||
.map_err(|e| format!("could not read output: {e}"))?
|
||||
.into_iter()
|
||||
.collect();
|
||||
if values.len() != TENSOR_LEN {
|
||||
return Err(format!("unexpected output length {}", values.len()));
|
||||
}
|
||||
Ok(values)
|
||||
}
|
||||
|
||||
// ---- STFT plumbing (pure DSP — no WinML below this line) -------------------
|
||||
|
||||
struct Stft {
|
||||
r2c: Arc<dyn RealToComplex<f32>>,
|
||||
c2r: Arc<dyn ComplexToReal<f32>>,
|
||||
/// Periodic Hann — torch.hann_window's default, what the model trained on.
|
||||
window: Vec<f32>,
|
||||
/// WOLA denominator Σ w² over the padded chunk layout. The 256 frames
|
||||
/// tile the padded buffer exactly, so it is identical for every chunk.
|
||||
envelope: Vec<f32>,
|
||||
}
|
||||
|
||||
impl Stft {
|
||||
fn new() -> Self {
|
||||
let mut planner = RealFftPlanner::<f32>::new();
|
||||
let r2c = planner.plan_fft_forward(N_FFT);
|
||||
let c2r = planner.plan_fft_inverse(N_FFT);
|
||||
let window: Vec<f32> = (0..N_FFT)
|
||||
.map(|n| 0.5 - 0.5 * (2.0 * std::f32::consts::PI * n as f32 / N_FFT as f32).cos())
|
||||
.collect();
|
||||
let mut envelope = vec![1e-8f32; CHUNK + 2 * TRIM];
|
||||
for t in 0..DIM_T {
|
||||
for (n, w) in window.iter().enumerate() {
|
||||
envelope[t * HOP + n] += w * w;
|
||||
}
|
||||
}
|
||||
Stft {
|
||||
r2c,
|
||||
c2r,
|
||||
window,
|
||||
envelope,
|
||||
}
|
||||
}
|
||||
|
||||
/// Centered STFT of one CHUNK-sized segment: reflect-pad n_fft/2 each
|
||||
/// side (torch.stft center=True), frames at the hop — frame-major
|
||||
/// [t * BINS + f].
|
||||
fn forward(&self, chunk: &[f32]) -> Result<Vec<Complex<f32>>, String> {
|
||||
debug_assert_eq!(chunk.len(), CHUNK);
|
||||
let mut padded = vec![0.0f32; CHUNK + 2 * TRIM];
|
||||
padded[TRIM..TRIM + CHUNK].copy_from_slice(chunk);
|
||||
for i in 1..=TRIM {
|
||||
padded[TRIM - i] = chunk[i];
|
||||
padded[TRIM + CHUNK - 1 + i] = chunk[CHUNK - 1 - i];
|
||||
}
|
||||
let mut spec = vec![Complex::default(); DIM_T * BINS];
|
||||
let mut frame = vec![0.0f32; N_FFT];
|
||||
for t in 0..DIM_T {
|
||||
for (n, w) in self.window.iter().enumerate() {
|
||||
frame[n] = padded[t * HOP + n] * w;
|
||||
}
|
||||
self.r2c
|
||||
.process(&mut frame, &mut spec[t * BINS..(t + 1) * BINS])
|
||||
.map_err(|e| format!("stft failed: {e}"))?;
|
||||
}
|
||||
Ok(spec)
|
||||
}
|
||||
|
||||
/// WOLA inverse of a frame-major spectrogram back to CHUNK samples with
|
||||
/// the centered pad trimmed off — torch.istft's window-envelope
|
||||
/// normalization. Consumes the spectrogram (the FFT scratches it).
|
||||
fn inverse(&self, spec: &mut [Complex<f32>]) -> Result<Vec<f32>, String> {
|
||||
let mut acc = vec![0.0f32; CHUNK + 2 * TRIM];
|
||||
let mut frame = vec![0.0f32; N_FFT];
|
||||
let scale = 1.0 / N_FFT as f32; // realfft's c2r is unnormalized
|
||||
for t in 0..DIM_T {
|
||||
let row = &mut spec[t * BINS..(t + 1) * BINS];
|
||||
// The model emits free re/im everywhere; a real inverse needs
|
||||
// pure-real DC and Nyquist. Nyquist sits in the cropped-zero
|
||||
// region already, DC gets clamped here.
|
||||
row[0].im = 0.0;
|
||||
row[BINS - 1].im = 0.0;
|
||||
self.c2r
|
||||
.process(row, &mut frame)
|
||||
.map_err(|e| format!("istft failed: {e}"))?;
|
||||
for (n, w) in self.window.iter().enumerate() {
|
||||
acc[t * HOP + n] += frame[n] * scale * w;
|
||||
}
|
||||
}
|
||||
Ok((TRIM..TRIM + CHUNK)
|
||||
.map(|i| acc[i] / self.envelope[i])
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
/// Pack two frame-major spectrograms into the model's flat input — channel
|
||||
/// order [L_re, L_im, R_re, R_im] (UVR's reshape/permute order), frequency
|
||||
/// cropped to DIM_F.
|
||||
fn pack_input(spec_l: &[Complex<f32>], spec_r: &[Complex<f32>]) -> Vec<f32> {
|
||||
let mut flat = vec![0.0f32; TENSOR_LEN];
|
||||
for (c, spec) in [spec_l, spec_r].into_iter().enumerate() {
|
||||
for f in 0..DIM_F {
|
||||
for t in 0..DIM_T {
|
||||
let z = spec[t * BINS + f];
|
||||
flat[((2 * c) * DIM_F + f) * DIM_T + t] = z.re;
|
||||
flat[((2 * c + 1) * DIM_F + f) * DIM_T + t] = z.im;
|
||||
}
|
||||
}
|
||||
}
|
||||
flat
|
||||
}
|
||||
|
||||
/// Unpack the model's flat output back into two frame-major spectrograms.
|
||||
/// Bins at and above DIM_F stay zero — the crop restored as silence, which
|
||||
/// for a vocal stem is honest (nothing sung lives above ~17.6 kHz).
|
||||
fn unpack_output(flat: &[f32]) -> (Vec<Complex<f32>>, Vec<Complex<f32>>) {
|
||||
let mut left = vec![Complex::default(); DIM_T * BINS];
|
||||
let mut right = vec![Complex::default(); DIM_T * BINS];
|
||||
for (c, spec) in [&mut left, &mut right].into_iter().enumerate() {
|
||||
for f in 0..DIM_F {
|
||||
for t in 0..DIM_T {
|
||||
spec[t * BINS + f] = Complex::new(
|
||||
flat[((2 * c) * DIM_F + f) * DIM_T + t],
|
||||
flat[((2 * c + 1) * DIM_F + f) * DIM_T + t],
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
(left, right)
|
||||
}
|
||||
|
||||
/// Chunk a stereo 44.1 kHz track through `run` (the [1,4,3072,256] →
|
||||
/// same-shape spectrogram transform) and reassemble the returned stem.
|
||||
/// UVR's non-overlap framing: TRIM zeros lead the track so chunk c's kept
|
||||
/// middle lands exactly at c·GEN, edge effects fall in the trimmed halves,
|
||||
/// and segments butt-join.
|
||||
fn process_chunks(
|
||||
stft: &Stft,
|
||||
left: &[f32],
|
||||
right: &[f32],
|
||||
mut run: impl FnMut(&[f32]) -> Result<Vec<f32>, String>,
|
||||
) -> Result<(Vec<f32>, Vec<f32>), String> {
|
||||
let len = left.len().min(right.len());
|
||||
if len < HOP {
|
||||
return Err("audio too short to separate".into());
|
||||
}
|
||||
let chunks = len.div_ceil(GEN);
|
||||
let padded_len = TRIM + chunks * GEN + TRIM;
|
||||
let mut mix_l = vec![0.0f32; padded_len];
|
||||
let mut mix_r = vec![0.0f32; padded_len];
|
||||
mix_l[TRIM..TRIM + len].copy_from_slice(&left[..len]);
|
||||
mix_r[TRIM..TRIM + len].copy_from_slice(&right[..len]);
|
||||
let mut stem_l = vec![0.0f32; chunks * GEN];
|
||||
let mut stem_r = vec![0.0f32; chunks * GEN];
|
||||
for c in 0..chunks {
|
||||
let s = c * GEN;
|
||||
let spec_l = stft.forward(&mix_l[s..s + CHUNK])?;
|
||||
let spec_r = stft.forward(&mix_r[s..s + CHUNK])?;
|
||||
let out = run(&pack_input(&spec_l, &spec_r))?;
|
||||
if out.len() != TENSOR_LEN {
|
||||
return Err(format!("unexpected chunk output length {}", out.len()));
|
||||
}
|
||||
let (mut vl, mut vr) = unpack_output(&out);
|
||||
let wl = stft.inverse(&mut vl)?;
|
||||
let wr = stft.inverse(&mut vr)?;
|
||||
stem_l[s..s + GEN].copy_from_slice(&wl[TRIM..TRIM + GEN]);
|
||||
stem_r[s..s + GEN].copy_from_slice(&wr[TRIM..TRIM + GEN]);
|
||||
// Politeness between chunks — this thread must never contend with
|
||||
// the audio path for sustained stretches (the passage loop's rule).
|
||||
std::thread::sleep(std::time::Duration::from_millis(3));
|
||||
}
|
||||
stem_l.truncate(len);
|
||||
stem_r.truncate(len);
|
||||
Ok((stem_l, stem_r))
|
||||
}
|
||||
|
||||
/// Whole-file vocal stem as 16 kHz mono for the CTC aligner: decode →
|
||||
/// chunked WinML separation → mid fold + anti-alias resample (player crate).
|
||||
/// Fails soft by design — the caller falls back to the band-passed mixture.
|
||||
pub fn separate_vocals_16k(
|
||||
app: &tauri::AppHandle,
|
||||
path: &std::path::Path,
|
||||
) -> Result<Vec<f32>, String> {
|
||||
if !model_present(app) {
|
||||
return Err("separation model not downloaded".into());
|
||||
}
|
||||
let [left, right] = player::vocals::separation_input_44k1(path)?;
|
||||
ensure_engine(app)?;
|
||||
let guard = ENGINE.lock().unwrap();
|
||||
let engine = guard.as_ref().ok_or("engine unavailable")?;
|
||||
let stft = Stft::new();
|
||||
let started = Instant::now();
|
||||
let (stem_l, stem_r) = process_chunks(&stft, &left, &right, |flat| {
|
||||
eval_chunk(&engine.session, flat)
|
||||
})?;
|
||||
log::info!(
|
||||
"vocal_sep: {:.0}s of audio separated in {:.1}s on {}",
|
||||
stem_l.len() as f32 / 44_100.0,
|
||||
started.elapsed().as_secs_f32(),
|
||||
engine.device
|
||||
);
|
||||
Ok(player::vocals::stem_pcm_16k(&stem_l, &stem_r))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// The WOLA plumbing alone — forward STFT straight back through the
|
||||
/// inverse, no crop — must be per-sample exact: the numerator and the
|
||||
/// envelope share the same window algebra, so any real error here is a
|
||||
/// tiling or normalization bug, not numerics.
|
||||
#[test]
|
||||
fn stft_wola_roundtrip_is_exact() {
|
||||
let chunk: Vec<f32> = (0..CHUNK)
|
||||
.map(|i| (2.0 * std::f32::consts::PI * 440.0 * i as f32 / 44_100.0).sin() * 0.5)
|
||||
.collect();
|
||||
let stft = Stft::new();
|
||||
let mut spec = stft.forward(&chunk).unwrap();
|
||||
let back = stft.inverse(&mut spec).unwrap();
|
||||
let worst = chunk
|
||||
.iter()
|
||||
.zip(&back)
|
||||
.map(|(a, b)| (a - b).abs())
|
||||
.fold(0.0f32, f32::max);
|
||||
assert!(worst < 1e-4, "un-cropped WOLA must be exact, worst {worst}");
|
||||
}
|
||||
|
||||
/// An identity model (output spectrogram = input spectrogram) must hand
|
||||
/// the track back through the whole STFT → crop → pack → unpack →
|
||||
/// zero-pad → iSTFT → trim → butt-join pipeline: proves the frame
|
||||
/// tiling, the chunk joins, and the layout round-trip all at once.
|
||||
/// Tolerance: the reflect pad puts a derivative kink at each chunk
|
||||
/// edge whose broadband energy reaches above bin 3072, and the model
|
||||
/// crop discards it — torch.stft + UVR's own framing lose exactly the
|
||||
/// same (~1% here); the model trained on that world.
|
||||
#[test]
|
||||
fn identity_model_reconstructs_across_chunk_joins() {
|
||||
let len = GEN + GEN / 2; // two chunks, one join, one truncation
|
||||
let track: Vec<f32> = (0..len)
|
||||
.map(|i| (2.0 * std::f32::consts::PI * 440.0 * i as f32 / 44_100.0).sin() * 0.5)
|
||||
.collect();
|
||||
let stft = Stft::new();
|
||||
let (out_l, out_r) =
|
||||
process_chunks(&stft, &track, &track, |flat| Ok(flat.to_vec())).unwrap();
|
||||
assert_eq!(out_l.len(), len);
|
||||
assert_eq!(out_r.len(), len);
|
||||
let mut worst = 0.0f32;
|
||||
for (a, b) in track.iter().zip(&out_l) {
|
||||
worst = worst.max((a - b).abs());
|
||||
}
|
||||
assert!(
|
||||
worst < 2e-2,
|
||||
"identity round-trip must reconstruct the tone, worst err {worst}"
|
||||
);
|
||||
// The join at GEN must not be a seam beyond that same crop error.
|
||||
let seam = (GEN - 50..GEN + 50)
|
||||
.map(|i| (track[i] - out_l[i]).abs())
|
||||
.fold(0.0f32, f32::max);
|
||||
assert!(seam < 2e-2, "butt-join must be seamless, err {seam}");
|
||||
}
|
||||
|
||||
/// The kept region of every chunk sits in full window-overlap coverage:
|
||||
/// the WOLA envelope there is flat and far from zero.
|
||||
#[test]
|
||||
fn envelope_is_stable_where_samples_are_kept() {
|
||||
let stft = Stft::new();
|
||||
let kept = &stft.envelope[2 * TRIM..2 * TRIM + GEN];
|
||||
let min = kept.iter().cloned().fold(f32::INFINITY, f32::min);
|
||||
let max = kept.iter().cloned().fold(0.0f32, f32::max);
|
||||
assert!(min > 2.0, "envelope must be well-conditioned, min {min}");
|
||||
assert!(
|
||||
max / min < 1.5,
|
||||
"envelope must be near-flat in the kept region, {min}..{max}"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user