transcriber/separate.py

"""Demucs source separation for full-song transcription."""

import subprocess
import sys
from pathlib import Path


def separate(input_path: str, output_dir: str, model: str = "htdemucs") -> dict:
    """Run Demucs source separation on an audio file.

    Args:
        input_path: Path to input audio file
        output_dir: Directory to write separated stems
        model: Demucs model name (default: htdemucs)

    Returns:
        dict with stem paths: {
            'vocals': str,
            'drums': str,
            'bass': str,
            'other': str,  # contains piano, guitar, synths, etc.
        }
    """
    input_path = Path(input_path)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    python = sys.executable

    cmd = [
        python, "-m", "demucs",
        "--name", model,
        "--out", str(output_dir),
        "--segment", "7",  # Process in 7s chunks (htdemucs max is 7.8)
        str(input_path),
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
    if result.returncode != 0:
        # Filter out download progress bars from stderr to find real error
        stderr_lines = result.stderr.splitlines()
        error_lines = [
            ln for ln in stderr_lines
            if not any(x in ln for x in ['%|', 'B/s]', 'Downloading', '━'])
        ]
        error_msg = '\n'.join(error_lines[-20:]).strip() or result.stderr[-500:]
        raise RuntimeError(f"Demucs failed: {error_msg}")

    # Demucs outputs to: output_dir/{model}/{filename_without_ext}/
    stems_dir = output_dir / model / input_path.stem

    if not stems_dir.exists():
        raise FileNotFoundError(f"Demucs output not found at {stems_dir}")

    stems = {}
    for name in ("vocals", "drums", "bass", "other"):
        path = stems_dir / f"{name}.wav"
        if not path.exists():
            raise FileNotFoundError(f"Missing stem: {path}")
        stems[name] = str(path)

    return stems