diff --git a/app/pipeline/analyze.py b/app/pipeline/analyze.py index 5af5f4f..c7cd467 100644 --- a/app/pipeline/analyze.py +++ b/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.""" diff --git a/app/pipeline/audio_stats.py b/app/pipeline/audio_stats.py new file mode 100644 index 0000000..432a6dd --- /dev/null +++ b/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 diff --git a/app/pipeline/collect.py b/app/pipeline/collect.py index 70e6757..5c46742 100644 --- a/app/pipeline/collect.py +++ b/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: diff --git a/app/pipeline/runner.py b/app/pipeline/runner.py index 0a7ad81..41ee9f5 100644 --- a/app/pipeline/runner.py +++ b/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) diff --git a/tests/test_pipeline_collect.py b/tests/test_pipeline_collect.py index f69810f..c5f56f0 100644 --- a/tests/test_pipeline_collect.py +++ b/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) diff --git a/tests/test_pipeline_runner.py b/tests/test_pipeline_runner.py index 69cb265..49a06c8 100644 --- a/tests/test_pipeline_runner.py +++ b/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}