#!/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 `_orig48k.wav` (naive-upsampled input, no model) is also written for an apples-to-apples A/B listen. """ 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 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): ck = torch.load(path, map_location="cpu") ch = ck.get("dec_channels", 3072) dec = dac.model.dac.Decoder(input_channel=1024, channels=ch, rates=[8, 5, 4, 3, 2]) dec.load_state_dict(ck["decoder"]) dec.to(device).eval() for q in dec.parameters(): q.requires_grad_(False) log(f"[dec] DAC decoder channels={ch} ({sum(p.numel() for p in dec.parameters())/1e6:.1f}M); " f"step {ck.get('step')}") return dec 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") @torch.no_grad() def restore_channel(wav16: np.ndarray, fe, dec, proc, device, chunk_s: float, bf16: bool): """wav16: 1-D float32 @16k (peak-normalized) -> restored 1-D float32 @48k. chunk_s <= 0 => single straight pass (recommended: w2v-BERT uses relative/rotary position embeddings and the DAC decoder is convolutional, so a full pass is length-invariant and cleaner). Windowing (chunk_s > 0) is only a memory fallback for very long audio (self-attention is O(T^2)); it crossfades outputs to avoid clicks.""" 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]] out = np.zeros(n * 3 + SR_OUT, dtype="float32") # 16k->48k is x3; 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] @48k y = y.squeeze().float().cpu().numpy() o0 = s * 3 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 * SR_OUT), 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 * 3] 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") ap.add_argument("--decoder", default="checkpoints/decoder_only.pt") ap.add_argument("--chunk", type=float, default=0.0, help="0 = single straight pass (default, recommended). >0 = window seconds " "(memory fallback for very long audio; crossfaded).") ap.add_argument("--mono", action="store_true", help="downmix to mono instead of per-channel") 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}") fe = load_fe(a.fe_adapter, dev) dec = load_decoder(a.decoder, dev) proc = AutoFeatureExtractor.from_pretrained(SSL_MODEL) 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) y48 = restore_channel(x16, fe, dec, proc, dev, a.chunk, bool(a.bf16)) chans.append(_peak_norm(y48, 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 orig = torchaudio.functional.resample(torch.from_numpy(data.T), sr, SR_OUT).T.numpy() sf.write(os.path.join(a.out_dir, f"{name}_orig48k.wav"), _peak_norm(orig, 0.97), SR_OUT) sf.write(os.path.join(a.out_dir, f"{name}_restored48k.wav"), restored, SR_OUT) dur = data.shape[0] / sr log(f"[ok] {name}: {nch}ch {sr}Hz {dur:.1f}s -> 48k 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()