apbwe_wrapper.py

"""
apbwe_wrapper.py — thin adapter between engine_ihya_v4.py's cascade and the
vendored AP-BWE model (apbwe_vendor/), replacing the dead SIDON slot.

AP-BWE (yxlu-0102/AP-BWE, MIT license): an amplitude+phase prediction
bandwidth-extension model. All-CNN (ConvNeXt blocks), no vocoder/diffusion
stage — it predicts wideband magnitude+phase spectra directly from the
narrowband STFT, then ISTFTs back to waveform. No internal language/text
model, nothing resembling word-level generation; the failure mode here is
spectral over/under-smoothing, not word substitution — much closer to the
masking-model risk profile than a full resynthesis pipeline.

CONFIG: expects checkpoints/<bucket>/{config.json, g_*} to exist, where
<bucket> matches the bandwidth pair the checkpoint was trained for (e.g.
"8kto48k" — Nyquist ~4kHz input, the right bucket for telephone-tier
audio). Checkpoints are NOT in the AP-BWE git repo (only a README pointing
at Google Drive) — download the matching folder yourself and place it at
apbwe_vendor/checkpoints/<bucket>/ before deploying.

This module never raises out to its caller — every failure path returns
False so the cascade (SIDON skip → here → VoiceFixer → DSP fallback)
degrades gracefully instead of crashing a job.
"""
import os
import glob
import json

try:
    import torch
    import torchaudio
    import torchaudio.functional as aF
    import soundfile as sf
    from apbwe_vendor.env import AttrDict
    from apbwe_vendor.datasets.dataset import amp_pha_stft, amp_pha_istft
    from apbwe_vendor.models.model import APNet_BWE_Model
    APBWE_DEPS_OK = True
    _DEPS_ERROR = None
except Exception as _e:
    APBWE_DEPS_OK = False
    _DEPS_ERROR = f"{type(_e).__name__}: {_e}"

_CHECKPOINT_BUCKET = os.environ.get("APBWE_CHECKPOINT_BUCKET", "8kto48k")
_BASE = os.path.dirname(os.path.abspath(__file__))
_CHECKPOINT_DIR = os.path.join(_BASE, "apbwe_vendor", "checkpoints", _CHECKPOINT_BUCKET)

_model = None
_h = None
_device = None
_load_error = None


def _load():
    """Lazy singleton load — first call pays the cost, every job after
    reuses the in-memory model. Records _load_error once so repeated calls
    don't re-attempt a load that's already known to fail."""
    global _model, _h, _device, _load_error
    if _model is not None or _load_error is not None:
        return

    if not APBWE_DEPS_OK:
        _load_error = f"imports failed — {_DEPS_ERROR}"
        return

    config_file = os.path.join(_CHECKPOINT_DIR, "config.json")
    if not os.path.isfile(config_file):
        _load_error = f"no config.json at {_CHECKPOINT_DIR} — checkpoint not deployed"
        return

    ckpt_candidates = sorted(glob.glob(os.path.join(_CHECKPOINT_DIR, "g_*")))
    if not ckpt_candidates:
        _load_error = f"no checkpoint (g_*) in {_CHECKPOINT_DIR}"
        return

    try:
        with open(config_file) as f:
            _h = AttrDict(json.loads(f.read()))
        _device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        _model = APNet_BWE_Model(_h).to(_device)
        state = torch.load(ckpt_candidates[-1], map_location=_device)
        _model.load_state_dict(state["generator"])
        _model.eval()
    except Exception as e:
        _model = None
        _load_error = f"checkpoint/model load failed: {e}"


def restore(wav_in: str, wav_out: str) -> bool:
    """Bandwidth-extend wav_in -> wav_out using AP-BWE. Returns True only
    if wav_out was actually written with content; False on any failure
    (including 'not deployed yet') so the caller's cascade falls through
    to the next stage rather than treating a missing checkpoint as a crash."""
    _load()
    if _model is None:
        if _load_error:
            print(f"[apbwe_wrapper] unavailable: {_load_error}")
        return False

    try:
        audio, orig_sr = torchaudio.load(wav_in)
        audio = audio.to(_device)
        if audio.shape[0] > 1:
            audio = audio.mean(dim=0, keepdim=True)  # AP-BWE is mono-only

        audio_hr = aF.resample(audio, orig_freq=orig_sr, new_freq=_h.hr_sampling_rate)
        audio_lr = aF.resample(audio, orig_freq=orig_sr, new_freq=_h.lr_sampling_rate)
        audio_lr = aF.resample(audio_lr, orig_freq=_h.lr_sampling_rate, new_freq=_h.hr_sampling_rate)
        audio_lr = audio_lr[:, : audio_hr.size(1)]

        with torch.no_grad():
            amp_nb, pha_nb, _com_nb = amp_pha_stft(audio_lr, _h.n_fft, _h.hop_size, _h.win_size)
            amp_wb_g, pha_wb_g, _com_wb_g = _model(amp_nb, pha_nb)
            audio_hr_g = amp_pha_istft(amp_wb_g, pha_wb_g, _h.n_fft, _h.hop_size, _h.win_size)

        sf.write(wav_out, audio_hr_g.squeeze().cpu().numpy(), _h.hr_sampling_rate, "PCM_16")
        return os.path.exists(wav_out) and os.path.getsize(wav_out) > 0
    except Exception as e:
        print(f"[apbwe_wrapper] restore failed: {e}")
        return False