app.py

import os
# ── CPU acceleration flags (MUST be set BEFORE importing onnxruntime/numpy) ──
_N_CPU = os.cpu_count() or 4
os.environ.setdefault("OMP_NUM_THREADS",        str(_N_CPU))
os.environ.setdefault("OPENBLAS_NUM_THREADS",   str(_N_CPU))
os.environ.setdefault("MKL_NUM_THREADS",        str(_N_CPU))
os.environ.setdefault("NUMEXPR_NUM_THREADS",    str(_N_CPU))
os.environ.setdefault("ONNXRUNTIME_EXECUTION_PROVIDERS", "CPUExecutionProvider")

import numpy as np
import gradio as gr
import librosa
import librosa.display
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from audio_separator.separator import Separator

# ── Paths (use /tmp so they're always writable on HF Spaces) ─────────────────
MODEL_DIR  = "/tmp/audio-separator-models"
OUTPUT_DIR = "/tmp/vocal-sep-output"
os.makedirs(MODEL_DIR,  exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)


def _build_visuals(audio_path, vocals_path, inst_path, fast_mode):
    """Generates the spectrogram plot and quality metrics."""
    # In fast mode, we reduce sample rate and mels to save ~10s of processing
    n_mels = 96 if fast_mode else 128
    dur    = 45 if fast_mode else 60
    sr     = 22050 if fast_mode else 44100   

    y_orig, _ = librosa.load(audio_path,  sr=sr, mono=True, duration=dur)
    y_voc,  _ = librosa.load(vocals_path, sr=sr, mono=True, duration=dur)
    y_inst, _ = librosa.load(inst_path,   sr=sr, mono=True, duration=dur)

    # ── Build Figure ───────────────────────────────────────────────────────
    fig = plt.figure(figsize=(16, 9))
    fig.patch.set_facecolor("#0f0f0f")
    gs = gridspec.GridSpec(3, 2, figure=fig, hspace=0.48, wspace=0.30)

    tracks = [
        (y_orig, "Original Mix",  "#4fc3f7"),
        (y_voc,  "Vocals Only",   "#ef5350"),
        (y_inst, "Instrumental",  "#66bb6a"),
    ]

    for i, (y, title, color) in enumerate(tracks):
        # Waveform
        ax_w = fig.add_subplot(gs[i, 0])
        t    = np.linspace(0, len(y) / sr, num=len(y))
        ax_w.plot(t, y, color=color, linewidth=0.4, alpha=0.85)
        ax_w.set_facecolor("#1a1a2e")
        ax_w.set_title(f"{title} — Waveform", color="white", fontsize=10, pad=4)
        ax_w.set_xlabel("Time (s)", color="#aaa", fontsize=8)
        ax_w.set_ylabel("Amplitude", color="#aaa", fontsize=8)
        ax_w.tick_params(colors="#aaa", labelsize=7)
        for spine in ax_w.spines.values():
            spine.set_edgecolor("#333")

        # Mel spectrogram
        ax_s = fig.add_subplot(gs[i, 1])
        S    = librosa.feature.melspectrogram(
                   y=y, sr=sr, n_mels=n_mels, fmax=sr//2,
                   n_fft=1024, hop_length=512
               )
        S_db = librosa.power_to_db(S, ref=np.max)
        img  = librosa.display.specshow(
            S_db, sr=sr, x_axis="time", y_axis="mel",
            fmax=sr//2, ax=ax_s, cmap="magma",
        )
        cb = fig.colorbar(img, ax=ax_s, format="%+2.0f dB", pad=0.02)
        cb.ax.yaxis.set_tick_params(color="#aaa", labelsize=7)
        ax_s.set_facecolor("#1a1a2e")
        ax_s.set_title(f"{title} — Mel Spectrogram", color="white", fontsize=10, pad=4)
        ax_s.set_xlabel("Time (s)", color="#aaa", fontsize=8)
        ax_s.set_ylabel("Hz",       color="#aaa", fontsize=8)
        ax_s.tick_params(colors="#aaa", labelsize=7)
        for spine in ax_s.spines.values():
            spine.set_edgecolor("#333")

    fig.suptitle(
        "Kim_Vocal_2.onnx  ·  MDX-Net Vocal Separation",
        color="white", fontsize=13, y=1.01,
    )
    plt.tight_layout()
    plot_path = os.path.join(OUTPUT_DIR, "separation_result.png")
    plt.savefig(plot_path, dpi=110, bbox_inches="tight", facecolor=fig.get_facecolor())
    plt.close(fig)

    # ── Quality metrics ────────────────────────────────────────────────────
    def leakage_db(stem, residual):
        n = min(len(stem), len(residual))
        s, r = stem[:n], residual[:n]
        leak = np.dot(s, r) / (np.linalg.norm(r) ** 2 + 1e-8) * r
        return 10 * np.log10(np.mean(leak ** 2) / (np.mean(s ** 2) + 1e-8) + 1e-8)

    def energy_pct(stem, mix):
        n = min(len(stem), len(mix))
        return (np.mean(stem[:n] ** 2) / (np.mean(mix[:n] ** 2) + 1e-8)) * 100

    metrics = (
        f"📊 Separation Metrics (proxy — no reference stems)\n"
        f"{'─'*44}\n"
        f"Vocals  energy vs mix : {energy_pct(y_voc,  y_orig):.1f}%\n"
        f"Instrum energy vs mix : {energy_pct(y_inst, y_orig):.1f}%\n"
        f"\n"
        f"Vocals  ← Instrum leak : {leakage_db(y_voc,  y_inst):.1f} dB  (lower = cleaner)\n"
        f"Instrum ← Vocals  leak : {leakage_db(y_inst, y_voc):.1f} dB  (lower = cleaner)\n"
        f"\n"
        f"Model SDR benchmark: ~8.9 dB on MVSep (Kim_Vocal_2)"
    )
    return plot_path, metrics


