restore(): plan, training side and handout for a studio-grade voice restorer

.work/todo.md

@@ -102,16 +102,15 @@ Missing: true harmonic resynthesis (`sbr` is a waveshaper exciter, not the voice
 - [ ] Measure against `sbr()` and the input: VoiceBank cut to 8 and 16 kHz, log-spectral distance, PESQ WB, DNSMOS; RX Spectral Recovery not hostable
 - [ ] Then the same for pitched solo instruments in a stem (voice, flute, violin); polyphony stays out
 
-### 3. Generative speech restoration: the Adobe Enhance class (1–2 weeks)
-- [ ] `restore()` on Sidon v0.1 (MIT; w2v-BERT 2.0's first 8 layers + LoRA, DAC decoder; 470 MB f16): export to ONNX on this laptop, int8, onnxruntime-node with CoreML, chunked as deepfilter
-- [ ] Measure against `deepfilter()` and `auto()`: VoiceBank+DEMAND, vbreverb, phone and laptop-mic takes; DNSMOS, UTMOS, WER (neural-asr), speaker similarity; listen blind
-- [ ] Strength knob: mix back into the input as deepfilter's mixback
-- [ ] `auto('speech')` gains it as the last repair stage when the package is installed
-
-### 4. A browser-size restorer (weeks, only if 3 proves the gap)
-- [ ] Student 2–20M params (10–30 MB int8): a net predicting band gains, deep-filter taps and harmonic amplitudes, the DSP from 2 rendering them
-- [ ] Trained on Sidon's outputs over our own speech data (distillation, MIT allows), PyTorch MPS on this laptop
-- [ ] Runs in onnxruntime-web, WebGPU, real time
+### 3. Generative speech restoration: the Adobe Enhance class ([restore.md](restore.md))
+- [x] Plan, prior art and data audit: `.work/restore.md` (FINALLY's recipe; Sidon MIT as the teacher; the open-license data that exists for English, Russian and singing, and what does not)
+- [x] Training side built and smoke-tested on CPU: `train/restore/` (degradation chain with codecs, mixed-bandwidth targets through one band-limited loss, ConvNeXt + attention generator in the STFT, MPD + MRD + WavLM losses, resume, ONNX export verified, chunked inference, bench scoring)
+- [ ] First run on a GPU (VCTK, LibriTTS-R, VocalSet, OpenAIR/BUT rooms, MUSAN/FSD50K): dim 384, 50k steps; `score.py` on VoiceBank+DEMAND and vbreverb against `deepfilter()`
+- [ ] Teacher pass (Sidon) over Golos crowd, RuLS and separated CC vocals with the quality gate (DNSMOS, WER, speaker similarity); full-size run
+- [ ] `restore()` on `fn/deepfilter.js`'s `enhancer()`, `@audio/neural-restore`, weights hosted; measured against `deepfilter()`, `auto()` and Adobe Enhance v2 by hand; `auto('speech')` gains it as the last repair stage when the package is installed
+
+### 4. A browser-size restorer
+- [ ] The same recipe at dim 384, 6 blocks (~8M, ~16 MB fp16); what it loses against full size; onnxruntime-web WebGPU and the RoFormer engine's kernels, real time in the page
 
 ### 5. Browser, measure what exists
 - [ ] **Fast, best-quality split on the GPU** (measured 2026-10-09 in Node, onnxruntime-node 1.30, machine under load: absolute times pessimistic, ratios hold)

train/restore/README.md

@@ -0,0 +1,39 @@
+# restore: a studio-grade voice from a bad recording
+
+The training side of `audio`'s coming `restore()` op: one network that takes a voice recorded in a room, in noise,
+through a cheap microphone, a codec or a phone, speech or singing, and gives back the voice as a studio would have
+recorded it, at 48 kHz. The plan, the data and what is known are in [`.work/restore.md`](../../.work/restore.md); the step-by-step runbook from an
+empty GPU machine to the published plugin, with go / no-go numbers, is [`HANDOUT.md`](HANDOUT.md).
+
+On a GPU machine, `bash run.sh all` does all of the below (about 20 GB of downloads: LibriTTS-R train-clean-100, VCTK, MUSAN, VocalSet). By hand:
+
+```
+pip install torch torchaudio soundfile scipy onnx onnxruntime onnxconverter-common     # ffmpeg on PATH for codecs
+python data.py scan en.json  /data/vctk /data/libritts-r --kind speech --lang en       # manifests, bandwidth per file
+python data.py scan ru.json  /data/ruls --kind speech --lang ru
+python data.py scan sing.json /data/vocalset --kind sing
+python train.py --out runs/a --clean en.json,ru.json,sing.json --rirs /data/rirs --noises /data/noise
+python export.py runs/a/net-0400000.pt --fp16 --verify                                   # → ~/.cache/audiojs/neural/restore/
+python infer.py runs/a/net-0400000.pt take.wav --out restored/
+python score.py --ref clean/ --out restored/                                             # PESQ, STOI, SI-SDR, DNSMOS, LSD
+python test.py                                                                           # smoke test, CPU, a minute
+```
+
+`score.py` also needs `pesq pyloudnorm` (`pesq` builds from source: on Python 3.13 it needs `setuptools<70` in the build, or Python 3.12).
+
+```
+```
+
+| File | What |
+|---|---|
+| `dsp.py` | the STFT (2048/512 at 48 kHz), levels, the band limit every loss passes through, mel, resampling, the bandwidth a file carries |
+| `degrade.py` | the on-the-fly degradation chain: room (measured or synthetic), noise, microphone colour, compressor, clipping, band limit, codec (ffmpeg), word length |
+| `data.py` | manifests and the pair sampler (an IterableDataset: crops by length, levels, the chain); `pairs` writes examples to listen to |
+| `model.py` | the generator (ConvNeXt over frames with self-attention, log-magnitude change + phase), the discriminators (multi-period, multi-resolution) |
+| `loss.py` | multi-resolution STFT, log-mel, WavLM feature distance, least-squares GAN, feature matching |
+| `train.py` | the run: regression warm-up, then adversarial; checkpoints, validation (LSD, SI-SDR, mel), samples |
+| `export.py` | ONNX between STFT and iSTFT, `[1, 2, 1025, T]` re/im in and out, fp16, verified against torch |
+| `infer.py` | files through a checkpoint or the ONNX, chunked |
+| `score.py` | `bench/speech.py`'s measures plus LSD on a directory of outputs against references |
+| `test.py` | the whole path on synthetic voices, CPU |
+| `run.sh` | one command on a GPU machine: fetch the open data, scan, train, export (`bash run.sh all`; `small` for the browser-size model) |

train/restore/data.py

@@ -0,0 +1,107 @@
+"""Clean voices, rooms and noises on disk → training pairs.
+
+    python data.py scan OUT.json DIR [DIR...] [--kind speech|sing] [--lang en|ru|..]
+        every audio file under the DIRs: its rate, length and the bandwidth it carries (dsp.bandwidth), one
+        manifest; a corpus at 16 or 24 kHz is kept, its band edge recorded, so the losses stop at it (dsp.band)
+    python data.py pairs OUT_DIR N --clean A.json[,B.json] [--rirs DIR] [--noises DIR] [--seg 4]
+        N training pairs as the trainer would draw them, <k>-in.wav and <k>-ref.wav: listen before training
+
+Pairs(manifests, rirs, noises, seg) draws a crop of `seg` seconds from a random file (by length), at 48 kHz mono,
+at a random level, through degrade.Chain, and yields (degraded, clean, bandwidth_hz). An IterableDataset: each
+worker seeds itself. Bank(dir) holds every file of a directory in memory at 48 kHz (rooms: aligned to the direct
+peak, 2.5 s at most, unit energy).
+"""
+import glob, json, os, random, sys
+import numpy as np
+import soundfile as sf
+import torch
+from torch.utils.data import IterableDataset, get_worker_info
+from dsp import SR, bandwidth, resample, rms_db
+from degrade import Chain
+
+EXT = ('.wav', '.flac', '.ogg', '.mp3', '.opus')
+
+def files(d):
+    return sorted(p for p in glob.glob(os.path.join(d, '**', '*'), recursive=True) if p.lower().endswith(EXT))
+
+def read(path, start=0, frames=-1):
+    x, sr = sf.read(path, start=start, frames=frames, dtype='float32', always_2d=True)
+    return torch.from_numpy(x.mean(1)), sr
+
+def scan(out, dirs, kind='speech', lang='en'):
+    items = json.load(open(out)) if os.path.exists(out) else []
+    seen = {i['path'] for i in items}
+    for d in dirs:
+        for p in files(d):
+            if p in seen: continue
+            try:
+                info = sf.info(p)
+                x, sr = read(p, frames=min(info.frames, info.samplerate * 30))
+            except Exception as e:
+                print('skip', p, e, file=sys.stderr); continue
+            bw = bandwidth(resample(x, sr), SR) if sr != SR else bandwidth(x, SR)
+            items.append(dict(path=p, sr=sr, frames=info.frames, bw=round(bw), kind=kind, lang=lang))
+            print(len(items), p, sr, round(info.frames / sr, 1), 's', round(bw), 'Hz', file=sys.stderr)
+            if len(items) % 200 == 0: json.dump(items, open(out, 'w'))
+    json.dump(items, open(out, 'w'))
+    return items
+
+class Bank:
+    def __init__(self, d, rir=False, limit=None):
+        self.items = []
+        for p in files(d)[:limit] if d else []:
+            x, sr = read(p)
+            x = resample(x, sr)
+            if rir:
+                k = int(x.abs().argmax()); x = x[max(0, k - int(0.002 * SR)): k + int(2.5 * SR)]
+                x = x / (x.pow(2).sum().sqrt() + 1e-9)
+            if len(x) > SR // 10: self.items.append(x)
+
+class Pairs(IterableDataset):
+    def __init__(self, manifests, rirs=(), noises=(), seg=4.0, p=None, level=(-35, -12), seed=0):
+        self.items = [i for m in manifests for i in json.load(open(m)) if i['frames'] >= i['sr'] * 0.5]
+        self.w = np.array([i['frames'] / i['sr'] for i in self.items]); self.w /= self.w.sum()
+        self.chain = Chain(rirs, noises, p)
+        self.n, self.level, self.seed = int(seg * SR), level, seed
+
+    def crop(self):
+        while True:
+            i = self.items[np.random.choice(len(self.items), p=self.w)]
+            n = int(self.n * i['sr'] / SR) + 1
+            start = random.randint(0, max(0, i['frames'] - n))
+            try: x, sr = read(i['path'], start, min(n, i['frames'] - start))
+            except Exception: continue
+            x = resample(x, sr)
+            if len(x) < self.n: x = torch.nn.functional.pad(x, (0, self.n - len(x)))
+            x = x[: self.n]
+            if rms_db(x) > -45: return x, float(min(i['bw'], sr / 2))
+
+    def __iter__(self):
+        w = get_worker_info()
+        s = self.seed + (w.id if w else 0) + 1000 * int.from_bytes(os.urandom(2), 'little')
+        random.seed(s); np.random.seed(s % 2 ** 32); torch.manual_seed(s)
+        while True:
+            x, bw = self.crop()
+            x = x * 10 ** ((random.uniform(*self.level) - float(rms_db(x))) / 20)
+            y = self.chain(x)
+            peak = y.abs().max()
+            if peak > 0.99: x, y = x * 0.99 / peak, y * 0.99 / peak
+            yield y, x, torch.tensor(bw)
+
+if __name__ == '__main__':
+    import argparse
+    a = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    a.add_argument('cmd', choices=['scan', 'pairs']); a.add_argument('out'); a.add_argument('dirs', nargs='*')
+    a.add_argument('--kind', default='speech', choices=['speech', 'sing']); a.add_argument('--lang', default='en')
+    a.add_argument('--clean'); a.add_argument('--rirs'); a.add_argument('--noises'); a.add_argument('--seg', type=float, default=4.0)
+    o = a.parse_args()
+    if o.cmd == 'scan':
+        scan(o.out, o.dirs, o.kind, o.lang)
+    else:
+        n = int(o.dirs[0]) if o.dirs else 10
+        pairs = Pairs(o.clean.split(','), Bank(o.rirs, rir=True).items if o.rirs else (), Bank(o.noises).items if o.noises else (), o.seg)
+        os.makedirs(o.out, exist_ok=True)
+        for k, (y, x, bw) in zip(range(n), pairs):
+            sf.write(os.path.join(o.out, f'{k}-in.wav'), y.numpy(), SR, subtype='FLOAT')
+            sf.write(os.path.join(o.out, f'{k}-ref.wav'), x.numpy(), SR, subtype='FLOAT')
+            print(k, f'bandwidth {int(bw)} Hz')

train/restore/degrade.py

@@ -0,0 +1,163 @@
+"""The take as a bad recording makes it: a clean 48 kHz voice through a random chain, on the fly, per item.
+
+Order is the physical one: the voice in a room, noise in the room, the microphone's colour and preamp, its
+compressor and clipping, the channel's band limit and codec, the file's word length. Each stage draws its own
+chance and strength from the ranges below; `Chain(bank).__call__(clean) → degraded`, same length and about the
+same level. CPU torch on (N,) tensors, cheap enough for data-loader workers; codecs through ffmpeg.
+"""
+import math, os, random, shutil, subprocess, tempfile
+import torch
+import torch.nn.functional as F
+import torchaudio.functional as AF
+from dsp import SR, band, resample
+
+FFMPEG = shutil.which('ffmpeg')
+CODECS = [  # (name, args, container, bitrates kb/s): codecs a voice reaches us through
+    ('mp3', ['-c:a', 'libmp3lame'], 'mp3', (24, 32, 48, 64, 96, 128)),
+    ('aac', ['-c:a', 'aac'], 'm4a', (24, 32, 48, 64, 96)),
+    ('opus', ['-c:a', 'libopus', '-application', 'voip'], 'ogg', (8, 12, 16, 24, 32, 48)),
+    ('gsm', ['-c:a', 'libgsm', '-ar', '8000'], 'gsm', (13,)),
+]
+
+def u(a, b): return random.uniform(a, b)
+def lu(a, b): return math.exp(u(math.log(a), math.log(b)))
+def chance(p): return random.random() < p
+
+def db(x): return 10 ** (x / 20)
+def power(x): return x.pow(2).mean() + 1e-12
+
+def spectral(x, curve_db):
+    """x through a zero-phase filter whose log-magnitude is `curve_db` over rfftfreq(n)."""
+    X = torch.fft.rfft(x)
+    return torch.fft.irfft(X * db(curve_db), len(x))
+
+def freqs(n): return torch.fft.rfftfreq(n, 1 / SR)
+
+# ---------------------------------------------------------------- stages
+
+def mic(x):
+    """A cheap microphone and its placement: tilt, two or three resonances and dips, a low cut; ±12 dB at most."""
+    f = freqs(len(x)); lf = torch.log2(f.clamp(min=20) / 1000)
+    c = u(-4, 4) * lf                                   # dB per octave about 1 kHz
+    for _ in range(random.randint(1, 3)):
+        fc, q, g = lu(150, 12000), lu(0.5, 4), u(-12, 9)
+        c = c + g * torch.exp(-((torch.log2(f.clamp(min=1) / fc)) * q) ** 2)
+    hp = lu(40, 300)                                    # proximity's absence, a windscreen, a phone
+    c = c + 12 * torch.log10(1 / (1 + (hp / f.clamp(min=1)) ** 4))
+    c = c - c[(f > 300) & (f < 3000)].mean()            # the voice band's level kept
+    return spectral(x, c.clamp(-30, 12))
+
+def synth_rir(t60=None):
+    """A room without a measurement: exponentially decaying Gaussian noise with a sparser first 20 ms and a
+    direct path `drr` dB over the tail; T60 0.15–2.5 s."""
+    t60 = t60 or lu(0.15, 2.5)
+    n = int(min(t60 * 1.2, 3.0) * SR)
+    t = torch.arange(n) / SR
+    h = torch.randn(n) * torch.exp(-6.908 * t / t60)
+    early = int(0.02 * SR)
+    h[:early] *= (torch.rand(early) < u(0.02, 0.2)).float() * 3
+    h[0] = 0
+    drr = u(-6, 12)
+    h = h * (h[1:].pow(2).sum() + 1e-9).rsqrt() * db(-drr)
+    h[0] = 1.0
+    if chance(0.5):  # air and walls take the top
+        h = spectral(h, (-u(2, 10)) * torch.log2(freqs(n).clamp(min=1000) / 1000))
+    return h
+
+def room(x, h):
+    """x in the room h, the direct path's level kept: y scaled so its energy equals x's."""
+    n = len(x) + len(h) - 1
+    y = torch.fft.irfft(torch.fft.rfft(x, n) * torch.fft.rfft(h, n), n)[: len(x)]
+    return y * (power(x) / power(y)).sqrt()
+
+def synth_noise(n):
+    """Coloured noise at a random slope, sometimes with mains hum and its harmonics."""
+    f = freqs(n)
+    y = spectral(torch.randn(n), -u(0, 12) * torch.log2(f.clamp(min=20) / 20))
+    if chance(0.3):
+        t = torch.arange(n) / SR; f0 = random.choice((50, 60))
+        for k in range(1, random.randint(2, 12)): y = y + torch.sin(2 * math.pi * f0 * k * t + u(0, 6.3)) * db(-u(0, 30)) / k
+    return y
+
+def noise(x, z, snr):
+    z = z * (power(x) / power(z) / db(snr) ** 2).sqrt()
+    return x + z
+
+def compress(x, ratio, thresh_db, attack=0.005, release=0.1):
+    """A preamp's or a phone's compressor: gain from a one-pole-smoothed envelope, over `thresh_db` reduced by ratio."""
+    env = x.abs()
+    a, r = math.exp(-1 / (attack * SR)), math.exp(-1 / (release * SR))
+    one = lambda c: AF.lfilter(env[None], torch.tensor([1.0, -c]), torch.tensor([1 - c, 0.0]), clamp=False)[0]
+    e = torch.maximum(one(r), one(a))
+    edb = 20 * torch.log10(e + 1e-6)
+    over = (edb - thresh_db).clamp(min=0)
+    g = db(-over * (1 - 1 / ratio))
+    y = x * g
+    return y * (power(x) / power(y)).sqrt()
+
+def clip(x):
+    peak = x.abs().max() + 1e-9
+    th = peak * u(0.15, 0.7)
+    if chance(0.5): return x.clamp(-th, th)
+    return torch.tanh(x / th) * th
+
+def bandlimit(x, hz):
+    return band(x[None], torch.tensor([hz]))[0]
+
+def codec(x):
+    """x through a real codec at a low bitrate: the file written, encoded, decoded back to 48 kHz by ffmpeg."""
+    name, args, ext, rates = random.choice(CODECS)
+    kb = random.choice(rates)
+    with tempfile.TemporaryDirectory() as d:
+        src, enc = os.path.join(d, 'in.wav'), os.path.join(d, 'enc.' + ext)
+        import soundfile as sf
+        sf.write(src, x.numpy(), SR, subtype='FLOAT')
+        q = ['-loglevel', 'error', '-y']
+        if subprocess.run([FFMPEG, *q, '-i', src, *args, '-b:a', f'{kb}k', enc]).returncode: return x
+        out = subprocess.run([FFMPEG, *q, '-i', enc, '-f', 'f32le', '-ac', '1', '-ar', str(SR), '-'], capture_output=True)
+        if out.returncode: return x
+        y = torch.frombuffer(bytearray(out.stdout), dtype=torch.float32)
+    # the codec's delay: the lag, under half a second, where y best meets x (cross-correlation by FFT)
+    k = min(len(x), len(y), SR * 2)
+    c = torch.fft.irfft(torch.fft.rfft(y[:k], 2 * k) * torch.fft.rfft(x[:k], 2 * k).conj(), 2 * k)[: SR // 2]
+    lag = int(c.argmax())
+    y = y[lag: lag + len(x)]
+    return F.pad(y, (0, len(x) - len(y))) if len(y) < len(x) else y
+
+def quantize(x):
+    bits = random.choice((8, 12, 16))
+    q = 2 ** (bits - 1)
+    return torch.round(x * q + (torch.rand_like(x) - 0.5)) / q
+
+# ---------------------------------------------------------------- the chain
+
+class Chain:
+    """P: the chance of each stage; `rirs`, `noises`: lists of (N,) tensors at 48 kHz, or empty for synthetic."""
+    P = dict(mic=0.7, room=0.6, noise=0.8, compress=0.3, clip=0.2, bandlimit=0.6, codec=0.3 if FFMPEG else 0, quantize=0.3)
+
+    def __init__(self, rirs=(), noises=(), p=None):
+        self.rirs, self.noises = list(rirs), list(noises)
+        self.p = {**self.P, **(p or {})}
+
+    def pick(self, bank, n):
+        z = random.choice(bank)
+        if len(z) < n: z = z.repeat(n // len(z) + 1)
+        at = random.randint(0, len(z) - n)
+        return z[at: at + n]
+
+    def __call__(self, x):
+        p = self.p
+        y = x
+        if chance(p['room']):
+            h = random.choice(self.rirs) if self.rirs and chance(0.7) else synth_rir()
+            y = room(y, h)
+        if chance(p['noise']):
+            z = self.pick(self.noises, len(y)) if self.noises and chance(0.85) else synth_noise(len(y))
+            y = noise(y, z, u(-5, 35))
+        if chance(p['mic']): y = mic(y)
+        if chance(p['compress']): y = compress(y, lu(2, 10), u(-35, -15))
+        if chance(p['clip']): y = clip(y)
+        if chance(p['bandlimit']): y = bandlimit(y, lu(2000, 20000))
+        if chance(p['codec']): y = codec(y)
+        if chance(p['quantize']): y = quantize(y)
+        return y