"""Demucs separation as a Gradio Space — the free audio backend for Stemwise.

Exposes one endpoint (/predict): audio file in, four mp3 stems out
(drums, bass, other, vocals). The Stemwise frontend calls it via
@gradio/client from the browser.
"""

import os
import subprocess
import sys
import tempfile
from pathlib import Path

import gradio as gr

# If the Space ever ends up on torch >= 2.6, torch.load defaults to
# weights_only=True and rejects demucs's pickled checkpoints. The official
# escape hatch restores the old behavior; it's a no-op on older torch.
os.environ.setdefault("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1")

MODEL = "htdemucs"  # 4 stems; fastest of the good models — free CPU is slow enough
STEMS = ["drums", "bass", "other", "vocals"]


def separate(audio_path):
    if not audio_path:
        raise gr.Error("Upload an audio file first.")
    out_dir = Path(tempfile.mkdtemp())
    proc = subprocess.run(
        [
            sys.executable, "-m", "demucs",
            "-n", MODEL,
            "--mp3", "--mp3-bitrate", "192",
            "--filename", "{stem}.{ext}",
            "-o", str(out_dir),
            str(audio_path),
        ],
        capture_output=True, text=True,
    )
    if proc.returncode != 0:
        raise gr.Error("demucs failed: " + (proc.stderr or proc.stdout)[-800:])
    model_dir = out_dir / MODEL
    return [str(model_dir / f"{s}.mp3") for s in STEMS]


demo = gr.Interface(
    fn=separate,
    inputs=gr.Audio(type="filepath", label="Track to separate"),
    outputs=[gr.Audio(type="filepath", label=s) for s in STEMS],
    title="Stemwise Demucs",
    description="Splits a song into drums / bass / other / vocals (htdemucs). "
                "Free CPU tier: expect 5–15 minutes per song.",
    flagging_mode="never",
)

demo.queue(max_size=10).launch()