From ea9648265c30b9ed4c53f80f9eba18bc4a1221c5 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 22:35:13 +0000 Subject: [PATCH 1/2] =?UTF-8?q?feat(backend):=20pluggable=20stem-separatio?= =?UTF-8?q?n=20backend=20(Demucs=20=E2=86=94=20MSST/BS-RoFormer)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a selectable separation backend so the Demucs stems stage can be swapped for the stronger MSST/BS-RoFormer models from SUC-DriverOld/MSST-WebUI, driving MSSeparator inference directly (no PySide6 GUI). Default stays Demucs; the stems JSON contract (vocals/bass/drums/other, 44.1 kHz) is unchanged. - separation_backend.py: ASA_SEPARATION_BACKEND dispatcher + model registry, mirroring loudness_backend.py (env-select, graceful degrade-to-Demucs on any failure). MSST runs as a subprocess under its own venv (ASA_MSST_PYTHON) so it works on the worker path (server.py hardcodes ./venv), keeps MSST stdout off analyze.py's JSON contract, and isolates MSST's conflicting deps. - scripts/msst_separate_runner.py: runs in the MSST venv; resamples to 44.1 kHz, transposes channels-last→first, maps stems by name (SCNet order varies), forces MSST logging to stderr, emits a one-line JSON manifest. - analyze.py: route all three separation sites (--separate, --pitch-note-only, --mt3-only) through separate_stems_backend. - separation_ab.py + scripts/ab_separation_backends.py: research-only A/B harness — synthetic gain-aligned SI-SDR smoke-test (labelled plumbing, not a quality ranking) plus optional real-track reference-free proxies and fair runtime. - requirements-msst.txt: separate-venv setup pointer (MSST brings its own torch). - tests/test_separation_backend.py: 16 tests (dispatch, fallback, registry, runner helpers) — no MSST install required. - Docs: CLAUDE.md env vars + core-files + scripts + dep-isolation; JSON_SCHEMA.md notes the backend is selectable and the schema backend-agnostic. Default path unchanged: test_analyze (147) and test_phase1_golden (11) green. https://claude.ai/code/session_01PCXgSHYKWnrD4PiRkwzJxR --- CLAUDE.md | 12 +- apps/backend/JSON_SCHEMA.md | 4 +- apps/backend/analyze.py | 10 +- apps/backend/requirements-msst.txt | 41 ++ .../backend/scripts/ab_separation_backends.py | 106 +++++ apps/backend/scripts/msst_separate_runner.py | 232 +++++++++++ apps/backend/separation_ab.py | 381 ++++++++++++++++++ apps/backend/separation_backend.py | 257 ++++++++++++ apps/backend/tests/test_separation_backend.py | 168 ++++++++ 9 files changed, 1204 insertions(+), 7 deletions(-) create mode 100644 apps/backend/requirements-msst.txt create mode 100644 apps/backend/scripts/ab_separation_backends.py create mode 100644 apps/backend/scripts/msst_separate_runner.py create mode 100644 apps/backend/separation_ab.py create mode 100644 apps/backend/separation_backend.py create mode 100644 apps/backend/tests/test_separation_backend.py diff --git a/CLAUDE.md b/CLAUDE.md index 9b1f23f1..268c1f86 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -191,10 +191,12 @@ Operational and one-shot scripts live in two places. They are not on the request `apps/backend/scripts/`: 1. `bootstrap.sh` — recreate the Python 3.11 venv (covered above). `dev.sh` here is a thin shim that `exec`s the repo-root `scripts/dev.sh` full-stack launcher. 2. `render_upload_limit_contract.py` — regenerate the operator-facing upload-limit contract text from `upload_limits.py`. Run after changing the canonical limits. -3. `evaluate_phase1.py`, `evaluate_polyphonic.py`, `evaluate_structure_sweep.py`, `evaluate_beats.py`, `evaluate_loudness_recs.py`, `evaluate_recommendations.py` — offline evaluation harnesses for the Phase 1 detector battery, the research polyphonic transcriber, structure-segmentation parameter sweeps, the beat/downbeat measurement gate, loudness-recommendation reachability, and the recommendation-quality scorer (`GOAL.md`'s recommendation-proof campaign — scores Phase 2 recs against the `tests/fixtures/recommendation_tracks/` answer-key corpus). Wired to `phase1_evaluation.py` / `polyphonic_evaluation.py` / `beat_evaluation.py` / `loudness_rec_evaluation.py` / `recommendation_evaluation.py`. `build_beat_manifest.py` assembles the beat-eval corpus manifest. `emit_deterministic_recs.ts` is the recommendation harness's deterministic-source bridge (Node 23+ native TS; wraps `apps/ui/src/data/abletonDevices.ts` to score the free path). All research-only. The beat gate's optional neural deps (`beat_this`, `mir_eval`) live in `apps/backend/requirements-eval.txt` — install into a separate venv, never the product venv. +3. `evaluate_phase1.py`, `evaluate_polyphonic.py`, `evaluate_structure_sweep.py`, `evaluate_beats.py`, `evaluate_loudness_recs.py`, `evaluate_recommendations.py` — offline evaluation harnesses for the Phase 1 detector battery, the research polyphonic transcriber, structure-segmentation parameter sweeps, the beat/downbeat measurement gate, loudness-recommendation reachability, and the recommendation-quality scorer (`GOAL.md`'s recommendation-proof campaign — scores Phase 2 recs against the `tests/fixtures/recommendation_tracks/` answer-key corpus). Wired to `phase1_evaluation.py` / `polyphonic_evaluation.py` / `beat_evaluation.py` / `loudness_rec_evaluation.py` / `recommendation_evaluation.py`. `build_beat_manifest.py` assembles the beat-eval corpus manifest. `emit_deterministic_recs.ts` is the recommendation harness's deterministic-source bridge (Node 23+ native TS; wraps `apps/ui/src/data/abletonDevices.ts` to score the free path). All research-only. The beat gate's optional neural deps (`beat_this`, `mir_eval`) live in `apps/backend/requirements-eval.txt` — install into a separate venv, never the product venv. The optional MSST separation backend follows the same isolation rule but more strictly: `apps/backend/requirements-msst.txt` is a setup-pointer (not a pinned list) because MSST brings its own conflicting torch/numpy/librosa — it lives in its own venv built from MSST-WebUI's own `requirements.txt`, reached only via the `scripts/msst_separate_runner.py` subprocess (see `separation_backend.py`), never imported into the product venv. 4. `audit_pass1.py`, `genre_check.py`, `replay_catalog_validation.py` — corpus auditing and Live 12 device-catalog validation. `genre_corpus.md` is the corpus manifest. 5. `import_midi_to_ground_truth.py`, `score_polyphonic_clip.py` — corpus-building helpers for the transcription/polyphonic ground-truth fixtures (`tests/fixtures/transcription_tracks/`, `tests/fixtures/polyphonic_tracks/`). Research-only; see those fixtures' READMEs and `docs/LAYER2_EVALUATION.md` / `docs/POLYPHONIC_TRANSCRIPTION_SPIKE.md`. 6. `parity_probe_synth_backends.py` — maintainer probe that renders the same `ClipPlan` through FluidSynth and `symusic.Synthesizer` and reports RMS/peak/spectral-centroid deltas before flipping the `ASA_SAMPLE_SYNTH_BACKEND` auto default. Research/operator-only; not invoked by the runtime. +7. `msst_separate_runner.py` — the MSST separation runner invoked by `separation_backend.py` as a subprocess under `ASA_MSST_PYTHON`. **Runs in the MSST venv, not the product venv** — it imports `MSSeparator` from a `SUC-DriverOld/MSST-WebUI` checkout, writes canonical `vocals/bass/drums/other` 44.1 kHz WAVs, and prints a single-line JSON manifest to stdout. Operator tooling; not a product-venv module. +8. `ab_separation_backends.py` — A/B harness CLI (wraps `separation_ab.py`) comparing the Demucs vs MSST separation backends on quality and runtime: an always-available synthetic SI-SDR smoke-test plus optional real-track reference-free proxies. Writes `.runtime/separation_ab/report.json`. Research-only; deleting `separation_ab.py` restores the product exactly. ## Architecture @@ -289,6 +291,8 @@ Artifact access goes through `artifact_storage.py` rather than direct disk paths 21. **`loudness_backend.py`**: Selectable Phase 1 loudness backend (default-off experiment). `ASA_LOUDNESS_BACKEND=wasm` overrides the four integrated/range/momentary-max/short-term-max LUFS scalars with readings from the native `measure-cli` binary (source-identical to `packages/loudness-spectro-wasm`). `truePeak` and `lufsCurve` always stay on Essentia. Any failure degrades back to Essentia. Default is `essentia` (no-op). +22. **`separation_backend.py`**: Selectable Phase 1 stem-separation backend (default-off experiment). `ASA_SEPARATION_BACKEND=msst` swaps torchaudio Hybrid Demucs (`analyze_audio_io.separate_stems`) for a stronger MSST/BS-RoFormer model from a `SUC-DriverOld/MSST-WebUI` checkout, while keeping the `{stem_name: wav_path}` contract (`vocals/bass/drums/other`, 44.1 kHz) unchanged. Because MSST pins deps that conflict with ASA's (`librosa==0.9.2`, `numpy<2`, its own torch), it runs in its **own** venv and ASA shells out to `scripts/msst_separate_runner.py` under `ASA_MSST_PYTHON` — mirroring how `loudness_backend.py` shells out to `measure-cli`. The subprocess boundary also keeps MSST's stdout/logging off `analyze.py`'s JSON contract (tripwire #1) and isolates its `utils`/`inference` packages. A small model registry (`ASA_MSST_MODEL`, default `scnet_4stem`) maps an id → `{model_type, config, checkpoint}`. Any failure (missing venv/checkout/checkpoint, non-zero exit, unparseable output) degrades back to Demucs. `separate_stems_backend` is the entry point called from `analyze.py`'s three separation sites (measurement `--separate`, `--pitch-note-only`, `--mt3-only`). + The subprocess isolation means `analyze.py` works as a standalone CLI. Check `apps/backend/JSON_SCHEMA.md` before adding new analyzer output fields. Check `apps/backend/ARCHITECTURE.md` for the full HTTP flow and contract details. **Phase 2 (`POST /api/phase2`, legacy compat):** Uploads audio to Gemini inline if ≤100 MiB, or via the Gemini Files API if larger. Phase 1 JSON is appended to the system prompt from `prompts/phase2_system.txt`. Also relevant: `prompts/stem_summary_system.txt` and `prompts/live12_device_catalog.json` (the *prompt-injected* device catalogue — distinct from the runtime-validation `data/live12_catalogue.json` consumed by `phase2_catalogue_gates.py`). Backend defense-in-depth: `server_phase2.py`'s `_validate_phase2_citation_paths` mirrors the frontend citation-existence check and emits `validationWarnings` when a recommendation cites a Phase 1 path that doesn't exist — it flags invented citations rather than failing, since Phase 1 stays authoritative (invariant #1). `phase2_catalogue_gates.apply_live12_catalogue_gates` runs immediately after for source-catalogue checks (see backend file #20 above). @@ -353,6 +357,12 @@ ASA_ENABLE_MT3="0" # set to "1" on the legacy analyze.py C ASA_SAMPLE_SYNTH_BACKEND="auto" # Phase 3 audition-sample synth backend: `auto` (default, prefers FluidSynth and falls back to symusic), `symusic`, or `fluidsynth` to pin one. See apps/backend/sample_synthesis.py. ASA_LOUDNESS_BACKEND="essentia" # Phase 1 loudness source for the LUFS scalars: `essentia` (default, authoritative) or `wasm` to override lufsIntegrated/lufsRange/lufsMomentaryMax/lufsShortTermMax with the asa-dsp (loudness-spectro-wasm) reading via the native measure-cli binary. truePeak + lufsCurve stay Essentia; degrades back to Essentia if measure-cli is unbuilt. Default-off experiment. See apps/backend/loudness_backend.py. ASA_MEASURE_CLI="" # optional absolute path to the measure-cli binary used by ASA_LOUDNESS_BACKEND=wasm; defaults to packages/loudness-spectro-wasm/target/release/measure-cli. +ASA_SEPARATION_BACKEND="demucs" # Phase 1 stem-separation backend: `demucs` (default, torchaudio Hybrid Demucs) or `msst` to drive a stronger MSST/BS-RoFormer model from a SUC-DriverOld/MSST-WebUI checkout. The stems contract (vocals/bass/drums/other WAVs) is unchanged. Runs MSST as a subprocess under its own venv (see ASA_MSST_*); any failure degrades back to Demucs. Default-off experiment. See apps/backend/separation_backend.py + requirements-msst.txt. +ASA_MSST_PYTHON="" # path to the MSST venv interpreter (built from MSST-WebUI's own requirements.txt) used to run scripts/msst_separate_runner.py. Required when ASA_SEPARATION_BACKEND=msst. +ASA_MSST_ROOT="" # path to the MSST-WebUI checkout (added to the runner's sys.path for `from inference.msst_infer import MSSeparator`). Required when ASA_SEPARATION_BACKEND=msst. +ASA_MSST_MODEL="scnet_4stem" # MSST model registry id in separation_backend._MSST_MODEL_REGISTRY. Default `scnet_4stem` (4-stem parity). `bs_roformer_vocals` is research/A-B-only (2-stem, leaves bass/drums empty). Unknown ids fall back to the default. +ASA_MSST_MODEL_DIR="" # optional config/checkpoint root (MSST configs/ + pretrain/ layout); defaults to ASA_MSST_ROOT. +ASA_MSST_DEVICE="" # optional device for MSST inference (auto|cpu|cuda|mps); defaults to auto. ``` Phase 2 is gated by `VITE_ENABLE_PHASE2_GEMINI`. `GEMINI_API_KEY` is backend-only. `SONIC_ANALYZER_ADMIN_KEY` is backend-only and never exposed to clients. diff --git a/apps/backend/JSON_SCHEMA.md b/apps/backend/JSON_SCHEMA.md index d9097464..b4f4a9ad 100644 --- a/apps/backend/JSON_SCHEMA.md +++ b/apps/backend/JSON_SCHEMA.md @@ -274,7 +274,9 @@ Sibling time-series partner for `spectralBalance`. Each row carries all seven ba Type: `object \| null` -Phase 1.B per-stem analytical surface — populated only when Demucs stem separation ran successfully (`--separate`). Null when separation wasn't requested or failed. Phase 2 cites individual stems for element-specific recommendations. +Phase 1.B per-stem analytical surface — populated only when stem separation ran successfully (`--separate`). Null when separation wasn't requested or failed. Phase 2 cites individual stems for element-specific recommendations. + +The separation backend is **selectable** (`ASA_SEPARATION_BACKEND`, default `demucs`; `msst` drives a stronger MSST/BS-RoFormer model — see `separation_backend.py`), but this schema is **backend-agnostic**: stems are always the canonical `drums`/`bass`/`other`/`vocals` at 44.1 kHz, so `stemAnalysis` is unchanged regardless of which backend produced them. For each available stem (`drums` / `bass` / `other` / `vocals`), the same high-value full-mix analyzers run on the stem's audio: diff --git a/apps/backend/analyze.py b/apps/backend/analyze.py index 5a388928..8050909d 100644 --- a/apps/backend/analyze.py +++ b/apps/backend/analyze.py @@ -61,10 +61,10 @@ _demucs_chunked_inference, _load_stem_mono, _load_stem_stereo, - separate_stems, analyze_crepe_pitch, cleanup_stems, ) +from separation_backend import separate_stems_backend # noqa: E402 from analyze_estimate import ( # noqa: E402 _format_duration_label, _estimate_stage_seconds, @@ -1162,7 +1162,7 @@ def _run_pitch_note_translation( if need_separation: temp_dir = stem_output_dir or tempfile.mkdtemp(prefix="asa_pitch_note_stems_") - separated = separate_stems(audio_path, output_dir=temp_dir) + separated = separate_stems_backend(audio_path, output_dir=temp_dir) if isinstance(separated, dict) and separated: stem_paths = separated @@ -1224,9 +1224,9 @@ def _run_mt3_transcription( # handover convention so both stages can short-circuit Demucs. if stems_dir_path is None and stem_output_dir is not None: os.makedirs(stem_output_dir, exist_ok=True) - separated = separate_stems(audio_path, output_dir=stem_output_dir) + separated = separate_stems_backend(audio_path, output_dir=stem_output_dir) if isinstance(separated, dict) and separated: - # separate_stems returns {"bass": "/path/...wav", ...}. Recover + # separate_stems_backend returns {"bass": "/path/...wav", ...}. Recover # the shared parent — all stem files live in one directory. for stem_path in separated.values(): if isinstance(stem_path, str) and os.path.isfile(stem_path): @@ -1555,7 +1555,7 @@ def main(): "Running source separation (this may take 30-60 seconds)...", file=sys.stderr, ) - stems = separate_stems(audio_path) + stems = separate_stems_backend(audio_path) print("@@SEPARATION_COMPLETE", file=sys.stderr) # Run torchcrepe pitch extraction on separated stems (if available) diff --git a/apps/backend/requirements-msst.txt b/apps/backend/requirements-msst.txt new file mode 100644 index 00000000..b5a3c7da --- /dev/null +++ b/apps/backend/requirements-msst.txt @@ -0,0 +1,41 @@ +# Optional MSST / BS-RoFormer separation backend. +# +# These deps are NOT part of the product baseline and must NOT be added to +# requirements.txt or installed into the product venv. The base ASA install +# runs separation on torchaudio Hybrid Demucs (analyze_audio_io.separate_stems); +# the MSST backend is a default-off experiment selected by ASA_SEPARATION_BACKEND=msst. +# +# Unlike requirements-mt3.txt / requirements-eval.txt, this file does NOT pin a +# package list. MSST-WebUI pins OLDER, conflicting versions of ASA's own deps +# (librosa==0.9.2 vs ASA 0.11.0, numpy<2 vs 2.4.3, protobuf==3.20.3 vs 7.34.0, +# demucs==4.0.0, transformers~=4.35, lightning, …) and brings its own torch. +# They cannot coexist with the product venv. So the MSST backend runs in a +# SEPARATE venv built from MSST-WebUI's OWN requirements.txt, and ASA shells out +# to it as a subprocess (apps/backend/scripts/msst_separate_runner.py) — exactly +# how loudness_backend.py shells out to measure-cli. ASA never imports MSST. +# +# Setup: +# +# # 1. Clone MSST-WebUI and create its own venv from ITS requirements. +# git clone https://github.com/SUC-DriverOld/MSST-WebUI.git /opt/MSST-WebUI +# python3.11 -m venv /opt/MSST-WebUI/venv +# /opt/MSST-WebUI/venv/bin/pip install -r /opt/MSST-WebUI/requirements.txt +# +# # 2. Download a 4-stem checkpoint + its config from the Sucial/MSST-WebUI +# # Hugging Face repo into the MSST layout (configs/ + pretrain/). The +# # default model in separation_backend._MSST_MODEL_REGISTRY ('scnet_4stem') +# # expects: +# # pretrain/multi_stem_models/model_scnet_sdr_9.3244.ckpt +# # configs/multi_stem_models/config_musdb18_scnet.yaml +# +# # 3. Point ASA at the checkout + interpreter and select the backend: +# export ASA_SEPARATION_BACKEND=msst +# export ASA_MSST_PYTHON=/opt/MSST-WebUI/venv/bin/python +# export ASA_MSST_ROOT=/opt/MSST-WebUI +# # optional: +# # ASA_MSST_MODEL_DIR (config/checkpoint root; defaults to ASA_MSST_ROOT) +# # ASA_MSST_MODEL (registry id; defaults to 'scnet_4stem') +# # ASA_MSST_DEVICE (auto|cpu|cuda|mps; default auto) +# +# Any failure (missing venv/checkout/checkpoint, non-zero exit, unparseable +# output) degrades gracefully back to Demucs — see separation_backend.py. diff --git a/apps/backend/scripts/ab_separation_backends.py b/apps/backend/scripts/ab_separation_backends.py new file mode 100644 index 00000000..4a95f232 --- /dev/null +++ b/apps/backend/scripts/ab_separation_backends.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""A/B separation backends CLI — EVAL / RESEARCH ONLY. + +Thin wrapper over ``separation_ab.run_ab`` (like ``scripts/evaluate_*.py`` wrap +their ``*_evaluation.py`` modules). Compares the default Demucs backend against +the optional MSST backend (``ASA_SEPARATION_BACKEND=msst``) on separation quality +and runtime, writes a JSON report, and prints a compact table. + +The synthetic smoke-test always runs. The MSST column is filled only when +``ASA_MSST_PYTHON`` + ``ASA_MSST_ROOT`` are configured (otherwise reported as +``skipped_no_msst``). This is observational, not a gate — exit code is 0 for any +completed/skipped run, 1 only on a hard error. + +Usage:: + + ./venv/bin/python scripts/ab_separation_backends.py \\ + [-i /dir/of/real/tracks] [-o .runtime/separation_ab] [--model scnet_4stem] +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO_DIR = Path(__file__).resolve().parent.parent +if str(REPO_DIR) not in sys.path: + sys.path.insert(0, str(REPO_DIR)) + +import separation_ab # noqa: E402 + + +def _fmt(value: object) -> str: + return "—" if value is None else str(value) + + +def _print_table(report: dict) -> None: + print("\nSeparation A/B — synthetic smoke-test (NOT a real-music quality ranking)") + print(f"{'backend':<10} {'status':<16} {'meanSI-SDR(dB)':>15} {'runtime(s)':>11} {'device':>8}") + synthetic = report.get("syntheticSmokeTest", {}).get("perBackend", {}) + for backend in ("demucs", "msst"): + block = synthetic.get(backend, {}) + status = block.get("status", "—") + quality = block.get("quality", {}) if isinstance(block.get("quality"), dict) else {} + mean = quality.get("meanSiSdrDb") + runtime = block.get("runtimeSeconds") + device = block.get("device") + print(f"{backend:<10} {status:<16} {_fmt(mean):>15} {_fmt(runtime):>11} {_fmt(device):>8}") + + real = report.get("realTracks", []) + if real: + print(f"\nReal-track reference-free proxies ({len(real)} track(s)):") + for entry in real: + print(f" {entry.get('track')}:") + for backend in ("demucs", "msst"): + block = entry.get("perBackend", {}).get(backend, {}) + status = block.get("status", "—") + runtime = block.get("runtimeSeconds") + residual = block.get("mixReconstructionResidualDb") + print( + f" {backend:<8} status={status:<16} " + f"runtime={_fmt(runtime)}s reconResidual={_fmt(residual)}dB" + ) + for caveat in report.get("caveats", []): + print(f"\n[caveat] {caveat}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "-i", "--input-dir", default=None, + help="Optional directory of real audio tracks for reference-free proxies.", + ) + parser.add_argument( + "-o", "--out-dir", type=Path, default=REPO_DIR / ".runtime" / "separation_ab", + help="Report output directory. Defaults to apps/backend/.runtime/separation_ab/.", + ) + parser.add_argument( + "--model", default=None, + help="MSST model registry id (sets ASA_MSST_MODEL). Defaults to scnet_4stem.", + ) + parser.add_argument("--repeats", type=int, default=2, help="Timed runs per backend (min wins).") + args = parser.parse_args() + + try: + report = separation_ab.run_ab( + input_dir=args.input_dir, + out_dir=str(args.out_dir), + model=args.model, + repeats=args.repeats, + ) + except Exception as exc: # noqa: BLE001 + print(json.dumps({"status": "error", "error": str(exc)}, indent=2)) + return 1 + + args.out_dir.mkdir(parents=True, exist_ok=True) + report_path = args.out_dir / "report.json" + report_path.write_text(json.dumps(report, indent=2), encoding="utf-8") + _print_table(report) + print(f"\nWrote {report_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/backend/scripts/msst_separate_runner.py b/apps/backend/scripts/msst_separate_runner.py new file mode 100644 index 00000000..28f5f2aa --- /dev/null +++ b/apps/backend/scripts/msst_separate_runner.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +"""MSST separation runner — RUNS UNDER THE MSST VENV, not ASA's product venv. + +This standalone script is invoked as a subprocess by +``apps/backend/separation_backend.py`` using the ``ASA_MSST_PYTHON`` interpreter. +It drives `SUC-DriverOld/MSST-WebUI `_'s +``MSSeparator`` (no PySide6 GUI) and writes canonical ``vocals/bass/drums/other`` +stems as 44.1 kHz PCM16 WAVs into ``--output-dir``. + +It imports ONLY MSST + numpy/soundfile/librosa/yaml (all present in the MSST +venv) — never ASA product modules (no essentia/torch-2.10 dependency), so the +two venvs stay fully isolated. + +Contract (load-bearing): the **only** thing written to stdout is a single line of +JSON: ``{"stems": {name: path}, "modelType": str, "loadSeconds": float, +"inferSeconds": float, "device": str}``. Every MSST progress bar / log line is +forced to stderr (the caller captures both streams; stdout must stay pure JSON). +A non-zero exit + stderr message signals failure, on which the caller falls back +to Demucs. + +Usage:: + + $ASA_MSST_PYTHON scripts/msst_separate_runner.py \\ + --input track.flac --output-dir /tmp/stems \\ + --msst-root /path/to/MSST-WebUI \\ + --model-type scnet --config --checkpoint \\ + [--device auto] +""" + +from __future__ import annotations + +import argparse +import contextlib +import json +import logging +import os +import sys +import time +import wave + +import numpy as np + +# Canonical ASA stem names + MSST aliases. Mapping is BY NAME (case-insensitive), +# never positional — MSST stem order varies by model (SCNet emits +# [drums, bass, other, vocals]). +_CANONICAL_STEMS = {"vocals", "bass", "drums", "other"} +_STEM_ALIASES = {"instrumental": "other", "instrum": "other"} + +_TARGET_SAMPLE_RATE_FALLBACK = 44100 + + +def _eprint(message: str) -> None: + print(message, file=sys.stderr) + + +def _stderr_logger() -> logging.Logger: + """A logger that writes to stderr only, so MSST never pollutes stdout.""" + logger = logging.getLogger("asa_msst_runner") + logger.handlers.clear() + handler = logging.StreamHandler(sys.stderr) + handler.setFormatter(logging.Formatter("[msst] %(message)s")) + logger.addHandler(handler) + logger.setLevel(logging.INFO) + logger.propagate = False + return logger + + +def _read_target_sample_rate(config_path: str) -> int: + """Best-effort read of ``audio.sample_rate`` from the MSST config YAML.""" + try: + import yaml + + with open(config_path, "r") as handle: + config = yaml.safe_load(handle) + audio = (config or {}).get("audio", {}) if isinstance(config, dict) else {} + sr = audio.get("sample_rate") + if isinstance(sr, int) and sr > 0: + return sr + except Exception as exc: # noqa: BLE001 + _eprint(f"[msst] could not read sample_rate from config ({exc}); using 44100.") + return _TARGET_SAMPLE_RATE_FALLBACK + + +def _load_mix(input_path: str, target_sr: int) -> np.ndarray: + """Load audio as channels-first stereo ``(2, N)`` float32 at ``target_sr``.""" + import soundfile as sf + + data, file_sr = sf.read(input_path, always_2d=True, dtype="float32") # (N, C) + mix = data.T # (C, N) + if mix.shape[0] == 1: + mix = np.vstack([mix, mix]) # mono -> stereo + elif mix.shape[0] > 2: + mix = mix[:2] # defensively clamp to stereo + + if file_sr != target_sr: + import librosa + + # Resample per channel — librosa's `axis` kwarg is not reliable across the + # 0.9.x line MSST pins, but 1-D resample works on every version. + channels = [ + librosa.resample(np.ascontiguousarray(ch), orig_sr=file_sr, target_sr=target_sr) + for ch in mix + ] + mix = np.stack(channels, axis=0) + return np.ascontiguousarray(mix, dtype=np.float32) + + +def _write_wav_pcm16(path: str, audio: np.ndarray, sample_rate: int) -> None: + """Write a channels-first ``(C, N)`` float waveform to PCM16 WAV. + + Matches ``analyze_audio_io._write_wav_pcm16`` encoding so the stems are + byte-compatible with the Demucs path (44.1 kHz stereo PCM16). + """ + data = np.asarray(audio, dtype=np.float32) + if data.ndim == 1: + data = data[np.newaxis, :] + data = np.clip(data, -1.0, 1.0) + interleaved = (data.T * 32767.0).astype(np.int16) # (N, C) + with wave.open(path, "wb") as wav_file: + wav_file.setnchannels(interleaved.shape[1] if interleaved.ndim == 2 else 1) + wav_file.setsampwidth(2) + wav_file.setframerate(int(sample_rate)) + wav_file.writeframes(interleaved.tobytes()) + + +def _canonical_name(stem_key: str) -> str | None: + """Map an MSST stem key to a canonical ASA stem name, or None to skip.""" + key = str(stem_key).strip().lower() + key = _STEM_ALIASES.get(key, key) + return key if key in _CANONICAL_STEMS else None + + +def _to_channels_first(stem: np.ndarray) -> np.ndarray: + """Coerce an MSST stem ``(N, C)`` (channels-last) to ``(C, N)``.""" + arr = np.asarray(stem, dtype=np.float32) + if arr.ndim == 1: + return arr[np.newaxis, :] + # MSSeparator.separate returns channels-last (N, C); pick the channel axis as + # the smaller dimension to be robust to either orientation. + if arr.shape[0] <= arr.shape[1]: + return arr # already (C, N) + return arr.T + + +def main() -> int: + parser = argparse.ArgumentParser(description="MSST separation runner (MSST venv)") + parser.add_argument("--input", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--msst-root", required=True) + parser.add_argument("--model-type", required=True) + parser.add_argument("--config", required=True) + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--device", default="auto") + args = parser.parse_args() + + # Append (not prepend) the MSST checkout so its top-level ``utils`` / + # ``inference`` packages don't shadow anything in this process unexpectedly. + if args.msst_root not in sys.path: + sys.path.append(args.msst_root) + + os.makedirs(args.output_dir, exist_ok=True) + target_sr = _read_target_sample_rate(args.config) + logger = _stderr_logger() + + # Force every MSST print()/tqdm/log onto stderr — stdout is the JSON contract. + with contextlib.redirect_stdout(sys.stderr): + try: + from inference.msst_infer import MSSeparator + except Exception as exc: # noqa: BLE001 + _eprint(f"[msst] could not import MSSeparator from {args.msst_root} ({exc}).") + return 2 + + separator = None + try: + load_start = time.perf_counter() + separator = MSSeparator( + model_type=args.model_type, + config_path=args.config, + model_path=args.checkpoint, + device=args.device, + output_format="wav", + store_dirs="", # we write stems ourselves; don't let MSST emit files + logger=logger, + ) + load_seconds = time.perf_counter() - load_start + + mix = _load_mix(args.input, target_sr) + + infer_start = time.perf_counter() + separated = separator.separate(mix) # {stem: ndarray (N, C)} + infer_seconds = time.perf_counter() - infer_start + except Exception as exc: # noqa: BLE001 + _eprint(f"[msst] separation failed ({exc}).") + return 3 + finally: + if separator is not None: + with contextlib.suppress(Exception): + separator.del_cache() + + if not isinstance(separated, dict) or not separated: + _eprint("[msst] separator returned no stems.") + return 4 + + device = getattr(separator, "device", args.device) + stems: dict[str, str] = {} + for raw_name, stem_audio in separated.items(): + canonical = _canonical_name(raw_name) + if canonical is None: + _eprint(f"[msst] skipping unrecognized stem '{raw_name}'.") + continue + out_path = os.path.join(args.output_dir, f"{canonical}.wav") + _write_wav_pcm16(out_path, _to_channels_first(stem_audio), target_sr) + stems[canonical] = out_path + + if not stems: + _eprint("[msst] no canonical stems written.") + return 5 + + manifest = { + "stems": stems, + "modelType": args.model_type, + "loadSeconds": round(load_seconds, 3), + "inferSeconds": round(infer_seconds, 3), + "device": str(device), + } + json.dump(manifest, sys.stdout) + sys.stdout.write("\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/backend/separation_ab.py b/apps/backend/separation_ab.py new file mode 100644 index 00000000..92452b34 --- /dev/null +++ b/apps/backend/separation_ab.py @@ -0,0 +1,381 @@ +"""A/B separation harness — EVAL / RESEARCH ONLY. + +Compares the default Demucs separation backend against the optional MSST backend +(``ASA_SEPARATION_BACKEND=msst``) on separation quality and runtime. This module +is intentionally separate from ``analyze.py`` and ``server.py`` — deleting it +restores the product exactly. + +Two complementary signals, because there is no honest SDR on real tracks without +ground-truth isolated stems (which we can't bundle — copyright): + +1. **Synthetic smoke-test (always available).** A deterministic synthetic + multitrack with KNOWN stems (sine/sub bass, click drums, chord pad -> other, + formant vocal) is summed to a mix, separated by both backends, and scored with + gain-aligned SI-SDR against the known stems. This is a PLUMBING / SANITY check + — "did each backend produce non-silent, correctly-named, roughly-correct-energy + stems on a known input" — NOT a real-music quality ranking. BS-RoFormer/SCNet + are trained on real spectra and can score poorly on toy synthetic sources, so + the synthetic SDR must not be read as "Demucs is better than MSST". The harness + carries this caveat in its report. + +2. **Real-track reference-free proxies (optional, ``-i``).** On real audio with no + reference stems we report per-stem RMS/peak/spectral-centroid, the + mix-reconstruction residual (sum of stems vs mix), and cross-backend deltas. + These measure self-consistency and plausibility, not correctness — but they're + the more trustworthy signal for a real Demucs-vs-MSST comparison. + +Runtime is reported as end-to-end wall-clock per backend (the cost the product +actually pays — neither backend caches the model across calls), after a warm-up +pass, taking the min over repeats, with the device recorded. +""" + +from __future__ import annotations + +import math +import os +import tempfile +import time +import wave +from typing import Any + +import numpy as np + +SAMPLE_RATE = 44100 +_CANONICAL_STEMS = ("vocals", "bass", "drums", "other") + + +# --------------------------------------------------------------------------- # +# Synthetic multitrack with known stems +# --------------------------------------------------------------------------- # +def build_synthetic_multitrack(seconds: float = 6.0) -> dict[str, np.ndarray]: + """Return KNOWN mono stems keyed by canonical name (deterministic). + + The sources are deliberately simple and well-separated in frequency so a + working separator can recover them; they are NOT a substitute for real music. + """ + n = int(seconds * SAMPLE_RATE) + t = np.arange(n, dtype=np.float64) / SAMPLE_RATE + rng = np.random.default_rng(20260603) + + # Bass: 55 Hz sine + 110 Hz harmonic. + bass = 0.5 * np.sin(2 * np.pi * 55.0 * t) + 0.2 * np.sin(2 * np.pi * 110.0 * t) + + # Drums: periodic noise bursts (a 2 Hz "kick" of decaying white noise). + drums = np.zeros(n) + burst_period = SAMPLE_RATE // 2 + burst_len = SAMPLE_RATE // 8 + env = np.exp(-np.linspace(0, 8, burst_len)) + for start in range(0, n - burst_len, burst_period): + drums[start : start + burst_len] += rng.standard_normal(burst_len) * env * 0.6 + + # Other (pad): a sustained major triad in the mid range. + pad = sum( + 0.18 * np.sin(2 * np.pi * f * t) for f in (261.63, 329.63, 392.0) + ) + + # Vocals: a formant-ish tone (440 Hz carrier with 5 Hz vibrato + 2nd formant). + vibrato = 1.0 + 0.01 * np.sin(2 * np.pi * 5.0 * t) + vocals = 0.35 * np.sin(2 * np.pi * 440.0 * t * vibrato) + 0.12 * np.sin( + 2 * np.pi * 1200.0 * t + ) + + return { + "bass": bass.astype(np.float32), + "drums": drums.astype(np.float32), + "other": pad.astype(np.float32), + "vocals": vocals.astype(np.float32), + } + + +def _write_stereo_wav(path: str, mono: np.ndarray, sample_rate: int = SAMPLE_RATE) -> None: + data = np.clip(np.asarray(mono, dtype=np.float32), -1.0, 1.0) + stereo = np.stack([data, data], axis=1) # (N, 2) + interleaved = (stereo * 32767.0).astype(np.int16) + with wave.open(path, "wb") as wav_file: + wav_file.setnchannels(2) + wav_file.setsampwidth(2) + wav_file.setframerate(int(sample_rate)) + wav_file.writeframes(interleaved.tobytes()) + + +def _load_wav_mono(path: str) -> np.ndarray | None: + try: + with wave.open(path, "rb") as wav_file: + channels = wav_file.getnchannels() + frames = wav_file.readframes(wav_file.getnframes()) + data = np.frombuffer(frames, dtype=np.int16).astype(np.float64) / 32768.0 + if channels > 1: + data = data.reshape(-1, channels).mean(axis=1) + return data + except Exception: + return None + + +# --------------------------------------------------------------------------- # +# Metrics +# --------------------------------------------------------------------------- # +def si_sdr(reference: np.ndarray, estimate: np.ndarray) -> float: + """Scale-invariant SDR (dB). Gain-aligns the estimate before the residual. + + SI-SDR = 10*log10( ||alpha*ref||^2 / ||alpha*ref - est||^2 ), + alpha = /. Returns -inf for a silent reference/estimate. + """ + ref = np.asarray(reference, dtype=np.float64) + est = np.asarray(estimate, dtype=np.float64) + length = min(ref.size, est.size) + if length == 0: + return float("-inf") + ref = ref[:length] + est = est[:length] + ref_energy = float(np.dot(ref, ref)) + if ref_energy <= 1e-12 or float(np.dot(est, est)) <= 1e-12: + return float("-inf") + alpha = float(np.dot(est, ref)) / ref_energy + projection = alpha * ref + noise = projection - est + noise_energy = float(np.dot(noise, noise)) + if noise_energy <= 1e-12: + return float("inf") + return 10.0 * math.log10(float(np.dot(projection, projection)) / noise_energy) + + +def buffer_stats(mono: np.ndarray) -> dict[str, Any]: + """Reference-free per-stem stats: RMS/peak dBFS + spectral centroid (Hz).""" + data = np.asarray(mono, dtype=np.float64).reshape(-1) + if data.size == 0: + return {"rmsDbfs": None, "peakDbfs": None, "spectralCentroidHz": None, "nonzero": False} + rms = float(np.sqrt(np.mean(data * data))) + peak = float(np.max(np.abs(data))) + if peak <= 0: + return {"rmsDbfs": None, "peakDbfs": None, "spectralCentroidHz": 0.0, "nonzero": False} + spectrum = np.abs(np.fft.rfft(data * np.hanning(data.size))) + freqs = np.fft.rfftfreq(data.size, d=1.0 / SAMPLE_RATE) + total = float(spectrum.sum()) + centroid = float((freqs * spectrum).sum() / total) if total > 0 else 0.0 + return { + "rmsDbfs": round(20.0 * math.log10(max(rms, 1e-12)), 3), + "peakDbfs": round(20.0 * math.log10(max(peak, 1e-12)), 3), + "spectralCentroidHz": round(centroid, 2), + "nonzero": True, + } + + +# --------------------------------------------------------------------------- # +# Backend execution +# --------------------------------------------------------------------------- # +def run_backend_separation( + backend: str, + mix_path: str, + out_dir: str, + *, + repeats: int = 2, +) -> dict[str, Any]: + """Separate ``mix_path`` with ``backend`` (demucs|msst); time it fairly. + + Does a warm-up pass (amortizes model download/load), then takes the min + wall-clock over ``repeats`` timed runs. Returns the stems + runtime + status. + """ + from separation_backend import separate_stems_backend + + prev = os.environ.get("ASA_SEPARATION_BACKEND") + os.environ["ASA_SEPARATION_BACKEND"] = backend + try: + device = "cuda" if _torch_cuda_available() else "cpu" + # Warm-up (untimed) — pulls weights / builds the graph. + warm = separate_stems_backend(mix_path, output_dir=os.path.join(out_dir, "warmup")) + if not warm: + return {"status": "error", "error": "separation returned no stems", "stems": {}} + + timings: list[float] = [] + stems: dict[str, str] = {} + for i in range(max(1, repeats)): + run_out = os.path.join(out_dir, f"run{i}") + start = time.perf_counter() + stems = separate_stems_backend(mix_path, output_dir=run_out) or {} + timings.append(time.perf_counter() - start) + if not stems: + return {"status": "error", "error": "separation returned no stems", "stems": {}} + + return { + "status": "completed", + "stems": stems, + "runtimeSeconds": round(min(timings), 3), + "device": device, + } + finally: + if prev is None: + os.environ.pop("ASA_SEPARATION_BACKEND", None) + else: + os.environ["ASA_SEPARATION_BACKEND"] = prev + + +def _torch_cuda_available() -> bool: + try: + import torch + + return bool(torch.cuda.is_available()) + except Exception: + return False + + +def _msst_configured() -> bool: + py = os.environ.get("ASA_MSST_PYTHON") + root = os.environ.get("ASA_MSST_ROOT") + return bool(py and os.path.exists(py) and root and os.path.isdir(root)) + + +# --------------------------------------------------------------------------- # +# Evaluation +# --------------------------------------------------------------------------- # +def _score_against_known( + stems: dict[str, str], + known: dict[str, np.ndarray], +) -> dict[str, Any]: + """SI-SDR per canonical stem against the known synthetic sources.""" + per_stem: dict[str, Any] = {} + sdrs: list[float] = [] + for name in _CANONICAL_STEMS: + path = stems.get(name) + est = _load_wav_mono(path) if isinstance(path, str) else None + if est is None or name not in known: + per_stem[name] = {"siSdrDb": None, "present": False} + continue + value = si_sdr(known[name], est) + per_stem[name] = { + "siSdrDb": round(value, 3) if math.isfinite(value) else None, + "present": True, + } + if math.isfinite(value): + sdrs.append(value) + per_stem["meanSiSdrDb"] = round(sum(sdrs) / len(sdrs), 3) if sdrs else None + return per_stem + + +def evaluate_synthetic(out_dir: str, *, repeats: int = 2) -> dict[str, Any]: + """Run the always-available synthetic smoke-test for both backends.""" + known = build_synthetic_multitrack() + mix = sum(known.values()) + peak = float(np.max(np.abs(mix))) or 1.0 + mix = (mix / peak * 0.9).astype(np.float32) + + work = tempfile.mkdtemp(prefix="asa_sep_ab_synth_", dir=out_dir) + mix_path = os.path.join(work, "mix.wav") + _write_stereo_wav(mix_path, mix) + + result: dict[str, Any] = {"perBackend": {}} + for backend in ("demucs", "msst"): + if backend == "msst" and not _msst_configured(): + result["perBackend"][backend] = {"status": "skipped_no_msst"} + continue + backend_dir = os.path.join(work, backend) + os.makedirs(backend_dir, exist_ok=True) + run = run_backend_separation(backend, mix_path, backend_dir, repeats=repeats) + if run["status"] != "completed": + result["perBackend"][backend] = run + continue + result["perBackend"][backend] = { + "status": "completed", + "runtimeSeconds": run["runtimeSeconds"], + "device": run["device"], + "quality": _score_against_known(run["stems"], known), + } + return result + + +def evaluate_real_track(audio_path: str, out_dir: str, *, repeats: int = 2) -> dict[str, Any]: + """Reference-free proxies on a real track for both backends.""" + work = tempfile.mkdtemp(prefix="asa_sep_ab_real_", dir=out_dir) + entry: dict[str, Any] = {"track": os.path.basename(audio_path), "perBackend": {}} + for backend in ("demucs", "msst"): + if backend == "msst" and not _msst_configured(): + entry["perBackend"][backend] = {"status": "skipped_no_msst"} + continue + backend_dir = os.path.join(work, backend) + os.makedirs(backend_dir, exist_ok=True) + run = run_backend_separation(backend, audio_path, backend_dir, repeats=repeats) + if run["status"] != "completed": + entry["perBackend"][backend] = run + continue + stems_stats = {} + recovered = None + for name, path in run["stems"].items(): + mono = _load_wav_mono(path) + if mono is None: + continue + stems_stats[name] = buffer_stats(mono) + recovered = mono if recovered is None else _sum_to(recovered, mono) + mix = _load_wav_mono(audio_path) + residual_db = _reconstruction_residual_db(mix, recovered) + entry["perBackend"][backend] = { + "status": "completed", + "runtimeSeconds": run["runtimeSeconds"], + "device": run["device"], + "stemStats": stems_stats, + "mixReconstructionResidualDb": residual_db, + } + return entry + + +def _sum_to(acc: np.ndarray, add: np.ndarray) -> np.ndarray: + length = min(acc.size, add.size) + return acc[:length] + add[:length] + + +def _reconstruction_residual_db(mix: np.ndarray | None, recovered: np.ndarray | None): + if mix is None or recovered is None: + return None + length = min(mix.size, recovered.size) + if length == 0: + return None + residual = mix[:length] - recovered[:length] + res_energy = float(np.dot(residual, residual)) + mix_energy = float(np.dot(mix[:length], mix[:length])) + if mix_energy <= 1e-12: + return None + return round(10.0 * math.log10(max(res_energy, 1e-12) / mix_energy), 3) + + +def run_ab( + *, + input_dir: str | None, + out_dir: str, + model: str | None = None, + repeats: int = 2, +) -> dict[str, Any]: + """Top-level A/B run: synthetic smoke-test + optional real-track proxies.""" + from datetime import datetime, timezone + + os.makedirs(out_dir, exist_ok=True) + if model: + os.environ["ASA_MSST_MODEL"] = model + + report: dict[str, Any] = { + "generatedAt": datetime.now(timezone.utc).isoformat(), + "researchOnly": True, + "status": "completed", + "config": { + "model": os.environ.get("ASA_MSST_MODEL", "scnet_4stem"), + "msstConfigured": _msst_configured(), + "repeats": repeats, + }, + "caveats": [ + "Synthetic SI-SDR is a PLUMBING smoke-test, not a real-music quality " + "ranking: RoFormer/SCNet are trained on real spectra and may score " + "poorly on toy synthetic sources. Do not read it as Demucs>MSST.", + "Real-track proxies measure self-consistency/plausibility, not " + "correctness (no ground-truth stems). Runtime is end-to-end wall-clock " + "(min of repeats after a warm-up); compare same-device rows only.", + ], + "syntheticSmokeTest": evaluate_synthetic(out_dir, repeats=repeats), + "realTracks": [], + } + + if input_dir and os.path.isdir(input_dir): + for name in sorted(os.listdir(input_dir)): + path = os.path.join(input_dir, name) + if os.path.isfile(path) and name.lower().endswith( + (".wav", ".flac", ".mp3", ".m4a", ".ogg", ".aiff", ".aif") + ): + report["realTracks"].append(evaluate_real_track(path, out_dir, repeats=repeats)) + + return report diff --git a/apps/backend/separation_backend.py b/apps/backend/separation_backend.py new file mode 100644 index 00000000..89740e36 --- /dev/null +++ b/apps/backend/separation_backend.py @@ -0,0 +1,257 @@ +"""Pluggable Phase 1 stem-separation backend (default-off experiment). + +ASA's authoritative separation is torchaudio Hybrid Demucs +(``analyze_audio_io.separate_stems``). This module lets the separation stage +optionally run a stronger MSST / BS-RoFormer model from +`SUC-DriverOld/MSST-WebUI `_, +selected at runtime by ``ASA_SEPARATION_BACKEND``: + +* ``demucs`` (default) — unchanged; this module delegates straight to + ``separate_stems`` and is a no-op wrapper. +* ``msst`` — drive MSST's ``MSSeparator`` inference (no PySide6 GUI) to produce + the same ``{stem_name: wav_path}`` contract. + +Design — subprocess, not in-process import (mirrors ``loudness_backend.py``): + +MSST pins older, conflicting versions of ASA's own deps (``librosa==0.9.2``, +``numpy<2``, ``demucs==4.0.0`` …) so it must live in its **own** virtualenv. We +therefore shell out to ``scripts/msst_separate_runner.py`` under +``ASA_MSST_PYTHON`` (the MSST venv interpreter) and capture its stdout. This: + +* lets MSST actually run on the staged worker path (``server.py`` always invokes + ``analyze.py`` under ``./venv/bin/python`` — an in-process import could never + see MSST's deps), +* keeps MSST's progress bars / logging off ``analyze.py``'s stdout, which is the + load-bearing JSON contract (tripwire #1), +* sidesteps MSST's ``utils`` / ``inference`` top-level package collision with + ASA's ``apps/backend/utils`` package, +* and isolates MSST's torch from ASA's ``torch==2.10.0``. + +Per PURPOSE.md invariant #1 the default stays Demucs and the authoritative +contract is unchanged. Any MSST failure — missing venv, missing checkout, +missing checkpoint, non-zero exit, unparseable output — degrades gracefully back +to Demucs with a ``[warn]`` on stderr. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +from analyze_audio_io import separate_stems + +# separation_backend.py lives at apps/backend/, so the runner is a sibling of +# this file under scripts/. +_RUNNER_SCRIPT = Path(__file__).resolve().parent / "scripts" / "msst_separate_runner.py" + +# Demucs writes its stems into a tempdir prefixed ``sonic_analyzer_demucs_`` and +# ``analyze_audio_io.cleanup_stems`` only reclaims a parent dir with that exact +# prefix. The MSST path reuses the same prefix so cleanup works unchanged and no +# temp directory leaks per measurement run. +_STEM_DIR_PREFIX = "sonic_analyzer_demucs_" + +# Model registry: ``ASA_MSST_MODEL`` -> the three coupled MSSeparator args. +# ``config_relpath`` / ``checkpoint_relpath`` resolve under ``ASA_MSST_MODEL_DIR`` +# (default: ``ASA_MSST_ROOT``), matching MSST-WebUI's own ``configs/`` + ``pretrain/`` +# layout. Checkpoints are NOT vendored — the operator downloads weights from the +# ``Sucial/MSST-WebUI`` Hugging Face repo into ``ASA_MSST_MODEL_DIR``. +_DEFAULT_MSST_MODEL = "scnet_4stem" +_MSST_MODEL_REGISTRY: dict[str, dict[str, str]] = { + # 4-stem parity model — fills vocals/bass/drums/other, so the product path + # (stemAnalysis overlay, pitch/note translation, MT3) behaves like Demucs. + "scnet_4stem": { + "model_type": "scnet", + "config_relpath": "configs/multi_stem_models/config_musdb18_scnet.yaml", + "checkpoint_relpath": "pretrain/multi_stem_models/model_scnet_sdr_9.3244.ckpt", + }, + # Strong 2-stem vocal BS-RoFormer (vocals + instrumental->other). RESEARCH / + # A-B ONLY: it leaves bass/drums absent, so pitch/note translation + MT3 fall + # back to the full mix. Never select this as a product-path default. + "bs_roformer_vocals": { + "model_type": "bs_roformer", + "config_relpath": "configs/vocal_models/config_vocals_bs_roformer.yaml", + "checkpoint_relpath": "pretrain/vocal_models/model_bs_roformer_ep_368_sdr_12.9628.ckpt", + }, +} + + +def separation_backend_name() -> str: + """Resolve the selected separation backend. Unknown values fall back to demucs.""" + name = (os.environ.get("ASA_SEPARATION_BACKEND") or "demucs").strip().lower() + return name if name in {"demucs", "msst"} else "demucs" + + +def _msst_model_entry() -> dict[str, str]: + """Resolve the MSST model registry entry. Unknown ids fall back to the default.""" + requested = (os.environ.get("ASA_MSST_MODEL") or _DEFAULT_MSST_MODEL).strip() + entry = _MSST_MODEL_REGISTRY.get(requested) + if entry is None: + print( + f"[warn] Unknown ASA_MSST_MODEL='{requested}'; valid values: " + f"{sorted(_MSST_MODEL_REGISTRY)}. Falling back to '{_DEFAULT_MSST_MODEL}'.", + file=sys.stderr, + ) + entry = _MSST_MODEL_REGISTRY[_DEFAULT_MSST_MODEL] + return entry + + +def separate_stems_backend(audio_path: str, output_dir: str | None = None): + """Run the selected separation backend, returning ``{stem_name: wav_path}``. + + The return shape is identical to ``separate_stems`` (Demucs) so every + downstream consumer — the ``stemAnalysis`` overlay, pitch/note translation, + MT3, and ``cleanup_stems`` — is byte-compatible. For ``ASA_SEPARATION_BACKEND`` + unset/``demucs`` this is a thin pass-through to ``separate_stems``. For + ``msst`` it drives MSST and, on any failure, degrades to Demucs. + """ + if separation_backend_name() != "msst": + return separate_stems(audio_path, output_dir=output_dir) + + result = _separate_via_msst_subprocess(audio_path, output_dir) + if result is not None: + return result + + print( + "[warn] MSST separation unavailable; falling back to Demucs.", + file=sys.stderr, + ) + return separate_stems(audio_path, output_dir=output_dir) + + +def _separate_via_msst_subprocess(audio_path: str, output_dir: str | None): + """Drive MSST separation via the runner subprocess under ``ASA_MSST_PYTHON``. + + Returns ``{stem_name: wav_path}`` on success, or ``None`` (warning on stderr) + on any failure so the caller can fall back to Demucs. All MSST stdout/stderr + is captured here and never reaches ``analyze.py``'s stdout JSON contract. + """ + interpreter = os.environ.get("ASA_MSST_PYTHON") + if not interpreter or not Path(interpreter).exists(): + print( + "[warn] ASA_SEPARATION_BACKEND=msst but ASA_MSST_PYTHON is unset or " + f"missing (got {interpreter!r}); cannot run the MSST venv.", + file=sys.stderr, + ) + return None + + msst_root = os.environ.get("ASA_MSST_ROOT") + if not msst_root or not Path(msst_root).is_dir(): + print( + "[warn] ASA_SEPARATION_BACKEND=msst but ASA_MSST_ROOT is unset or not a " + f"directory (got {msst_root!r}); cannot find the MSST-WebUI checkout.", + file=sys.stderr, + ) + return None + + if not _RUNNER_SCRIPT.exists(): + print(f"[warn] MSST runner script missing at {_RUNNER_SCRIPT}.", file=sys.stderr) + return None + + model_dir = Path(os.environ.get("ASA_MSST_MODEL_DIR") or msst_root) + entry = _msst_model_entry() + config_path = model_dir / entry["config_relpath"] + checkpoint_path = model_dir / entry["checkpoint_relpath"] + for label, candidate in (("config", config_path), ("checkpoint", checkpoint_path)): + if not candidate.exists(): + print( + f"[warn] MSST {label} not found at {candidate} (model " + f"'{os.environ.get('ASA_MSST_MODEL') or _DEFAULT_MSST_MODEL}').", + file=sys.stderr, + ) + return None + + # Reuse the Demucs tempdir prefix so cleanup_stems reclaims the directory. + if output_dir is None: + output_dir = tempfile.mkdtemp(prefix=_STEM_DIR_PREFIX) + else: + os.makedirs(output_dir, exist_ok=True) + + command = [ + interpreter, + str(_RUNNER_SCRIPT), + "--input", + audio_path, + "--output-dir", + output_dir, + "--msst-root", + msst_root, + "--model-type", + entry["model_type"], + "--config", + str(config_path), + "--checkpoint", + str(checkpoint_path), + ] + device = os.environ.get("ASA_MSST_DEVICE") + if device: + command.extend(["--device", device]) + + try: + proc = subprocess.run( + command, + capture_output=True, + text=True, + timeout=1800, + check=False, + ) + except Exception as exc: # noqa: BLE001 - degrade to Demucs, never crash analysis + print(f"[warn] MSST runner failed to launch ({exc}).", file=sys.stderr) + return None + + if proc.stderr: + # Surface the runner's diagnostics on our stderr (never stdout). + print(proc.stderr.rstrip(), file=sys.stderr) + if proc.returncode != 0: + print( + f"[warn] MSST runner exited {proc.returncode}; falling back to Demucs.", + file=sys.stderr, + ) + return None + + return _parse_runner_manifest(proc.stdout) + + +def _parse_runner_manifest(stdout: str): + """Parse the runner's single-line JSON manifest into ``{stem_name: path}``.""" + try: + payload = json.loads(stdout.strip()) + except Exception as exc: # noqa: BLE001 + print(f"[warn] MSST runner produced unparseable output ({exc}).", file=sys.stderr) + return None + + stems = payload.get("stems") if isinstance(payload, dict) else None + if not isinstance(stems, dict) or not stems: + print("[warn] MSST runner reported no stems.", file=sys.stderr) + return None + + resolved: dict[str, str] = {} + for name, path in stems.items(): + if isinstance(path, str) and os.path.isfile(path): + resolved[name] = path + if not resolved: + print("[warn] MSST runner stems do not exist on disk.", file=sys.stderr) + return None + + load_s = payload.get("loadSeconds") + infer_s = payload.get("inferSeconds") + device = payload.get("device") + print( + f"[info] separation backend: MSST ({entry_label(payload)}) produced " + f"{sorted(resolved)} on {device} " + f"(load {load_s}s, infer {infer_s}s).", + file=sys.stderr, + ) + return resolved + + +def entry_label(payload: dict) -> str: + """Best-effort model label for the info line (model_type from the manifest).""" + if isinstance(payload, dict): + model_type = payload.get("modelType") + if isinstance(model_type, str) and model_type: + return model_type + return os.environ.get("ASA_MSST_MODEL") or _DEFAULT_MSST_MODEL diff --git a/apps/backend/tests/test_separation_backend.py b/apps/backend/tests/test_separation_backend.py new file mode 100644 index 00000000..9422f1a2 --- /dev/null +++ b/apps/backend/tests/test_separation_backend.py @@ -0,0 +1,168 @@ +"""Tests for the pluggable separation backend (Demucs <-> MSST). + +The MSST path shells out to a runner under a separate venv that isn't installed +in CI — so these mock the subprocess and the env wiring and assert the +select / dispatch / fall-back logic in pure Python. No MSST install required. +""" + +import importlib.util +import json +import os +import unittest +from pathlib import Path +from unittest import mock + +import numpy as np + +import separation_backend as sb + + +def _load_runner_module(): + """Load the MSST runner from scripts/ by path (it is not on sys.path).""" + runner_path = Path(__file__).resolve().parent.parent / "scripts" / "msst_separate_runner.py" + spec = importlib.util.spec_from_file_location("msst_separate_runner", runner_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_runner = _load_runner_module() + + +_SENTINEL_STEMS = {"vocals": "/tmp/v.wav", "bass": "/tmp/b.wav"} + + +def _runner_proc(stdout: str, returncode: int = 0, stderr: str = "") -> mock.Mock: + return mock.Mock(returncode=returncode, stdout=stdout, stderr=stderr) + + +class SeparationBackendNameTests(unittest.TestCase): + def test_default_is_demucs(self) -> None: + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop("ASA_SEPARATION_BACKEND", None) + self.assertEqual(sb.separation_backend_name(), "demucs") + + def test_unknown_value_falls_back_to_demucs(self) -> None: + with mock.patch.dict(os.environ, {"ASA_SEPARATION_BACKEND": "bogus"}): + self.assertEqual(sb.separation_backend_name(), "demucs") + + def test_msst_is_recognized_case_insensitively(self) -> None: + with mock.patch.dict(os.environ, {"ASA_SEPARATION_BACKEND": "MSST"}): + self.assertEqual(sb.separation_backend_name(), "msst") + + +class ModelRegistryTests(unittest.TestCase): + def test_default_model_resolves(self) -> None: + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop("ASA_MSST_MODEL", None) + entry = sb._msst_model_entry() + self.assertEqual(entry["model_type"], "scnet") + self.assertIn("checkpoint_relpath", entry) + + def test_unknown_model_falls_back_to_default(self) -> None: + with mock.patch.dict(os.environ, {"ASA_MSST_MODEL": "does_not_exist"}): + entry = sb._msst_model_entry() + self.assertEqual(entry, sb._MSST_MODEL_REGISTRY[sb._DEFAULT_MSST_MODEL]) + + def test_two_stem_entry_is_not_the_default(self) -> None: + # Guard: the contract-incomplete 2-stem entry must never be the default. + self.assertNotEqual(sb._DEFAULT_MSST_MODEL, "bs_roformer_vocals") + self.assertIn("bs_roformer_vocals", sb._MSST_MODEL_REGISTRY) + + +class DispatchTests(unittest.TestCase): + def test_demucs_default_delegates_to_separate_stems(self) -> None: + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop("ASA_SEPARATION_BACKEND", None) + with mock.patch.object( + sb, "separate_stems", return_value=_SENTINEL_STEMS + ) as demucs: + out = sb.separate_stems_backend("/tmp/track.flac", output_dir="/tmp/out") + demucs.assert_called_once_with("/tmp/track.flac", output_dir="/tmp/out") + self.assertEqual(out, _SENTINEL_STEMS) + + def test_msst_without_python_env_falls_back_to_demucs(self) -> None: + with mock.patch.dict(os.environ, {"ASA_SEPARATION_BACKEND": "msst"}, clear=False): + os.environ.pop("ASA_MSST_PYTHON", None) + with mock.patch.object( + sb, "separate_stems", return_value=_SENTINEL_STEMS + ) as demucs: + out = sb.separate_stems_backend("/tmp/track.flac") + demucs.assert_called_once() + self.assertEqual(out, _SENTINEL_STEMS) + + def test_msst_runner_nonzero_exit_falls_back_to_demucs(self) -> None: + env = { + "ASA_SEPARATION_BACKEND": "msst", + "ASA_MSST_PYTHON": __file__, # any existing file path + "ASA_MSST_ROOT": str(Path(__file__).parent), # any existing dir + } + with mock.patch.dict(os.environ, env, clear=False), \ + mock.patch.object(Path, "exists", return_value=True), \ + mock.patch("separation_backend.subprocess.run", + return_value=_runner_proc("", returncode=3, stderr="boom")), \ + mock.patch.object(sb, "separate_stems", return_value=_SENTINEL_STEMS) as demucs: + out = sb.separate_stems_backend("/tmp/track.flac", output_dir="/tmp/out") + demucs.assert_called_once() + self.assertEqual(out, _SENTINEL_STEMS) + + +class RunnerManifestParseTests(unittest.TestCase): + def test_valid_manifest_returns_existing_stems(self) -> None: + manifest = json.dumps( + { + "stems": {"vocals": "/tmp/v.wav", "bass": "/tmp/b.wav"}, + "modelType": "scnet", + "loadSeconds": 1.0, + "inferSeconds": 2.0, + "device": "cpu", + } + ) + with mock.patch("separation_backend.os.path.isfile", return_value=True): + out = sb._parse_runner_manifest(manifest) + self.assertEqual(out, {"vocals": "/tmp/v.wav", "bass": "/tmp/b.wav"}) + + def test_unparseable_output_returns_none(self) -> None: + self.assertIsNone(sb._parse_runner_manifest("not json")) + + def test_missing_files_on_disk_returns_none(self) -> None: + manifest = json.dumps({"stems": {"vocals": "/nope/v.wav"}}) + with mock.patch("separation_backend.os.path.isfile", return_value=False): + self.assertIsNone(sb._parse_runner_manifest(manifest)) + + +class RunnerHelperTests(unittest.TestCase): + def test_canonical_name_maps_by_name_not_position(self) -> None: + # SCNet emits [drums, bass, other, vocals]; mapping is by name. + self.assertEqual(_runner._canonical_name("Vocals"), "vocals") + self.assertEqual(_runner._canonical_name("BASS"), "bass") + self.assertEqual(_runner._canonical_name("instrumental"), "other") + self.assertIsNone(_runner._canonical_name("piano")) + + def test_to_channels_first_transposes_channels_last(self) -> None: + # MSSeparator.separate returns channels-last (N, C); writer needs (C, N). + stem = np.zeros((44100, 2), dtype=np.float32) + out = _runner._to_channels_first(stem) + self.assertEqual(out.shape, (2, 44100)) + + def test_to_channels_first_keeps_channels_first(self) -> None: + stem = np.zeros((2, 44100), dtype=np.float32) + out = _runner._to_channels_first(stem) + self.assertEqual(out.shape, (2, 44100)) + + def test_write_wav_pcm16_roundtrip(self) -> None: + import tempfile + import wave + + audio = np.stack([np.linspace(-0.5, 0.5, 100), np.linspace(0.5, -0.5, 100)]) + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "stem.wav") + _runner._write_wav_pcm16(path, audio, 44100) + with wave.open(path, "rb") as handle: + self.assertEqual(handle.getnchannels(), 2) + self.assertEqual(handle.getframerate(), 44100) + self.assertEqual(handle.getnframes(), 100) + + +if __name__ == "__main__": + unittest.main() From 08b4dec405e750727334b60ea15595010999ad8e Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 22:41:44 +0000 Subject: [PATCH 2/2] test(backend): cover the MSST happy-path dispatch chain Add an integration test exercising the full success path: separate_stems_backend -> _separate_via_msst_subprocess -> subprocess exits 0 with a valid manifest -> stems returned with no Demucs fallback. The pieces were tested in isolation; this closes the end-to-end gap (PR #141 review). https://claude.ai/code/session_01PCXgSHYKWnrD4PiRkwzJxR --- apps/backend/tests/test_separation_backend.py | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/apps/backend/tests/test_separation_backend.py b/apps/backend/tests/test_separation_backend.py index 9422f1a2..146cdd89 100644 --- a/apps/backend/tests/test_separation_backend.py +++ b/apps/backend/tests/test_separation_backend.py @@ -91,6 +91,48 @@ def test_msst_without_python_env_falls_back_to_demucs(self) -> None: demucs.assert_called_once() self.assertEqual(out, _SENTINEL_STEMS) + def test_msst_happy_path_returns_runner_stems(self) -> None: + # Full chain: separate_stems_backend -> _separate_via_msst_subprocess -> + # subprocess exits 0 with a valid manifest -> stems returned (no fallback). + env = { + "ASA_SEPARATION_BACKEND": "msst", + "ASA_MSST_PYTHON": __file__, + "ASA_MSST_ROOT": str(Path(__file__).parent), + } + manifest = json.dumps( + { + "stems": { + "vocals": "/tmp/v.wav", + "bass": "/tmp/b.wav", + "drums": "/tmp/d.wav", + "other": "/tmp/o.wav", + }, + "modelType": "scnet", + "loadSeconds": 3.2, + "inferSeconds": 5.1, + "device": "cpu", + } + ) + with mock.patch.dict(os.environ, env, clear=False), \ + mock.patch.object(Path, "exists", return_value=True), \ + mock.patch.object(Path, "is_dir", return_value=True), \ + mock.patch("separation_backend.os.makedirs"), \ + mock.patch("separation_backend.os.path.isfile", return_value=True), \ + mock.patch("separation_backend.subprocess.run", + return_value=_runner_proc(manifest, returncode=0)), \ + mock.patch.object(sb, "separate_stems") as demucs: + out = sb.separate_stems_backend("/tmp/track.flac", output_dir="/tmp/out") + demucs.assert_not_called() # MSST succeeded — no Demucs fallback + self.assertEqual( + out, + { + "vocals": "/tmp/v.wav", + "bass": "/tmp/b.wav", + "drums": "/tmp/d.wav", + "other": "/tmp/o.wav", + }, + ) + def test_msst_runner_nonzero_exit_falls_back_to_demucs(self) -> None: env = { "ASA_SEPARATION_BACKEND": "msst",