From 2c9c87a1088472f8e8dc97f57d4766d1cdee3c9d Mon Sep 17 00:00:00 2001 From: audiojs Date: Wed, 14 Oct 2026 21:18:41 +0000 Subject: [PATCH 1/4] restore(): the plan and the training side of a studio-grade voice restorer (.work/restore.md, train/restore/) Prior art checked (FINALLY's recipe, Sidon MIT as the teacher, AnyEnhance, UniPASE, Vocos BWE, URGENT 2026), the open-license data that exists for English, Russian and singing and what does not, the design: one pass in the STFT, ConvNeXt over frames with self-attention for the room's tail, a change of log magnitude and a phase, trained as a regressor then adversarially with WavLM's features in the loss; every loss through a band limit at the target's own bandwidth, so 24 and 16 kHz corpora train beside full-band ones in one run. train/restore: the on-the-fly degradation chain (rooms measured or synthetic, noise, microphone, compressor, clipping, band limit, codecs through ffmpeg, word length), manifests with each file's bandwidth, the generator and discriminators, losses, the run with resume and validation, ONNX export in neural-separate's contract, chunked inference, scoring on bench/speech.py's measures plus LSD, and a CPU smoke test. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01MUVeWzmg33J7gc9R3MWMYh --- .work/restore.md | 193 +++++++++++++++++++++++++++++++++++++++ .work/todo.md | 19 ++-- train/restore/README.md | 30 ++++++ train/restore/data.py | 95 +++++++++++++++++++ train/restore/degrade.py | 163 +++++++++++++++++++++++++++++++++ train/restore/dsp.py | 80 ++++++++++++++++ train/restore/export.py | 49 ++++++++++ train/restore/infer.py | 66 +++++++++++++ train/restore/loss.py | 52 +++++++++++ train/restore/model.py | 137 +++++++++++++++++++++++++++ train/restore/score.py | 44 +++++++++ train/restore/test.py | 76 +++++++++++++++ train/restore/train.py | 144 +++++++++++++++++++++++++++++ 13 files changed, 1138 insertions(+), 10 deletions(-) create mode 100644 .work/restore.md create mode 100644 train/restore/README.md create mode 100644 train/restore/data.py create mode 100644 train/restore/degrade.py create mode 100644 train/restore/dsp.py create mode 100644 train/restore/export.py create mode 100644 train/restore/infer.py create mode 100644 train/restore/loss.py create mode 100644 train/restore/model.py create mode 100644 train/restore/score.py create mode 100644 train/restore/test.py create mode 100644 train/restore/train.py diff --git a/.work/restore.md b/.work/restore.md new file mode 100644 index 0000000..9d326d5 --- /dev/null +++ b/.work/restore.md @@ -0,0 +1,193 @@ +# restore(): a studio-grade voice from a bad recording + +Plan and findings, 2026-10-14. The training side lives in [`train/restore/`](../train/restore/README.md); the op and +its package follow once weights exist. "Verified" marks what was checked at the source this session; "recalled" +marks training knowledge not re-checked (the container could not fetch arxiv, openreview or most paper mirrors). + +## Goal + +An offline alternative to Adobe Podcast Enhance Speech (v2): a voice recorded in a room, in noise, on a laptop or +phone microphone, through a codec, speech or singing, English or Russian, comes back as a studio would have +recorded it, 48 kHz, full band. Small (20 MB or less in fp16), one pass on the GPU (no diffusion steps, no +autoregression), real-time factor under 0.05 on a consumer GPU, runnable in the page on WebGPU and in Node. + +What this is not: a denoiser. `deepfilter()` (DeepFilterNet3, 8 MB) already beats iZotope RX 12 Dialogue Isolate +on VoiceBank+DEMAND, and `dereverb()` (WPE) beats RX De-reverb (docs/rx.md). Masking keeps what is there and takes +what is not; it cannot give back a band a phone never recorded, a consonant a codec smeared, or a voice under a +long tail. That is generation, and it needs an adversarially trained model (or a diffusion / token model, which +is slow). Adobe's is that class: on real takes its output has bandwidth and texture the input never had. + +## What exists, measured (reused here) + +- `@audio/neural-denoise`: DeepFilterNet3 through `@audio/neural-runtime` (ONNX Runtime, wasm / WebGPU / Node), + 10 s chunks, the input mixed back at a limit (`fn/deepfilter.js`'s `enhancer()`): the op shape `restore()` takes. +- `@audio/neural-separate`: the STFT-outside-the-graph contract (`[1, 2, F, T]` re/im in and out, `export-roformer.py`), + the Mel-RoFormer WebGPU engine (subgroup matrices, 1.5–1.8 s an 8 s segment) and the fp16 hosting on Hugging Face + (`audiojs/mel-roformer`). The restorer exports to the same contract, so the JS side is the same code. +- `bench/speech.py` and `bench/speech.mjs`: PESQ, STOI, SI-SDR, DNSMOS at 16 kHz on VoiceBank+DEMAND, MIT IR Survey + rooms (`vbreverb`), ten narrations; `bench/rx/*.mjs` the same against RX 12. `train/restore/score.py` reads them + and adds LSD at 48 kHz, the bandwidth measure. +- `dereverb()`'s room sets (`bench/rx/dereverb.py`: MIT IR Survey, DEMAND, VocalSet) and `auto()`'s multi-defect + takes (`bench/rx/assistant.mjs`) are evaluation sets ready to run. + +## Prior art, checked 2026-10-14 + +| System | What | Verified | Use here | +|---|---|---|---| +| FINALLY (Samsung, NeurIPS 2024) | HiFi++ generator with a WavLM encoder, GAN over MS-STFT discriminators, LMOS (WavLM-conv features + STFT L1), UTMOS and PESQ losses in the last stage, 48 kHz out; three stages: content on LibriTTS-R at 16 kHz, then an upsampler and adversarial training on DAPS-clean at 48 kHz; noise from DNS; claims studio-like quality on real takes, one pass | paper found, training data and staging confirmed from the proceedings text; no code, weights or license found (demo page only) | the recipe (a speech model's features as the loss that carries quality through GAN training; regression first, adversarial after). DAPS is CC BY-NC-SA: not for our targets | +| Sidon v0.1 (U. Tokyo, 2025, MIT) | w2v-BERT 2.0 predictor + DAC decoder, 16 kHz in → 48 kHz out, multilingual (dozens of languages), 470 MB, 500× real time on a GPU; made for dataset cleansing | repo MIT, two models released; corpora not named; weights' own license not stated | the teacher: cleans Russian audiobooks and separated vocals into 48 kHz targets; a quality ceiling to measure against; too big to ship | +| Resemble Enhance (2023, MIT) | UNet denoiser + latent conditional flow matching enhancer, 44.1 kHz | MIT confirmed from several sources | second teacher / baseline; multi-step, not the shape we ship | +| AnyEnhance (CUHK-SZ / Huawei, 2025) | masked generative model on DAC tokens for speech and singing (denoise, dereverb, declip, super-resolution, target speaker) | code MIT, a 45.7M baseline checkpoint released (trained on 2×A800, 200k steps), ~200 h paired data | the one open restorer that names singing; token models are slower and hallucinate more: a baseline, not the design | +| UniPASE (2026) | generative universal SE "with low hallucinations": a Vocos-style iSTFT vocoder at 16 kHz and a 48 kHz PostNet for bandwidth | abstract and architecture summary only | confirms the Vocos backbone as an enhancement generator | +| "Fast and flexible audio bandwidth extension via Vocos" (Sharma, 2026) | one Vocos network for 8–48 kHz input, a refiner blends the low band back; RTF 0.0001 on an A100, 0.0053 on an 8-core CPU | abstract only, code link not retrieved | same backbone, the speed figure we aim at | +| URGENT 2026 (ICASSP) | universal SE across distortions and rates; 2025's finding: purely generative models show language dependency; 2026 adds whispered and singing voices, ages, accents | overview papers found; Russian not confirmed in its data | the finding is the reason for Russian targets; its data lists are a source map | +| Adobe Enhance Speech v2 | web app, March 2026 release added source separation; reviewers: much better than v1, adds bass, still flattens voices and damages consonants at times; RX 11 and Waves Clarity Vx preferred for natural tone in several tests | reviews only | the reference set: our test takes through it, by hand (no API) | + +Recalled, not re-checked: VoiceFixer (2021, mel + TFGAN vocoder, 44.1 kHz, slow), Miipher / Miipher-2 (Google, not +open; LibriTTS-R is Miipher's output), MaskSR / Genhancer (token models), VoiceRestore (flow matching on mel + +BigVGAN), UNIVERSE++ (Dolby, score-based). + +## Design + +One network, one pass, in the STFT domain, trained first as a regressor and then adversarially with a speech +model's features in the loss. Reasons, against the alternatives: + +- Not DeepFilterNet3 → Vocos in series. DFN3 is a masker trained to erase; what it erases (breath, held-note texture, + Russian fricatives it takes for hiss) a second net must invent back. Two models, two error surfaces, two + downloads, and the second cannot see what the first threw away. One net trained end to end on the same + degradations learns to keep what is signal and to make what is missing. (DFN3 stays the cheap tier for noise alone.) +- Not flow matching or tokens (Resemble Enhance, AnyEnhance, UniPASE's generative core). They give the best + worst-case restoration but need 8–32 passes or an autoregressive loop, hallucinate words and speaker identity + at low SNR (URGENT's finding), and the token codecs cap singing. A GAN-trained one-pass model is what FINALLY + shows reaches "studio-like" on real takes, and what fits a browser. +- Backbone: Vocos's ConvNeXt stack over frames (dim 512, 8 blocks, kernel 7, ×3 MLP) with two self-attention + layers (rotary positions) after blocks 3 and 6: the convolutions see ~0.5 s, the attention the whole 10 s chunk, + which a room's tail needs. ~17M parameters; no recurrence, so the GPU spreads it over time (SCNet's LSTMs are + what made our split slow, .work/todo.md). STFT 2048/512 at 48 kHz (1025 bins, 93.75 frames/s). +- Output: a change of log magnitude added to the input's (unbounded, so a band can be attenuated 60 dB or raised + from nothing) and a phase as a unit vector (Vocos's head without the atan2). The STFT and iSTFT stay outside the + graph (neural-separate's contract), so the export is the network alone. +- Level: the net hears every input at −20 dBFS RMS (as `deepfilter()` has DFN3 hear it) and the target is scaled + by the same gain: the output keeps the input's level, which `mix` and `auto()` need; loudness is the next stage's. +- Losses: log-magnitude L1 + spectral convergence over 512/1024/2048/4096 windows, log-mel L1 (×45), WavLM Base+ + conv features and layer 6 L1 (FINALLY's LMOS made of a 95M MIT model that torchaudio ships), then least-squares + GAN over HiFi-GAN's multi-period (2, 3, 5, 7, 11) and UnivNet's multi-resolution discriminators with L1 + feature matching (×2). Pitch loss: none; the period discriminators and WavLM carry f0 (CREPE-feature losses + exist in vocoder work; add only if vibrato measurably drifts). +- Mixed-bandwidth targets, the novelty worth keeping: every loss and discriminator sees the target and the + estimate through the same low-pass at the target's measured bandwidth (`dsp.band`, `dsp.bandwidth`). LibriTTS-R + is 24 kHz, Russian corpora 16 kHz, VCTK and VocalSet full band: a 24 kHz target judges nothing above 12 kHz and + the full-band ones teach the top. FINALLY needed three stages and two corpora for this; here it is one run. +- Degradations on the fly, in the physical order (`degrade.py`): room (measured IRs 70 %, else an exponential + tail with a sparse early part, T60 0.15–2.5 s, DRR −6…12 dB), noise (banks or coloured noise and mains hum, + SNR −5…35 dB), microphone (tilt ±4 dB/oct, 1–3 resonances, low cut 40–300 Hz), compressor, clipping, band + limit (2–20 kHz, log-uniform), codec (MP3, AAC, Opus, GSM through ffmpeg at 8–128 kb/s, the codec's lag removed), + 8/12/16-bit words with dither; chances 0.2–0.8 a stage. Chunks of 4 s, levels −35…−12 dBFS. + +## Data + +Targets must be studio-clean and, for the top octave, full band; the license must allow hosting weights for anyone +(the audiojs policy: audited before any promise). Checked this session: + +| Source | Hours | Rate | License | Status | +|---|---|---|---|---| +| VCTK | 44 | 48 kHz | CC BY 4.0 | use; its room and noise are mild but real: teacher-clean it (Sidon) or filter by DNSMOS | +| LibriTTS-R | 585 | 24 kHz | CC BY 4.0 per OpenSLR (one arXiv entry says CC BY-NC-ND: read the download page) | use, bandwidth 12 kHz; Miipher-restored, so already a teacher's output | +| EARS | 100 | 48 kHz | CC BY-NC 4.0 (secondhand) | not for hosted weights; evaluation only | +| DAPS | 4.5 | 44.1 kHz | CC BY-NC-SA | evaluation only (FINALLY's 48 kHz targets) | +| Hi-Fi TTS (NVIDIA) | 292 | 44.1 kHz | not checked | check: LibriVox-derived, likely CC BY-SA-ish | +| Golos (Sber, Russian) | 1240 | 16 kHz | "close to CC BY-SA", research and commercial use allowed (SberDevices), ShareAlike | use after teacher cleaning: crowd and far-field, not studio; ShareAlike may bind the weights' license: read `license/en_us.pdf` | +| RuLS (Russian LibriSpeech, OpenSLR 96) | 98 | 16 kHz | public domain in the USA (OpenSLR), CC BY 4.0 on a mirror | use after teacher cleaning (LibriVox amateurs' rooms and microphones) | +| M-AILABS Russian | ~46 | 16 kHz | not retrieved | check | +| RUSLAN, Natasha (Russian TTS corpora) | 31, ? | ? | RUSLAN recalled CC BY-NC-SA | likely evaluation only | +| Balalaika (Russian, 2025/26) | 1000 | ? | CC BY-NC-SA / NC-ND (sources disagree) | not for hosted weights | +| Dialogs (Russian studio, 2026) | 20.6 | 44.1 kHz | OpenRAIL | read its use restrictions before use | +| VocalSet | 10.1 | 44.1 kHz | CC BY 4.0 (the paper; confirm on Zenodo) | use: the one full-band singing set with a permissive license | +| GTSinger | 80.6 | ? | CC BY-NC-SA 4.0 | not for hosted weights (it has Russian singing: evaluation set for exactly that) | +| MUSDB18-HQ vocals, M4Singer, OpenSinger, Opencpop | — | — | research-only / NC / agreement | evaluation only | +| CC BY and CC0 music (Jamendo, FMA, ccMixter) → `vocals({ model: 'mel-roformer' })` | unbounded | 44.1 kHz | per track | the singing data factory: separate, teacher-clean, keep the clips a quality gate passes | +| MIT IR Survey | 271 IRs | 32 kHz | unspecified (a mirror says unknown) | already used in benches; for training prefer the clearly CC BY sets below | +| OpenAIR, BUT ReverbDB (1300+ and 2325 IRs) | — | — | CC BY / CC BY 4.0 | use; plus synthetic rooms (unlimited) | +| DEMAND | 18 scenes | 48 kHz | CC BY-SA 3.0 (a mirror) | use for training noise (ShareAlike again), test noise as benches do | +| MUSAN | 109 | 16 kHz | CC BY 4.0 overall, per-file LICENSE files, commercial use allowed by selection | use (noise and music beds) | +| FSD50K | 108 | 44.1 kHz | CC BY 4.0 overall; 84.7 % of clips CC0 or CC BY, the rest NC | use the CC0 + CC BY clips | +| WHAM!, DNS Challenge noise and IRs | — | — | not retrieved (WHAM! recalled CC BY-NC) | check before use | + +What this means in hours, after cleaning and gating: English ~300–600 h (LibriTTS-R, VCTK, maybe Hi-Fi TTS), of +it ~50 h full band; Russian ~100–300 h at 16 kHz through the teacher (its top octave then is the teacher's); +singing 10 h full band (VocalSet) plus whatever the separation factory yields (aim 50–100 h). The pasted +strategy's "150 h of studio Russian" and "100–200 h of singing" do not exist under an open license; GTSinger, +M4Singer, OpenSinger, MUSDB18-HQ, EARS and RUSLAN are all non-commercial, and scraping LitRes is a copyright +breach. The teacher route (Sidon, MIT) is how both gaps close, and it is what Google did for LibriTTS-R. + +Teacher cleaning, concretely: run Sidon over Golos's crowd subset, RuLS and the separated vocals; keep a clip when +DNSMOS OVRL ≥ 3.3, the teacher's output stays within 0.5 of the input's WER (neural-asr, so no words are +invented) and speaker similarity ≥ 0.85 (neural-diarize's WavLM embeddings); drop the rest. The student then +trains on teacher outputs the way LibriTTS-R is Miipher's: distillation by data, no teacher at training time. + +## Training and compute + +- Steps: 30k regression warm-up, then adversarial to 400k; batch 16 × 4 s at 48 kHz. Memory with the WavLM loss + and 8 discriminators: ~18 GB at bf16. A 24 GB card (RTX 4090 / 3090, A5000) does it; a 16 GB card at batch 8. +- Time: the generator is ~17M parameters over 375 frames a clip, the discriminators and WavLM dominate. Estimate + 2.5–3.5 steps/s on a 4090 → 400k steps in 1.5–2 days; on an A100 about the same; on the M4 Max (MPS) 5–10× + slower and `torch.stft` / `istft` on MPS have had gaps: debug there, train on a rented card. Cloud: a 4090 at + $0.4–0.7/h on Vast or RunPod → $20–60 a run; budget five runs (losses, data mix, size) → under $300. +- The teacher pass: Sidon at 500× real time cleans 1000 h in a few GPU hours. Separation of music for singing + targets: mel-roformer at RTF 0.44 on the laptop's WebGPU, 100 h of music a day; on a CUDA card far less. +- Smaller sibling for the page: dim 384, 6 blocks (~8M, ~16 MB fp16) trained the same way; measure what it loses. +- Checked to work before any GPU is paid for: `train/restore/test.py` runs the chain, a training run with resume, + ONNX export matching torch, inference in chunks, on CPU with synthetic voices. + +## Evaluation + +The sets exist in `bench/`: VoiceBank+DEMAND (noise), vbreverb (MIT rooms), the ten narrations, `bench/rx/ +assistant.mjs`'s multi-defect takes, `dereverb.py`'s sung and music rooms. Add: phone and laptop-mic takes (record +them: no open set), Russian (GTSinger's Russian songs and Balalaika's or RUSLAN's clips are fine as test material), +and the same takes through Adobe Enhance v2 by hand, as `bench/rx` does with RX. Measures: PESQ, STOI, SI-SDR, +DNSMOS (there already), LSD for bandwidth, WER through `neural-asr` (a restorer that invents consonants raises it), +speaker similarity, UTMOS; and a blind listen, which decides. The policy from `deepfilter()` holds: DNSMOS SIG +must not fall on clean input, VocalSet's held notes must not go down, music alone must pass. + +## Shipping + +- Export: `export.py` → `restore.onnx` (fp32) and `restore.fp16.onnx` (~34 MB), `[1, 2, 1025, T]` in and out, + verified against torch; hosted as `audiojs/restore` on Hugging Face, pinned by revision as neural-separate pins. +- Package: `@audio/neural-restore` in audiojs/neural: STFT 2048/512 (JS, as neural-separate's), the net through + `neural-runtime` (WebGPU where there is one, Node's providers else), 10 s chunks with 1 s crossfades, the input at + −20 dBFS RMS and back. The Mel-RoFormer WGSL engine's matmul and attention kernels apply if ONNX Runtime is too + slow; the ConvNeXt blocks add depthwise convolutions and LayerNorm. +- Op: `restore(strength?, { model?, weights?, device? })` in `fn/restore.js` on `fn/deepfilter.js`'s `enhancer()`: + the whole output made ahead of rendering, memoized on the input; `strength` 0–1 mixes the input back (1: the + model's output; as `deepfilter`'s limit). `auto('speech')` gains it as the last repair stage when installed. + Playground: a "Restore" entry beside Enhance; the CLI the same. + +## What the pasted strategy got right and wrong + +Right: the problem class (generative restoration and bandwidth extension, not masking), synthetic degradation of +clean targets, 48 kHz, multi-resolution STFT + mel + MPD/MSD losses, 15–40M parameters, one pass, fine-tune +cost in the tens of dollars. Wrong or unsupported: the data list (GTSinger, M4Singer, OpenSinger, MUSDB18-HQ, EARS, +RUSLAN are non-commercial; LitRes scraping is not a source); "150 h studio Russian" (none open; Golos is crowd +and far-field, 16 kHz); DFN3 → Vocos as two models (see Design); "Vocos is orders of magnitude faster than +wave-domain GANs" (true against HiFi-GAN's transposed convolutions, and why we use its backbone, but FINALLY's +HiFi++ is also one pass); a pitch loss as a fix for vibrato (unmeasured; MPD and WavLM already hold f0); and the +most important ingredient was missing: the speech-model feature loss (LMOS) that FINALLY shows turns adversarial +training from "clean" into "studio". The Russian fricative worry is plausible but unmeasured: test it (WER and a +listen on щ / ц / ч words) before building for it. + +## Risks + +- GAN training can collapse or ring: the warm-up and feature matching are the usual guards; watch `samples/`. +- Hallucination at low SNR (the generative tax): the WavLM loss and the WER gate in evaluation; `strength` mixes + the input back where the user prefers damage to invention. +- Singing: VocalSet alone is 20 singers; the separation factory decides whether singing is good or merely not broken. +- Russian top octave comes from the teacher, not from recordings: listen for a "Sidon accent" in sibilants. +- Licenses: three "verify" rows above before any weights are hosted; ShareAlike (Golos, DEMAND) may bind the + weights' license to CC BY-SA, which is still hostable. + +## Next + +1. On a GPU machine: fetch VCTK, LibriTTS-R, VocalSet, OpenAIR + BUT IRs, MUSAN + FSD50K CC0/BY; `data.py scan`; + a 50k-step run at dim 384 to validate the recipe on VoiceBank+DEMAND and vbreverb (`score.py`); then full size. +2. Sidon installed; the Russian and separated-vocal teacher pass with the quality gate; rerun with them. +3. Export, host, `@audio/neural-restore`, `restore()`; benches against `deepfilter()`, `auto()`, Adobe v2. diff --git a/.work/todo.md b/.work/todo.md index 8645fe4..ba61e2b 100644 --- a/.work/todo.md +++ b/.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) diff --git a/train/restore/README.md b/train/restore/README.md new file mode 100644 index 0000000..82149a9 --- /dev/null +++ b/train/restore/README.md @@ -0,0 +1,30 @@ +# 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). + +``` +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 +``` + +| 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) | +| `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 | diff --git a/train/restore/data.py b/train/restore/data.py new file mode 100644 index 0000000..249aec6 --- /dev/null +++ b/train/restore/data.py @@ -0,0 +1,95 @@ +"""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) + +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']); a.add_argument('out'); a.add_argument('dirs', nargs='+') + a.add_argument('--kind', default='speech', choices=['speech', 'sing']); a.add_argument('--lang', default='en') + o = a.parse_args() + scan(o.out, o.dirs, o.kind, o.lang) diff --git a/train/restore/degrade.py b/train/restore/degrade.py new file mode 100644 index 0000000..0426697 --- /dev/null +++ b/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 diff --git a/train/restore/dsp.py b/train/restore/dsp.py new file mode 100644 index 0000000..fca7ce1 --- /dev/null +++ b/train/restore/dsp.py @@ -0,0 +1,80 @@ +"""Shared signal parts: the STFT the model lives between, levels, band limits, mel, resampling. + +Everything is torch, batched, device-agnostic. Tensors are (B, N) waveforms unless said. +""" +import math +import torch +import torch.nn.functional as F +import torchaudio.functional as AF + +SR = 48000 # the model's rate: everything in and out at 48 kHz +N_FFT, HOP = 2048, 512 # 1025 bins of 23.4 Hz, 93.75 frames a second +BINS = N_FFT // 2 + 1 +LEVEL = -20.0 # dBFS RMS the model hears its input at +EPS = 1e-5 + +_win = {} +def window(device): + if device not in _win: _win[device] = torch.hann_window(N_FFT, device=device) + return _win[device] + +def stft(x): + """(B, N) → complex (B, F, T), centered, reflect-padded: torch's convention, which the JS side mirrors.""" + return torch.stft(x, N_FFT, HOP, N_FFT, window(x.device), center=True, pad_mode='reflect', return_complex=True) + +def istft(X, length): + return torch.istft(X, N_FFT, HOP, N_FFT, window(X.device), center=True, length=length) + +def rms_db(x): + return 10 * torch.log10(x.pow(2).mean(-1) + 1e-12) + +def to_level(x, db=LEVEL): + """x scaled to `db` dBFS RMS; returns (y, gain) so the result can be undone.""" + g = 10 ** ((db - rms_db(x)) / 20) + return x * g[..., None], g + +def band(x, hz): + """x low-passed at `hz` (B,) with a 1/6-octave raised-cosine edge, zero-phase (FFT over the whole segment). + hz ≥ SR/2 leaves x as it is. Both target and estimate pass through it: a target recorded at 24 kHz holds + nothing above 12 kHz, so what the model invents there must be kept out of every loss.""" + n = x.shape[-1] + X = torch.fft.rfft(x, dim=-1) + f = torch.fft.rfftfreq(n, 1 / SR, device=x.device) + lo, hi = hz[:, None] * 2 ** (-1 / 12), hz[:, None] * 2 ** (1 / 12) + m = ((hi - f) / (hi - lo)).clamp(0, 1) + m = 0.5 - 0.5 * torch.cos(math.pi * m) + m = torch.where(hz[:, None] >= SR / 2, torch.ones_like(m), m) + return torch.fft.irfft(X * m, n, dim=-1) + +def bandwidth(x, sr=SR, floor_db=55.0): + """The highest frequency a recording carries, Hz: where the long-term spectrum has fallen `floor_db` under + its 1–4 kHz level and stays there. Reads a 24 kHz file as ≤ 12 kHz and a band-limited 48 kHz one as its edge.""" + x = x.reshape(-1) + n = 4096 + if len(x) < n: return sr / 2 + frames = x[: len(x) // n * n].reshape(-1, n) * torch.hann_window(n) + psd = torch.fft.rfft(frames).abs().pow(2).mean(0) + db = 10 * torch.log10(psd + 1e-20) + f = torch.fft.rfftfreq(n, 1 / sr) + ref = db[(f >= 1000) & (f <= 4000)].mean() + k = 12 # bins smoothed, 140 Hz at 48 kHz + sm = F.avg_pool1d(db[None, None], k, 1, k // 2)[0, 0, : len(db)] + below = (sm < ref - floor_db).int().flip(0).cumprod(0).flip(0).bool() # true from the edge to the top + idx = below.nonzero() + return float(f[idx[0, 0]]) if len(idx) else sr / 2 + +_mel = {} +def mel(x, n_mels=128): + """log mel power of (B, N) at 48 kHz, (B, n_mels, T).""" + key = (n_mels, x.device) + if key not in _mel: + _mel[key] = AF.melscale_fbanks(BINS, 0, SR / 2, n_mels, SR, norm='slaney', mel_scale='slaney').to(x.device) + S = stft(x).abs().pow(2) + return torch.log(torch.matmul(S.transpose(1, 2), _mel[key]).transpose(1, 2) + EPS) + +def resample(x, sr, to=SR): + if sr == to: return x + return AF.resample(x, sr, to, lowpass_filter_width=64, rolloff=0.9475937167399596, resampling_method='sinc_interp_kaiser', beta=14.769656459379492) + +def fade(n, device): + return 0.5 - 0.5 * torch.cos(math.pi * torch.arange(n, device=device) / n) diff --git a/train/restore/export.py b/train/restore/export.py new file mode 100644 index 0000000..18a475f --- /dev/null +++ b/train/restore/export.py @@ -0,0 +1,49 @@ +"""The trained net as ONNX, between the STFT and the iSTFT the JS side computes. + + python export.py runs/a/net-0400000.pt --out ~/.cache/audiojs/neural/restore [--fp16] [--verify] + + input spec [1, 2, 1025, T] STFT of one channel at 48 kHz (n_fft 2048, hop 512, Hann, centered, reflect), re and im + output spec [1, 2, 1025, T] the restored STFT +The frame axis is dynamic; the model hears its input at −20 dBFS RMS (dsp.to_level), the JS side undoes the gain. +--fp16 writes a second file with float16 weights (half the download; computed in float32 by onnxruntime's +WebGPU as neural-separate's mel-roformer is). --verify runs onnxruntime on a random segment against torch. +""" +import argparse, os +import numpy as np +import torch +from model import Net, Graph, params +from dsp import BINS + +def load(path, dev='cpu'): + ck = torch.load(path, map_location=dev) + net = Net(**ck['cfg']); net.load_state_dict(ck['net']); return net.eval() + +def main(): + a = argparse.ArgumentParser() + a.add_argument('ckpt'); a.add_argument('--out', default=os.path.join(os.environ.get('AUDIO_NEURAL_CACHE') or os.path.expanduser('~/.cache/audiojs/neural'), 'restore')) + a.add_argument('--fp16', action='store_true'); a.add_argument('--verify', action='store_true'); a.add_argument('--frames', type=int, default=376) + o = a.parse_args() + os.makedirs(o.out, exist_ok=True) + net = load(o.ckpt); g = Graph(net).eval() + print(f'{params(net) / 1e6:.2f}M parameters') + x = torch.randn(1, 2, BINS, o.frames) * 0.01 + path = os.path.join(o.out, 'restore.onnx') + torch.onnx.export(g, (x,), path, input_names=['spec'], output_names=['spec_out'], dynamic_axes={'spec': {3: 'T'}, 'spec_out': {3: 'T'}}, opset_version=17, dynamo=False) + print('wrote', path, os.path.getsize(path) // 1024, 'KB') + if o.fp16: + import onnx + from onnxconverter_common import float16 + m = float16.convert_float_to_float16(onnx.load(path), keep_io_types=True) + p16 = os.path.join(o.out, 'restore.fp16.onnx'); onnx.save(m, p16); print('wrote', p16, os.path.getsize(p16) // 1024, 'KB') + if o.verify: + import onnxruntime as ort + s = ort.InferenceSession(path, providers=['CPUExecutionProvider']) + for T in (o.frames, 97): + x = torch.randn(1, 2, BINS, T) * 0.01 + with torch.no_grad(): ref = g(x).numpy() + out = s.run(None, {'spec': x.numpy()})[0] + err = np.abs(out - ref).max(); snr = 10 * np.log10((ref ** 2).sum() / ((out - ref) ** 2).sum()) + print(f'T={T}: max |Δ| {err:.2e}, {snr:.1f} dB from torch') + +if __name__ == '__main__': + main() diff --git a/train/restore/infer.py b/train/restore/infer.py new file mode 100644 index 0000000..12d62e7 --- /dev/null +++ b/train/restore/infer.py @@ -0,0 +1,66 @@ +"""Restore files with a checkpoint or the ONNX export. + + python infer.py MODEL IN.wav [IN2.wav ...] --out DIR MODEL: net-*.pt or restore.onnx + +Each channel apart: resampled to 48 kHz, brought to −20 dBFS RMS, run in 10 s chunks overlapping by 1 s +(raised-cosine crossfade), the gain undone, resampled back to the input's rate, written as float WAV. +""" +import argparse, os +import numpy as np +import soundfile as sf +import torch +from dsp import SR, stft, istft, to_level, resample, fade + +CHUNK, OVERLAP = 10 * SR, SR + +class Model: + def __init__(self, path, dev): + self.dev = dev + if path.endswith('.onnx'): + import onnxruntime as ort + self.s = ort.InferenceSession(path, providers=['CUDAExecutionProvider', 'CoreMLExecutionProvider', 'CPUExecutionProvider']) + self.net = None + else: + from export import load + self.net = load(path, dev).to(dev) + def __call__(self, x): + X = stft(x[None].to(self.dev)) + if self.net: + with torch.no_grad(): Y = self.net.spec(X) + else: + spec = torch.stack([X.real, X.imag], 1).cpu().numpy() + o = self.s.run(None, {'spec': spec})[0] + Y = torch.complex(torch.from_numpy(o[:, 0]), torch.from_numpy(o[:, 1])).to(self.dev) + return istft(Y, x.shape[-1])[0].cpu() + +def restore(model, x): + xin, g = to_level(x[None]); xin = xin[0] + n = len(xin) + if n <= CHUNK: return model(xin) / g + y, w = torch.zeros(n), torch.zeros(n) + step = CHUNK - OVERLAP; f = fade(OVERLAP, 'cpu') + for at in range(0, n, step): + seg = xin[at: at + CHUNK]; out = model(seg); win = torch.ones(len(seg)) + if at > 0: win[:OVERLAP] = f[: len(seg)] + if at + CHUNK < n: win[-OVERLAP:] = 1 - f + y[at: at + len(seg)] += out * win; w[at: at + len(seg)] += win + if at + CHUNK >= n: break + return y / w.clamp(min=1e-6) / g + +def main(): + a = argparse.ArgumentParser(); a.add_argument('model'); a.add_argument('files', nargs='+'); a.add_argument('--out', required=True); a.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu') + o = a.parse_args(); os.makedirs(o.out, exist_ok=True) + m = Model(o.model, torch.device(o.device)) + for p in o.files: + x, sr = sf.read(p, dtype='float32', always_2d=True) + chans = [] + for c in range(x.shape[1]): + v = resample(torch.from_numpy(x[:, c]), sr) + v = restore(m, v) + chans.append(resample(v, SR, sr)[: len(x)].numpy()) + y = np.stack(chans, 1) + if len(y) < len(x): y = np.pad(y, ((0, len(x) - len(y)), (0, 0))) + sf.write(os.path.join(o.out, os.path.basename(p)), y, sr, subtype='FLOAT'); print(p) + +if __name__ == '__main__': + main() diff --git a/train/restore/loss.py b/train/restore/loss.py new file mode 100644 index 0000000..0dab6c4 --- /dev/null +++ b/train/restore/loss.py @@ -0,0 +1,52 @@ +"""What the restorer is trained toward. + +Regression: log-magnitude L1 and spectral convergence over four STFT resolutions, log-mel L1, and (where its +weights are present) the distance of WavLM's representations (FINALLY's finding, Babaev et al., NeurIPS 2024: +a speech model's features, not a waveform or spectrum distance, carry "studio" quality through adversarial +training). Adversarial: least-squares GAN (HiFi-GAN) over the discriminators' logits and L1 feature matching. + +Every loss sees the estimate and the target through the same band limit (dsp.band at the target's bandwidth): +a target that stops at 12 kHz judges nothing above it. +""" +import torch +import torch.nn.functional as F +import torchaudio +from dsp import SR, mel, resample + +RES = ((512, 128), (1024, 256), (2048, 512), (4096, 1024)) + +def mag(x, n_fft, hop): + w = torch.hann_window(n_fft, device=x.device) + return torch.stft(x, n_fft, hop, n_fft, w, center=True, return_complex=True).abs() + +def mrstft(y, yh): + """Σ over resolutions of L1(log|Y|, log|Ŷ|) + ‖|Y| − |Ŷ|‖_F / ‖|Y|‖_F""" + loss = 0 + for n, h in RES: + Y, Yh = mag(y, n, h), mag(yh, n, h) + loss = loss + F.l1_loss(torch.log(Yh + 1e-5), torch.log(Y + 1e-5)) + (Y - Yh).norm() / (Y.norm() + 1e-5) + return loss / len(RES) + +def mel_l1(y, yh): return F.l1_loss(mel(yh), mel(y)) + +def fm(real, fake): + return sum(F.l1_loss(f, r.detach()) for rs, fs in zip(real, fake) for r, f in zip(rs, fs)) / len(real) + +def gan_g(fake): return sum(((1 - l) ** 2).mean() for l in fake) / len(fake) +def gan_d(real, fake): return sum(((1 - r) ** 2).mean() + (f ** 2).mean() for r, f in zip(real, fake)) / len(real) + +class SSL: + """L1 between WavLM Base+'s convolutional features and its layer `layer` on y and ŷ, at 16 kHz.""" + def __init__(self, device, layer=6): + bundle = torchaudio.pipelines.WAVLM_BASE_PLUS + self.m = bundle.get_model().to(device).eval() + for p in self.m.parameters(): p.requires_grad_(False) + self.layer = layer + def feats(self, x): + x = resample(x, SR, 16000) + c, _ = self.m.feature_extractor(x, None) + return c, self.m.encoder.extract_features(c, None, self.layer)[-1] + def __call__(self, y, yh): + with torch.no_grad(): c, t = self.feats(y) + ch, th = self.feats(yh) + return F.l1_loss(ch, c) + F.l1_loss(th, t) diff --git a/train/restore/model.py b/train/restore/model.py new file mode 100644 index 0000000..3f877f4 --- /dev/null +++ b/train/restore/model.py @@ -0,0 +1,137 @@ +"""The restorer: one pass over the STFT, from a degraded voice's spectrum to a clean one's. + +Between torch's STFT and iSTFT (dsp.py) a ConvNeXt stack over frames (Vocos's backbone: Siuzdak, ICLR 2024) +with a few self-attention layers in it for a room's long memory, reading each frame's log magnitude and phase +and writing a change of log magnitude (so it attenuates what is noise and raises what is missing, both +unbounded) and a new phase as a unit vector. `Net.spec(X) → Y` is the whole network on complex (B, F, T); +`Net(x) → y` wraps it in the STFT for training; `Graph` is the ONNX view, re and im as channels. + +Discriminators for the adversarial phase: HiFi-GAN's multi-period (Kong et al., 2020) and a multi-resolution +spectrogram discriminator (UnivNet, Jang et al., 2021), each returning logits and feature maps. +""" +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.utils.parametrizations import weight_norm +from dsp import BINS, stft, istft, EPS + +class Block(nn.Module): + """ConvNeXt block over (B, T, C): depthwise conv 7, LayerNorm, MLP ×3, GELU, layer scale, residual.""" + def __init__(self, dim, mult=3, scale=1e-6): + super().__init__() + self.dw = nn.Conv1d(dim, dim, 7, padding=3, groups=dim) + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.up, self.down = nn.Linear(dim, dim * mult), nn.Linear(dim * mult, dim) + self.gamma = nn.Parameter(scale * torch.ones(dim)) + def forward(self, h): + r = h + h = self.dw(h.transpose(1, 2)).transpose(1, 2) + h = self.down(F.gelu(self.up(self.norm(h)))) + return r + self.gamma * h + +class Attn(nn.Module): + """Self-attention over frames with rotary positions, pre-norm, residual; the room's whole tail in view.""" + def __init__(self, dim, heads=8, scale=1e-6): + super().__init__() + self.h, self.d = heads, dim // heads + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.qkv, self.out = nn.Linear(dim, dim * 3, bias=False), nn.Linear(dim, dim) + self.gamma = nn.Parameter(scale * torch.ones(dim)) + def rope(self, x): + B, H, T, D = x.shape + pos = torch.arange(T, device=x.device, dtype=torch.float32) + inv = 1.0 / (10000 ** (torch.arange(0, D, 2, device=x.device, dtype=torch.float32) / D)) + ang = pos[:, None] * inv[None] + c, s = ang.cos()[None, None], ang.sin()[None, None] + x1, x2 = x[..., : D // 2], x[..., D // 2:] + return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], -1) + def forward(self, h): + B, T, C = h.shape + q, k, v = self.qkv(self.norm(h)).view(B, T, 3, self.h, self.d).permute(2, 0, 3, 1, 4) + q, k = self.rope(q), self.rope(k) + a = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(self.d), -1) @ v + return h + self.gamma * self.out(a.transpose(1, 2).reshape(B, T, C)) + +class Net(nn.Module): + def __init__(self, dim=512, blocks=8, attn=(3, 6), heads=8): + super().__init__() + self.inp = nn.Linear(3 * BINS, dim) + self.norm = nn.LayerNorm(dim, eps=1e-6) + layers = [] + for i in range(blocks): + layers.append(Block(dim)) + if i + 1 in attn: layers.append(Attn(dim, heads)) + self.layers = nn.Sequential(*layers) + self.final = nn.LayerNorm(dim, eps=1e-6) + self.out = nn.Linear(dim, 3 * BINS) + self.cfg = dict(dim=dim, blocks=blocks, attn=list(attn), heads=heads) + + def spec(self, X): + """complex (B, F, T) → complex (B, F, T)""" + return torch.complex(*self.parts(X.real, X.imag, torch.log(X.abs() + EPS))) + + def parts(self, re, im, logmag): + mag = torch.sqrt(re * re + im * im) + EPS + f = torch.cat([logmag, re / mag, im / mag], 1).transpose(1, 2) # (B, T, 3F) + h = self.final(self.layers(self.norm(self.inp(f)))) + d, px, py = self.out(h).transpose(1, 2).chunk(3, 1) + m = torch.exp((logmag + d).clamp(max=6.0)) + n = torch.sqrt(px * px + py * py) + EPS + return m * px / n, m * py / n + + def forward(self, x): + return istft(self.spec(stft(x)), x.shape[-1]) + +class Graph(nn.Module): + """ONNX contract, as @audio/neural-separate's 'complex' models: in (1, 2, F, T) re/im, out the same.""" + def __init__(self, net): super().__init__(); self.net = net + def forward(self, spec): + re, im = spec[:, 0], spec[:, 1] + return torch.stack(self.net.parts(re, im, 0.5 * torch.log(re * re + im * im + EPS * EPS)), 1) + +# ---------------------------------------------------------------- discriminators + +class Period(nn.Module): + def __init__(self, p, ch=(32, 128, 512, 1024)): + super().__init__() + self.p = p; c = 1; layers = [] + for i, o in enumerate(ch): + layers.append(weight_norm(nn.Conv2d(c, o, (5, 1), (3, 1), (2, 0)))); c = o + layers.append(weight_norm(nn.Conv2d(c, 1024, (5, 1), 1, (2, 0)))) + self.layers = nn.ModuleList(layers) + self.post = weight_norm(nn.Conv2d(1024, 1, (3, 1), 1, (1, 0))) + def forward(self, x): + B, N = x.shape + if N % self.p: x = F.pad(x, (0, self.p - N % self.p), 'reflect') + h = x.view(B, 1, -1, self.p); fm = [] + for l in self.layers: h = F.leaky_relu(l(h), 0.1); fm.append(h) + h = self.post(h); fm.append(h) + return h.flatten(1), fm + +class Resolution(nn.Module): + def __init__(self, n_fft, hop, ch=32): + super().__init__() + self.n_fft, self.hop = n_fft, hop + k, s, p = (3, 9), (1, 2), (1, 4) + self.layers = nn.ModuleList([weight_norm(nn.Conv2d(1, ch, k, 1, p)), weight_norm(nn.Conv2d(ch, ch, k, s, p)), + weight_norm(nn.Conv2d(ch, ch, k, s, p)), weight_norm(nn.Conv2d(ch, ch, k, s, p)), + weight_norm(nn.Conv2d(ch, ch, (3, 3), 1, (1, 1)))]) + self.post = weight_norm(nn.Conv2d(ch, 1, (3, 3), 1, (1, 1))) + def forward(self, x): + w = torch.hann_window(self.n_fft, device=x.device) + S = torch.stft(x, self.n_fft, self.hop, self.n_fft, w, center=True, return_complex=True).abs() + h = torch.log(S + 1e-5)[:, None]; fm = [] + for l in self.layers: h = F.leaky_relu(l(h), 0.1); fm.append(h) + h = self.post(h); fm.append(h) + return h.flatten(1), fm + +class Discriminators(nn.Module): + def __init__(self, periods=(2, 3, 5, 7, 11), resolutions=((512, 128), (1024, 256), (2048, 512))): + super().__init__() + self.d = nn.ModuleList([Period(p) for p in periods] + [Resolution(n, h) for n, h in resolutions]) + def forward(self, x): + out = [d(x) for d in self.d] + return [o[0] for o in out], [o[1] for o in out] + +def params(m): return sum(p.numel() for p in m.parameters()) diff --git a/train/restore/score.py b/train/restore/score.py new file mode 100644 index 0000000..c9583e0 --- /dev/null +++ b/train/restore/score.py @@ -0,0 +1,44 @@ +"""Scores of restored files against their references: bench/speech.py's measures (PESQ, STOI, SI-SDR, DNSMOS) +and LSD at 48 kHz (the bandwidth measure), each file's by name, the mean last. + + python score.py --ref DIR --out DIR [--json OUT.json] +Needs bench/speech.py's dependencies (pesq pystoi onnxruntime pyloudnorm) and its DNSMOS model in the neural cache. +""" +import argparse, json, os, sys, importlib.util +import numpy as np, soundfile as sf +from scipy.signal import resample_poly + +HERE = os.path.dirname(os.path.abspath(__file__)) +spec = importlib.util.spec_from_file_location('speech', os.path.join(HERE, '..', '..', 'bench', 'speech.py')) +speech = importlib.util.module_from_spec(spec); spec.loader.exec_module(speech) + +def lsd(ref, out): + n, h = 2048, 512 + def S(x): + f = np.lib.stride_tricks.sliding_window_view(np.pad(x, n // 2), n)[::h] * np.hanning(n) + return np.log10(np.abs(np.fft.rfft(f)) ** 2 + 1e-9) + a, b = S(ref), S(out); m = min(len(a), len(b)) + return float(np.sqrt(((a[:m] - b[:m]) ** 2).mean(1)).mean()) + +def main(): + a = argparse.ArgumentParser(); a.add_argument('--ref', required=True); a.add_argument('--out', required=True); a.add_argument('--json') + o = a.parse_args(); rows = {} + for name in sorted(os.listdir(o.out)): + if not name.endswith('.wav') or not os.path.exists(os.path.join(o.ref, name)): continue + r, sr = sf.read(os.path.join(o.ref, name), dtype='float64', always_2d=True); y, sr2 = sf.read(os.path.join(o.out, name), dtype='float64', always_2d=True) + r, y = r[:, 0], y[:, 0] + if sr2 != sr: y = resample_poly(y, sr, sr2) + m = min(len(r), len(y)); r, y = r[:m], y[:m] + r16, y16 = speech.to16(r, sr), speech.to16(y, sr) + from pesq import pesq; from pystoi import stoi + row = dict(pesq=pesq(16000, r16, y16, 'wb'), stoi=stoi(r16, y16, 16000), sisdr=speech.sisdr(r16, y16), lsd=lsd(r, y)) + try: row['sig'], row['bak'], row['ovrl'] = speech.dnsmos(y, sr) + except Exception as e: row['dnsmos_error'] = str(e) + rows[name] = row; print(name, json.dumps({k: round(v, 3) for k, v in row.items() if isinstance(v, float)})) + keys = [k for k in next(iter(rows.values())) if isinstance(rows[next(iter(rows))][k], float)] if rows else [] + mean = {k: float(np.mean([r[k] for r in rows.values() if k in r])) for k in keys} + print('mean', json.dumps({k: round(v, 3) for k, v in mean.items()})) + if o.json: json.dump(dict(files=rows, mean=mean), open(o.json, 'w'), indent=1) + +if __name__ == '__main__': + main() diff --git a/train/restore/test.py b/train/restore/test.py new file mode 100644 index 0000000..11124ef --- /dev/null +++ b/train/restore/test.py @@ -0,0 +1,76 @@ +"""Smoke test on synthetic voices, CPU: the chain degrades, the net trains a few steps, exports to ONNX and +matches torch. python test.py [--steps 6]""" +import os, sys, json, subprocess, tempfile, math +import numpy as np, soundfile as sf, torch +HERE = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, HERE) +from dsp import SR, bandwidth, band, stft, istft, to_level +from degrade import Chain, synth_rir, codec, FFMPEG +from model import Net, Graph, Discriminators, params +import loss as L + +def voice(sec, f0=140.0, bw=None, sr=SR): + """a sung vowel: harmonics under a formant envelope, vibrato, breath noise; low-passed at bw""" + t = torch.arange(int(sec * sr)) / sr + f = f0 * (1 + 0.02 * torch.sin(2 * math.pi * 5.5 * t)) + ph = 2 * math.pi * torch.cumsum(f, 0) / sr + x = sum(torch.sin(k * ph) * math.exp(-k / 12) * (1 + 0.5 * math.cos(k * 0.7)) for k in range(1, 120)) + x = x + 0.01 * torch.randn(len(t)); x = x * (0.5 - 0.5 * torch.cos(2 * math.pi * t / sec)) + if bw: x = band(x[None], torch.tensor([bw]))[0] + return x / x.abs().max() * 0.3 + +steps = int(sys.argv[sys.argv.index('--steps') + 1]) if '--steps' in sys.argv else 6 +torch.manual_seed(0) +with tempfile.TemporaryDirectory() as d: + # data: a few files at 48 and 24 kHz + os.makedirs(f'{d}/clean'); os.makedirs(f'{d}/rir'); os.makedirs(f'{d}/noise') + for i in range(4): + sf.write(f'{d}/clean/a{i}.wav', voice(5, 110 + 40 * i).numpy(), SR, subtype='FLOAT') + sf.write(f'{d}/clean/b{i}.wav', voice(5, 200 + 30 * i, sr=24000).numpy(), 24000, subtype='FLOAT') + sf.write(f'{d}/rir/r.wav', synth_rir(0.8).numpy(), SR, subtype='FLOAT') + sf.write(f'{d}/noise/n.wav', (torch.randn(SR * 3) * 0.05).numpy(), SR, subtype='FLOAT') + x24, _ = sf.read(f'{d}/clean/b0.wav'); bw24 = bandwidth(torch.from_numpy(x24).float(), 24000) + assert bw24 <= 12000, bw24 + x48 = voice(3, 150); assert bandwidth(x48) > 15000, bandwidth(x48) + x8 = voice(3, 150, bw=4000); assert 3000 < bandwidth(x8) < 5000, bandwidth(x8) + print(f'bandwidth: 24 kHz file {bw24:.0f} Hz, 48 kHz voice {bandwidth(x48):.0f}, low-passed at 4 kHz {bandwidth(x8):.0f}') + # degradation chain, every stage forced + ch = Chain(p={k: 1.0 for k in Chain.P}) + y = ch(x48); assert y.shape == x48.shape and torch.isfinite(y).all() + print(f'chain: {20 * math.log10((y - x48).pow(2).mean().sqrt() / x48.pow(2).mean().sqrt()):.1f} dB of change, codecs {"on" if FFMPEG else "off"}') + if FFMPEG: + import degrade + every = list(degrade.CODECS) + for c in every: # one codec at a time; a build without an encoder passes the take through + degrade.CODECS[:] = [c] + z = codec(x48); assert z.shape == x48.shape and torch.isfinite(z).all() + print(f'codec {c[0]}: {"changed the take" if not torch.equal(z, x48) else "unavailable, passed through"}') + degrade.CODECS[:] = every + # manifests and a training run + subprocess.check_call([sys.executable, f'{HERE}/data.py', 'scan', f'{d}/clean.json', f'{d}/clean', '--kind', 'sing'], stderr=subprocess.DEVNULL) + m = json.load(open(f'{d}/clean.json')); assert len(m) == 8 and any(i['sr'] == 24000 and i['bw'] <= 12000 for i in m) + out = f'{d}/run' + subprocess.check_call([sys.executable, f'{HERE}/train.py', '--out', out, '--clean', f'{d}/clean.json', '--rirs', f'{d}/rir', '--noises', f'{d}/noise', + '--dim', '64', '--blocks', '2', '--attn', '1', '--heads', '4', '--seg', '1', '--batch', '2', '--steps', str(steps), '--warmup', str(steps // 2), + '--every', str(steps), '--workers', '0', '--ssl', '0', '--val-batches', '1', '--device', 'cpu']) + log = [json.loads(l) for l in open(f'{out}/log.jsonl')] + assert any('val' in r for r in log), log + assert os.path.exists(f'{out}/last.pt') and os.path.exists(f'{out}/samples') + # resume + subprocess.check_call([sys.executable, f'{HERE}/train.py', '--out', out, '--clean', f'{d}/clean.json', '--dim', '64', '--blocks', '2', '--attn', '1', '--heads', '4', + '--seg', '1', '--batch', '2', '--steps', str(steps + 2), '--warmup', '0', '--every', '1000', '--workers', '0', '--ssl', '0', '--val-batches', '1', '--device', 'cpu', '--resume']) + # export and verify, then infer with both + ck = sorted(f for f in os.listdir(out) if f.startswith('net-'))[-1] + subprocess.check_call([sys.executable, f'{HERE}/export.py', f'{out}/{ck}', '--out', f'{d}/onnx', '--verify', '--frames', '50']) + subprocess.check_call([sys.executable, f'{HERE}/infer.py', f'{out}/{ck}', f'{d}/clean/a0.wav', f'{d}/clean/b0.wav', '--out', f'{d}/pt', '--device', 'cpu']) + subprocess.check_call([sys.executable, f'{HERE}/infer.py', f'{d}/onnx/restore.onnx', f'{d}/clean/a0.wav', '--out', f'{d}/onnx-out', '--device', 'cpu']) + a, _ = sf.read(f'{d}/pt/a0.wav'); b, _ = sf.read(f'{d}/onnx-out/a0.wav'); o, _ = sf.read(f'{d}/clean/a0.wav') + assert len(a) == len(o) and len(b) == len(o) + snr = 10 * np.log10((a ** 2).sum() / ((a - b) ** 2).sum()) + print(f'infer: torch and onnx within {snr:.1f} dB; outputs {len(a)} samples as the input') + # the long path: a 25 s file in chunks + sf.write(f'{d}/clean/long.wav', voice(25, 130).numpy(), SR, subtype='FLOAT') + subprocess.check_call([sys.executable, f'{HERE}/infer.py', f'{out}/{ck}', f'{d}/clean/long.wav', '--out', f'{d}/pt', '--device', 'cpu']) + l, _ = sf.read(f'{d}/pt/long.wav'); assert len(l) == 25 * SR and np.isfinite(l).all() + # the full-size net's cost + net = Net(); print(f'full net {params(net) / 1e6:.1f}M parameters, discriminators {params(Discriminators()) / 1e6:.1f}M') +print('ok') diff --git a/train/restore/train.py b/train/restore/train.py new file mode 100644 index 0000000..30c7af0 --- /dev/null +++ b/train/restore/train.py @@ -0,0 +1,144 @@ +"""Train the restorer. + + python train.py --out runs/a --clean speech-en.json,speech-ru.json,sing.json [--rirs DIR] [--noises DIR] + [--dim 512 --blocks 8] [--seg 4] [--batch 16] [--steps 400000] [--warmup 30000] [--ssl 1] [--resume] + +Two phases in one run. To `warmup`: regression alone (mrstft, mel, ssl). After: the discriminators join +(adversarial, feature matching), the regression terms kept (HiFi-GAN's weights: mel 45, fm 2, adv 1). Both +nets AdamW (0.8, 0.99), lr 2e-4 decayed ×0.999 every 1000 steps, bf16 autocast on CUDA. Every `every` steps a +checkpoint (`last.pt`: nets, optimizers, step) and a validation on `--val` pairs (LSD, SI-SDR, log-mel distance) +with a few restored examples written as wav. Logs a JSON line per 50 steps to /log.jsonl. + +The model hears its input at −20 dBFS RMS (dsp.to_level); the target is scaled by the same gain, so a level +is preserved, not imposed. +""" +import argparse, json, os, time, glob +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader +import soundfile as sf +from dsp import SR, to_level, band, mel +from data import Pairs, Bank +from model import Net, Discriminators, params +import loss as L + +def device(): + if torch.cuda.is_available(): return torch.device('cuda') + if getattr(torch.backends, 'mps', None) and torch.backends.mps.is_available(): return torch.device('mps') + return torch.device('cpu') + +def lsd(y, yh): + """log-spectral distance, dB: the bandwidth measure (Vocos / NVSR convention: 2048-point, log10 power)""" + Y, Yh = L.mag(y, 2048, 512).pow(2), L.mag(yh, 2048, 512).pow(2) + return torch.sqrt(((torch.log10(Yh + 1e-9) - torch.log10(Y + 1e-9)) ** 2).mean(1)).mean().item() + +def sisdr(y, yh): + y, yh = y - y.mean(-1, keepdim=True), yh - yh.mean(-1, keepdim=True) + s = (yh * y).sum(-1, keepdim=True) * y / (y.pow(2).sum(-1, keepdim=True) + 1e-9) + return (10 * torch.log10(s.pow(2).sum(-1) / ((yh - s).pow(2).sum(-1) + 1e-9) + 1e-9)).mean().item() + +def validate(net, val, dev, out, step): + net.eval(); m = dict(lsd=0, sisdr=0, mel=0); n = 0 + with torch.no_grad(): + for i, (x, y, bw) in enumerate(val): + x, y = x.to(dev), y.to(dev) + xin, g = to_level(x); yh = net(xin) / g[:, None] + y, yh = band(y, bw.to(dev)), band(yh, bw.to(dev)) + m['lsd'] += lsd(y, yh); m['sisdr'] += sisdr(y, yh); m['mel'] += F.l1_loss(mel(yh), mel(y)).item(); n += 1 + if i == 0: + d = os.path.join(out, 'samples'); os.makedirs(d, exist_ok=True) + for k in range(min(4, len(x))): + for name, v in (('in', x), ('out', yh), ('ref', y)): + sf.write(os.path.join(d, f'{step:07d}-{k}-{name}.wav'), v[k].cpu().numpy(), SR, subtype='FLOAT') + net.train() + return {k: v / max(n, 1) for k, v in m.items()} + +def main(): + a = argparse.ArgumentParser() + a.add_argument('--out', required=True); a.add_argument('--clean', required=True) + a.add_argument('--val', help='manifests for validation pairs (else 2% of a fixed draw from --clean)') + a.add_argument('--rirs'); a.add_argument('--noises') + a.add_argument('--dim', type=int, default=512); a.add_argument('--blocks', type=int, default=8) + a.add_argument('--attn', default='3,6'); a.add_argument('--heads', type=int, default=8) + a.add_argument('--seg', type=float, default=4.0); a.add_argument('--batch', type=int, default=16) + a.add_argument('--steps', type=int, default=400000); a.add_argument('--warmup', type=int, default=30000) + a.add_argument('--lr', type=float, default=2e-4); a.add_argument('--every', type=int, default=2000) + a.add_argument('--workers', type=int, default=6); a.add_argument('--ssl', type=float, default=1.0) + a.add_argument('--w-mel', type=float, default=45.0); a.add_argument('--w-stft', type=float, default=10.0) + a.add_argument('--w-fm', type=float, default=2.0); a.add_argument('--w-adv', type=float, default=1.0) + a.add_argument('--val-batches', type=int, default=8); a.add_argument('--resume', action='store_true') + a.add_argument('--device') + o = a.parse_args() + dev = torch.device(o.device) if o.device else device() + os.makedirs(o.out, exist_ok=True) + rirs = Bank(o.rirs, rir=True).items if o.rirs else () + noises = Bank(o.noises).items if o.noises else () + manifests = o.clean.split(',') + data = Pairs(manifests, rirs, noises, o.seg) + loader = DataLoader(data, o.batch, num_workers=o.workers, pin_memory=dev.type == 'cuda', persistent_workers=o.workers > 0) + vdata = Pairs(o.val.split(',') if o.val else manifests, rirs, noises, o.seg, seed=12345) + val = [] + for i, b in enumerate(DataLoader(vdata, o.batch, num_workers=0)): + val.append(b) + if i + 1 >= o.val_batches: break + + net = Net(o.dim, o.blocks, tuple(int(v) for v in o.attn.split(',') if v), o.heads).to(dev) + disc = Discriminators().to(dev) + print(f'generator {params(net) / 1e6:.1f}M, discriminators {params(disc) / 1e6:.1f}M, {dev}') + og = torch.optim.AdamW(net.parameters(), o.lr, (0.8, 0.99), weight_decay=0.01) + od = torch.optim.AdamW(disc.parameters(), o.lr, (0.8, 0.99), weight_decay=0.01) + sg = torch.optim.lr_scheduler.StepLR(og, 1000, 0.999); sd = torch.optim.lr_scheduler.StepLR(od, 1000, 0.999) + step = 0 + last = os.path.join(o.out, 'last.pt') + if o.resume and os.path.exists(last): + ck = torch.load(last, map_location=dev) + net.load_state_dict(ck['net']); disc.load_state_dict(ck['disc']); og.load_state_dict(ck['og']); od.load_state_dict(ck['od']) + sg.load_state_dict(ck['sg']); sd.load_state_dict(ck['sd']); step = ck['step'] + print('resumed at', step) + ssl = None + if o.ssl > 0: + try: ssl = L.SSL(dev) + except Exception as e: print('no SSL loss:', e) + json.dump(dict(vars(o), cfg=net.cfg), open(os.path.join(o.out, 'config.json'), 'w'), indent=1) + log = open(os.path.join(o.out, 'log.jsonl'), 'a') + amp = dev.type == 'cuda' + t0 = time.time() + for x, y, bw in loader: + if step >= o.steps: break + x, y, bw = x.to(dev, non_blocking=True), y.to(dev, non_blocking=True), bw.to(dev) + xin, g = to_level(x); y = y * g[:, None] + adv = step >= o.warmup + with torch.autocast(dev.type, torch.bfloat16, enabled=amp): + yh = net(xin) + yh = yh.float() + y, yhb = band(y, bw), band(yh, bw) + if adv: + od.zero_grad(set_to_none=True) + lr_, _ = disc(y); lf_, _ = disc(yhb.detach()) + ld = L.gan_d(lr_, lf_); ld.backward(); od.step() + og.zero_grad(set_to_none=True) + terms = dict(stft=L.mrstft(y, yhb), mel=L.mel_l1(y, yhb)) + if ssl: terms['ssl'] = ssl(y, yhb) + lg = o.w_stft * terms['stft'] + o.w_mel * terms['mel'] + o.ssl * terms.get('ssl', 0) + if adv: + lr_, fr = disc(y); lf_, ff = disc(yhb) + terms['adv'], terms['fm'] = L.gan_g(lf_), L.fm(fr, ff) + lg = lg + o.w_adv * terms['adv'] + o.w_fm * terms['fm'] + lg.backward() + torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0) + og.step(); sg.step() + if adv: sd.step() + step += 1 + if step % 50 == 0: + rec = dict(step=step, loss=round(lg.item(), 4), **{k: round(float(v), 4) for k, v in terms.items()}, s=round(time.time() - t0, 1)) + if adv: rec['d'] = round(ld.item(), 4) + print(json.dumps(rec)); log.write(json.dumps(rec) + '\n'); log.flush() + if step % o.every == 0 or step == o.steps: + m = validate(net, val, dev, o.out, step) + print(json.dumps(dict(step=step, val=m))); log.write(json.dumps(dict(step=step, val=m)) + '\n'); log.flush() + torch.save(dict(net=net.state_dict(), disc=disc.state_dict(), og=og.state_dict(), od=od.state_dict(), sg=sg.state_dict(), sd=sd.state_dict(), step=step, cfg=net.cfg), last + '.part') + os.replace(last + '.part', last) + torch.save(dict(net=net.state_dict(), cfg=net.cfg, step=step), os.path.join(o.out, f'net-{step:07d}.pt')) + +if __name__ == '__main__': + main() From c743750ebe6bd49541352464775b1d0cacbd3a89 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 14 Oct 2026 21:53:46 +0000 Subject: [PATCH 2/4] restore: the smoke test's findings noted (codecs, export, the WavLM path), pesq's build on 3.13 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01MUVeWzmg33J7gc9R3MWMYh --- .work/restore.md | 6 ++++-- train/restore/README.md | 5 +++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/.work/restore.md b/.work/restore.md index 9d326d5..2536c8b 100644 --- a/.work/restore.md +++ b/.work/restore.md @@ -136,8 +136,10 @@ trains on teacher outputs the way LibriTTS-R is Miipher's: distillation by data, - The teacher pass: Sidon at 500× real time cleans 1000 h in a few GPU hours. Separation of music for singing targets: mel-roformer at RTF 0.44 on the laptop's WebGPU, 100 h of music a day; on a CUDA card far less. - Smaller sibling for the page: dim 384, 6 blocks (~8M, ~16 MB fp16) trained the same way; measure what it loses. -- Checked to work before any GPU is paid for: `train/restore/test.py` runs the chain, a training run with resume, - ONNX export matching torch, inference in chunks, on CPU with synthetic voices. +- Checked to work before any GPU is paid for (torch 2.14, CPU, this session): `train/restore/test.py` runs the chain + with all four codecs, a training run with resume, ONNX export within 123 dB of torch at two lengths, inference in + chunks on a 25 s file; the WavLM loss path by shape on an untrained model of the same class (the weights' host + was unreachable here). Untested: CUDA / MPS, real data at scale, the adversarial phase beyond a few steps. ## Evaluation diff --git a/train/restore/README.md b/train/restore/README.md index 82149a9..9d9be11 100644 --- a/train/restore/README.md +++ b/train/restore/README.md @@ -16,6 +16,11 @@ python score.py --ref clean/ --out restored/ 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 | From 7174009273c304b2d6e7f339bf27be27ceda24ba Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 01:19:19 +0000 Subject: [PATCH 3/4] restore: run.sh, the run as one command on a GPU machine (fetch the open data, scan, train, export); measured here: no GPU, no route to the corpora, the full model past 15 GB on CPU Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01MUVeWzmg33J7gc9R3MWMYh --- train/restore/README.md | 3 ++ train/restore/run.sh | 77 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+) create mode 100755 train/restore/run.sh diff --git a/train/restore/README.md b/train/restore/README.md index 9d9be11..c0489ed 100644 --- a/train/restore/README.md +++ b/train/restore/README.md @@ -4,6 +4,8 @@ The training side of `audio`'s coming `restore()` op: one network that takes a v 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). +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 @@ -33,3 +35,4 @@ python test.py | `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) | diff --git a/train/restore/run.sh b/train/restore/run.sh new file mode 100755 index 0000000..254f3c3 --- /dev/null +++ b/train/restore/run.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# One command on a GPU machine (CUDA, 24 GB; or an Apple Silicon Mac for the small model): the open data fetched, +# scanned, the model trained, exported and scored. +# +# bash run.sh fetch [DATA] LibriTTS-R train-clean-100 (CC BY 4.0, 24 kHz), VCTK (CC BY 4.0, 48 kHz), MUSAN noise +# (CC BY 4.0), VocalSet (CC BY 4.0, 44.1 kHz) into DATA (default ~/.cache/audiojs/data/restore) +# bash run.sh scan [DATA] manifests with each file's bandwidth +# bash run.sh train [DATA] [RUN] the run (RUN default runs/restore); resumes if RUN exists +# bash run.sh small [DATA] [RUN] the browser-size model (dim 384, 6 blocks), batch 8 +# bash run.sh export RUN ONNX fp32 and fp16 into the neural cache, verified +# bash run.sh all [DATA] fetch, scan, train, export +# +# Rooms: no single archive serves OpenAIR or BUT ReverbDB; put any impulse responses (wav) under DATA/rirs and the run +# takes them (70 % measured, 30 % synthetic); without them every room is synthetic. Python 3.12 is used where found +# (pesq, which score.py needs, builds on it; the trainer runs on 3.13 too). +set -euo pipefail +HERE=$(cd "$(dirname "$0")" && pwd) +CMD=${1:-all}; DATA=${2:-$HOME/.cache/audiojs/data/restore}; RUN=${3:-$HERE/runs/restore} +PY=$(command -v python3.12 || command -v python3) + +venv() { + [ -d "$HERE/.venv" ] || { "$PY" -m venv "$HERE/.venv"; "$HERE/.venv/bin/pip" -q install -U pip; } + "$HERE/.venv/bin/pip" -q install torch torchaudio soundfile scipy pystoi onnx onnxruntime onnxconverter-common pyloudnorm \ + && ("$HERE/.venv/bin/pip" -q install pesq || echo "pesq did not build: score.py's PESQ needs it, the rest runs") + PY="$HERE/.venv/bin/python" +} + +get() { # get URL DIR: fetched once (resumable), unpacked once + local url=$1 dir=$2 f; f="$dir/$(basename "${url%%\?*}")" + mkdir -p "$dir" + if [ ! -e "$f" ]; then curl -fL --retry 5 -C - -o "$f.part" "$url" && mv -f "$f.part" "$f"; fi + [ -e "$f" ] || { echo "could not fetch $url" >&2; return 1; } + if [ ! -e "$f.unpacked" ]; then + case "$f" in *.tar.gz|*.tgz) tar xzf "$f" -C "$dir";; *.zip) unzip -qo "$f" -d "$dir";; esac + touch "$f.unpacked" + fi +} + +fetch() { + get https://www.openslr.org/resources/141/train_clean_100.tar.gz "$DATA/libritts-r" + get https://www.openslr.org/resources/17/musan.tar.gz "$DATA/musan" + get https://datashare.ed.ac.uk/bitstream/handle/10283/3443/VCTK-Corpus-0.92.zip "$DATA/vctk" + get "https://zenodo.org/records/1193957/files/VocalSet1-2.zip?download=1" "$DATA/vocalset" || echo "VocalSet: download VocalSet1-2.zip from zenodo.org/records/1193957 into $DATA/vocalset and unzip it" + mkdir -p "$DATA/rirs" "$DATA/noise" + [ -d "$DATA/musan/musan/noise" ] && ln -sfn "$DATA/musan/musan/noise" "$DATA/noise/musan" +} + +scan() { + cd "$HERE" + [ -d "$DATA/libritts-r/LibriTTS_R" ] && "$PY" data.py scan "$DATA/en-libritts.json" "$DATA/libritts-r/LibriTTS_R" --kind speech --lang en + [ -d "$DATA/vctk/wav48_silence_trimmed" ] && "$PY" data.py scan "$DATA/en-vctk.json" "$DATA/vctk/wav48_silence_trimmed" --kind speech --lang en + [ -d "$DATA/vocalset" ] && "$PY" data.py scan "$DATA/sing.json" "$DATA/vocalset" --kind sing + ls "$DATA"/*.json +} + +manifests() { ls "$DATA"/*.json | paste -sd, -; } + +train() { + cd "$HERE"; local resume=""; [ -e "$RUN/last.pt" ] && resume=--resume + "$PY" train.py --out "$RUN" --clean "$(manifests)" --rirs "$DATA/rirs" --noises "$DATA/noise" $resume "$@" 2>&1 | tee -a "$RUN.log" +} + +export_() { + cd "$HERE"; local ck; ck=$(ls "$RUN"/net-*.pt | tail -1) + "$PY" export.py "$ck" --fp16 --verify +} + +venv +case "$CMD" in + fetch) fetch;; + scan) scan;; + train) mkdir -p "$RUN"; train;; + small) mkdir -p "$RUN"; train --dim 384 --blocks 6 --attn 2,4 --batch 8;; + export) export_;; + all) fetch; scan; mkdir -p "$RUN"; train; export_;; + *) sed -n 2,20p "$0"; exit 1;; +esac From 2501f5ff150222f71fea2e5be6f78c4b50ca2463 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 01:44:46 +0000 Subject: [PATCH 4/4] restore: the handout, from an empty GPU machine to @audio/neural-restore and restore() (data by license, the teacher pass and its gate, stage A go/no-go, the run, trouble, evaluation sets and baselines, export and hosting, the plugin's contract, the op); data.py pairs writes examples to listen to Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01MUVeWzmg33J7gc9R3MWMYh --- train/restore/HANDOUT.md | 292 +++++++++++++++++++++++++++++++++++++++ train/restore/README.md | 5 +- train/restore/data.py | 16 ++- 3 files changed, 309 insertions(+), 4 deletions(-) create mode 100644 train/restore/HANDOUT.md diff --git a/train/restore/HANDOUT.md b/train/restore/HANDOUT.md new file mode 100644 index 0000000..51c401d --- /dev/null +++ b/train/restore/HANDOUT.md @@ -0,0 +1,292 @@ +# Handout: from an empty GPU box to `@audio/neural-restore` + +The steps, in order, to train the restorer and ship it as a plugin of the audiojs ecosystem, with the numbers +that say go or stop at each one. Why each choice was made is in [`.work/restore.md`](../../.work/restore.md); +this is the how. Everything in `train/restore/` moves into the plugin's repository as `scripts/` when the first +weights exist (the convention of `audiojs/neural`: `packages/neural-pitch/scripts/train.py`, +`packages/neural-separate/scripts/export-*.py`). + +What "done" means: a `restore()` op in `audio` that takes a voice recorded in a room, in noise, on a phone or +through a codec, speech or singing, English or Russian, and gives back a 48 kHz full-band studio take; one +model under 40 MB in fp16, one pass, faster than real time on WebGPU; measured against `deepfilter()`, +`auto()` and Adobe Enhance v2 on the sets below; weights hosted under a license that allows anyone to use them. + +Budget, if nothing goes wrong: a week of your attention spread over three weeks of wall time, and under $300 +of rented GPU. Data preparation is most of the work; training is the cheap part. + +## 1. Machine + +| Choice | Fits | Notes | +|---|---|---| +| Rented RTX 4090 / A5000 / L4 (24 GB), Linux | the full model, batch 16 × 4 s, ~2.5–3.5 steps/s | RunPod, Vast, Lambda: $0.4–0.9/h. 150 GB disk for data. Keep the data on a persistent volume: the run is restartable, the downloads are the slow part | +| A100 40/80 GB | same, batch 32 | not faster per step than a 4090 for this size; buy it for batch, not speed | +| M4 Max 36 GB (MPS) | the small model (dim 384, batch 8), 5–10× slower | for debugging the recipe on real data, not the production run. `torch.stft` / `istft` on MPS have had gaps between versions: if they fail, `--device cpu` for the smoke test and the GPU for the run | +| 16 GB card | small model, batch 8 × 4 s, or full model batch 6 × 3 s | adjust `--batch` and `--seg` until it fits | + +Setup on the box: + +``` +sudo apt install -y ffmpeg unzip # ffmpeg with libmp3lame, libopus, libgsm, aac: check `ffmpeg -encoders` +python3.12 -m venv .venv && . .venv/bin/activate # 3.12: pesq builds; 3.13 runs the trainer but not pesq +pip install torch torchaudio soundfile scipy pystoi onnx onnxruntime onnxconverter-common pyloudnorm pesq +python test.py # ~1 minute: the whole path on synthetic voices must print ok +``` + +If `test.py` prints every codec as "changed the take" and ends with `ok`, the machine is ready. + +## 2. Data + +Targets must be clean and, for the top octave, full band; licenses must allow hosting weights for anyone. +Sizes are the downloads; keep ~100 GB free after unpacking. + +### 2.1 English speech (the bulk) + +| Corpus | Get | Use | +|---|---|---| +| LibriTTS-R train-clean-100 and train-clean-360 (CC BY 4.0, 24 kHz, ~7 + 27 GB) | `openslr.org/141` | all of it; `scan` records 12 kHz bandwidth, so it trains nothing above | +| VCTK 0.92 (CC BY 4.0, 48 kHz, 11 GB) | `datashare.ed.ac.uk/handle/10283/3443` | `wav48_silence_trimmed`, mic1 only (`*_mic1.flac`): drop mic2. Its room is mild but real: see 2.5 | + +Skip: DAPS, EARS (non-commercial), anything from YouTube. + +### 2.2 Russian speech + +Nothing studio-grade and open exists at 48 kHz. The route is 16 kHz corpora through a teacher (2.5): + +| Corpus | Get | Use | +|---|---|---| +| Golos crowd (SberDevices; research and commercial use allowed, ShareAlike-like: read `license/en_us.pdf` in `github.com/sberdevices/golos`) | `openslr.org/114` | the crowd part only (far-field is a test set); teacher-cleaned and gated | +| RuLS, Russian LibriSpeech (public domain in the USA / CC BY 4.0 on mirrors, 98 h) | `openslr.org/96` | teacher-cleaned and gated | +| M-AILABS Russian | search "M-AILABS speech dataset" (the original host is gone; mirrors exist) | read its license file first; teacher-cleaned | + +Not for training weights you host: RUSLAN, Natasha, Balalaika (non-commercial), Common Voice (quality), LitRes or +any audiobook you do not own the rights to. + +### 2.3 Singing + +| Source | Get | Use | +|---|---|---| +| VocalSet (CC BY 4.0, 44.1 kHz, 10 h, 20 singers) | `zenodo.org/records/1193957` | all; the only full-band open singing set | +| CC BY / CC0 music → vocals | Jamendo (filter by license), Free Music Archive, ccMixter | the factory: `npx audio song.mp3 --vocals isolate --model mel-roformer -o vocal.wav` (your own separator), then the teacher and the gate (2.5). Aim at 50–100 h. Keep a CSV of track, artist, license: it goes into the model card | + +Not for training: GTSinger, M4Singer, OpenSinger, Opencpop, MUSDB18-HQ (non-commercial or by agreement). GTSinger's +Russian songs are a fine *test* set. + +### 2.4 Rooms and noise + +| Source | Get | Use | +|---|---|---| +| OpenAIR (CC BY) | `openairlib.net`, per room | every mono / omni IR into `rirs/`; B-format: take the W channel | +| BUT ReverbDB (CC BY 4.0, 1300+ IRs in 8 rooms) | `speech.fit.vutbr.cz/software/but-speech-fit-reverb-database` | all `RIR` wavs into `rirs/` | +| Your own | a balloon pop or a clap in bathrooms, stairwells, offices, cars, recorded on a phone | the rooms real users have; 50 of them are worth more than 500 concert halls | +| MUSAN noise (CC BY 4.0) | `openslr.org/17` | `musan/noise` into `noise/` | +| FSD50K (CC BY 4.0; per-clip licenses) | `zenodo.org/records/4060432` | keep clips whose license in `FSD50K.metadata/collection/*.csv` is CC0 or CC BY (85 %); drop `Speech`, `Singing`, `Music` labels (they are targets, not noise) | +| DEMAND (CC BY-SA 3.0) | `zenodo.org/records/1227121` | **test only**: VoiceBank+DEMAND's noises are from it, and the bench must not see training noise | + +Without `rirs/` every room is synthetic (`degrade.synth_rir`): the model still learns dereverberation, weaker on +real tails. 30 % of rooms stay synthetic even with a bank, by design. + +### 2.5 The teacher pass (Russian, separated vocals, VCTK) + +Install Sidon (`github.com/sarulab-speech/Sidon`, MIT; 16 kHz in, 48 kHz out, ~500× real time on a GPU). Run it +over Golos crowd, RuLS, the separated vocals and VCTK mic1. Then gate each clip; keep it only if all hold: + +| Gate | Tool | Threshold | +|---|---|---| +| quality | DNSMOS OVRL on the teacher's output (`bench/speech.py`'s `dnsmos`) | ≥ 3.3 (speech); singing: ≥ 3.0 and SIG ≥ 3.3 | +| no invented words | WER of `neural-asr` (whisper) on the output against the input's transcript (Golos and RuLS have transcripts; else against the input's own ASR) | relative increase ≤ 10 % | +| same person | cosine of `neural-diarize`'s WavLM speaker embeddings, input vs output | ≥ 0.85 | +| full band | `dsp.bandwidth` of the output | ≥ 18 kHz (else the teacher gave up on the top: drop) | + +Expect to keep 30–60 % of Golos crowd, 60–80 % of RuLS and VCTK, 20–50 % of separated vocals. The student trains +on the teacher's outputs as ordinary targets (that is what LibriTTS-R is to Miipher); the teacher is not needed at +training time. Listen to 20 kept Russian clips for a "teacher accent" on с/ш/щ before trusting the set. + +### 2.6 Scan and look + +``` +python data.py scan en.json $DATA/libritts-r $DATA/vctk-clean --kind speech --lang en +python data.py scan ru.json $DATA/ru-clean --kind speech --lang ru +python data.py scan sing.json $DATA/vocalset $DATA/vocals-clean --kind sing +python data.py pairs listen 20 --clean en.json,ru.json,sing.json --rirs $DATA/rirs --noises $DATA/noise +``` + +Open `listen/` in the playground. Each `k-in.wav` should sound like a recording someone would actually send you: +a room, a laptop mic, a phone, a Zoom call, a cassette. If they all sound like the same kind of broken, change +the chances in `degrade.Chain.P`; if an input is unrecognizable, the SNR floor (−5 dB) or the clip range is too +hard for your use: raise them. This half hour decides more than any hyperparameter. + +Check the manifests: `python -c "import json;m=json.load(open('en.json'));print(sum(i['frames']/i['sr'] for i in m)/3600,'h')"`. +Target before the full run: ≥ 300 h English, ≥ 100 h Russian, ≥ 30 h singing, of it ≥ 60 h full band in total. + +## 3. Training + +### 3.1 Stage A: does the recipe learn on real data? (half a day) + +``` +python train.py --out runs/a --clean en.json --rirs $DATA/rirs --noises $DATA/noise \ + --dim 384 --blocks 6 --attn 2,4 --batch 12 --steps 60000 --warmup 20000 --every 2000 +``` + +Watch `runs/a/log.jsonl` and `runs/a/samples/`. Go / no-go at 60k steps, on the validation lines (`val`): + +| Measure | Go | Stop and look | +|---|---|---| +| `lsd` | falls below 1.0 by 20k and keeps falling | flat above 1.3: the chain or the data is broken, not the model | +| `sisdr` | above the input's (compute the input's once: `score.py` on `listen/`), typically 8–14 dB | below the input's: the level or alignment is off (a codec lag the search missed) | +| `mel` | under 0.5 | — | +| samples at 40k+ | the room gone, consonants crisp, no buzzing in sustained vowels, the top octave present on inputs cut at 4–8 kHz | metallic ringing: `--w-adv 0.5`; muffled: the adversarial phase is not biting, check `d` hovers 0.3–0.7 and `adv` 0.2–1.0 | + +A clean input must come out clean: run `infer.py` on three VCTK files and compare by ear and with `score.py` against +the originals (PESQ ≥ 4.0, LSD ≤ 0.6). If stage A passes, nothing in the recipe changes for stage B. + +### 3.2 Stage B: the model (2 days of GPU) + +``` +python train.py --out runs/b --clean en.json,ru.json,sing.json --rirs $DATA/rirs --noises $DATA/noise \ + --batch 16 --steps 400000 --warmup 30000 --every 2000 +``` + +Defaults are the full model (dim 512, 8 blocks, attention after 3 and 6, ~18M). Resume after any interruption with +the same command plus `--resume`. Checkpoints land every 2000 steps as `net-NNNNNNN.pt` (keep every 20th; delete +the rest). Expect: regression phase 2–3 hours, then `d` settles near 0.5, `adv` near 0.5–1.0, `fm` falling slowly +for the whole run; `lsd` under 0.8 and `sisdr` 10–16 dB by the end. If `val` numbers plateau by 250k and the +samples are good, stop there: the last 150k steps buy texture, not intelligibility. + +Knobs worth touching, in order: the data mix (duplicate `sing.json` in `--clean` to weight singing ×2); `--seg 6` +if rooms over 1 s stay (longer context); `--w-fm 4` if the top octave sounds sprayed; `--ssl 0.5` if speech comes +out overly smooth. Not worth touching: learning rates, betas, the discriminator set, the STFT. + +### 3.3 Stage C: the browser-size sibling (1 day) + +The same command with `--dim 384 --blocks 6 --attn 2,4` (~8M, ~16 MB fp16). Measure what it loses in section 4; +ship both and let `restore({ model })` choose, the small one where there is no GPU. + +### 3.4 Trouble + +| Symptom | Cause | Fix | +|---|---|---| +| `nan` in the log | a silent crop (RMS gate at −45 dBFS failed on a file of zeros) or fp16 overflow | `--device cuda` uses bf16 already; find the file: `data.py pairs` with the same seed; drop it | +| GPU under 60 % busy | the loaders: codecs through ffmpeg are the slow stage | `--workers 12`, or `Chain.P['codec'] = 0.15` | +| `d` → 0, `adv` → 2+, samples buzz | the discriminators won | restart from the last good checkpoint with `--w-adv 0.5 --w-fm 4`; never raise the generator's lr | +| output 2–3 dB quieter than the input | expected on noisy input (the noise is gone); on clean input it is a bug | check `to_level` is applied in both `train.py` and `infer.py` the same way | +| Russian щ/ц come out as hiss | too little Russian in the mix, or the teacher's accent | weight `ru.json` ×2; listen to the teacher's outputs again | +| held sung notes lose level | the singing share is too small | VocalSet ×3 in `--clean`; the factory of 2.3 | +| 24 kHz-only output on everything | every target was band-limited (a scan bug) | `python -c` over the manifest: `bw` must be 24000 for VCTK and VocalSet | + +## 4. Evaluation + +Run before claiming anything. The sets exist in `audio`'s bench or download in minutes: + +| Set | What it measures | Baseline to beat | +|---|---|---| +| VoiceBank+DEMAND test (824 files, `datashare.ed.ac.uk/handle/10283/2791`) | noise; PESQ, STOI, SI-SDR, DNSMOS | `deepfilter()` PESQ 3.05, `deepfilter(0)` 3.15, RX 12 Dialogue Isolate 2.72 (`bench/rx/isolate.mjs`). A generative model may land near 3.0 PESQ while winning DNSMOS and the listen: report both | +| vbreverb (its clean speech in MIT IR Survey rooms: `bench/speech.mjs` builds it) | rooms | `dereverb()`, `deepfilter()` + `dereverb()`; PESQ 2.52 for WPE | +| the ten narrations (`bench/speech.mjs rooms`) | real recordings, no reference: DNSMOS, ACX floor | `auto('speech')` OVRL 3.12 | +| band limits: VCTK test speakers low-passed at 4 and 8 kHz, and through Opus 12 kb/s (`degrade.bandlimit`, `degrade.codec`) | recovery; LSD, DNSMOS | the input itself (LSD 2.5–3.5 at 4 kHz); target ≤ 1.0 | +| GTSinger's Russian songs, VocalSet held-out singers, with rooms and noise added | singing | the input; DNSMOS SIG must not fall on clean singing | +| 20 real takes of your own (phone, laptop, car, Zoom, cassette), the same files through Adobe Enhance v2 by hand | the comparison that matters | blind A/B, you and two others, which of each pair they would publish | +| WER: `npx audio take.wav --asr` on input, output and reference (`neural-asr`) | invented or lost words | relative WER increase ≤ 10 % over the reference's; Adobe's and FINALLY's class loses words at −5 dB SNR too: measure, do not assume | + +``` +python infer.py runs/b/net-0400000.pt $VB/noisy_testset_wav/*.wav --out out/vb +python score.py --ref $VB/clean_testset_wav --out out/vb --json out/vb.json +``` + +Publish only numbers from `score.py` and the bench scripts, with the commit that produced them, as every README in +the ecosystem does (the `Measured` sections of `neural-denoise`, `neural-separate`). + +## 5. Export and host + +``` +python export.py runs/b/net-0400000.pt --fp16 --verify # ~/.cache/audiojs/neural/restore/restore{,.fp16}.onnx +``` + +`--verify` must print over 100 dB from torch at both lengths. Host `restore.fp16.onnx` (and the small model's) in a +Hugging Face model repo `audiojs/restore`, and pin its revision hash in the plugin (as `neural-separate`'s +`REVISIONS` does): a URL without a hash is a model that can change under your users. The model card lists every +corpus, its license and attribution, the teacher, the bench numbers, and the weights' own license: CC BY 4.0 if +every source is CC BY or public domain, CC BY-SA 4.0 if Golos's ShareAlike applies (read its PDF; if in doubt, train +the hosted weights without Golos and keep a Golos variant for yourself). + +## 6. The plugin: `@audio/neural-restore` + +A package in `audiojs/neural/packages/neural-restore`, the shape of `neural-denoise`: + +``` +neural-restore/ + package.json "@audio/neural-restore", type module, peer: @audio/neural-runtime + index.js export default restore(pcm, { sampleRate, model, strength }) ; export { load, MODEL, REVISIONS } + restore.js the pipeline below + stft.js or @audio/stft if it does centered reflect padding with a periodic Hann window + scripts/ this directory: train.py, data.py, degrade.py, model.py, loss.py, export.py, infer.py, score.py, test.py, run.sh + test.js fixtures/in.wav through the JS pipeline against fixtures/ref.wav (Python's infer.py output): ≥ 60 dB + README.md What, How, Measured, Weights (license, sources), as the other packages + LICENSE, NOTICE MIT; the data attributions +``` + +The pipeline `restore.js` must reproduce `infer.py` exactly; the contract is small: + +1. Per channel, resample to 48 kHz (`@audio/resample-sinc`); remember the gain `g` that brings the chunk to −20 dBFS RMS + (`dsp.to_level`: `g = 10^((−20 − rms_db) / 20)`), apply it. +2. Chunks of 10 s overlapping by 1 s; each chunk: STFT with `n_fft` 2048, hop 512, **periodic** Hann + (`0.5 − 0.5 cos(2πn/2048)`, torch's default), centered with 1024 samples of reflect padding at both ends, no + normalization → 1025 bins × T frames; feed `[1, 2, 1025, T]` re/im float32 through `neural-runtime` (`load(url, + { backend })`, `webgpu` where `backends()` lists it, else `wasm` / Node); iSTFT the output with the same window, + overlap-add divided by the summed squared window (torch's `istft`), cut to the chunk's length. +3. Raised-cosine crossfade over the 1 s overlaps; divide by `g`; resample back to the input's rate; the output is + exactly the input's length. +4. `strength` 0–1: `out = strength · y + (1 − strength) · x` (the same mixback `deepfilter()` uses for its limit). + +Verify the JS STFT against torch once with a fixture (`scripts/export.py --verify` writes none: add +`scripts/reference.py` that dumps one chunk's spectrum, as `neural-separate/scripts/reference.py` does); then the +whole pipeline against `infer.py` on `fixtures/in.wav`. Speed: measure RTF on onnxruntime-web WebGPU first; if it +is over 0.2, the Mel-RoFormer engine (`neural-separate/roformer.js`) has the matmul and attention kernels, and the +ConvNeXt blocks add depthwise convolution and LayerNorm shaders. + +## 7. The op in `audio` + +`fn/deepfilter.js`'s `enhancer(name, id, finish, again)` already does what `restore()` needs: the whole input made +once ahead of rendering, memoized on the input's fingerprint and the options, channels apart, live sources waited +for, playback from the kept output. Two things in it are DeepFilterNet3's and become parameters: the package it +imports (`@audio/neural-denoise`) and the run (`removal()`, DFN3's unlimited output). Give `enhancer` a fifth +argument `{ pkg, run }` with the current values as defaults, and `fn/restore.js` is: + +```js +import audio from '../core.js' +import { enhancer } from './deepfilter.js' + +// restore(strength = 1, { model = 'restore', weights, device }): the take as a studio would have recorded it, 48 kHz full band +audio.op('restore', { + params: ['strength'], + ...enhancer('restore', o => { + let s = o.strength ?? 1 + if (typeof s !== 'number' || !(s >= 0 && s <= 1)) throw new TypeError(`restore: strength is 0 to 1, not ${s}`) + return `${s}:${o.model ?? 'restore'}` + }, (x, y, o) => { // the model's output, the input mixed back by 1 − strength + let s = o.strength ?? 1, out = new Float32Array(y.length) + for (let i = 0; i < y.length; i++) out[i] = s * y[i] + (1 - s) * x[i] + return out + }, null, { + pkg: '@audio/neural-restore', + run: (neural, pcm, sampleRate, o) => neural.default(pcm, { sampleRate, model: o.model, weights: o.weights, device: o.device }), + }), +}) +``` + +`restore()` is a full-band op: unlike `deepfilter`, no `music: 'pass'` guard (a song is a valid input), so `music` +is not read. Then: `fn/` is auto-registered by name, so a row in README's plugin table (what it is, the measured +numbers, ≡ Adobe Podcast Enhance Speech, iZotope RX Dialogue Isolate + Spectral Recovery), `audio.d.ts`, the +playground's Enhance menu ("Restore"), the CLI (`--restore`), `auto('speech')` taking it as the last repair stage +when the package is installed (as it takes `deepfilter` through `@audio/chain`), a test in `test/` with a +fixture, and `bench/speech.mjs` gaining a `restore` system so the recipe benches and `bench/rx` compare it. + +## 8. Ship list + +- [ ] `test.py` ok on the GPU box; `listen/` pairs sound like real bad recordings +- [ ] Stage A go on every row of 3.1 +- [ ] Stage B: `val` plateaued, samples clean, a clean input stays clean +- [ ] Section 4 numbers, with the commit hash, and the blind A/B against Adobe v2 +- [ ] `export.py --verify` over 100 dB; weights on `audiojs/restore` with a pinned revision and a complete model card +- [ ] `@audio/neural-restore` tests pass in Node and the browser; RTF measured on WebGPU, wasm, Node +- [ ] `restore()` in `audio`: README row, d.ts, playground, CLI, `auto()`; `bench/rx` numbers against RX and Adobe +- [ ] Licenses: every corpus in NOTICE; the weights' license decided (CC BY or CC BY-SA); nothing non-commercial touched the hosted weights diff --git a/train/restore/README.md b/train/restore/README.md index c0489ed..e8fa9bf 100644 --- a/train/restore/README.md +++ b/train/restore/README.md @@ -2,7 +2,8 @@ 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). +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: @@ -27,7 +28,7 @@ python test.py |---|---| | `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) | +| `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 | diff --git a/train/restore/data.py b/train/restore/data.py index 249aec6..b53da64 100644 --- a/train/restore/data.py +++ b/train/restore/data.py @@ -3,6 +3,8 @@ 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, -in.wav and -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 @@ -89,7 +91,17 @@ def __iter__(self): if __name__ == '__main__': import argparse a = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - a.add_argument('cmd', choices=['scan']); a.add_argument('out'); a.add_argument('dirs', nargs='+') + 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() - scan(o.out, o.dirs, o.kind, o.lang) + 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')