path: root/lib/tests/test_engines.py

import os
import sys
import types
import warnings

import numpy as np
import pytest

AUDIO_META_MSG = (
    "`torchaudio.backend.common.AudioMetaData` has been moved to "
    "`torchaudio.AudioMetaData`. Please update the import path."
)


class _FakeTensor:
    def __init__(self, arr):
        self.arr = np.asarray(arr)

    def unsqueeze(self, axis):
        return _FakeTensor(self.arr[None, ...])

    def detach(self):
        return self

    def cpu(self):
        return self

    def numpy(self):
        return self.arr


def _install_torch_stub(monkeypatch):
    fake = types.ModuleType("torch")
    fake.Tensor = _FakeTensor
    fake.from_numpy = lambda a: _FakeTensor(a)
    fake.cuda = types.SimpleNamespace(is_available=lambda: False)
    monkeypatch.setitem(sys.modules, "torch", fake)


def _install_df_enhance_stub(monkeypatch, transform, calls, records=None):
    pkg = types.ModuleType("df")

    class _State:
        def sr(self):
            return 44100

    class _Model:
        def to(self, dev):
            return self

    def init_df(*_args, **kwargs):
        if records is not None:
            records["init_df"] = kwargs
        return _Model(), _State(), "fake"

    def enhance(_model, _state, audio, **kwargs):
        calls.append(audio.arr.shape[-1])
        if records is not None:
            records["enhance"] = kwargs
        return _FakeTensor(transform(audio.arr.copy()))

    pkg.enhance = types.ModuleType("df.enhance")
    pkg.enhance.init_df = init_df
    pkg.enhance.enhance = enhance
    for sub in ("df.io", "df.logger", "df.utils"):
        mod = types.ModuleType(sub)
        setattr(pkg, sub.split(".")[1], mod)
        monkeypatch.setitem(sys.modules, sub, mod)
    monkeypatch.setitem(sys.modules, "df", pkg)
    monkeypatch.setitem(sys.modules, "df.enhance", pkg.enhance)


def test_dfn3_chunked_engine_stitches_full_length(monkeypatch, sr, noisy_speech):
    from producer.engines import denoise_dfn

    monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None)
    monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None)
    _install_torch_stub(monkeypatch)
    calls: list[int] = []
    _install_df_enhance_stub(monkeypatch, lambda a: a * 0.5 + 0.001, calls)

    x = noisy_speech[: sr * 4]
    y, eng, dev = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=1.0, overlap_s=0.1)
    assert "dfn" in eng
    assert "cpu" in dev
    assert len(calls) > 2
    assert y.shape == x.shape
    np.testing.assert_allclose(y, x * 0.5 + 0.001, atol=1e-6)


def test_dfn3_whole_file_mode_single_call(monkeypatch, sr, noisy_speech):
    from producer.engines import denoise_dfn

    monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None)
    monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None)
    _install_torch_stub(monkeypatch)
    calls: list[int] = []
    _install_df_enhance_stub(monkeypatch, lambda a: a * 0.5, calls)

    x = noisy_speech[: sr * 2]
    y, _eng, _dev = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=0.0)
    assert calls == [x.size]
    np.testing.assert_allclose(y, x * 0.5, atol=1e-7)


def test_dfn3_post_filter_opt_in_no_atten_lim(monkeypatch, sr, noisy_speech):
    from producer.engines import denoise_dfn

    monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None)
    monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None)
    _install_torch_stub(monkeypatch)
    calls: list[int] = []
    records: dict = {}
    _install_df_enhance_stub(monkeypatch, lambda a: a * 0.5, calls, records)

    x = noisy_speech[: sr * 2]
    denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=0.0)
    assert records["init_df"]["post_filter"] is False
    assert records["enhance"] == {}  # stock df_enhance call, no atten-lim override

    calls.clear()
    records.clear()
    denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=0.0, post_filter=True)
    assert records["init_df"]["post_filter"] is True


