from __future__ import annotations import json import logging import shutil import subprocess import time from pathlib import Path import numpy as np import soundfile as sf from app.core.config import ( DEMUCS_MODEL, JOB_TTL_SECONDS, STEM_NAMES, TIMEOUT_FFMPEG, ffmpeg_executable, ) from app.core.models import Job from app.core.registry import all_jobs as registry_all from app.core.registry import persist as registry_persist from app.core.registry import remove as registry_remove from app.core.registry import set_proc logger = logging.getLogger("stemdeck.collect") def _rmtree(path: Path) -> None: try: shutil.rmtree(path) except FileNotFoundError: pass except Exception: logger.warning("failed to remove %s", path, exc_info=True) def _run_ffmpeg(job: Job, cmd: list[str]) -> bool: """Run an ffmpeg command, registering the subprocess with the job registry so POST /api/jobs/{id}/cancel can terminate it. Returns True on success, False on failure or external termination. Without registering the proc, an in-flight ffmpeg amix would block cancellation for up to its 300s timeout -- the cancel flag is set but the runner can't see it until subprocess.run returns. With set_proc, the cancel API can call proc.terminate() directly and communicate() returns within ~1s with a non-zero returncode.""" proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) set_proc(job.id, proc) try: try: _, stderr = proc.communicate(timeout=TIMEOUT_FFMPEG) except subprocess.TimeoutExpired: proc.kill() proc.communicate() logger.warning("ffmpeg timed out for job %s", job.id) return False if proc.returncode != 0: tail = (stderr or b"").decode(errors="replace").splitlines()[-3:] logger.warning( "ffmpeg exit %s for job %s: %s", proc.returncode, job.id, " | ".join(tail) or "(no stderr)", ) return False return True finally: set_proc(job.id, None) _TERMINAL = frozenset(("done", "error", "cancelled")) def collect(job: Job, stems_root: Path, job_dir: Path) -> list[str]: """Move Demucs-emitted stems into the job's stems/ dir and clean up the demucs intermediate dir. Does NOT delete the source download -- cleanup_source() is called by the runner after any post-processing that needs to re-encode the source (e.g. building original.wav).""" target_dir = job_dir / "stems" target_dir.mkdir(exist_ok=True) found: list[str] = [] for name in STEM_NAMES: src = stems_root / f"{name}.wav" if src.exists(): shutil.move(str(src), target_dir / f"{name}.wav") found.append(name) _rmtree(job_dir / DEMUCS_MODEL) if not found: raise RuntimeError("no stems produced by demucs") return found def cleanup_source(job_dir: Path) -> None: """Delete the source audio file. Called after collect AND after any post-processing that re-encodes the source (make_original_track). The source is 100-300 MB, so getting rid of it is the bulk of disk reclaim per job; only the stems remain.""" for f in job_dir.glob("source.*"): f.unlink(missing_ok=True) def make_original_track(job: Job, job_dir: Path, stems_dir: Path) -> Path | None: """Build the "Original" backing track at stems/original.wav as the sum of the stems the user did NOT select. This way the studio can play (original + each selected stem) and reconstruct the full song without doubling the selected stems -- which is what would happen if "original" were the raw source download (drum hits in original + isolated drums.wav = drums at 2x amplitude). Skipped when the user kept all 6 stems (no complement to mix) or when none of the unselected stem WAVs are on disk.""" unselected = [s for s in STEM_NAMES if s not in job.selected_stems] inputs = [stems_dir / f"{name}.wav" for name in unselected] inputs = [p for p in inputs if p.exists()] if not inputs: return None out = stems_dir / "original.wav" cmd: list[str] = [ ffmpeg_executable(), "-y", "-nostdin", "-loglevel", "error", ] for p in inputs: cmd += ["-i", str(p)] if len(inputs) == 1: # Single complement stem -- copy as-is so we still produce a # canonical mix.wav-shaped output without invoking amix on a # 1-input graph (which is a no-op anyway). cmd += ["-c:a", "pcm_s16le", str(out)] else: filter_inputs = "".join(f"[{i}:a]" for i in range(len(inputs))) cmd += [ "-filter_complex", f"{filter_inputs}amix=inputs={len(inputs)}:normalize=0", "-c:a", "pcm_s16le", str(out), ] return out if _run_ffmpeg(job, cmd) else None def make_selected_mix(job: Job, stems_dir: Path, found: list[str]) -> Path | None: """If the user picked a strict subset of stems at submit time, sum those stems with ffmpeg amix into mix.wav. Returns the output path on success, or None when there's nothing to mix. Returns the existing single stem path (no ffmpeg) if exactly one stem was selected -- copying it to mix.wav would be 30 MB of duplicate data. The caller uses the returned path's name for the download URL, so a single-stem selection points the Download Mix button directly at the existing stem file. amix normalize=0 keeps stem amplitudes as-is. Demucs separations sum back to (close to) the original signal, so a 2-stem subset fits comfortably below 0 dBFS without normalization headroom.""" selected = [s for s in job.selected_stems if s in found] if not selected: return None if len(selected) == 1: return stems_dir / f"{selected[0]}.wav" inputs = [stems_dir / f"{name}.wav" for name in selected] out = stems_dir / "mix.wav" cmd: list[str] = [ ffmpeg_executable(), "-y", "-nostdin", "-loglevel", "error", ] for p in inputs: cmd += ["-i", str(p)] filter_inputs = "".join(f"[{i}:a]" for i in range(len(inputs))) cmd += [ "-filter_complex", f"{filter_inputs}amix=inputs={len(inputs)}:normalize=0", "-c:a", "pcm_s16le", str(out), ] return out if _run_ffmpeg(job, cmd) else None _PEAK_POINTS = 1500 # matches OVERVIEW_WAVE_POINTS in player.js def compute_stem_peaks(stems_dir: Path, stem_names: list[str]) -> None: """Compute and cache [min, max] waveform peaks for each stem. Failure is non-fatal — missing peaks.json degrades to client-side decode.""" peaks: dict[str, list[list[float]]] = {} for name in stem_names: path = stems_dir / f"{name}.wav" if not path.is_file(): continue try: data, _ = sf.read(path, dtype="float32", always_2d=True) ch = data[:, 0] n = len(ch) if n == 0: continue chunk = max(1, n // _PEAK_POINTS) result: list[list[float]] = [] for i in range(0, n, chunk): block = ch[i : i + chunk] result.append([float(np.min(block)), float(np.max(block))]) peaks[name] = result[:_PEAK_POINTS] except Exception: logger.warning("could not compute peaks for %s/%s", stems_dir.name, name, exc_info=True) if not peaks: return try: tmp = stems_dir / "peaks.json.tmp" tmp.write_text(json.dumps(peaks), encoding="utf-8") tmp.replace(stems_dir / "peaks.json") except Exception: logger.warning("could not write peaks.json for %s", stems_dir.name, exc_info=True) def sweep_old_jobs(jobs_dir: Path) -> None: """Delete job directories older than JOB_TTL_SECONDS and remove them from the in-memory registry. Called hourly from the background sweep loop started at app startup. Prefers Job.created_at over directory mtime (which can be touched by unrelated filesystem events), and never deletes the directory of an active (non-terminal) registered job even if its timestamp looks old. Falls back to mtime for orphan directories left over from a previous server run, since the registry is in-memory only.""" cutoff = time.time() - JOB_TTL_SECONDS if not jobs_dir.is_dir(): return jobs = registry_all() removed = False for d in jobs_dir.iterdir(): if not d.is_dir(): continue job = jobs.get(d.name) if job is not None: if job.status not in _TERMINAL: continue # never delete an active job's working dir if job.created_at >= cutoff: continue elif d.stat().st_mtime >= cutoff: continue _rmtree(d) registry_remove(d.name) removed = True if removed: registry_persist(jobs_dir)