#!/usr/bin/env python3 """Timed extraction pipeline used by the FastAPI app and benchmarks.""" from __future__ import annotations import hashlib import importlib.util import json import os import shutil import subprocess import sys import tempfile import time from contextlib import contextmanager from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any, Callable import librosa import numpy as np import soundfile as sf from sample_extractor import ( DEMUCS_MODELS, DEMUCS_STEMS, build_archive, classify_hits, cluster_hits, cluster_hits_online, detect_bpm, detect_onsets, export_midi, extract_stem, render_midi_with_samples, sample_quality_score, select_best, synthesize_from_cluster, ) ProgressCallback = Callable[[dict[str, Any]], None] SPLEETER_MODELS = ["spleeter:4stems", "spleeter:2stems", "spleeter:5stems"] SPLEETER_STEMS = { "spleeter:2stems": ["vocals", "accompaniment"], "spleeter:4stems": ["vocals", "drums", "bass", "other"], "spleeter:5stems": ["vocals", "drums", "bass", "piano", "other"], } SEPARATION_BACKENDS = ["spleeter", "demucs", "none"] def _module_available(name: str) -> bool: return importlib.util.find_spec(name) is not None def separation_runtime_status() -> dict[str, Any]: """Return lightweight runtime diagnostics for optional separation engines. This is intentionally import-spec/CLI based and does not instantiate models or download weights. It is safe to call on page load. """ spleeter_available = _module_available("spleeter") or shutil.which("spleeter") is not None demucs_available = _module_available("demucs") return { "backends": { "spleeter": { "available": bool(spleeter_available), "default": True, "install_hint": "Install requirements-spleeter.txt or use the full-mix fallback.", }, "demucs": { "available": bool(demucs_available), "default": False, "install_hint": "Install demucs and model dependencies for quality fallback separation.", }, "none": { "available": True, "default": False, "install_hint": "Always available; uses the uploaded mix without source separation.", }, } } @dataclass class PipelineParams: stem: str = "drums" separation_backend: str = "spleeter" spleeter_model: str = "spleeter:4stems" demucs_model: str = "htdemucs_ft" demucs_shifts: int = 1 demucs_overlap: float = 0.25 onset_mode: str = "auto" onset_delta: float = 0.12 energy_threshold_db: float = -35.0 pre_pad: float = 0.003 min_dur: float = 0.02 max_dur: float = 1.5 min_gap: float = 0.03 ncc_threshold: float = 0.80 attack_ms: float = 25.0 mel_threshold: float = 0.75 linkage: str = "average" clustering_mode: str = "batch_quality" target_min: int = 5 target_max: int = 20 synthesize: bool = True quantize_midi: bool = True subdivision: int = 16 device: str = "cpu" auto_tune: bool = True use_disk_cache: bool = True allow_backend_fallback: bool = True @classmethod def from_mapping(cls, data: dict[str, Any] | None) -> "PipelineParams": data = dict(data or {}) allowed = {field.name for field in cls.__dataclass_fields__.values()} unknown = sorted(set(data) - allowed) if unknown: raise ValueError(f"Unknown pipeline parameter(s): {', '.join(unknown)}") int_fields = {"demucs_shifts", "target_min", "target_max", "subdivision"} float_fields = { "demucs_overlap", "onset_delta", "energy_threshold_db", "pre_pad", "min_dur", "max_dur", "min_gap", "ncc_threshold", "attack_ms", "mel_threshold", } bool_fields = {"synthesize", "quantize_midi", "auto_tune", "use_disk_cache", "allow_backend_fallback"} def coerce_bool(name: str, value: Any) -> bool: if isinstance(value, bool): return value if isinstance(value, (int, float)) and value in {0, 1}: return bool(value) if isinstance(value, str): normalized = value.strip().lower() if normalized in {"true", "1", "yes", "on"}: return True if normalized in {"false", "0", "no", "off"}: return False raise ValueError(f"{name} must be a boolean") coerced: dict[str, Any] = {} for name, value in data.items(): if value is None: coerced[name] = value elif name in int_fields: try: number = float(value) if isinstance(value, str) else value if int(number) != number: raise ValueError coerced[name] = int(number) except (TypeError, ValueError): raise ValueError(f"{name} must be an integer") from None elif name in float_fields: try: coerced[name] = float(value) except (TypeError, ValueError): raise ValueError(f"{name} must be a number") from None elif name in bool_fields: coerced[name] = coerce_bool(name, value) else: coerced[name] = value params = cls(**coerced) params.validate() return params def validate(self) -> None: if self.separation_backend not in set(SEPARATION_BACKENDS): raise ValueError(f"Unsupported separation backend: {self.separation_backend}") if self.spleeter_model not in SPLEETER_MODELS: raise ValueError(f"Unsupported Spleeter model: {self.spleeter_model}") if self.demucs_model not in DEMUCS_MODELS: raise ValueError(f"Unsupported Demucs model: {self.demucs_model}") if self.separation_backend == "demucs": allowed_stems = set(DEMUCS_STEMS.get(self.demucs_model, [])) | {"all"} backend_label = self.demucs_model elif self.separation_backend == "spleeter": allowed_stems = set(SPLEETER_STEMS.get(self.spleeter_model, [])) | {"all"} backend_label = self.spleeter_model else: allowed_stems = {"all"} backend_label = "full mix" if self.stem not in allowed_stems: raise ValueError(f"Stem '{self.stem}' is not available for {backend_label}") if self.onset_mode not in {"auto", "percussive", "harmonic", "broadband"}: raise ValueError(f"Unsupported onset mode: {self.onset_mode}") if self.linkage not in {"average", "complete", "single"}: raise ValueError(f"Unsupported clustering linkage: {self.linkage}") if self.clustering_mode not in {"batch_quality", "online_preview"}: raise ValueError(f"Unsupported clustering mode: {self.clustering_mode}") if not 0 <= self.demucs_shifts <= 8: raise ValueError("demucs_shifts must be between 0 and 8") if not 0.0 <= self.demucs_overlap <= 0.9: raise ValueError("demucs_overlap must be between 0.0 and 0.9") if not 0.001 <= self.onset_delta <= 1.0: raise ValueError("onset_delta must be between 0.001 and 1.0") if not -100.0 <= self.energy_threshold_db <= 0.0: raise ValueError("energy_threshold_db must be between -100 and 0 dB") if not 0.0 <= self.pre_pad <= 0.25: raise ValueError("pre_pad must be between 0 and 0.25 seconds") if not 0.001 <= self.min_dur <= self.max_dur <= 10.0: raise ValueError("duration bounds must satisfy 0.001 <= min_dur <= max_dur <= 10") if not 0.001 <= self.min_gap <= 1.0: raise ValueError("min_gap must be between 0.001 and 1.0 seconds") if not 0.0 <= self.ncc_threshold <= 1.0: raise ValueError("ncc_threshold must be between 0 and 1") if not 1.0 <= self.attack_ms <= 250.0: raise ValueError("attack_ms must be between 1 and 250 ms") if not 0.0 <= self.mel_threshold <= 1.0: raise ValueError("mel_threshold must be between 0 and 1") if self.target_min < 0 or self.target_max < 0: raise ValueError("target_min and target_max must be non-negative") if self.target_max and self.target_min and self.target_min > self.target_max: raise ValueError("target_min cannot be greater than target_max") if self.subdivision not in {4, 8, 16, 32, 64}: raise ValueError("subdivision must be one of 4, 8, 16, 32, 64") @dataclass class StageTiming: key: str label: str duration_sec: float = 0.0 status: str = "pending" detail: str = "" progress: float = 0.0 work_done: int | None = None work_total: int | None = None @dataclass class PipelineResult: params: dict[str, Any] duration_sec: float audio_duration_sec: float realtime_factor: float bpm: float | None sample_rate: int hit_count: int cluster_count: int stages: list[dict[str, Any]] samples: list[dict[str, Any]] hits: list[dict[str, Any]] overview: dict[str, Any] files: dict[str, str] STAGE_DEFS = [ ("stem", "Stem separation / source load"), ("auto_tune", "Automatic parameter tuning"), ("bpm", "Tempo detection"), ("onsets", "Onset detection + slicing"), ("classification", "Spectral rule classification"), ("clustering", "Mel fingerprint + transient NCC clustering"), ("selection", "Best representative scoring"), ("synthesis", "Optional sample synthesis"), ("export", "MIDI, reproduced mix, WAV, ZIP export"), ] def initial_stages() -> list[dict[str, Any]]: return [asdict(StageTiming(key=key, label=label)) for key, label in STAGE_DEFS] def _notify(cb: ProgressCallback | None, payload: dict[str, Any]) -> None: if cb: cb(payload) def _progress_payload(stages: list[StageTiming]) -> dict[str, Any]: running = next((stage for stage in stages if stage.status == "running"), None) running_progress = float(np.clip(running.progress, 0.0, 1.0)) if running else 0.0 total = 0.0 completed_units = 0.0 completed = 0 for stage in stages: units = float(stage.work_total or 1) total += units if stage.status == "done": completed += 1 completed_units += units elif stage.status == "running": completed_units += float(stage.work_done) if stage.work_total and stage.work_done is not None else running_progress * units fraction = float(np.clip(completed_units / max(total, 1.0), 0.0, 1.0)) return { "fraction": round(fraction, 6), "completed_steps": completed, "total_steps": len(stages), "completed_units": round(completed_units, 6), "total_units": round(total, 6), "stage_key": running.key if running else None, "stage_label": running.label if running else None, "stage_fraction": round(running_progress, 6), "stage_work_done": running.work_done if running else None, "stage_work_total": running.work_total if running else None, "basis": "exact completed work units: Demucs chunks when available; Spleeter and non-instrumented stages advance only at real stage boundaries; no time-based estimates", } def _notify_stages(cb: ProgressCallback | None, stages: list[StageTiming], event_type: str, stage: StageTiming) -> None: _notify(cb, { "type": event_type, "stage": asdict(stage), "stages": [asdict(item) for item in stages], "progress": _progress_payload(stages), }) @contextmanager def _timed_stage(stages: list[StageTiming], key: str, cb: ProgressCallback | None = None): stage = next(stage for stage in stages if stage.key == key) stage.status = "running" stage.progress = 0.0 stage.work_done = None stage.work_total = None _notify_stages(cb, stages, "stage", stage) started = time.perf_counter() try: yield stage except Exception as exc: stage.duration_sec = time.perf_counter() - started stage.status = "error" stage.detail = str(exc) _notify_stages(cb, stages, "stage", stage) raise else: stage.duration_sec = time.perf_counter() - started stage.status = "done" stage.progress = 1.0 _notify_stages(cb, stages, "stage", stage) def _update_stage_progress( stages: list[StageTiming], key: str, cb: ProgressCallback | None, *, fraction: float, detail: str | None = None, work_done: int | None = None, work_total: int | None = None, ) -> None: stage = next(stage for stage in stages if stage.key == key) if stage.status != "running": return stage.progress = float(np.clip(fraction, 0.0, 1.0)) stage.work_done = work_done stage.work_total = work_total if detail: stage.detail = detail _notify_stages(cb, stages, "stage_progress", stage) def _normalise_audio(audio: np.ndarray) -> np.ndarray: audio = audio.astype(np.float32) if audio.ndim > 1: audio = audio.mean(axis=1) peak = float(np.max(np.abs(audio))) if audio.size else 0.0 if peak > 0: audio = audio / peak return audio.astype(np.float32) def _mono(audio: np.ndarray) -> np.ndarray: audio = np.asarray(audio, dtype=np.float32) if audio.ndim > 1: audio = audio.mean(axis=1) return audio.astype(np.float32) def _pad_or_trim(audio: np.ndarray, length: int) -> np.ndarray: audio = _mono(audio) if len(audio) == length: return audio if len(audio) > length: return audio[:length] return np.pad(audio, (0, max(0, length - len(audio)))).astype(np.float32) def _load_source_mix(audio_path: str | os.PathLike[str], sr: int) -> np.ndarray: audio, _ = librosa.load(audio_path, sr=sr, mono=True) return _mono(audio) def _common_gain(reference: np.ndarray, fallback: np.ndarray) -> float: peak = float(np.max(np.abs(reference))) if reference.size else 0.0 if peak <= 1e-8: peak = float(np.max(np.abs(fallback))) if fallback.size else 0.0 return peak if peak > 1e-8 else 1.0 def _rms(audio: np.ndarray) -> float: audio = _mono(audio) return float(np.sqrt(np.mean(np.square(audio, dtype=np.float64)))) if audio.size else 0.0 def _match_rms(rendered: np.ndarray, reference: np.ndarray, *, min_gain: float = 0.05, max_gain: float = 8.0) -> np.ndarray: rendered_rms = _rms(rendered) reference_rms = _rms(reference) if rendered_rms <= 1e-10 or reference_rms <= 1e-10: return _mono(rendered) gain = float(np.clip(reference_rms / rendered_rms, min_gain, max_gain)) return (_mono(rendered) * gain).astype(np.float32) def _soft_limit(audio: np.ndarray, ceiling: float = 0.98) -> np.ndarray: audio = _mono(audio).astype(np.float32) peak = float(np.max(np.abs(audio))) if audio.size else 0.0 if peak > ceiling > 0: audio = audio * (ceiling / peak) return audio.astype(np.float32) def _make_reproduction_mix(target_reconstruction: np.ndarray, context_bed: np.ndarray, length: int) -> np.ndarray: target = _pad_or_trim(target_reconstruction, length) context = _pad_or_trim(context_bed, length) return _soft_limit(context + target) MODULE_ROOT = Path(__file__).resolve().parent CACHE_DIR = Path(os.environ["DSE_CACHE_DIR"]) if os.environ.get("DSE_CACHE_DIR") else MODULE_ROOT / ".cache" STEM_CACHE_DIR = CACHE_DIR / "stems" CACHE_VERSION = "dse-cache-v4-runtime-fallback" def _write_audio(path: Path, audio: np.ndarray, sr: int, subtype: str = "PCM_24") -> None: path.parent.mkdir(parents=True, exist_ok=True) sf.write(path, audio.astype(np.float32), sr, subtype=subtype) def _sha256_file(path: str | os.PathLike[str]) -> str: h = hashlib.sha256() with Path(path).open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): h.update(chunk) return h.hexdigest() def _stem_cache_path(audio_path: str | os.PathLike[str], params: PipelineParams) -> Path: key_payload = { "version": CACHE_VERSION, "source_sha256": _sha256_file(audio_path), "stem": params.stem, "separation_backend": params.separation_backend, "spleeter_model": params.spleeter_model, "demucs_model": params.demucs_model, "demucs_shifts": params.demucs_shifts, "demucs_overlap": params.demucs_overlap, "device": params.device if params.stem != "all" else "decode", } key = hashlib.sha256(json.dumps(key_payload, sort_keys=True).encode("utf-8")).hexdigest() return STEM_CACHE_DIR / f"{key}.wav" def _context_cache_path(stem_cache_path: Path) -> Path: return stem_cache_path.with_name(f"{stem_cache_path.stem}.context.wav") def clear_disk_cache() -> None: if CACHE_DIR.exists(): shutil.rmtree(CACHE_DIR) def _load_spleeter_audio(path: Path, sr: int | None = None) -> tuple[np.ndarray, int]: audio, loaded_sr = librosa.load(path, sr=sr, mono=True) return _mono(audio), int(loaded_sr) def _spleeter_output_file(root: Path, source_stem: str, stem: str) -> Path | None: candidates = [root / source_stem / f"{stem}.wav", root / f"{stem}.wav"] candidates.extend(root.glob(f"**/{stem}.wav")) for candidate in candidates: if candidate.exists(): return candidate return None def _extract_spleeter_separation( audio_path: str | os.PathLike[str], params: PipelineParams, progress_cb: Callable[[dict[str, Any]], None] | None = None, ) -> tuple[np.ndarray, int, np.ndarray | None, str]: """Run Spleeter and return target stem plus the sum of non-target stems. Progress is deliberately coarse because Spleeter does not expose reliable chunk callbacks through the CLI/Python API. We report start/completion only; the global UI therefore never interpolates fake progress. """ if progress_cb: progress_cb({"fraction": 0.0, "completed_units": 0, "total_units": 1, "detail": f"Spleeter {params.spleeter_model} starting"}) with tempfile.TemporaryDirectory(prefix="dse_spleeter_") as tmp: tmpdir = Path(tmp) commands = [ [ sys.executable, "-m", "spleeter", "separate", "-p", params.spleeter_model, "-o", str(tmpdir), str(audio_path), ], [ "spleeter", "separate", "-p", params.spleeter_model, "-o", str(tmpdir), str(audio_path), ], ] failures: list[str] = [] completed = None for cmd in commands: try: completed = subprocess.run(cmd, capture_output=True, text=True, check=False, timeout=60 * 30) except FileNotFoundError as exc: failures.append(str(exc)) continue if completed.returncode == 0: break failures.append((completed.stderr or completed.stdout or "Spleeter failed").strip()[-1200:]) if completed is None or completed.returncode != 0: raise RuntimeError("; ".join(part for part in failures if part) or "Spleeter failed") source_stem = Path(audio_path).stem stems = SPLEETER_STEMS[params.spleeter_model] paths = {stem: _spleeter_output_file(tmpdir, source_stem, stem) for stem in stems} missing = [stem for stem, path in paths.items() if path is None] if missing: raise RuntimeError(f"Spleeter finished but did not write expected stem(s): {', '.join(missing)}") if params.stem not in paths: raise RuntimeError(f"Spleeter model {params.spleeter_model} does not provide stem '{params.stem}'") target, sr = _load_spleeter_audio(paths[params.stem]) # type: ignore[arg-type] context_parts: list[np.ndarray] = [] for name, path in paths.items(): if name == params.stem or path is None: continue part, _ = _load_spleeter_audio(path, sr=sr) context_parts.append(_pad_or_trim(part, len(target))) context = np.sum(np.stack(context_parts), axis=0).astype(np.float32) if context_parts else None if progress_cb: progress_cb({"fraction": 1.0, "completed_units": 1, "total_units": 1, "detail": f"Spleeter {params.spleeter_model} complete"}) return target.astype(np.float32), sr, context, f"{params.stem} via {params.spleeter_model}" def _extract_demucs_separation( audio_path: str | os.PathLike[str], params: PipelineParams, progress_cb: Callable[[dict[str, Any]], None] | None = None, ) -> tuple[np.ndarray, int, np.ndarray | None, str]: audio, sr = extract_stem( str(audio_path), stem=params.stem, device=params.device, model_name=params.demucs_model, shifts=int(params.demucs_shifts), overlap=float(params.demucs_overlap), progress_cb=progress_cb, ) return audio, sr, None, f"{params.stem} via Demucs {params.demucs_model}" def _load_full_mix_separation( audio_path: str | os.PathLike[str], progress_cb: Callable[[dict[str, Any]], None] | None = None, detail: str = "loaded full mix", ) -> tuple[np.ndarray, int, str, np.ndarray | None]: audio, sr = librosa.load(audio_path, sr=44100, mono=True) if progress_cb: progress_cb({"fraction": 1.0, "completed_units": 1, "total_units": 1, "detail": detail}) return _mono(audio), int(sr), detail, None def _fallback_to_full_mix( audio_path: str | os.PathLike[str], progress_cb: Callable[[dict[str, Any]], None] | None, reason: str, ) -> tuple[np.ndarray, int, str, np.ndarray | None]: return _load_full_mix_separation( audio_path, progress_cb=progress_cb, detail=f"fallback full mix because separation failed: {reason}", ) def _load_or_extract_separation(audio_path: str | os.PathLike[str], params: PipelineParams, progress_cb: Callable[[dict[str, Any]], None] | None = None) -> tuple[np.ndarray, int, str, np.ndarray | None]: if params.stem == "all" or params.separation_backend == "none": return _load_full_mix_separation(audio_path, progress_cb=progress_cb, detail="loaded full mix") cache_path = _stem_cache_path(audio_path, params) context_cache = _context_cache_path(cache_path) if params.use_disk_cache and cache_path.exists(): audio, sr = sf.read(cache_path, dtype="float32", always_2d=False) context = None if context_cache.exists(): context_audio, context_sr = sf.read(context_cache, dtype="float32", always_2d=False) context = _pad_or_trim(_mono(context_audio), len(_mono(audio))) if int(context_sr) == int(sr) else _mono(librosa.resample(_mono(context_audio), orig_sr=int(context_sr), target_sr=int(sr))) if progress_cb: progress_cb({"fraction": 1.0, "completed_units": 1, "total_units": 1, "detail": f"{params.stem} disk-cache hit"}) return np.asarray(audio, dtype=np.float32), int(sr), f"{params.stem} disk-cache hit", context context: np.ndarray | None = None if params.separation_backend == "spleeter": try: audio, sr, context, detail = _extract_spleeter_separation(audio_path, params, progress_cb=progress_cb) except Exception as spleeter_exc: if not params.allow_backend_fallback: raise if progress_cb: progress_cb({"fraction": 0.0, "completed_units": 0, "total_units": 1, "detail": f"Spleeter unavailable; using full-mix fallback: {spleeter_exc}"}) audio, sr, detail, context = _fallback_to_full_mix(audio_path, progress_cb, f"Spleeter failed ({spleeter_exc})") elif params.separation_backend == "demucs": try: audio, sr, context, detail = _extract_demucs_separation(audio_path, params, progress_cb=progress_cb) except Exception as demucs_exc: if not params.allow_backend_fallback: raise audio, sr, detail, context = _fallback_to_full_mix(audio_path, progress_cb, f"Demucs failed ({demucs_exc})") else: raise ValueError(f"Unsupported separation backend: {params.separation_backend}") if params.use_disk_cache: _write_audio(cache_path, audio, sr, subtype="PCM_16") if context is not None: _write_audio(context_cache, context, sr, subtype="PCM_16") detail += " · cached" return audio, sr, detail, context def _build_onset_envelope(audio: np.ndarray, sr: int) -> np.ndarray: """Build one onset envelope for automatic tuning; candidate deltas reuse it.""" y = _mono(audio) if y.size == 0: return np.zeros(0, dtype=np.float32) hop_length = 220 n_fft = 1024 try: yh, _ = librosa.effects.hpss(y) def _sf(part: np.ndarray, fmin: float | None = None, fmax: float | None = None, ms: int = 3, lag: int = 1) -> np.ndarray: S = librosa.feature.melspectrogram( y=part, sr=sr, n_fft=n_fft, hop_length=hop_length, n_mels=96, fmin=fmin or 20, fmax=fmax or sr // 2, power=1, ) return librosa.onset.onset_strength( S=librosa.power_to_db(S, ref=np.max), sr=sr, hop_length=hop_length, lag=lag, max_size=ms, ) def _norm(values: np.ndarray) -> np.ndarray: peak = float(np.max(values)) if values.size else 0.0 return values / peak if peak > 0 else values return np.maximum.reduce([ _norm(_sf(y, 20, 300)), _norm(_sf(y, 300, 4000)), _norm(_sf(y, 4000, 16000)), _norm(_sf(yh, lag=2, ms=5)), ]).astype(np.float32) except Exception: return np.zeros(0, dtype=np.float32) def _count_onsets_from_envelope(envelope: np.ndarray, sr: int, delta: float) -> int: if envelope.size == 0: return 0 hop_length = 220 frames = librosa.onset.onset_detect( onset_envelope=envelope, sr=sr, hop_length=hop_length, wait=max(1, int(0.03 * sr / hop_length)), pre_avg=3, post_avg=3, pre_max=3, post_max=5, delta=float(delta), backtrack=True, units="frames", ) return int(len(frames)) def _auto_tune_params(audio: np.ndarray, sr: int, params: PipelineParams) -> dict[str, Any]: """Tune the common detection/grouping parameters from the separated target stem.""" duration = max(0.001, len(audio) / max(sr, 1)) # The target density keeps enough repeated hits for clustering without drowning the UI in bleed. desired_min = max(4, int(duration * 1.1)) desired_max = max(desired_min + 2, int(duration * 4.8)) candidates = [0.06, 0.075, 0.09, 0.105, 0.12, 0.145, 0.17, 0.21, 0.26, 0.32] envelope = _build_onset_envelope(audio, sr) counts = [(delta, _count_onsets_from_envelope(envelope, sr, delta)) for delta in candidates] viable = [(delta, count) for delta, count in counts if desired_min <= count <= desired_max] if viable: chosen_delta, chosen_count = min(viable, key=lambda item: abs(item[1] - ((desired_min + desired_max) / 2))) else: chosen_delta, chosen_count = min(counts, key=lambda item: min(abs(item[1] - desired_min), abs(item[1] - desired_max))) type_budget = max(6, min(48, int(round((chosen_count or desired_min) ** 0.5 * 3)))) params.onset_delta = float(chosen_delta) params.target_min = max(2, min(int(params.target_min or 0) or 5, type_budget)) params.target_max = max(params.target_min, type_budget) if params.clustering_mode == "batch_quality" and duration <= 180: # Keep quality for normal songs; online preview remains user/preset controlled. params.clustering_mode = "batch_quality" return { "chosen_onset_delta": round(float(chosen_delta), 4), "estimated_onsets": int(chosen_count), "desired_onset_range": [int(desired_min), int(desired_max)], "target_min": int(params.target_min), "target_max": int(params.target_max), "candidate_counts": [{"onset_delta": round(float(delta), 4), "count": int(count)} for delta, count in counts], } def _make_overview(audio: np.ndarray, sr: int, hits: list[Any], max_points: int = 1600) -> dict[str, Any]: if len(audio) == 0: return {"sample_rate": sr, "duration_sec": 0, "envelope": [], "onsets": []} frame = max(1, int(np.ceil(len(audio) / max_points))) usable = (len(audio) // frame) * frame if usable == 0: envelope = [float(np.max(np.abs(audio)))] else: envelope = np.max(np.abs(audio[:usable].reshape(-1, frame)), axis=1).astype(float).tolist() return { "sample_rate": sr, "duration_sec": round(len(audio) / sr, 6), "frame_duration_sec": round(frame / sr, 6), "envelope": [round(float(x), 6) for x in envelope], "onsets": [ { "index": int(getattr(h, "index", -1)), "time_sec": round(float(h.onset_time), 6), "duration_sec": round(float(h.duration), 6), "label": h.label, "energy": round(float(h.rms_energy), 6), "cluster_id": int(getattr(h, "cluster_id", -1)), } for h in hits ], } def _copy_temp_file(src: str | os.PathLike[str], dst: Path) -> str: dst.parent.mkdir(parents=True, exist_ok=True) shutil.copyfile(src, dst) return str(dst) def _safe_file_component(value: str) -> str: cleaned = "".join(ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in value.lower()) while "__" in cleaned: cleaned = cleaned.replace("__", "_") return cleaned.strip("_") or "item" def run_extraction_pipeline( audio_path: str | os.PathLike[str], output_dir: str | os.PathLike[str], params: PipelineParams | dict[str, Any] | None = None, progress_cb: ProgressCallback | None = None, ) -> PipelineResult: """Run extraction and write all runtime artifacts into output_dir.""" if not isinstance(params, PipelineParams): params = PipelineParams.from_mapping(params) out = Path(output_dir) out.mkdir(parents=True, exist_ok=True) samples_dir = out / "samples" samples_dir.mkdir(parents=True, exist_ok=True) stages = [StageTiming(key=key, label=label) for key, label in STAGE_DEFS] started_total = time.perf_counter() bpm: float | None = None stem_audio: np.ndarray source_audio: np.ndarray context_bed: np.ndarray stem_sr: int hits: list[Any] = [] clusters: list[Any] = [] target_rendered: np.ndarray | None = None reproduced: np.ndarray | None = None _notify(progress_cb, {"type": "start", "stages": [asdict(s) for s in stages], "progress": _progress_payload(stages)}) with _timed_stage(stages, "stem", progress_cb) as stage: def _stem_progress(event: dict[str, Any]) -> None: _update_stage_progress( stages, "stem", progress_cb, fraction=float(event.get("fraction", 0.0)), detail=str(event.get("detail") or stage.detail or "stem extraction"), work_done=event.get("completed_units"), work_total=event.get("total_units"), ) raw_stem_audio, stem_sr, stem_detail, separated_context = _load_or_extract_separation(audio_path, params, progress_cb=_stem_progress) source_raw = _load_source_mix(audio_path, stem_sr) length = max(len(raw_stem_audio), len(source_raw)) raw_stem_audio = _pad_or_trim(raw_stem_audio, length) source_raw = _pad_or_trim(source_raw, length) gain = _common_gain(raw_stem_audio if params.stem != "all" else source_raw, source_raw) stem_audio = (raw_stem_audio / gain).astype(np.float32) source_audio = (source_raw / gain).astype(np.float32) context_bed = np.zeros_like(source_audio) if params.stem == "all" else (_pad_or_trim(separated_context, length) / gain if separated_context is not None else (source_audio - stem_audio)).astype(np.float32) context_detail = "full mix" if params.stem == "all" else ("separated non-target stems" if separated_context is not None else "residual non-target stems") stage.detail = stem_detail + f" · reproduction uses {context_detail}" _write_audio(out / "source.wav", _soft_limit(source_audio), stem_sr, subtype="PCM_16") _write_audio(out / "stem.wav", _soft_limit(stem_audio), stem_sr, subtype="PCM_16") _write_audio(out / "context_bed.wav", _soft_limit(context_bed), stem_sr, subtype="PCM_16") audio_duration_sec = len(source_audio) / stem_sr if stem_sr else 0.0 with _timed_stage(stages, "auto_tune", progress_cb) as stage: if params.auto_tune: tuning = _auto_tune_params(stem_audio, stem_sr, params) stage.detail = ( f"onset_delta={tuning['chosen_onset_delta']} · " f"estimated {tuning['estimated_onsets']} hits · " f"groups {tuning['target_min']}-{tuning['target_max']}" ) else: stage.detail = "disabled" with _timed_stage(stages, "bpm", progress_cb) as stage: bpm = detect_bpm(stem_audio, stem_sr) stage.detail = f"{bpm} BPM" with _timed_stage(stages, "onsets", progress_cb) as stage: hits = detect_onsets( stem_audio, stem_sr, mode=params.onset_mode, onset_delta=float(params.onset_delta), energy_threshold_db=float(params.energy_threshold_db), pre_pad=float(params.pre_pad), min_dur=float(params.min_dur), max_dur=float(params.max_dur), min_gap=float(params.min_gap), ) stage.detail = f"{len(hits)} hits" if hits: with _timed_stage(stages, "classification", progress_cb) as stage: hits = classify_hits(hits) counts: dict[str, int] = {} for hit in hits: counts[hit.label] = counts.get(hit.label, 0) + 1 stage.detail = ", ".join(f"{key}:{value}" for key, value in sorted(counts.items())) with _timed_stage(stages, "clustering", progress_cb) as stage: if params.clustering_mode == "online_preview": clusters = cluster_hits_online( hits, audio=stem_audio, sr=stem_sr, ncc_threshold=float(params.ncc_threshold), attack_ms=float(params.attack_ms), mel_threshold=float(params.mel_threshold), target_min=int(params.target_min), target_max=int(params.target_max), ) stage.detail = f"{len(clusters)} clusters · online preview" else: clusters = cluster_hits( hits, audio=stem_audio, sr=stem_sr, ncc_threshold=float(params.ncc_threshold), attack_ms=float(params.attack_ms), mel_threshold=float(params.mel_threshold), target_min=int(params.target_min), target_max=int(params.target_max), linkage=params.linkage, ) stage.detail = f"{len(clusters)} clusters · batch quality" for cluster in clusters: for hit in cluster.hits: hit.cluster_id = cluster.cluster_id with _timed_stage(stages, "selection", progress_cb) as stage: select_best(clusters) stage.detail = "quality-scored representatives" with _timed_stage(stages, "synthesis", progress_cb) as stage: if params.synthesize: synth_count = 0 for cluster in clusters: if cluster.count >= 2: cluster.synthesized = synthesize_from_cluster(cluster) synth_count += int(cluster.synthesized is not None) stage.detail = f"{synth_count} synthesized alternates" else: stage.detail = "disabled" else: for key, detail in [ ("classification", "skipped: no hits"), ("clustering", "skipped: no hits"), ("selection", "skipped: no hits"), ("synthesis", "skipped: no hits"), ]: stage = next(s for s in stages if s.key == key) stage.status = "done" stage.detail = detail sample_rows: list[dict[str, Any]] = [] hit_rows: list[dict[str, Any]] = [] files: dict[str, str] = {"source": "source.wav", "stem": "stem.wav", "context_bed": "context_bed.wav"} with _timed_stage(stages, "export", progress_cb) as stage: midi_path = out / "reconstruction.mid" if clusters: export_midi( clusters, str(midi_path), bpm=bpm or 120.0, quantize=bool(params.quantize_midi), subdivision=int(params.subdivision), ) target_rendered = render_midi_with_samples(clusters, sr=stem_sr) target_rendered = _match_rms(target_rendered, stem_audio) else: target_rendered = np.zeros_like(stem_audio) midi_path.write_bytes(b"") reproduced = _make_reproduction_mix(target_rendered, context_bed, max(len(source_audio), len(target_rendered))) _write_audio(out / "target_reconstruction.wav", _soft_limit(target_rendered), stem_sr, subtype="PCM_16") _write_audio(out / "reconstruction.wav", reproduced, stem_sr, subtype="PCM_16") files["target_reconstruction"] = "target_reconstruction.wav" files["reconstruction"] = "reconstruction.wav" files["midi"] = "reconstruction.mid" cluster_labels = {int(cluster.cluster_id): cluster.label for cluster in clusters} representative_ids = {id(cluster.best_hit) for cluster in clusters} review_hits_dir = out / "review" / "hits" if hits: review_hits_dir.mkdir(parents=True, exist_ok=True) for hit in sorted(hits, key=lambda item: item.index): safe_label = _safe_file_component(hit.label or "hit") file_name = f"hit_{int(hit.index):05d}_{safe_label}.wav" rel_file = f"review/hits/{file_name}" hit.save(str(out / rel_file)) cluster_id = int(getattr(hit, "cluster_id", -1)) hit_rows.append( { "index": int(hit.index), "label": hit.label, "cluster_id": cluster_id, "cluster_label": cluster_labels.get(cluster_id, "unclustered"), "is_representative": id(hit) in representative_ids, "onset_sec": round(float(hit.onset_time), 6), "duration_ms": round(float(hit.duration * 1000), 1), "rms_energy": round(float(hit.rms_energy), 6), "spectral_centroid_hz": round(float(hit.spectral_centroid), 1), "file": rel_file, } ) for cluster in sorted(clusters, key=lambda item: item.count, reverse=True): best = cluster.best_hit sample_path = samples_dir / f"{cluster.label}.wav" best.save(str(sample_path)) quality = sample_quality_score(best.audio, best.sr, cluster.label.rsplit("_", 1)[0]) sample_row = { "label": cluster.label, "classification": cluster.label.rsplit("_", 1)[0], "hits": int(cluster.count), "midi_note": int(cluster.midi_note), "score": round(float(quality["total"]), 2), "cleanness": round(float(quality["cleanness"]), 4), "completeness": round(float(quality["completeness"]), 4), "duration_ms": round(float(best.duration * 1000), 1), "first_onset_sec": round(float(min(hit.onset_time for hit in cluster.hits)), 4), "representative_hit_index": int(best.index), "cluster_id": int(cluster.cluster_id), "file": f"samples/{cluster.label}.wav", } sample_rows.append(sample_row) _notify(progress_cb, {"type": "sample", "sample": sample_row, "samples": list(sample_rows)}) if cluster.synthesized is not None: synth_path = samples_dir / f"{cluster.label}__synth.wav" _write_audio(synth_path, cluster.synthesized, stem_sr) archive_tmp = build_archive( clusters, bpm or 120.0, stem_sr, midi_path=str(midi_path), rendered_audio=reproduced, target_rendered_audio=target_rendered, ) files["archive"] = _copy_temp_file(archive_tmp, out / "sample-pack.zip") files["archive"] = "sample-pack.zip" try: os.unlink(archive_tmp) except OSError: pass stage.detail = f"{len(sample_rows)} samples + {len(hit_rows)} review hits + MIDI + full-context reproduction + ZIP" duration_sec = time.perf_counter() - started_total result = PipelineResult( params=asdict(params), duration_sec=round(duration_sec, 6), audio_duration_sec=round(audio_duration_sec, 6), realtime_factor=round(duration_sec / max(audio_duration_sec, 1e-9), 6), bpm=bpm, sample_rate=stem_sr, hit_count=len(hits), cluster_count=len(clusters), stages=[asdict(stage) for stage in stages], samples=sample_rows, hits=hit_rows, overview=_make_overview(stem_audio, stem_sr, hits), files=files, ) (out / "manifest.json").write_text(json.dumps(asdict(result), indent=2), encoding="utf-8") _notify(progress_cb, {"type": "complete", "result": asdict(result), "stages": result.stages, "progress": {**_progress_payload(stages), "fraction": 1.0}}) return result