Spaces:
Sleeping
Sleeping
| 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() |