Blame for infer_callcentre.py

#!/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 `<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

# 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()