| |
| """Sidon call-centre restoration — inference. |
| |
| Restores telephony / call-centre audio (narrowband, codec'd, noisy) to clean |
| 48 kHz, using the two trained stages: |
| |
| input audio --resample 16k--> (call-centre FE: 24L w2v-BERT + LoRA) --features[T,1024]--> |
| (DAC decoder, 188M) --> 48 kHz waveform |
| |
| The FE LoRA adapter is **merged into the base weights** here (W_eff = W + (alpha/r)*B@A, |
| bias = trained bias), so inference needs NO `peft` — just transformers + descript-audio-codec. |
| |
| Usage: |
| python runpod/infer_callcentre.py --input audio --out-dir audio/out \ |
| --fe-adapter checkpoints/fe_adapter_full.pt --decoder checkpoints/decoder_only.pt |
| |
| Stereo inputs (e.g. agent/customer on separate channels) are restored per-channel |
| and recombined to stereo. A `<name>_orig<khz>.wav` (naive-upsampled input, no model) |
| is also written for an apples-to-apples A/B listen. |
| |
| The output rate comes from the DECODER CHECKPOINT (`out_sr`/`dec_rates`), so a 24 kHz |
| enhancer-lite decoder needs no extra flag; 48 kHz models keep the `_restored48k.wav` |
| filenames the eval scripts already expect. |
| """ |
| from __future__ import annotations |
|
|
| import os |
|
|
| os.environ.setdefault("OMP_NUM_THREADS", "4") |
|
|
| import argparse |
| import glob |
| import time |
|
|
| import numpy as np |
| import soundfile as sf |
| import torch |
| import torchaudio |
| from transformers import AutoFeatureExtractor, Wav2Vec2BertModel |
|
|
| import dac |
|
|
| |
| |
| |
| if os.environ.get("FAST_CONV1D", "0") == "1": |
| import fast_conv1d as _fc |
| if _fc.patch(): |
| print("[fast_conv1d] conv1d/conv_transpose1d via GEMM (sm_10x cuDNN workaround)", flush=True) |
|
|
|
|
| SSL_MODEL = "facebook/w2v-bert-2.0" |
| FE_SR = 16000 |
| SR_OUT = 48000 |
| AUDIO_EXTS = (".wav", ".flac", ".mp3", ".ogg", ".opus", ".m4a") |
|
|
|
|
| def log(m: str) -> None: |
| print(m, flush=True) |
|
|
|
|
| def load_fe(adapter_path: str, device: torch.device) -> Wav2Vec2BertModel: |
| """Build the 24L w2v-BERT base and merge the trained LoRA adapter into it.""" |
| ck = torch.load(adapter_path, map_location="cpu") |
| ad = ck["adapter"] |
| scaling = ck["lora_alpha"] / ck["r"] |
| layers = ck.get("layers", 24) |
| model = Wav2Vec2BertModel.from_pretrained(SSL_MODEL, num_hidden_layers=layers, layerdrop=0.0) |
| sd = model.state_dict() |
| prefixes = sorted({k[: -len(".lora_A.default.weight")] |
| for k in ad if k.endswith(".lora_A.default.weight")}) |
| merged = 0 |
| for p in prefixes: |
| A = ad[p + ".lora_A.default.weight"].float() |
| B = ad[p + ".lora_B.default.weight"].float() |
| delta = scaling * (B @ A) |
| wkey = p + ".weight" |
| sd[wkey] = sd[wkey].float() + delta.to(sd[wkey].dtype) |
| bkey = p + ".base_layer.bias" |
| if bkey in ad: |
| sd[p + ".bias"] = ad[bkey].to(sd[p + ".bias"].dtype) |
| merged += 1 |
| model.load_state_dict(sd) |
| model.to(device).eval() |
| for q in model.parameters(): |
| q.requires_grad_(False) |
| log(f"[fe] merged LoRA into {merged} output_dense layers (scaling={scaling}); " |
| f"step {ck.get('step')}") |
| return model |
|
|
|
|
| def load_decoder(path: str, device: torch.device, input_channel: int = 1024): |
| """Returns (decoder, out_sr). The graph is reconstructed from the checkpoint, so a |
| 24 kHz decoder (`rates=[8,5,4,3]`) loads without any CLI flag; checkpoints written |
| before --out-sr existed have no `out_sr` key and were all 48 kHz.""" |
| ck = torch.load(path, map_location="cpu") |
| ch = ck.get("dec_channels", 3072) |
| extra = ck.get("extra_res", "") |
| out_sr = int(ck.get("out_sr", SR_OUT)) |
| rates = ([int(r) for r in ck["dec_rates"].split(",")] if ck.get("dec_rates") |
| else [8, 5, 4, 3, 2]) |
| |
| |
| arch = ck.get("dec_arch", "dac") |
| if arch == "hifigan": |
| from hifigan_decoder import build_hifigan_decoder |
| dec = build_hifigan_decoder(input_channel, ch, rates=tuple(rates), |
| activation=ck.get("dec_activation", "lrelu")) |
| else: |
| from expand_decoder import build_decoder, parse_extra_res |
| dec = build_decoder(input_channel, ch, parse_extra_res(extra), rates=rates) |
| dec.load_state_dict(ck["decoder"]) |
| dec.to(device).eval() |
| for q in dec.parameters(): |
| q.requires_grad_(False) |
| label = "HiFi-GAN(MRF)" if arch == "hifigan" else "DAC" |
| log(f"[dec] {label} decoder channels={ch}{' +' + extra if extra else ''} rates={rates} " |
| f"-> {out_sr} Hz ({sum(p.numel() for p in dec.parameters())/1e6:.1f}M); " |
| f"step {ck.get('step')}") |
| return dec, out_sr |
|
|
|
|
| def _peak_norm(x: np.ndarray, peak: float = 0.95) -> np.ndarray: |
| m = float(np.abs(x).max()) |
| return (x / m * peak).astype("float32") if m > 1e-6 else x.astype("float32") |
|
|
|
|
| def _lowpass(x: np.ndarray, sr: int, fc: float) -> np.ndarray: |
| """Zero-phase low-pass of the decoder output (see --out-lowpass for the measurement |
| that justifies the 8 kHz default). scipy is a hard dep of the training stack; if it |
| is somehow absent, pass the audio through rather than fail an inference run.""" |
| try: |
| from scipy.signal import butter, sosfiltfilt |
| except Exception: |
| log("[lp] scipy unavailable — skipping --out-lowpass") |
| return x |
| sos = butter(8, fc / (sr / 2), btype="low", output="sos") |
| return sosfiltfilt(sos, x).astype("float32") |
|
|
|
|
| @torch.no_grad() |
| def restore_channel(wav16: np.ndarray, fe, dec, proc, device, chunk_s: float, bf16: bool, |
| out_sr: int = SR_OUT): |
| """wav16: 1-D float32 @16k (peak-normalized) -> restored 1-D float32 @out_sr. |
| chunk_s > 0 => WINDOWED. The DEFAULT IS ARCH-AWARE (resolved in main): 4.0 s for w2vbert, |
| 0 (straight pass) for distilhubert. Both are measured, and they differ because only one of |
| the two feature extractors is length-dependent: |
| w2vbert SeamlessM4TFeatureExtractor normalises PER UTTERANCE -> the same audio |
| encodes differently depending on how much follows it (max|diff| 0 -> 10.4 |
| from 6 s -> 29 s). Training used 6 s windows, so whole-file is OOD. |
| distilhubert its INPUT features are bit-identical across length (do_normalize=False on |
| raw waveform), but its HIDDEN STATES are NOT -- the transformer attends over |
| the whole utterance, so the same 6 s drifts 18.9-39.5 % as more audio follows. |
| Window it too. (An earlier note here claimed invariance from the input test |
| alone; that inference was wrong.) |
| |
| The old docstring claimed a full pass is "length-invariant and cleaner" because w2v-BERT |
| uses rotary embeddings and the decoder is convolutional. **That is wrong, measured |
| 2026-08-13.** `SeamlessM4TFeatureExtractor` normalises each utterance to zero-mean / |
| unit-variance, so the SAME audio gets different features depending on how much audio |
| follows it: on one 30 s clip the first 6 s came out as ~1.61x + 0.37 relative to encoding |
| those 6 s alone. Training only ever used 6 s windows, so whole-file inference is |
| out-of-distribution and the decoder collapses to silence on parts of long inputs. |
| Measured on 20 worst-case clips: whole-pass 82.7 % of speech muted with 46 dropouts >3 s, |
| vs 33.6 % / 20 at chunk=4. On a random 40 it is a wash (14.9 -> 14.3 %), so windowing |
| rescues the pathological tail rather than improving the average -- but the tail is 20 s |
| holes in customer calls. Sweep: 4 s best, then 6 / 8 / 12, monotone in length. |
| COST: the 2 s overlap means chunk=4 is ~2x the compute of a straight pass (chunk=6 is |
| ~1.5x) -- relevant for the CPU-realtime lite path, not for GPU.""" |
| n = len(wav16) |
| win = int(chunk_s * FE_SR) |
| if chunk_s <= 0 or n <= win: |
| bounds = [(0, n)] |
| else: |
| ov = int(2.0 * FE_SR) |
| hop = win - ov |
| bounds = [(s, min(s + win, n)) for s in range(0, n, hop)] |
| bounds = [b for b in bounds if b[1] > b[0]] |
| |
| |
| |
| |
| |
| MIN_SEG = int(0.6 * FE_SR) |
| while len(bounds) > 1 and (bounds[-1][1] - bounds[-1][0]) < MIN_SEG: |
| bounds[-2] = (bounds[-2][0], bounds[-1][1]) |
| bounds.pop() |
| |
| |
| up_num, up_den = out_sr, FE_SR |
| out = np.zeros(n * up_num // up_den + out_sr, dtype="float32") |
| wsum = np.zeros_like(out) |
| amp = torch.bfloat16 if (bf16 and device.type == "cuda") else torch.float32 |
| for s, e in bounds: |
| seg = np.pad(wav16[s:e], (40, 40)) |
| feats_in = proc(seg, sampling_rate=FE_SR, return_tensors="pt") |
| feats_in = {k: v.to(device) for k, v in feats_in.items()} |
| with torch.autocast(device.type, dtype=amp, enabled=(amp == torch.bfloat16)): |
| h = fe(**feats_in).last_hidden_state |
| y = dec(h.transpose(1, 2)) |
| y = y.squeeze().float().cpu().numpy() |
| o0 = s * up_num // up_den |
| L = min(len(y), len(out) - o0) |
| ramp = np.ones(L, dtype="float32") |
| if len(bounds) > 1: |
| r = min(int(2.0 * out_sr), L // 2) |
| if r > 0: |
| ramp[:r] = np.linspace(0, 1, r) |
| ramp[-r:] = np.linspace(1, 0, r) |
| out[o0:o0 + L] += y[:L] * ramp |
| wsum[o0:o0 + L] += ramp |
| valid = wsum > 1e-6 |
| out[valid] /= wsum[valid] |
| return out[: n * up_num // up_den] |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--input", default="audio", help="audio file or directory") |
| ap.add_argument("--out-dir", default="audio/out") |
| ap.add_argument("--fe-adapter", default="checkpoints/fe_adapter_full.pt", |
| help="w2vbert: LoRA adapter (merged into the base at load)") |
| ap.add_argument("--decoder", default="checkpoints/decoder_only.pt") |
| ap.add_argument("--fe-arch", default="w2vbert", choices=["w2vbert", "distilhubert"], |
| help="distilhubert = the ~37M realtime FE (enhancer-lite); pass the " |
| "stage-1 checkpoint via --fe-adapter (it holds the full FE)") |
| ap.add_argument("--chunk", type=float, default=None, |
| help="window seconds, crossfaded with 2 s overlap. DEFAULT IS ARCH-AWARE: " |
| "4.0 for w2vbert (its SeamlessM4T fbank normalises PER UTTERANCE, so a " |
| "whole-file pass is out-of-distribution and drops speech), 0 for " |
| "distilhubert (raw waveform, do_normalize=False -> verified BIT-IDENTICAL " |
| "at 6/10/20/29 s, so windowing only costs ~2x compute). Pass a value to " |
| "override; 0 = single straight pass.") |
| ap.add_argument("--mono", action="store_true", help="downmix to mono instead of per-channel") |
| ap.add_argument("--out-lowpass", type=float, default=0.0, |
| help="low-pass the decoder output at this Hz (0 = off, the DEFAULT since 2026-10-03). From 2026-08-07 to " |
| "2026-10-03 the default was 8000: measured on the 58-clip EMGS benchmark as a small gain on both ASRs " |
| "(WER 56.18 vs 57.63 whisper-large-v3; 50.00 vs 51.77 Malaysian-whisper-turbo) and on UTMOS (2.230 vs " |
| "2.205). It also removed the whole 8-24 kHz band of a 48 kHz model, so every full-band measurement read " |
| "a 10 kHz ceiling that was this filter, not the decoder (ABLATION.md 2026-10-03). The user removed it: " |
| "the shipped output is the decoder's native band. Pass --out-lowpass 8000 to reproduce the old outputs.") |
| ap.add_argument("--bf16", type=int, default=1) |
| ap.add_argument("--device", default="cuda") |
| a = ap.parse_args() |
| dev = torch.device(a.device if torch.cuda.is_available() else "cpu") |
| torch.set_float32_matmul_precision("medium") |
| os.makedirs(a.out_dir, exist_ok=True) |
|
|
| if os.path.isdir(a.input): |
| files = sorted(f for f in glob.glob(os.path.join(a.input, "*")) |
| if f.lower().endswith(AUDIO_EXTS)) |
| else: |
| files = [a.input] |
| if not files: |
| raise SystemExit(f"no audio under {a.input}") |
|
|
| if a.fe_arch == "distilhubert": |
| import fe_distilhubert as fedh |
| fe = fedh.load_fe(a.fe_adapter, dev) |
| proc = fedh.make_processor() |
| dec, out_sr = load_decoder(a.decoder, dev, input_channel=fe.out_dim) |
| |
| |
| |
| |
| |
| |
| |
| |
| if a.chunk is None: |
| a.chunk = 4.0 |
| else: |
| fe = load_fe(a.fe_adapter, dev) |
| dec, out_sr = load_decoder(a.decoder, dev) |
| proc = AutoFeatureExtractor.from_pretrained(SSL_MODEL) |
| |
| |
| |
| |
| if a.chunk is None: |
| a.chunk = 4.0 |
| khz = f"{out_sr // 1000}k" |
|
|
| for path in files: |
| name = os.path.splitext(os.path.basename(path))[0] |
| data, sr = sf.read(path, always_2d=True, dtype="float32") |
| if a.mono: |
| data = data.mean(axis=1, keepdims=True) |
| nch = data.shape[1] |
| t0 = time.time() |
| chans = [] |
| for c in range(nch): |
| x = data[:, c] |
| x16 = (torchaudio.functional.resample(torch.from_numpy(x)[None], sr, FE_SR)[0].numpy() |
| if sr != FE_SR else x) |
| x16 = _peak_norm(x16, 0.95) |
| yout = restore_channel(x16, fe, dec, proc, dev, a.chunk, bool(a.bf16), out_sr) |
| if a.out_lowpass and a.out_lowpass < out_sr / 2: |
| yout = _lowpass(yout, out_sr, a.out_lowpass) |
| chans.append(_peak_norm(yout, 0.97)) |
| L = max(len(c) for c in chans) |
| chans = [np.pad(c, (0, L - len(c))) for c in chans] |
| restored = np.stack(chans, axis=1) |
|
|
| |
| orig = torchaudio.functional.resample(torch.from_numpy(data.T), sr, out_sr).T.numpy() |
| sf.write(os.path.join(a.out_dir, f"{name}_orig{khz}.wav"), _peak_norm(orig, 0.97), out_sr) |
| sf.write(os.path.join(a.out_dir, f"{name}_restored{khz}.wav"), restored, out_sr) |
| dur = data.shape[0] / sr |
| log(f"[ok] {name}: {nch}ch {sr}Hz {dur:.1f}s -> {khz} in {time.time()-t0:.1f}s " |
| f"(RTF {(time.time()-t0)/dur:.2f})") |
|
|
| log(f"[done] outputs in {a.out_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
