dj sets + lyrics: a mix that names its own tracks, and words pulled off the separated voice

This commit is contained in:
KeelrLeek committed 2026-08-24 10:49:10 -05:00
1 parent 4efe85b6bf
commit 1dd8573808
41 files changed
+9288 -647

No files matched your search

+499
View File
@@ -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}"
);
}
}