def test_dfn3_chunked_feeds_context_padding(monkeypatch, sr, noisy_speech):
    from producer.engines import denoise_dfn

    monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None)
    monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None)
    _install_torch_stub(monkeypatch)
    calls: list[int] = []
    _install_df_enhance_stub(monkeypatch, lambda a: a * 0.5, calls)

    x = noisy_speech[: sr * 10]
    y, _eng, _dev = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=3.0, overlap_s=0.5)
    # chunks are widened with context, then trimmed back to the spans
    assert max(calls) > 3.0 * sr
    assert min(calls) >= 3.0 * sr
    np.testing.assert_allclose(y, x * 0.5, atol=1e-7)


_Z_MODEL_REPO = "iic/speech_zipenhancer_ans_multiloss_16k_base"


def _install_zipenhancer_stub(monkeypatch, calls):
    fake = types.ModuleType("zipenhancer")
    fake.MODEL_ZIPENHANCER = _Z_MODEL_REPO

    def denoise(chunk, sample_rate, model=_Z_MODEL_REPO, normalize=True, strength=1.0, **_kw):
        calls.append(
            {"n": chunk.size, "model": model, "normalize": normalize, "strength": strength}
        )
        scale = 0.1 if len(calls) % 2 == 1 else 1.0
        return (chunk * scale, 0.0, chunk.size / sample_rate)

    fake.denoise = denoise
    monkeypatch.setitem(sys.modules, "zipenhancer", fake)


