feat: chunking, zipenhancer denoising

1 files changed, 250 insertions, 0 deletions

diff --git a/lib/tests/test_chunking.py b/lib/tests/test_chunking.py
new file mode 100644
index 0000000..f00c04e
--- /dev/null
+++ b/lib/tests/test_chunking.py
@@ -0,0 +1,250 @@
+import numpy as np
+import pytest
+
+from producer.engines.base import blend
+from producer.engines.chunking import apply_chunked, plan_chunks, stitch
+
+
+def test_plan_chunks_basic():
+    assert plan_chunks(0, 5, 2) == []
+    assert plan_chunks(4, 10, 2) == [(0, 4)]
+    assert plan_chunks(10, 5, 0) == [(0, 5), (5, 10)]
+    assert plan_chunks(10, 5, 2) == [(0, 5), (3, 8), (6, 10)]
+    assert plan_chunks(7, 5, 2) == [(0, 5), (3, 7)]
+
+
+def test_plan_chunks_last_span_covers_tail():
+    spans = plan_chunks(11, 5, 4)
+    assert spans[-1][1] == 11
+    assert spans[-1][1] - spans[-1][0] >= 5
+
+
+def test_plan_chunks_rejects_overlap_ge_chunk():
+    with pytest.raises(ValueError):
+        plan_chunks(100, 5, 5)
+    with pytest.raises(ValueError):
+        plan_chunks(100, 5, 6)
+
+
+def test_stitch_identity_pieces_reconstruct_signal():
+    rng = np.random.default_rng(0)
+    x = rng.standard_normal(10_000).astype(np.float32)
+    spans = plan_chunks(x.size, 2500, 300)
+    y = stitch(spans, x.size, (x[a:b] for a, b in spans))
+    assert y.dtype == np.float32
+    assert y.shape == x.shape
+    np.testing.assert_allclose(y, x, atol=1e-5)
+
+
+def test_stitch_single_span_passthrough():
+    x = np.arange(100, dtype=np.float32)
+    y = stitch([(0, 100)], 100, [x])
+    np.testing.assert_array_equal(y, x)
+
+
+def test_stitch_crossfade_of_complementary_pieces():
+    # piece A is silence, piece B is a full-scale ramp; the crossfade zone
+    # must be a smooth blend, not a jump.
+    spans = [(0, 10), (5, 15)]
+    pieces = [np.zeros(10, dtype=np.float32), np.ones(10, dtype=np.float32)]
+    y = stitch(spans, 15, pieces)
+    mid = y[7]  # 50% through the overlap
+    assert 0.4 < mid < 0.6
+    assert y[0] == 0.0
+    assert y[14] == 1.0
+
+
+def test_stitch_pads_short_pieces():
+    spans = [(0, 10), (5, 15)]
+    pieces = [np.arange(10, dtype=np.float32), np.arange(8, dtype=np.float32)]
+    y = stitch(spans, 15, pieces)
+    assert y.shape == (15,)
+    assert np.all(np.isfinite(y))
+
+
+def test_apply_chunked_identity_matches_whole_file():
+    rng = np.random.default_rng(1)
+    x = rng.standard_normal(50_000).astype(np.float32)
+    sr = 8000
+    calls = []
+
+    def fn(chunk):
+        calls.append(chunk.size)
+        return chunk
+
+    y = apply_chunked(x, sr, 7.0, 0.05, fn)
+    assert calls == [50_000]
+    np.testing.assert_array_equal(y, x)
+
+    calls.clear()
+    y = apply_chunked(x, sr, 0.5, 0.05, fn)
+    assert len(calls) >= 10
+    assert all(c <= 4000 for c in calls)
+    np.testing.assert_allclose(y, x, atol=1e-5)
+
+
+def test_apply_chunked_zero_disables_chunking():
+    x = np.zeros(10, dtype=np.float32)
+    calls = []
+
+    def fn(chunk):
+        calls.append(chunk.size)
+        return chunk
+
+    apply_chunked(x, 8000, 0.0, 0.5, fn)
+    assert calls == [10]
+
+
+def test_apply_chunked_context_reconstructs_identity():
+    rng = np.random.default_rng(3)
+    x = rng.standard_normal(30_000).astype(np.float32)
+    y = apply_chunked(x, 8000, 1.5, 0.25, lambda c: c, context_s=0.5)
+    np.testing.assert_allclose(y, x, atol=1e-5)
+
+
+def test_apply_chunked_context_feeds_padded_chunks():
+    x = np.ones(5_000, dtype=np.float32)
+    sizes = []
+
+    def fn(chunk):
+        sizes.append(chunk.size)
+        return chunk
+
+    apply_chunked(x, 1000, 2.0, 0.0, fn, context_s=0.5)
+    # plan_chunks(5000, 2000, 0) -> (0,2000), (2000,4000), (4000,5000)
+    assert sizes == [2500, 3000, 1500]
+
+
+def test_apply_chunked_context_trims_cold_start_artifacts():
+    x = np.ones(6_000, dtype=np.float32)
+
+    def fn(chunk):
+        out = chunk.copy()
+        out[0] = 0.0  # cold-start artifact at the start of every model call
+        return out
+
+    broken = apply_chunked(x, 1000, 2.0, 0.0, fn)
+    assert int(np.sum(broken == 0.0)) > 1
+
+    y = apply_chunked(x, 1000, 2.0, 0.0, fn, context_s=0.5)
+    assert y[0] == 0.0  # only the true file start stays degraded
+    assert np.all(y[1:] == 1.0)
+
+
+def test_apply_chunked_context_survives_oom_retry():
+    x = np.ones(8_000, dtype=np.float32)
+    attempts = []
+
+    def fn(chunk):
+        attempts.append(chunk.size)
+        if chunk.size > 3000:
+            raise MemoryError("simulated oom")
+        out = chunk.copy()
+        out[0] = 0.0
+        return out
+
+    y = apply_chunked(x, 1000, 6.0, 0.0, fn, min_chunk_s=1.0, context_s=0.5)
+    assert max(attempts) > 3000
+    assert y[0] == 0.0
+    assert np.all(y[1:] == 1.0)
+
+
+def test_apply_chunked_retries_smaller_on_oom():
+    x = np.random.default_rng(2).standard_normal(40_000).astype(np.float32)
+    sr = 8000
+    attempts = []
+
+    def fn(chunk):
+        attempts.append(chunk.size)
+        if chunk.size > 3000:
+            raise MemoryError("simulated oom")
+        return chunk
+
+    y = apply_chunked(x, sr, 5.0, 0.0, fn, min_chunk_s=0.25)
+    assert attempts[0] == 40_000
+    assert attempts[-1] == 2500
+    assert max(a for a in attempts if a <= 3000) == 2500
+    np.testing.assert_allclose(y, x, atol=1e-5)
+
+
+def test_apply_chunked_oom_gives_up_at_min_chunk():
+    x = np.zeros(40_000, dtype=np.float32)
+
+    def fn(chunk):
+        raise MemoryError("always")
+
+    with pytest.raises(MemoryError):
+        apply_chunked(x, 8000, 1.0, 0.0, fn, min_chunk_s=0.5)
+
+
+def test_apply_chunked_progress_callback():
+    x = np.random.default_rng(3).standard_normal(20_000).astype(np.float32)
+    sr = 8000
+    seen: list[tuple[int, int]] = []
+    y = apply_chunked(x, sr, 0.5, 0.0, lambda c: c, on_progress=lambda d, t: seen.append((d, t)))
+    total = len(plan_chunks(x.size, round(0.5 * sr), 0))
+    assert seen[0] == (0, total)
+    assert seen[-1] == (total, total)
+    assert [d for d, _t in seen[1:]] == list(range(1, total + 1))
+    np.testing.assert_allclose(y, x, atol=1e-6)
+
+
+def test_apply_chunked_progress_resets_after_oom_retry():
+    x = np.random.default_rng(4).standard_normal(40_000).astype(np.float32)
+    sr = 8000
+    seen: list[tuple[int, int]] = []
+
+    def fn(chunk):
+        if chunk.size > 3000:
+            raise MemoryError("simulated oom")
+        return chunk
+
+    y = apply_chunked(
+        x, sr, 5.0, 0.0, fn, min_chunk_s=0.25, on_progress=lambda d, t: seen.append((d, t))
+    )
+    assert seen[0] == (0, 1)  # first attempt is one whole-file chunk
+    done, total = seen[-1]
+    assert done == total and total > 2  # retried into smaller chunks
+    np.testing.assert_allclose(y, x, atol=1e-5)
+
+
+def test_apply_chunked_progress_absent_when_disabled():
+    x = np.zeros(10_000, dtype=np.float32)
+    y = apply_chunked(x, 8000, 0.5, 0.0, lambda c: c, on_progress=None)
+    assert y.shape == x.shape
+
+
+def test_apply_chunked_whole_file_mode_falls_back_on_oom():
+    x = np.random.default_rng(5).standard_normal(10_000).astype(np.float32)
+    sr = 1000
+
+    def fn(chunk):
+        if chunk.size > 1500:
+            raise MemoryError("simulated oom")
+        return chunk * 2
+
+    y = apply_chunked(x, sr, 0.0, 0.0, fn, min_chunk_s=0.5)
+    np.testing.assert_allclose(y, x * 2, atol=1e-6)
+
+
+def test_blend_float32_no_float64_temporaries():
+    x = np.full(1000, 0.25, dtype=np.float32)
+    y = np.full(1000, 0.75, dtype=np.float32)
+    out = blend(x, y, 0.5)
+    assert out.dtype == np.float32
+    np.testing.assert_allclose(out, 0.5, atol=1e-7)
+    assert blend(x, y, 1.0) is y
+    np.testing.assert_allclose(blend(x, y, 0.0), x, atol=1e-7)
+
+
+def test_blend_realigns_resample_drift():
+    # a 48k -> 16k -> 48k round trip can come back a sample or two long
+    # (resample_poly emits ceil(n * up/down) per hop); blend must cope
+    x = np.ones(5, dtype=np.float32)
+    long_y = np.full(7, 0.5, dtype=np.float32)
+    np.testing.assert_allclose(blend(x, long_y, 0.5), np.full(5, 0.75), atol=1e-7)
+
+    short_y = np.full(4, 0.5, dtype=np.float32)
+    out = blend(x, short_y, 0.5)
+    assert out.shape == x.shape
+    np.testing.assert_allclose(out, [0.75, 0.75, 0.75, 0.75, 0.5], atol=1e-7)