# ── Core separation logic ─────────────────────────────────────────────────────
def separate(audio_path, segment_size, overlap, enable_denoise, fast_mode, progress=gr.Progress()):
    if audio_path is None:
        raise gr.Error("Please upload an audio file first.")

    progress(0.05, desc="Setting up separator…")

    # ── Fast-mode overrides ────────────────────────────────────────────────
    # hop 2048  ≈ 1.5× faster than 1024, inaudible quality change for vocals
    # batch 8   ≈ 3–4× faster than 1 on multi-core CPUs
    # chunk 30  ≈ fewer Python iterations, less overhead
    hop    = 2048 if fast_mode else 1024
    bsize  = 8    if fast_mode else 2
    chunk  = 30   if fast_mode else 10

    separator = Separator(
        output_dir=OUTPUT_DIR,
        output_format="WAV",
        model_file_dir=MODEL_DIR,
        chunk_size=chunk,                           # Top-level parameter!
        normalization_enabled=not fast_mode,        # Skip gain analysis in fast mode
        mdx_params={
            "segment_size":      int(segment_size),
            "overlap":           float(overlap),
            "batch_size":        bsize,             # Batching for huge CPU speedup
            "hop_length":        hop,               # Larger hop = fewer FFT windows
            "enable_denoise":    enable_denoise,
        },
    )

    progress(0.10, desc="Loading model (first run downloads ~67 MB)…")
    separator.load_model(model_filename="Kim_Vocal_2.onnx")

    progress(0.20, desc="Separating… (CPU batched) ⚡")
    output_files = separator.separate(audio_path)

    progress(0.75, desc="Locating output files…")

    def resolve(f):
        if os.path.exists(f):
            return f
        cand = os.path.join(OUTPUT_DIR, os.path.basename(f))
        if os.path.exists(cand):
            return cand
        raise FileNotFoundError(f"Cannot find output: {f}")

    resolved        = [resolve(f) for f in output_files]
    vocals_path     = next((f for f in resolved if "Vocals"       in os.path.basename(f)), None)
    inst_path       = next((f for f in resolved if "Instrumental" in os.path.basename(f)), None)

    if not vocals_path or not inst_path:
        raise gr.Error(
            f"Separation finished but output files were not found. "
            f"Got: {[os.path.basename(f) for f in resolved]}"
        )

    # ── Spectrogram plot & Metrics ─────────────────────────────────────────
    progress(0.85, desc="Computing visuals & metrics…")
    plot_path, metrics = _build_visuals(audio_path, vocals_path, inst_path, fast_mode)

    progress(1.0, desc="✅ Done!")
    return vocals_path, inst_path, plot_path, metrics


