perf(pipeline): single-pass streamed peaks + presence

app/pipeline/analyze.py

@@ -214,38 +214,6 @@ def _load_audio_ffmpeg(
     return y, sr
 
 
-def compute_stem_presence(stems_dir: Path, selected_stems: list[str]) -> dict[str, int]:
-    """Load each extracted stem WAV, compute mean absolute amplitude, normalize
-    to 0-100. Only the stems that were selected (and therefore extracted) are
-    measured; the rest are omitted from the returned dict."""
-    import numpy as np
-
-    result: dict[str, int] = {}
-    rms_values: dict[str, float] = {}
-
-    for name in selected_stems:
-        wav_path = stems_dir / f"{name}.wav"
-        if not wav_path.is_file():
-            continue
-        loaded = _load_audio_ffmpeg(wav_path, sr=22050, duration=180.0)
-        if loaded is None:
-            continue
-        y, _ = loaded
-        rms_values[name] = float(np.sqrt(np.mean(y**2)))
-
-    if not rms_values:
-        return result
-
-    max_rms = max(rms_values.values())
-    if max_rms < 1e-9:
-        return {name: 0 for name in rms_values}
-
-    for name, rms in rms_values.items():
-        result[name] = max(0, min(100, round(rms / max_rms * 100)))
-
-    return result
-
-
 def analyze(job: Job, source: Path) -> tuple[int | None, str | None]:
     """Best-effort BPM and key detection. On failure, returns (None, None)
     and leaves job fields untouched -- the chips stay as placeholders."""

app/pipeline/audio_stats.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+
+
+def scan_stem(path: Path, buckets: int = 1500) -> tuple[list[list[float]], float]:
+    """One streamed pass over the WAV at `path`: per-bucket [min, max] over
+    channel 0 (for the waveform display) and RMS over channel 0 (for stem
+    presence) -- both derived from the same blocks, so a stem is only
+    decoded once instead of twice.
+
+    Constant memory via sf.blocks() -- a block is
+    ~frames/buckets * channels * 4 bytes, a few MB even for a 20-minute
+    stereo stem -- instead of sf.read()'s full-file load (#286, ~420 MB for
+    the same file)."""
+    info = sf.info(str(path))
+    frames = info.frames
+    if frames == 0:
+        return [], 0.0
+
+    # Floor division, matching the old sf.read()-then-chunk implementation's
+    # `n // buckets` exactly: sequential fixed-size blocks with the leftover
+    # remainder folded into one final partial block, so peaks are bit-for-bit
+    # identical to before, just computed one block at a time instead of after
+    # loading the whole file.
+    blocksize = max(1, frames // buckets)
+    result: list[list[float]] = []
+    sumsq = 0.0
+    n = 0
+    for block in sf.blocks(str(path), blocksize=blocksize, dtype="float32", always_2d=True):
+        ch = block[:, 0]
+        if ch.size == 0:
+            continue
+        result.append([float(np.min(ch)), float(np.max(ch))])
+        sumsq += float(np.sum(ch.astype(np.float64) ** 2))
+        n += ch.size
+
+    rms = math.sqrt(sumsq / n) if n else 0.0
+    return result[:buckets], rms

app/pipeline/collect.py

@@ -7,9 +7,6 @@
 import time
 from pathlib import Path
 
-import numpy as np
-import soundfile as sf
-
 from app.core.config import (
     DEMUCS_MODEL,
     FAILED_TTL_SECONDS,
@@ -23,6 +20,7 @@
 from app.core.registry import persist as registry_persist
 from app.core.registry import remove as registry_remove
 from app.core.registry import set_proc
+from app.pipeline.audio_stats import scan_stem
 
 logger = logging.getLogger("stemdeck.collect")
 
@@ -187,38 +185,37 @@ def make_selected_mix(job: Job, stems_dir: Path, found: list[str]) -> Path | Non
 _PEAK_POINTS = 1500  # matches OVERVIEW_WAVE_POINTS in player.js
 
 
-def compute_stem_peaks(stems_dir: Path, stem_names: list[str]) -> None:
-    """Compute and cache [min, max] waveform peaks for each stem.
-    Failure is non-fatal — missing peaks.json degrades to client-side decode."""
+def compute_stem_peaks(stems_dir: Path, stem_names: list[str]) -> dict[str, float]:
+    """Compute and cache [min, max] waveform peaks for each stem via a single
+    streamed pass per file (#286), and return each stem's RMS from that same
+    pass so the caller can derive stem presence without a second decode
+    (#287). Peaks failure is non-fatal — missing peaks.json degrades to
+    client-side decode; a stem missing from the returned dict is simply
+    excluded from presence."""
     peaks: dict[str, list[list[float]]] = {}
+    rms_values: dict[str, float] = {}
     for name in stem_names:
         path = stems_dir / f"{name}.wav"
         if not path.is_file():
             continue
         try:
-            data, _ = sf.read(path, dtype="float32", always_2d=True)
-            ch = data[:, 0]
-            n = len(ch)
-            if n == 0:
+            result, rms = scan_stem(path, _PEAK_POINTS)
+            if not result:
                 continue
-            chunk = max(1, n // _PEAK_POINTS)
-            result: list[list[float]] = []
-            for i in range(0, n, chunk):
-                block = ch[i : i + chunk]
-                result.append([float(np.min(block)), float(np.max(block))])
-            peaks[name] = result[:_PEAK_POINTS]
+            peaks[name] = result
+            rms_values[name] = rms
         except Exception:
             logger.warning("could not compute peaks for %s/%s", stems_dir.name, name, exc_info=True)
 
-    if not peaks:
-        return
+    if peaks:
+        try:
+            tmp = stems_dir / "peaks.json.tmp"
+            tmp.write_text(json.dumps(peaks), encoding="utf-8")
+            tmp.replace(stems_dir / "peaks.json")
+        except Exception:
+            logger.warning("could not write peaks.json for %s", stems_dir.name, exc_info=True)
 
-    try:
-        tmp = stems_dir / "peaks.json.tmp"
-        tmp.write_text(json.dumps(peaks), encoding="utf-8")
-        tmp.replace(stems_dir / "peaks.json")
-    except Exception:
-        logger.warning("could not write peaks.json for %s", stems_dir.name, exc_info=True)
+    return rms_values
 
 
 def sweep_old_jobs(jobs_dir: Path) -> None:

app/pipeline/runner.py

@@ -12,7 +12,7 @@
 from app.core.config import DEMUCS_MODEL, TIMEOUT_FFMPEG
 from app.core.models import Job, JobCancelled, _set
 from app.core.registry import persist as persist_registry
-from app.pipeline.analyze import analyze, compute_stem_presence
+from app.pipeline.analyze import analyze
 from app.pipeline.collect import (
     cleanup_source,
     collect,
@@ -133,6 +133,19 @@ def _lap(job: Job, name: str, start: float) -> float:
     return now
 
 
+def _presence_from_rms(rms_values: dict[str, float]) -> dict[str, int]:
+    """Normalize per-stem RMS to 0-100 relative to the loudest stem -- exact
+    logic the old analyze.compute_stem_presence used, now fed by the single
+    streamed pass in compute_stem_peaks (#287) instead of a second full
+    decode of every stem."""
+    if not rms_values:
+        return {}
+    max_rms = max(rms_values.values())
+    if max_rms < 1e-9:
+        return {name: 0 for name in rms_values}
+    return {name: max(0, min(100, round(rms / max_rms * 100))) for name, rms in rms_values.items()}
+
+
 def _run_common(job: Job, source: Path, job_dir: Path) -> None:
     """Analyze → separate → collect → mix. Shared by both YouTube and local
     upload pipelines after their respective source acquisition steps."""
@@ -145,7 +158,6 @@ def _run_common(job: Job, source: Path, job_dir: Path) -> None:
     mark = _lap(job, "separate", mark)
     found = collect(job, stems_root, job_dir)
     stems_dir = job_dir / "stems"
-    job.stem_presence = compute_stem_presence(stems_dir, found)
     # Source (100-300 MB or the local upload) is no longer needed after
     # collect; delete it before the ffmpeg amix steps in case scratch space
     # is tight.
@@ -171,7 +183,13 @@ def _run_common(job: Job, source: Path, job_dir: Path) -> None:
     all_stem_names = [s["name"] for s in job.stems]
     if mix_path is not None and mix_path.stem not in all_stem_names:
         all_stem_names.append(mix_path.stem)
-    compute_stem_peaks(stems_dir, all_stem_names)
+    # "original"/"mix" get peaks (all_stem_names) but are excluded from
+    # presence (found -- the demucs-produced stems only), matching the old
+    # two-pass behavior.
+    rms_values = compute_stem_peaks(stems_dir, all_stem_names)
+    job.stem_presence = _presence_from_rms(
+        {name: rms for name, rms in rms_values.items() if name in found}
+    )
     _lap(job, "post", mark)
 
 

tests/test_pipeline_collect.py

@@ -6,6 +6,8 @@
 from pathlib import Path
 
 import numpy as np
+import pytest
+import soundfile as sf
 
 from app.pipeline.collect import _PEAK_POINTS, compute_stem_peaks
 
@@ -99,8 +101,70 @@ def test_non_fatal_on_corrupt_wav(tmp_path):
     _write_wav(stems_dir / "drums.wav", [0.1, -0.1])
 
     # Should not raise; drums should still be computed
-    compute_stem_peaks(stems_dir, ["vocals", "drums"])
+    rms_values = compute_stem_peaks(stems_dir, ["vocals", "drums"])
 
     data = json.loads((stems_dir / "peaks.json").read_text())
     assert "drums" in data
     assert "vocals" not in data
+    assert "drums" in rms_values
+    assert "vocals" not in rms_values
+
+
+# ─── #287: RMS returned from the same streamed pass ──────────────────────────
+
+
+def test_returns_rms_matching_full_load_reference(tmp_path):
+    stems_dir = tmp_path / "stems"
+    stems_dir.mkdir()
+    sr = 44100
+    t = np.linspace(0, 2, sr * 2, endpoint=False)
+    samples = (np.sin(2 * np.pi * 440 * t) * 0.6).tolist()
+    _write_wav(stems_dir / "vocals.wav", samples, sr)
+
+    rms_values = compute_stem_peaks(stems_dir, ["vocals"])
+
+    reference, _ = sf.read(stems_dir / "vocals.wav", dtype="float32", always_2d=True)
+    expected_rms = float(np.sqrt(np.mean(reference[:, 0].astype(np.float64) ** 2)))
+    assert rms_values["vocals"] == pytest.approx(expected_rms, rel=1e-3)
+
+
+def test_missing_stem_excluded_from_rms(tmp_path):
+    stems_dir = tmp_path / "stems"
+    stems_dir.mkdir()
+    _write_wav(stems_dir / "drums.wav", [0.1, -0.1])
+
+    rms_values = compute_stem_peaks(stems_dir, ["vocals", "drums"])
+
+    assert "drums" in rms_values
+    assert "vocals" not in rms_values
+
+
+def test_peaks_match_full_load_reference(tmp_path):
+    """Golden test: the streamed implementation's peaks must match the old
+    full-load (sf.read + manual chunking) implementation within float
+    tolerance for a multi-tone signal."""
+    stems_dir = tmp_path / "stems"
+    stems_dir.mkdir()
+    sr = 44100
+    t = np.linspace(0, 3, sr * 3, endpoint=False)
+    samples = (0.5 * np.sin(2 * np.pi * 220 * t) + 0.3 * np.sin(2 * np.pi * 1760 * t)).tolist()
+    _write_wav(stems_dir / "vocals.wav", samples, sr)
+
+    compute_stem_peaks(stems_dir, ["vocals"])
+    actual = json.loads((stems_dir / "peaks.json").read_text())["vocals"]
+
+    # Reference: the old sf.read()-then-chunk implementation.
+    data, _ = sf.read(stems_dir / "vocals.wav", dtype="float32", always_2d=True)
+    ch = data[:, 0]
+    n = len(ch)
+    chunk = max(1, n // _PEAK_POINTS)
+    expected = []
+    for i in range(0, n, chunk):
+        block = ch[i : i + chunk]
+        expected.append([float(np.min(block)), float(np.max(block))])
+    expected = expected[:_PEAK_POINTS]
+
+    assert len(actual) == len(expected)
+    for (a_min, a_max), (e_min, e_max) in zip(actual, expected, strict=True):
+        assert a_min == pytest.approx(e_min, abs=1e-4)
+        assert a_max == pytest.approx(e_max, abs=1e-4)

tests/test_pipeline_runner.py

@@ -7,7 +7,12 @@
 
 from app.core.models import Job, JobCancelled
 from app.core.registry import _jobs
-from app.pipeline.runner import _extract_video_track, run_local_pipeline, run_pipeline
+from app.pipeline.runner import (
+    _extract_video_track,
+    _presence_from_rms,
+    run_local_pipeline,
+    run_pipeline,
+)
 
 
 def _ffmpeg_available() -> bool:
@@ -289,3 +294,19 @@ def test_extract_video_track_audio_only_mp4(tmp_path: Path):
 
     assert job.has_video is False
     assert not (job_dir / "video.mp4").exists()
+
+
+# ─── #287: presence normalization (moved from analyze.compute_stem_presence) ─
+
+
+def test_presence_from_rms_normalizes_to_loudest_stem():
+    result = _presence_from_rms({"vocals": 0.5, "drums": 0.25, "bass": 0.0})
+    assert result == {"vocals": 100, "drums": 50, "bass": 0}
+
+
+def test_presence_from_rms_empty_input():
+    assert _presence_from_rms({}) == {}
+
+
+def test_presence_from_rms_all_silent():
+    assert _presence_from_rms({"vocals": 0.0, "drums": 0.0}) == {"vocals": 0, "drums": 0}