#!/usr/bin/env python3 """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 `_orig.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 # tm-b300 (B300 = sm_103): torch 2.8's cuDNN has no tuned 1-D conv kernels there (32 TFLOP/s in every precision), # so FAST_CONV1D=1 routes conv1d / conv_transpose1d through unfold + GEMM (runpod/fast_conv1d.py, verified # against the reference to 1e-7). Off by default: other boxes keep cuDNN byte for byte. if os.environ.get("FAST_CONV1D", "0") == "1": import fast_conv1d as _fc if _fc.patch(): # only on Blackwell-class GPUs; elsewhere cuDNN stays and nothing is printed 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: # p e.g. encoder.layers.0.ffn1.output_dense A = ad[p + ".lora_A.default.weight"].float() # (r, in) B = ad[p + ".lora_B.default.weight"].float() # (out, r) delta = scaling * (B @ A) # (out, in) wkey = p + ".weight" sd[wkey] = sd[wkey].float() + delta.to(sd[wkey].dtype) bkey = p + ".base_layer.bias" # trained (lora_only) 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", "") # depth-expanded ckpts, see expand_decoder.py 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]) # `dec_arch` only exists on ckpts written after the HiFi-GAN port; everything older # is a DAC decoder, so the default keeps every existing checkpoint loading unchanged. 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: # noqa: BLE001 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) # 2 s overlap 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]] # A trailing sliver CRASHES the stack: the FE conv frontend strides 320, so <~0.5 s # yields 1-2 frames and the decoder's conv stack raises "Kernel size can't be greater # than actual input size". Latent until chunking became the default (it killed a # full-set run after 244/1289 clips). Absorb any short tail into the previous window # instead of emitting it — the overlap-add already covers that audio. 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() # 16k -> out_sr upsample ratio (x3 at 48 kHz, x1.5 at 24 kHz) — integer math only, so # a 24 kHz model doesn't silently round its overlap-add offsets. up_num, up_den = out_sr, FE_SR out = np.zeros(n * up_num // up_den + out_sr, dtype="float32") # + pad slack 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 # [1,T,1024] y = dec(h.transpose(1, 2)) # [1,1,L] @out_sr 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: # taper edges for crossfade 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() # raw waveform in (no fbank) dec, out_sr = load_decoder(a.decoder, dev, input_channel=fe.out_dim) # WINDOW THIS TOO. A 2026-08-14 conclusion that distilhubert is "length-invariant" # was WRONG: it measured the FEATURE EXTRACTOR (do_normalize=False on raw waveform -> # input features genuinely bit-identical at 6/10/20/29 s) and then wrongly inferred # the MODEL was invariant. The FE is a TRANSFORMER, so attention spans the whole # utterance: re-measured 2026-08-18, the hidden states for the SAME first 6 s drift # 18.9 % / 38.6 % / 39.5 % (mean rel err) when 4 / 14 / 23 s follow. Whole-pass # inference then produced 24 s dropouts and ~29k mute events on the 1,289-clip set, # vs 6.6k at chunk=4. Input invariance does NOT imply representation invariance. 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) # w2v-BERT: SeamlessM4TFeatureExtractor normalises PER UTTERANCE, so the same audio # encodes differently depending on how much follows it (max|diff| 0 -> 10.4 from # 6 s -> 29 s; ~1.61x + 0.37 affine). Training used 6 s windows, so a whole-file pass # is out-of-distribution. Window by default. if a.chunk is None: a.chunk = 4.0 khz = f"{out_sr // 1000}k" # filenames follow the model, not a constant for path in files: name = os.path.splitext(os.path.basename(path))[0] data, sr = sf.read(path, always_2d=True, dtype="float32") # [N, C] 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) # [L, C] # naive-upsampled input (no model) for an A/B reference, at the model's rate 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()