feat: chunking, zipenhancer denoising

1 files changed, 396 insertions, 0 deletions

diff --git a/lib/tests/test_engines.py b/lib/tests/test_engines.py
index 781b46e..9881859 100644
--- a/lib/tests/test_engines.py
+++ b/lib/tests/test_engines.py
@@ -12,6 +12,365 @@ AUDIO_META_MSG = (
 )
 
 
+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 = []
@@ -84,6 +443,43 @@ def test_dfn3_reduces_noise(sr, noisy_speech):
 
 
 @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")