aboutsummaryrefslogtreecommitdiff
path: root/lib/tests/test_engines.py
blob: 781b46ebfe2846d5fce8b4a2789075e3c9a0656d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
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."
)


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_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))