def test_zipenhancer_chunked_normalizes_once(monkeypatch, sr):
    from producer.engines import denoise_zip

    seq: list[str] = []
    monkeypatch.setattr(
        "producer.lazy.ensure",
        lambda pkgs, **_k: seq.append("ensure:" + ",".join(str(p) for p in pkgs)),
    )
    monkeypatch.setattr("producer.lazy.ensure_torch", lambda: seq.append("torch"))
    monkeypatch.setattr(
        "producer.lazy.ensure_import", lambda mod, **_k: seq.append("import:" + mod)
    )
    monkeypatch.setattr("producer.lazy.ensure_call", lambda fn, **_k: (seq.append("call"), fn())[1])
    calls: list[dict] = []
    _install_zipenhancer_stub(monkeypatch, calls)

    n = 4 * 44100
    t = np.arange(n) / sr
    x = (0.5 * np.sin(2 * np.pi * 160.0 * t)).astype(np.float32)
    y, eng, _dev = denoise_zip.denoise(x, sr, 1.0, "cpu", chunk_s=1.0, overlap_s=0.0)
    assert "zipenhancer" in eng
    # torch is pinned before package installs; the undeclared modelscope
    # import ships with the engine; import and call probes run after both
    assert seq == [
        "torch",
        "ensure:zipenhancer==0.3.2,modelscope",
        "import:zipenhancer",
        "call",
        "call",
        "call",
        "call",
    ]
    assert len(calls) == 4
    # the library API takes the full modelscope repo id; the short name
    # would be treated as a repo id and fail with modelscope E3021
    assert all(c["model"] == _Z_MODEL_REPO for c in calls)
    assert all(c["normalize"] is False for c in calls)
    assert y.shape == x.shape
    peak = float(np.max(np.abs(y)))
    assert abs(peak - 10 ** (-3.0 / 20.0)) < 0.01
    even_rms = float(np.sqrt(np.mean(y[: n // 4].astype(np.float64) ** 2)))
    odd_rms = float(np.sqrt(np.mean(y[n // 4 : n // 2].astype(np.float64) ** 2)))
    assert 8.0 < odd_rms / even_rms < 12.0


def test_zipenhancer_resample_roundtrip_length_realigned(monkeypatch):
    # 120s @ 48k round-tripped through the 16k engine can come back a sample
    # or two long (resample_poly emits ceil(n * up/down) per hop); the blend
    # used to crash on the mismatch instead of realigning
    from producer.engines import denoise_zip

    monkeypatch.setattr("producer.lazy.ensure_torch", lambda: None)
    monkeypatch.setattr("producer.lazy.ensure", lambda *_a, **_k: None)
    calls: list[dict] = []
    _install_zipenhancer_stub(monkeypatch, calls)

    sr = 48000
    n = 100001  # 48k -> 16k -> 48k drifts +1 for this length
    x = (0.3 * np.sin(2 * np.pi * 220.0 * np.arange(n) / sr)).astype(np.float32)
    y, _eng, _dev = denoise_zip.denoise(x, sr, 0.5, "cpu")
    assert y.shape == x.shape


def test_mossformer_clearvoice_probe(monkeypatch, sr, noisy_speech):
    from producer.engines import enhance_mossformer

    seq: list[str] = []
    monkeypatch.setattr("producer.lazy.ensure_torch", lambda: seq.append("torch"))
    monkeypatch.setattr(
        "producer.lazy.ensure",
        lambda pkgs, **_k: seq.append("ensure:" + ",".join(str(p) for p in pkgs)),
    )
    monkeypatch.setattr(
        "producer.lazy.ensure_import", lambda mod, **_k: seq.append("import:" + mod)
    )

    fake = types.ModuleType("clearvoice")

    class _FakeCV:
        def __init__(self, task=None, model_name=None):
            pass

        def __call__(self, chunk):
            return (chunk * 0.5, 48000)

    fake.ClearVoice = _FakeCV
    monkeypatch.setitem(sys.modules, "clearvoice", fake)

    x = noisy_speech[: sr * 4]
    y, eng, _dev = enhance_mossformer.enhance(x, 48000, 1.0, "cpu", chunk_s=2.0)
    assert "mossformer2" in eng
    assert seq == ["torch", "ensure:clearvoice==0.1.2", "import:clearvoice"]
    assert y.shape == x.shape
    np.testing.assert_allclose(y, x * 0.5, atol=1e-6)


def _hissy_speech(sr, dur=12.0, noise_db=-42.0, seed=3):
    from conftest import speechish

    x = speechish(dur, sr, level_dbfs=-20.0, seed=seed)
    rng = np.random.default_rng(seed)
    noise = rng.standard_normal(x.size)
    noise *= (10 ** (noise_db / 20.0)) / np.sqrt(np.mean(np.square(noise)))
    return (x + noise).astype(np.float32)


def test_spectral_reduces_steady_noise(sr):
    from producer import meters
    from producer.engines import denoise_spectral

    x = _hissy_speech(sr)
    y, eng, dev = denoise_spectral.denoise(x, sr, 1.0, "cpu")
    assert "spectral" in eng and dev == "cpu"
    assert y.shape == x.shape
    assert np.all(np.isfinite(y))
    assert meters.noise_floor_db(y, sr) < meters.noise_floor_db(x, sr) - 8.0
    assert np.corrcoef(x.astype(np.float64), y.astype(np.float64))[0, 1] > 0.8


def test_spectral_speech_level_flat(sr):
    from producer.engines import denoise_spectral

    x = _hissy_speech(sr)
    y, _eng, _dev = denoise_spectral.denoise(x, sr, 1.0, "cpu")
    frame = int(0.03 * sr)
    nf = x.size // frame

    def frms_db(z):
        rms = np.sqrt(np.mean(z[: nf * frame].reshape(nf, frame).astype(np.float64) ** 2, axis=1))
        return 20.0 * np.log10(rms + 1e-12)

    xdb, ydb = frms_db(x), frms_db(y)
    speech = xdb > np.percentile(xdb, 10) + 12.0
    assert speech.sum() > 20
    # the deterministic engine must not pump the speech level (dfn3's failure mode)
    swing = np.abs(ydb[speech] - xdb[speech])
    assert np.percentile(swing, 95) < 2.0


def test_spectral_strength_zero_is_identity(sr):
    from producer.engines import denoise_spectral

    x = _hissy_speech(sr)
    y, _eng, _dev = denoise_spectral.denoise(x, sr, 0.0, "cpu")
    np.testing.assert_array_equal(y, x)


def test_spectral_sample_aligned(sr):
    from producer.engines import denoise_spectral

    x = _hissy_speech(sr)
    y, _eng, _dev = denoise_spectral.denoise(x, sr, 1.0, "cpu")
    lo, hi = int(sr * 2.0), int(sr * 9.0)
    corrs = {k: float(np.corrcoef(x[lo:hi], y[lo + k : hi + k])[0, 1]) for k in range(-3, 4)}
    assert max(corrs, key=corrs.get) == 0
    assert corrs[0] > 0.8


def test_spectral_suppression_capped(sr):
    from producer import meters
    from producer.engines import denoise_spectral

    noise = (0.008 * np.random.default_rng(5).standard_normal(sr * 6)).astype(np.float32)
    y, _eng, _dev = denoise_spectral.denoise(noise, sr, 1.0, "cpu")
    drop = meters.noise_floor_db(noise, sr) - meters.noise_floor_db(y, sr)
    # bounded: deep enough to matter, never gated to digital silence
    assert 20.0 < drop < 40.0


def test_spectral_chunked_matches_whole_file(sr):
    from producer.engines import denoise_spectral

    x = _hissy_speech(sr, dur=20.0)
    y1, _e, _d = denoise_spectral.denoise(x, sr, 1.0, "cpu", chunk_s=60.0, overlap_s=0.5)
    y2, _e, _d = denoise_spectral.denoise(x, sr, 1.0, "cpu", chunk_s=8.0, overlap_s=0.5)
    assert y1.shape == y2.shape == x.shape
    corr = np.corrcoef(y1.astype(np.float64), y2.astype(np.float64))[0, 1]
    assert corr > 0.99


def test_spectral_profile_global_not_per_chunk(sr):
    """Tail hiss must get the same suppression with or without speech up front.

    The old per-chunk percentile leaked speech into the noise estimate and
    under-suppressed exactly where it matters (between sentences).
    """
    from conftest import speechish

    from producer import meters
    from producer.engines import denoise_spectral

    speech = speechish(8.0, sr, level_dbfs=-20.0, seed=11)
    rng = np.random.default_rng(9)
    noise = rng.standard_normal(sr * 16).astype(np.float32)
    noise *= (10 ** (-40.0 / 20.0)) / np.sqrt(np.mean(np.square(noise)))
    tail_lo, tail_hi = sr * 10, sr * 16
    with_speech = np.concatenate([speech + noise[: speech.size], noise[speech.size :]]).astype(
        np.float32
    )
    hiss_only = noise.copy()
    y1, _e, _d = denoise_spectral.denoise(with_speech, sr, 1.0, "cpu")
    y2, _e, _d = denoise_spectral.denoise(hiss_only, sr, 1.0, "cpu")
    drop1 = meters.noise_floor_db(with_speech[tail_lo:tail_hi], sr) - meters.noise_floor_db(
        y1[tail_lo:tail_hi], sr
    )
    drop2 = meters.noise_floor_db(hiss_only[tail_lo:tail_hi], sr) - meters.noise_floor_db(
        y2[tail_lo:tail_hi], sr
    )
    assert drop1 > 14.0
    assert drop1 > drop2 - 4.0


def _stub_df_modules(calls: list[str]) -> tuple[types.ModuleType, list[types.ModuleType]]:
    pkg = types.ModuleType("df")
    mods = []
    for full in ("df.utils", "df.logger", "df.io"):
        mod = types.ModuleType(full)

        def probe(name: str):
            def fn(*_args):
                calls.append(name)
                return "deadbeef"

            return fn

        for fn_name in ("get_git_root", "get_commit_hash", "get_branch_name"):
            setattr(mod, fn_name, probe(f"{full}.{fn_name}"))
        setattr(pkg, full.split(".")[1], mod)
        mods.append(mod)
    return pkg, mods


def test_dfn_shim_neutralizes_git_probes(monkeypatch):
    from producer.engines import denoise_dfn

    calls: list[str] = []
    pkg, mods = _stub_df_modules(calls)
    for name, mod in zip(("df", "df.utils", "df.logger", "df.io"), (pkg, *mods), strict=True):
        monkeypatch.setitem(sys.modules, name, mod)
    assert pkg.logger.get_commit_hash() == "deadbeef"
    calls.clear()
    denoise_dfn._shim_df_git()
    for mod in mods:
        for fn_name in ("get_git_root", "get_commit_hash", "get_branch_name"):
            assert mod.__dict__[fn_name]() is None
    assert calls == []


def test_dfn_shim_silences_torchaudio_warning():
    from producer.engines import denoise_dfn

    code = "import warnings\nwarnings.warn(MESSAGE, UserWarning)"
    denoise_dfn._shim_torchaudio_backend()
    with warnings.catch_warnings(record=True) as caught:
        exec(compile(code, "df/io.py", "exec"), {"__name__": "df.io", "MESSAGE": AUDIO_META_MSG})
    assert caught == []
    with warnings.catch_warnings(record=True) as caught:
        exec(
            compile(code, "other/mod.py", "exec"),
            {"__name__": "other.mod", "MESSAGE": AUDIO_META_MSG},
        )
    assert len(caught) == 1


def _metrics_floor(x, sr):
    from producer import meters

    return meters.noise_floor_db(x, sr)


@pytest.mark.slow
def test_dfn3_reduces_noise(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("df")
    from producer.engines import denoise_dfn

    x = noisy_speech[: sr * 4]
    y, eng, _dev = denoise_dfn.denoise(x, sr, 1.0, "cpu")
    assert "dfn" in eng
    assert _metrics_floor(y, sr) < _metrics_floor(x, sr) - 5.0
    assert np.corrcoef(x, y.astype(np.float64))[0, 1] > 0.9


@pytest.mark.slow
def test_dfn3_speech_level_flat(sr, noisy_speech):
    """Guard against dfn3's reported failure mode: volume wobble in sentences."""
    pytest.importorskip("torch")
    pytest.importorskip("df")
    from producer.engines import denoise_dfn

    x = noisy_speech[: sr * 8]
    y, _eng, _dev = denoise_dfn.denoise(x, sr, 0.9, "cpu")
    frame = int(0.03 * sr)
    nf = x.size // frame

    def frms_db(z):
        rms = np.sqrt(np.mean(z[: nf * frame].reshape(nf, frame).astype(np.float64) ** 2, axis=1))
        return 20.0 * np.log10(rms + 1e-12)

    xdb, ydb = frms_db(x), frms_db(y)
    speech = xdb > np.percentile(xdb, 10) + 12.0
    swing = np.abs(ydb[speech] - xdb[speech])
    assert np.percentile(swing, 95) < 2.5


@pytest.mark.slow
def test_dfn3_chunked_matches_whole_file(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("df")
    from producer.engines import denoise_dfn

    x = noisy_speech[: sr * 60]
    y_full, _, _ = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=0.0)
    y_chunk, _, _ = denoise_dfn.denoise(x, sr, 1.0, "cpu", chunk_s=15.0, overlap_s=0.5)
    assert y_full.shape == y_chunk.shape == x.shape
    corr = np.corrcoef(y_full.astype(np.float64), y_chunk.astype(np.float64))[0, 1]
    assert corr > 0.99
    assert float(np.max(np.abs(y_full - y_chunk))) < 0.1


@pytest.mark.slow
def test_zipenhancer_reduces_noise(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("zipenhancer")
    from producer.engines import denoise_zip

    x = noisy_speech[: sr * 4]
    y, eng, _ = denoise_zip.denoise(x, sr, 1.0, "cpu")
    assert "zipenhancer" in eng
    assert _metrics_floor(y, sr) < _metrics_floor(x, sr) - 5.0


@pytest.mark.slow
def test_mossformer2_enhances(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("clearvoice")
    from producer.engines import enhance_mossformer

    x = noisy_speech[: sr * 4]
    y, eng, _ = enhance_mossformer.enhance(x, sr, 1.0, "cpu")
    assert "mossformer2" in eng
    assert np.all(np.isfinite(y))


@pytest.mark.slow
def test_resemble_enhance(sr, noisy_speech):
    if not os.environ.get("PRODUCER_TEST_RESEMBLE"):
        pytest.skip("set PRODUCER_TEST_RESEMBLE=1 to run the isolated-venv generative engine")
    from producer.engines import enhance_resemble

    x = noisy_speech[: sr * 4]
    y, eng, _ = enhance_resemble.enhance(x, sr, 1.0, "cpu")
    assert "resemble" in eng
    assert np.all(np.isfinite(y))