feat: demucs 6-stem separation
Radix-2 FFT and torch-compatible stft/istft (hann, center, reflect, normalized). HTDemucs-6s ONNX separator mirroring demucs apply_model: whole-mix normalisation, 343980-sample centred segments, 25% overlap with triangular weights, _spec/_ispec around the graph, spectral + time branch sum. Brief's guitar/bass > 3 dB criterion is unmet by the model on t1b/1 and is kept as test.failing pending a ruling. Co-Authored-By: Claude Opus 5.5 <[email protected]>
This commit is contained in:
Lucas Winther
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
6bd8ab9cee
commit
86416f7637
+173
@@ -0,0 +1,173 @@
|
||||
// Radix-2 FFT and torch-compatible STFT / iSTFT (hann window, center=True, reflect pad, normalized=True).
|
||||
// Pure: no I/O, browser-portable.
|
||||
|
||||
type Buf = Float32Array | Float64Array;
|
||||
|
||||
const twiddleCache = new Map<number, { cos: Float64Array; sin: Float64Array; rev: Uint32Array }>();
|
||||
|
||||
function plan(n: number) {
|
||||
let p = twiddleCache.get(n);
|
||||
if (p) return p;
|
||||
if (n < 1 || (n & (n - 1)) !== 0) throw new Error(`fft: length ${n} is not a power of two`);
|
||||
const half = n >> 1;
|
||||
const cos = new Float64Array(half);
|
||||
const sin = new Float64Array(half);
|
||||
for (let k = 0; k < half; k++) {
|
||||
cos[k] = Math.cos((2 * Math.PI * k) / n);
|
||||
sin[k] = -Math.sin((2 * Math.PI * k) / n);
|
||||
}
|
||||
const bits = Math.log2(n);
|
||||
const rev = new Uint32Array(n);
|
||||
for (let i = 0; i < n; i++) {
|
||||
let r = 0;
|
||||
for (let b = 0; b < bits; b++) r = (r << 1) | ((i >> b) & 1);
|
||||
rev[i] = r;
|
||||
}
|
||||
p = { cos, sin, rev };
|
||||
twiddleCache.set(n, p);
|
||||
return p;
|
||||
}
|
||||
|
||||
/** In-place forward DFT (no scaling): X[k] = sum_t x[t] e^{-2πikt/N}. N must be a power of two. */
|
||||
export function fft(re: Buf, im: Buf): void {
|
||||
const n = re.length;
|
||||
if (im.length !== n) throw new Error("fft: re/im length mismatch");
|
||||
if (n <= 1) return;
|
||||
const { cos, sin, rev } = plan(n);
|
||||
for (let i = 0; i < n; i++) {
|
||||
const j = rev[i]!;
|
||||
if (j > i) {
|
||||
let t = re[i]!; re[i] = re[j]!; re[j] = t;
|
||||
t = im[i]!; im[i] = im[j]!; im[j] = t;
|
||||
}
|
||||
}
|
||||
for (let size = 2; size <= n; size <<= 1) {
|
||||
const halfSize = size >> 1;
|
||||
const step = n / size;
|
||||
for (let start = 0; start < n; start += size) {
|
||||
for (let k = 0, tw = 0; k < halfSize; k++, tw += step) {
|
||||
const wr = cos[tw]!, wi = sin[tw]!;
|
||||
const a = start + k, b = a + halfSize;
|
||||
const xr = re[b]! * wr - im[b]! * wi;
|
||||
const xi = re[b]! * wi + im[b]! * wr;
|
||||
re[b] = re[a]! - xr; im[b] = im[a]! - xi;
|
||||
re[a] = re[a]! + xr; im[a] = im[a]! + xi;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const hannCache = new Map<number, Float64Array>();
|
||||
/** Periodic hann window (torch.hann_window default). */
|
||||
export function hann(n: number): Float64Array {
|
||||
let w = hannCache.get(n);
|
||||
if (!w) {
|
||||
w = new Float64Array(n);
|
||||
for (let i = 0; i < n; i++) w[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / n);
|
||||
hannCache.set(n, w);
|
||||
}
|
||||
return w;
|
||||
}
|
||||
|
||||
/** Reflect index (torch 'reflect': edge sample not repeated) into [0, len). */
|
||||
function reflectIndex(i: number, len: number): number {
|
||||
if (len === 1) return 0;
|
||||
const period = 2 * (len - 1);
|
||||
let m = i % period;
|
||||
if (m < 0) m += period;
|
||||
return m < len ? m : period - m;
|
||||
}
|
||||
|
||||
export interface Spectrogram { re: Float32Array[]; im: Float32Array[] }
|
||||
|
||||
/**
|
||||
* torch.stft(x, nFft, hop, window=hann(nFft), center=True, pad_mode="reflect", normalized=True, onesided=True).
|
||||
* Returns frames × bins (nFft/2 + 1 bins, 1 + floor(len/hop) frames).
|
||||
*/
|
||||
export function stft(x: Float32Array, nFft: number, hop: number): Spectrogram {
|
||||
const w = hann(nFft);
|
||||
const half = nFft >> 1;
|
||||
if (x.length <= half) throw new Error(`stft: input length ${x.length} too short for reflect pad ${half}`);
|
||||
const nFrames = 1 + Math.floor(x.length / hop);
|
||||
const bins = half + 1;
|
||||
const scale = 1 / Math.sqrt(nFft);
|
||||
const re: Float32Array[] = [];
|
||||
const im: Float32Array[] = [];
|
||||
const br = new Float64Array(nFft);
|
||||
const bi = new Float64Array(nFft);
|
||||
for (let f = 0; f < nFrames; f++) {
|
||||
const start = f * hop - half; // index in the original signal
|
||||
for (let i = 0; i < nFft; i++) {
|
||||
const j = start + i;
|
||||
const v = j >= 0 && j < x.length ? x[j]! : x[reflectIndex(j, x.length)]!;
|
||||
br[i] = v * w[i]!;
|
||||
bi[i] = 0;
|
||||
}
|
||||
fft(br, bi);
|
||||
const fr = new Float32Array(bins);
|
||||
const fi = new Float32Array(bins);
|
||||
for (let k = 0; k < bins; k++) {
|
||||
fr[k] = br[k]! * scale;
|
||||
fi[k] = bi[k]! * scale;
|
||||
}
|
||||
re.push(fr);
|
||||
im.push(fi);
|
||||
}
|
||||
return { re, im };
|
||||
}
|
||||
|
||||
/**
|
||||
* torch.istft(spec, nFft, hop, window=hann(nFft), center=True, normalized=True, onesided=True, length).
|
||||
* Overlap-add of windowed inverse frames divided by the summed squared window; output is exactly `length` samples
|
||||
* (zero-padded if the frames do not cover it).
|
||||
*/
|
||||
export function istft(spec: Spectrogram, nFft: number, hop: number, length: number): Float32Array {
|
||||
const w = hann(nFft);
|
||||
const half = nFft >> 1;
|
||||
const bins = half + 1;
|
||||
const nFrames = spec.re.length;
|
||||
const total = nFft + hop * (nFrames - 1);
|
||||
const acc = new Float64Array(total);
|
||||
const env = new Float64Array(total);
|
||||
const scale = Math.sqrt(nFft) / nFft; // undo normalized=True, then 1/N of the inverse DFT
|
||||
const br = new Float64Array(nFft);
|
||||
const bi = new Float64Array(nFft);
|
||||
// Two real frames per complex IFFT: z = X_a + i X_b -> ifft(z) = x_a + i x_b.
|
||||
for (let f = 0; f < nFrames; f += 2) {
|
||||
const ar = spec.re[f]!, ai = spec.im[f]!;
|
||||
const hasB = f + 1 < nFrames;
|
||||
const cr = hasB ? spec.re[f + 1]! : null, ci = hasB ? spec.im[f + 1]! : null;
|
||||
for (let k = 0; k < bins; k++) {
|
||||
// Onesided -> Hermitian. torch ignores imag of DC/Nyquist (C2R), so zero them.
|
||||
const edge = k === 0 || k === half;
|
||||
const xar = ar[k]!, xai = edge ? 0 : ai[k]!;
|
||||
const xbr = cr ? cr[k]! : 0, xbi = ci && !edge ? ci[k]! : 0;
|
||||
// Conjugate input for inverse-via-forward: ifft(Z) = conj(fft(conj(Z))) / N
|
||||
// Z[k] = Xa[k] + i Xb[k]; Z[N-k] = conj(Xa[k]) + i conj(Xb[k])
|
||||
br[k] = xar - xbi;
|
||||
bi[k] = -(xai + xbr);
|
||||
if (!edge) {
|
||||
br[nFft - k] = xar + xbi;
|
||||
bi[nFft - k] = -(-xai + xbr);
|
||||
}
|
||||
}
|
||||
fft(br, bi);
|
||||
const off = f * hop;
|
||||
for (let i = 0; i < nFft; i++) {
|
||||
const wi = w[i]!;
|
||||
acc[off + i]! += br[i]! * scale * wi;
|
||||
env[off + i]! += wi * wi;
|
||||
if (hasB) {
|
||||
acc[off + hop + i]! += -bi[i]! * scale * wi;
|
||||
env[off + hop + i]! += wi * wi;
|
||||
}
|
||||
}
|
||||
}
|
||||
const out = new Float32Array(length);
|
||||
const n = Math.min(length, total - half);
|
||||
for (let i = 0; i < n; i++) {
|
||||
const e = env[i + half]!;
|
||||
out[i] = e > 1e-11 ? acc[i + half]! / e : 0;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
// HTDemucs 6-stem separation via ONNX Runtime. The exported graph takes the time-domain segment (`input`) and its
|
||||
// complex-as-channels spectrogram (`x`); STFT/iSTFT are done here, mirroring demucs htdemucs.py / apply.py (MIT).
|
||||
// Pure: no fs / Bun APIs; only ORT APIs that also exist in onnxruntime-web (InferenceSession.create, Tensor, run).
|
||||
import * as ort from "onnxruntime-node";
|
||||
import type { Audio } from "../core/types.ts";
|
||||
import { stft, istft, type Spectrogram } from "../dsp/fft.ts";
|
||||
|
||||
export type Stem = "drums" | "bass" | "other" | "vocals" | "guitar" | "piano";
|
||||
/** Model output order (models/demucs/manifest.json). */
|
||||
export const STEMS: readonly Stem[] = ["drums", "bass", "other", "vocals", "guitar", "piano"];
|
||||
|
||||
const SR = 44100;
|
||||
const SEGMENT = 343980; // int(model.segment * samplerate)
|
||||
const OVERLAP = 0.25;
|
||||
const NFFT = 4096;
|
||||
const HOP = 1024;
|
||||
const PAD = (HOP / 2) * 3; // 1536
|
||||
const LE = Math.ceil(SEGMENT / HOP); // 336 frames
|
||||
const BINS = NFFT / 2; // 2048 (last bin dropped)
|
||||
|
||||
/** Reflect-pad (torch 'reflect') by `left`/`right` samples. Requires x.length > max(left, right). */
|
||||
function reflectPad(x: Float32Array, left: number, right: number): Float32Array {
|
||||
const n = x.length;
|
||||
const out = new Float32Array(n + left + right);
|
||||
out.set(x, left);
|
||||
for (let i = 1; i <= left; i++) out[left - i] = x[i]!;
|
||||
for (let i = 1; i <= right; i++) out[left + n - 1 + i] = x[n - 1 - i]!;
|
||||
return out;
|
||||
}
|
||||
|
||||
/** htdemucs._spec for one channel: frames [2, 2+LE), bins [0, 2048). */
|
||||
function specChannel(x: Float32Array): Spectrogram {
|
||||
const padded = reflectPad(x, PAD, PAD + LE * HOP - x.length);
|
||||
const z = stft(padded, NFFT, HOP);
|
||||
if (z.re.length !== LE + 4) throw new Error(`demucs _spec: expected ${LE + 4} frames, got ${z.re.length}`);
|
||||
return { re: z.re.slice(2, 2 + LE), im: z.im.slice(2, 2 + LE) };
|
||||
}
|
||||
|
||||
/** htdemucs._ispec for one channel given [bins=2048][frames=LE] planes read through accessors. */
|
||||
function ispecChannel(get: (part: 0 | 1, bin: number, frame: number) => number): Float32Array {
|
||||
const frames = LE + 4;
|
||||
const re: Float32Array[] = [];
|
||||
const im: Float32Array[] = [];
|
||||
for (let f = 0; f < frames; f++) {
|
||||
const fr = new Float32Array(BINS + 1);
|
||||
const fi = new Float32Array(BINS + 1);
|
||||
const src = f - 2;
|
||||
if (src >= 0 && src < LE) {
|
||||
for (let k = 0; k < BINS; k++) {
|
||||
fr[k] = get(0, k, src);
|
||||
fi[k] = get(1, k, src);
|
||||
}
|
||||
}
|
||||
re.push(fr);
|
||||
im.push(fi);
|
||||
}
|
||||
const le = HOP * LE + 2 * PAD;
|
||||
return istft({ re, im }, NFFT, HOP, le).subarray(PAD, PAD + SEGMENT);
|
||||
}
|
||||
|
||||
/** Triangular overlap-add weight from apply_model (transition_power = 1). */
|
||||
function triangleWeight(): Float32Array {
|
||||
const w = new Float32Array(SEGMENT);
|
||||
const h = SEGMENT >> 1;
|
||||
for (let i = 0; i < h; i++) w[i] = i + 1;
|
||||
for (let i = h; i < SEGMENT; i++) w[i] = SEGMENT - i;
|
||||
let max = 0;
|
||||
for (const v of w) max = Math.max(max, v);
|
||||
for (let i = 0; i < SEGMENT; i++) w[i] = w[i]! / max;
|
||||
return w;
|
||||
}
|
||||
|
||||
export interface SeparatorOptions {
|
||||
/** ORT intra-op threads (CPU EP). Default: ORT's choice. 8 was fastest on a 16-core host (~9 s per minute of audio). */
|
||||
intraOpNumThreads?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Load HTDemucs-6s and return a separator. Mirrors demucs `apply_model(shifts=0, split=True, overlap=0.25)`:
|
||||
* whole-mix normalisation, 343980-sample segments centred/zero-padded like `TensorChunk.padded`, triangular
|
||||
* overlap-add weights. Input must be 44.1 kHz mono or stereo; every stem is stereo with the input's length.
|
||||
* The graph normalises `input` and `x` internally (verified: ReduceMean/Sub/Div on both inputs) and its outputs are
|
||||
* already de-normalised, so no per-segment normalisation is applied here.
|
||||
*/
|
||||
export async function createSeparator(
|
||||
modelBytes: Uint8Array,
|
||||
opts: SeparatorOptions = {},
|
||||
): Promise<(mix: Audio) => Promise<Record<Stem, Audio>>> {
|
||||
const session = await ort.InferenceSession.create(modelBytes, {
|
||||
executionProviders: ["cpu"],
|
||||
graphOptimizationLevel: "all",
|
||||
...(opts.intraOpNumThreads ? { intraOpNumThreads: opts.intraOpNumThreads } : {}),
|
||||
});
|
||||
const weight = triangleWeight();
|
||||
|
||||
/** Model on one exactly-SEGMENT-long stereo chunk → [stem][channel] time signals of length SEGMENT. */
|
||||
async function runSegment(seg: [Float32Array, Float32Array]): Promise<Float32Array[][]> {
|
||||
const input = new Float32Array(2 * SEGMENT);
|
||||
input.set(seg[0], 0);
|
||||
input.set(seg[1], SEGMENT);
|
||||
const x = new Float32Array(4 * BINS * LE);
|
||||
for (let c = 0; c < 2; c++) {
|
||||
const z = specChannel(seg[c]!);
|
||||
const reBase = (2 * c) * BINS * LE, imBase = (2 * c + 1) * BINS * LE;
|
||||
for (let f = 0; f < LE; f++) {
|
||||
const zr = z.re[f]!, zi = z.im[f]!;
|
||||
for (let k = 0; k < BINS; k++) {
|
||||
x[reBase + k * LE + f] = zr[k]!;
|
||||
x[imBase + k * LE + f] = zi[k]!;
|
||||
}
|
||||
}
|
||||
}
|
||||
const res = await session.run({
|
||||
input: new ort.Tensor("float32", input, [1, 2, SEGMENT]),
|
||||
x: new ort.Tensor("float32", x, [1, 4, BINS, LE]),
|
||||
});
|
||||
const spec = res["output"]!.data as Float32Array; // [1, 6, 4, 2048, 336]
|
||||
const time = res["add_67"]!.data as Float32Array; // [1, 6, 2, SEGMENT]
|
||||
const out: Float32Array[][] = [];
|
||||
for (let s = 0; s < STEMS.length; s++) {
|
||||
const chans: Float32Array[] = [];
|
||||
for (let c = 0; c < 2; c++) {
|
||||
const base = (s * 4 + 2 * c) * BINS * LE;
|
||||
const y = ispecChannel((part, k, f) => spec[base + part * BINS * LE + k * LE + f]!);
|
||||
const t = time.subarray((s * 2 + c) * SEGMENT, (s * 2 + c + 1) * SEGMENT);
|
||||
const sum = new Float32Array(SEGMENT);
|
||||
for (let i = 0; i < SEGMENT; i++) sum[i] = y[i]! + t[i]!;
|
||||
chans.push(sum);
|
||||
}
|
||||
out.push(chans);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
return async (mix: Audio): Promise<Record<Stem, Audio>> => {
|
||||
if (mix.sampleRate !== SR) throw new Error(`separate: sample rate must be ${SR}, got ${mix.sampleRate}`);
|
||||
if (mix.channels.length < 1 || mix.channels.length > 2) {
|
||||
throw new Error(`separate: expected 1 or 2 channels, got ${mix.channels.length}`);
|
||||
}
|
||||
const src = mix.channels.length === 1 ? [mix.channels[0]!, mix.channels[0]!] : mix.channels;
|
||||
const length = src[0]!.length;
|
||||
if (src[1]!.length !== length) throw new Error("separate: channel lengths differ");
|
||||
|
||||
// Whole-mix normalisation on the mono reference (demucs separate: (wav - ref.mean()) / ref.std()).
|
||||
let mean = 0;
|
||||
for (let i = 0; i < length; i++) mean += (src[0]![i]! + src[1]![i]!) / 2;
|
||||
mean /= Math.max(length, 1);
|
||||
let v = 0;
|
||||
for (let i = 0; i < length; i++) v += ((src[0]![i]! + src[1]![i]!) / 2 - mean) ** 2;
|
||||
const std = Math.sqrt(v / Math.max(length - 1, 1)) + 1e-8; // torch.std is unbiased
|
||||
const norm = src.map((ch) => {
|
||||
const o = new Float32Array(length);
|
||||
for (let i = 0; i < length; i++) o[i] = (ch[i]! - mean) / std;
|
||||
return o;
|
||||
});
|
||||
|
||||
const acc = STEMS.map(() => [new Float64Array(length), new Float64Array(length)]);
|
||||
const sumWeight = new Float64Array(length);
|
||||
const stride = Math.floor((1 - OVERLAP) * SEGMENT);
|
||||
for (let offset = 0; offset < length; offset += stride) {
|
||||
const chunkLen = Math.min(SEGMENT, length - offset);
|
||||
// TensorChunk.padded: centre the chunk in a SEGMENT window, reading neighbouring audio, zeros outside.
|
||||
const start = offset - Math.floor((SEGMENT - chunkLen) / 2);
|
||||
const seg: [Float32Array, Float32Array] = [new Float32Array(SEGMENT), new Float32Array(SEGMENT)];
|
||||
const lo = Math.max(0, start), hi = Math.min(length, start + SEGMENT);
|
||||
for (let c = 0; c < 2; c++) seg[c]!.set(norm[c]!.subarray(lo, hi), lo - start);
|
||||
const est = await runSegment(seg);
|
||||
const trim = offset - start; // center_trim back to chunkLen
|
||||
for (let s = 0; s < STEMS.length; s++) {
|
||||
for (let c = 0; c < 2; c++) {
|
||||
const e = est[s]![c]!, a = acc[s]![c]!;
|
||||
for (let i = 0; i < chunkLen; i++) a[offset + i]! += weight[i]! * e[trim + i]!;
|
||||
}
|
||||
}
|
||||
for (let i = 0; i < chunkLen; i++) sumWeight[offset + i]! += weight[i]!;
|
||||
}
|
||||
|
||||
const result = {} as Record<Stem, Audio>;
|
||||
STEMS.forEach((stem, s) => {
|
||||
const channels = acc[s]!.map((a) => {
|
||||
const o = new Float32Array(length);
|
||||
for (let i = 0; i < length; i++) o[i] = (a[i]! / sumWeight[i]!) * std + mean;
|
||||
return o;
|
||||
});
|
||||
result[stem] = { sampleRate: SR, channels };
|
||||
});
|
||||
return result;
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
import { test, expect } from "bun:test";
|
||||
import { fft, stft, istft } from "../../src/dsp/fft.ts";
|
||||
import { createRng } from "../../src/core/rng.ts";
|
||||
|
||||
test("fft of an impulse has all-ones magnitude", () => {
|
||||
const n = 1024;
|
||||
const re = new Float32Array(n);
|
||||
const im = new Float32Array(n);
|
||||
re[0] = 1;
|
||||
fft(re, im);
|
||||
for (let k = 0; k < n; k++) expect(Math.hypot(re[k]!, im[k]!)).toBeCloseTo(1, 6);
|
||||
});
|
||||
|
||||
test("fft matches a direct DFT on a random signal", () => {
|
||||
const n = 64;
|
||||
const rng = createRng(3);
|
||||
const x = Array.from({ length: n }, () => rng.next() * 2 - 1);
|
||||
const re = Float32Array.from(x);
|
||||
const im = new Float32Array(n);
|
||||
fft(re, im);
|
||||
for (let k = 0; k < n; k++) {
|
||||
let sr = 0, si = 0;
|
||||
for (let t = 0; t < n; t++) {
|
||||
sr += x[t]! * Math.cos((-2 * Math.PI * k * t) / n);
|
||||
si += x[t]! * Math.sin((-2 * Math.PI * k * t) / n);
|
||||
}
|
||||
expect(re[k]!).toBeCloseTo(sr, 4);
|
||||
expect(im[k]!).toBeCloseTo(si, 4);
|
||||
}
|
||||
});
|
||||
|
||||
test("stft follows torch semantics: frame count, bins, normalized=True scaling", () => {
|
||||
const nFft = 4096, hop = 1024, len = 50_000;
|
||||
const x = new Float32Array(len).fill(1); // DC: bin 0 of an interior frame = sum(hann)/sqrt(nFft) = nFft/2/sqrt(nFft)
|
||||
const s = stft(x, nFft, hop);
|
||||
expect(s.re.length).toBe(1 + Math.floor(len / hop));
|
||||
expect(s.re[0]!.length).toBe(nFft / 2 + 1);
|
||||
expect(s.re[10]![0]!).toBeCloseTo(nFft / 2 / Math.sqrt(nFft), 3);
|
||||
expect(Math.abs(s.re[10]![1]!)).toBeCloseTo(nFft / 4 / Math.sqrt(nFft), 3);
|
||||
});
|
||||
|
||||
test("istft(stft(x)) reconstructs x (max err < 1e-4)", () => {
|
||||
const nFft = 4096, hop = 1024, len = 50_000;
|
||||
const rng = createRng(7);
|
||||
const x = new Float32Array(len);
|
||||
for (let i = 0; i < len; i++) x[i] = rng.next() * 2 - 1;
|
||||
const y = istft(stft(x, nFft, hop), nFft, hop, len);
|
||||
expect(y.length).toBe(len);
|
||||
let maxErr = 0;
|
||||
for (let i = 0; i < len; i++) maxErr = Math.max(maxErr, Math.abs(y[i]! - x[i]!));
|
||||
expect(maxErr).toBeLessThan(1e-4);
|
||||
});
|
||||
@@ -0,0 +1,117 @@
|
||||
import { test, expect, beforeAll } from "bun:test";
|
||||
import { join } from "node:path";
|
||||
import { buildSong, loadSongDir } from "../../src/bench/dataset.ts";
|
||||
import { createSeparator, STEMS, type Stem } from "../../src/pipeline/separate.ts";
|
||||
import type { Audio } from "../../src/core/types.ts";
|
||||
|
||||
const MODEL = join(import.meta.dir, "../../models/demucs/htdemucs_6s.onnx");
|
||||
const SLOW = 300_000;
|
||||
|
||||
let separate: (mix: Audio) => Promise<Record<Stem, Audio>>;
|
||||
let data: Awaited<ReturnType<typeof loadSongDir>>;
|
||||
|
||||
beforeAll(async () => {
|
||||
const dir = await buildSong("t1b", 1);
|
||||
data = await loadSongDir(dir);
|
||||
separate = await createSeparator(new Uint8Array(await Bun.file(MODEL).arrayBuffer()), { intraOpNumThreads: 8 });
|
||||
}, SLOW);
|
||||
|
||||
/** Sum of oracle stems for every track of the given role. */
|
||||
function oracle(role: string): Audio {
|
||||
const len = data.mix.channels[0]!.length;
|
||||
const ch = [new Float32Array(len), new Float32Array(len)];
|
||||
for (const t of data.song.tracks) {
|
||||
if (t.role !== role) continue;
|
||||
const s = data.stems.get(t.name)!;
|
||||
for (let c = 0; c < 2; c++) for (let i = 0; i < len; i++) ch[c]![i]! += s.channels[c]![i]!;
|
||||
}
|
||||
return { sampleRate: data.mix.sampleRate, channels: ch };
|
||||
}
|
||||
|
||||
/** SDR in dB over all channels: 10 log10(|ref|^2 / |ref - est|^2). */
|
||||
function sdr(ref: Audio, est: Audio): number {
|
||||
let num = 0, den = 0;
|
||||
for (let c = 0; c < ref.channels.length; c++) {
|
||||
const r = ref.channels[c]!, e = est.channels[c]!;
|
||||
for (let i = 0; i < r.length; i++) {
|
||||
num += r[i]! * r[i]!;
|
||||
den += (r[i]! - e[i]!) ** 2;
|
||||
}
|
||||
}
|
||||
return 10 * Math.log10(num / Math.max(den, 1e-20));
|
||||
}
|
||||
|
||||
let out: Record<Stem, Audio>;
|
||||
let secsPerMin = 0;
|
||||
const sdrs: Record<string, number> = {};
|
||||
|
||||
function sum(...as: Audio[]): Audio {
|
||||
const len = as[0]!.channels[0]!.length;
|
||||
return { sampleRate: 44100, channels: [0, 1].map((c) => {
|
||||
const o = new Float32Array(len);
|
||||
for (const a of as) for (let i = 0; i < len; i++) o[i]! += a.channels[c]![i]!;
|
||||
return o;
|
||||
}) };
|
||||
}
|
||||
|
||||
beforeAll(async () => {
|
||||
const t0 = performance.now();
|
||||
out = await separate(data.mix);
|
||||
const minutes = data.mix.channels[0]!.length / data.mix.sampleRate / 60;
|
||||
secsPerMin = (performance.now() - t0) / 1000 / minutes;
|
||||
for (const [stem, role] of [["guitar", "guitar"], ["bass", "bass"], ["drums", "drums"]] as const) {
|
||||
sdrs[stem] = sdr(oracle(role), out[stem]);
|
||||
sdrs[`${stem}_mixBaseline`] = sdr(oracle(role), data.mix);
|
||||
}
|
||||
sdrs.guitarPlusPiano = sdr(oracle("guitar"), sum(out.guitar, out.piano));
|
||||
console.log(`separation t1b/1: ${secsPerMin.toFixed(1)} s per minute of audio`, sdrs);
|
||||
}, SLOW);
|
||||
|
||||
test("separates t1b/1: every stem stereo 44.1 kHz with exactly the input length", () => {
|
||||
const len = data.mix.channels[0]!.length;
|
||||
for (const s of STEMS) {
|
||||
expect(out[s].sampleRate).toBe(44100);
|
||||
expect(out[s].channels.length).toBe(2);
|
||||
for (const ch of out[s].channels) expect(ch.length).toBe(len);
|
||||
}
|
||||
});
|
||||
|
||||
test("separates t1b/1: drums > 3 dB, bass beats the mix, guitar(+piano) > 3 dB and beats the mix", () => {
|
||||
// Measured (2026-10-30): drums 5.4 dB (mix -13.7), bass 1.67 dB (mix -2.8), guitar 1.85 dB (mix 1.87),
|
||||
// guitar+piano 4.7 dB. HTDemucs-6s routes ~30% of the GM steel guitar into its piano stem.
|
||||
expect(sdrs.drums!).toBeGreaterThan(3);
|
||||
expect(sdrs.drums!).toBeGreaterThan(sdrs.drums_mixBaseline!);
|
||||
expect(sdrs.bass!).toBeGreaterThan(sdrs.bass_mixBaseline!);
|
||||
expect(sdrs.guitarPlusPiano!).toBeGreaterThan(3);
|
||||
expect(sdrs.guitarPlusPiano!).toBeGreaterThan(sdrs.guitar_mixBaseline!);
|
||||
});
|
||||
|
||||
// Brief's literal criterion. Unmet by the model on this song (guitar 1.85 dB, bass 1.67 dB); kept as an
|
||||
// expected failure so it flips loudly if separation improves. Needs a controller ruling.
|
||||
test.failing("brief criterion: guitar and bass stems each > 3 dB and > mix baseline", () => {
|
||||
expect(sdrs.guitar!).toBeGreaterThan(3);
|
||||
expect(sdrs.guitar!).toBeGreaterThan(sdrs.guitar_mixBaseline!);
|
||||
expect(sdrs.bass!).toBeGreaterThan(3);
|
||||
expect(sdrs.bass!).toBeGreaterThan(sdrs.bass_mixBaseline!);
|
||||
});
|
||||
|
||||
test("a 3-second slice (shorter than one segment) returns 3 s stems", async () => {
|
||||
const n = 3 * 44100;
|
||||
const slice: Audio = { sampleRate: 44100, channels: data.mix.channels.map((c) => c.slice(44100, 44100 + n)) };
|
||||
const out = await separate(slice);
|
||||
for (const s of STEMS) for (const ch of out[s].channels) expect(ch.length).toBe(n);
|
||||
}, SLOW);
|
||||
|
||||
test("mono input is accepted and returns stereo stems of the same length", async () => {
|
||||
const n = 44100;
|
||||
const mono: Audio = { sampleRate: 44100, channels: [data.mix.channels[0]!.slice(0, n)] };
|
||||
const out = await separate(mono);
|
||||
for (const s of STEMS) {
|
||||
expect(out[s].channels.length).toBe(2);
|
||||
expect(out[s].channels[0]!.length).toBe(n);
|
||||
}
|
||||
}, SLOW);
|
||||
|
||||
test("rejects sample rates other than 44100", async () => {
|
||||
await expect(separate({ sampleRate: 22050, channels: [new Float32Array(100)] })).rejects.toThrow(/44100/);
|
||||
});
|
||||
Reference in New Issue
Block a user