diff --git a/.work/restore.md b/.work/restore.md new file mode 100644 index 0000000..2536c8b --- /dev/null +++ b/.work/restore.md @@ -0,0 +1,195 @@ +# 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 (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 + +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/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 new file mode 100644 index 0000000..e8fa9bf --- /dev/null +++ b/train/restore/README.md @@ -0,0 +1,39 @@ +# restore: a studio-grade voice from a bad recording + +The training side of `audio`'s coming `restore()` op: one network that takes a voice recorded in a room, in noise, +through a cheap microphone, a codec or a phone, speech or singing, and gives back the voice as a studio would have +recorded it, at 48 kHz. The plan, the data and what is known are in [`.work/restore.md`](../../.work/restore.md); the step-by-step runbook from an +empty GPU machine to the published plugin, with go / no-go numbers, is [`HANDOUT.md`](HANDOUT.md). + +On a GPU machine, `bash run.sh all` does all of the below (about 20 GB of downloads: LibriTTS-R train-clean-100, VCTK, MUSAN, VocalSet). By hand: + +``` +pip install torch torchaudio soundfile scipy onnx onnxruntime onnxconverter-common # ffmpeg on PATH for codecs +python data.py scan en.json /data/vctk /data/libritts-r --kind speech --lang en # manifests, bandwidth per file +python data.py scan ru.json /data/ruls --kind speech --lang ru +python data.py scan sing.json /data/vocalset --kind sing +python train.py --out runs/a --clean en.json,ru.json,sing.json --rirs /data/rirs --noises /data/noise +python export.py runs/a/net-0400000.pt --fp16 --verify # → ~/.cache/audiojs/neural/restore/ +python infer.py runs/a/net-0400000.pt take.wav --out restored/ +python score.py --ref clean/ --out restored/ # PESQ, STOI, SI-SDR, DNSMOS, LSD +python test.py # smoke test, CPU, a minute +``` + +`score.py` also needs `pesq pyloudnorm` (`pesq` builds from source: on Python 3.13 it needs `setuptools<70` in the build, or Python 3.12). + +``` +``` + +| File | What | +|---|---| +| `dsp.py` | the STFT (2048/512 at 48 kHz), levels, the band limit every loss passes through, mel, resampling, the bandwidth a file carries | +| `degrade.py` | the on-the-fly degradation chain: room (measured or synthetic), noise, microphone colour, compressor, clipping, band limit, codec (ffmpeg), word length | +| `data.py` | manifests and the pair sampler (an IterableDataset: crops by length, levels, the chain); `pairs` writes examples to listen to | +| `model.py` | the generator (ConvNeXt over frames with self-attention, log-magnitude change + phase), the discriminators (multi-period, multi-resolution) | +| `loss.py` | multi-resolution STFT, log-mel, WavLM feature distance, least-squares GAN, feature matching | +| `train.py` | the run: regression warm-up, then adversarial; checkpoints, validation (LSD, SI-SDR, mel), samples | +| `export.py` | ONNX between STFT and iSTFT, `[1, 2, 1025, T]` re/im in and out, fp16, verified against torch | +| `infer.py` | files through a checkpoint or the ONNX, chunked | +| `score.py` | `bench/speech.py`'s measures plus LSD on a directory of outputs against references | +| `test.py` | the whole path on synthetic voices, CPU | +| `run.sh` | one command on a GPU machine: fetch the open data, scan, train, export (`bash run.sh all`; `small` for the browser-size model) | diff --git a/train/restore/data.py b/train/restore/data.py new file mode 100644 index 0000000..b53da64 --- /dev/null +++ b/train/restore/data.py @@ -0,0 +1,107 @@ +"""Clean voices, rooms and noises on disk → training pairs. + + python data.py scan OUT.json DIR [DIR...] [--kind speech|sing] [--lang en|ru|..] + every audio file under the DIRs: its rate, length and the bandwidth it carries (dsp.bandwidth), one + manifest; a corpus at 16 or 24 kHz is kept, its band edge recorded, so the losses stop at it (dsp.band) + python data.py pairs OUT_DIR N --clean A.json[,B.json] [--rirs DIR] [--noises DIR] [--seg 4] + N training pairs as the trainer would draw them, -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 +worker seeds itself. Bank(dir) holds every file of a directory in memory at 48 kHz (rooms: aligned to the direct +peak, 2.5 s at most, unit energy). +""" +import glob, json, os, random, sys +import numpy as np +import soundfile as sf +import torch +from torch.utils.data import IterableDataset, get_worker_info +from dsp import SR, bandwidth, resample, rms_db +from degrade import Chain + +EXT = ('.wav', '.flac', '.ogg', '.mp3', '.opus') + +def files(d): + return sorted(p for p in glob.glob(os.path.join(d, '**', '*'), recursive=True) if p.lower().endswith(EXT)) + +def read(path, start=0, frames=-1): + x, sr = sf.read(path, start=start, frames=frames, dtype='float32', always_2d=True) + return torch.from_numpy(x.mean(1)), sr + +def scan(out, dirs, kind='speech', lang='en'): + items = json.load(open(out)) if os.path.exists(out) else [] + seen = {i['path'] for i in items} + for d in dirs: + for p in files(d): + if p in seen: continue + try: + info = sf.info(p) + x, sr = read(p, frames=min(info.frames, info.samplerate * 30)) + except Exception as e: + print('skip', p, e, file=sys.stderr); continue + bw = bandwidth(resample(x, sr), SR) if sr != SR else bandwidth(x, SR) + items.append(dict(path=p, sr=sr, frames=info.frames, bw=round(bw), kind=kind, lang=lang)) + print(len(items), p, sr, round(info.frames / sr, 1), 's', round(bw), 'Hz', file=sys.stderr) + if len(items) % 200 == 0: json.dump(items, open(out, 'w')) + json.dump(items, open(out, 'w')) + return items + +class Bank: + def __init__(self, d, rir=False, limit=None): + self.items = [] + for p in files(d)[:limit] if d else []: + x, sr = read(p) + x = resample(x, sr) + if rir: + k = int(x.abs().argmax()); x = x[max(0, k - int(0.002 * SR)): k + int(2.5 * SR)] + x = x / (x.pow(2).sum().sqrt() + 1e-9) + if len(x) > SR // 10: self.items.append(x) + +class Pairs(IterableDataset): + def __init__(self, manifests, rirs=(), noises=(), seg=4.0, p=None, level=(-35, -12), seed=0): + self.items = [i for m in manifests for i in json.load(open(m)) if i['frames'] >= i['sr'] * 0.5] + self.w = np.array([i['frames'] / i['sr'] for i in self.items]); self.w /= self.w.sum() + self.chain = Chain(rirs, noises, p) + self.n, self.level, self.seed = int(seg * SR), level, seed + + def crop(self): + while True: + i = self.items[np.random.choice(len(self.items), p=self.w)] + n = int(self.n * i['sr'] / SR) + 1 + start = random.randint(0, max(0, i['frames'] - n)) + try: x, sr = read(i['path'], start, min(n, i['frames'] - start)) + except Exception: continue + x = resample(x, sr) + if len(x) < self.n: x = torch.nn.functional.pad(x, (0, self.n - len(x))) + x = x[: self.n] + if rms_db(x) > -45: return x, float(min(i['bw'], sr / 2)) + + def __iter__(self): + w = get_worker_info() + s = self.seed + (w.id if w else 0) + 1000 * int.from_bytes(os.urandom(2), 'little') + random.seed(s); np.random.seed(s % 2 ** 32); torch.manual_seed(s) + while True: + x, bw = self.crop() + x = x * 10 ** ((random.uniform(*self.level) - float(rms_db(x))) / 20) + y = self.chain(x) + peak = y.abs().max() + if peak > 0.99: x, y = x * 0.99 / peak, y * 0.99 / peak + yield y, x, torch.tensor(bw) + +if __name__ == '__main__': + import argparse + a = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + a.add_argument('cmd', choices=['scan', 'pairs']); a.add_argument('out'); a.add_argument('dirs', nargs='*') + a.add_argument('--kind', default='speech', choices=['speech', 'sing']); a.add_argument('--lang', default='en') + a.add_argument('--clean'); a.add_argument('--rirs'); a.add_argument('--noises'); a.add_argument('--seg', type=float, default=4.0) + o = a.parse_args() + if o.cmd == 'scan': + scan(o.out, o.dirs, o.kind, o.lang) + else: + n = int(o.dirs[0]) if o.dirs else 10 + pairs = Pairs(o.clean.split(','), Bank(o.rirs, rir=True).items if o.rirs else (), Bank(o.noises).items if o.noises else (), o.seg) + os.makedirs(o.out, exist_ok=True) + for k, (y, x, bw) in zip(range(n), pairs): + sf.write(os.path.join(o.out, f'{k}-in.wav'), y.numpy(), SR, subtype='FLOAT') + sf.write(os.path.join(o.out, f'{k}-ref.wav'), x.numpy(), SR, subtype='FLOAT') + print(k, f'bandwidth {int(bw)} Hz') 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/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 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()