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 <noreply@anthropic.com>
This commit is contained in:
Lucas WintherandClaude Opus 5.5 committed
1 parent
6bd8ab9cee
commit
86416f7637
4 files changed
+531
No files matched your search
+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;
|
||||
}
|
||||
Reference in new issue
Block a user