diff options
Diffstat (limited to 'lib/tests/test_engines.py')
| -rw-r--r-- | lib/tests/test_engines.py | 515 |
1 files changed, 0 insertions, 515 deletions
diff --git a/lib/tests/test_engines.py b/lib/tests/test_engines.py deleted file mode 100644 index 9881859..0000000 --- a/lib/tests/test_engines.py +++ /dev/null @@ -1,515 +0,0 @@ -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)) |