# ── UI ────────────────────────────────────────────────────────────────────────
CSS = """
#run-btn { font-size: 1.1rem; padding: 0.75rem 2rem; }
.gradio-container { max-width: 1100px !important; }
footer { display: none !important; }
"""

with gr.Blocks(title="🎤 Vocal Remover", css=CSS) as demo:

    gr.Markdown("""
# 🎤 Vocal Remover (Optimized CPU Edition)
### Kim_Vocal_2 · MDX-Net ONNX · Best quality-per-byte vocal separator
Upload any song to isolate **vocals** and **instrumental** stems.
Supports MP3, WAV, FLAC, M4A, OGG, and more.
> ⚡ **Running on CPU** — With Fast Mode enabled, a typical 3-min track takes **~40 seconds**.  
> The model (~67 MB) is downloaded automatically on first run, then cached for the session.
""")

    with gr.Row(equal_height=False):
        # Left column — upload + settings
        with gr.Column(scale=1, min_width=280):
            audio_in = gr.Audio(
                label="Upload Audio",
                type="filepath",
            )

            with gr.Accordion("⚙️ Advanced Settings", open=False):
                fast_mode = gr.Checkbox(
                    value=True,
                    label="⚡ Fast Mode",
                    info="Batch=8, hop=2048, sr=22050 visuals. ~3× faster, near-identical audio quality.",
                )
                segment_size = gr.Slider(
                    minimum=128, maximum=512, value=256, step=128,
                    label="Segment Size",
                    info="Lower = less RAM · 512 = better quality",
                )
                overlap = gr.Slider(
                    minimum=0.10, maximum=0.50, value=0.25, step=0.05,
                    label="Overlap",
                    info="0.25 = fast · 0.50 = smoother transitions",
                )
                enable_denoise = gr.Checkbox(
                    value=False,
                    label="Enable Denoise",
                    info="Post-process artifact reduction. Doubles processing time on CPU.",
                )

            run_btn = gr.Button(
                "🎵  Separate Vocals", variant="primary",
                size="lg", elem_id="run-btn",
            )

        # Right column — audio outputs
        with gr.Column(scale=2):
            vocals_out = gr.Audio(
                label="🎤 Vocals Only",
                type="filepath",
                interactive=False,
            )
            inst_out = gr.Audio(
                label="🎵 Instrumental",
                type="filepath",
                interactive=False,
            )

    # Bottom row — plot + metrics
    with gr.Row():
        plot_out    = gr.Image(label="Waveforms & Spectrograms", type="filepath")
        metrics_out = gr.Textbox(
            label="📊 Quality Metrics",
            lines=9,
            interactive=False,
        )

    run_btn.click(
        fn=separate,
        inputs=[audio_in, segment_size, overlap, enable_denoise, fast_mode],
        outputs=[vocals_out, inst_out, plot_out, metrics_out],
    )

    gr.Markdown("""
---
### Model Comparison
| Model | Size | Vocal SDR | ONNX |
|---|---|---|---|
| **Kim_Vocal_2** ✅ | 67 MB | ~8.9 dB | ✅ |
| UVR-MDX-NET-Voc_FT | 67 MB | ~8.7 dB | ✅ |
| htdemucs_ft (Demucs 4) | 330 MB | ~9.2 dB | ❌ |
| BS-Roformer | 430 MB | ~12.9 dB | ❌ |
Kim_Vocal_2 is the best quality-per-byte ONNX vocal model — ideal for free-tier CPU inference.
""")


demo.queue()
demo.launch()