| """Demucs Audio Stem Separator - Separate audio into vocals, drums, bass, other.""" | |
| import os | |
| import tempfile | |
| import gradio as gr | |
| import torch | |
| import soundfile as sf | |
| import numpy as np | |
| import librosa | |
| from demucs.pretrained import get_model | |
| from demucs.apply import apply_model | |
| device = torch.device("cpu") | |
| _model_cache = {} | |
| def load_model(model_name: str): | |
| """Load and cache Demucs model.""" | |
| if model_name not in _model_cache: | |
| model = get_model(model_name) | |
| model.to(device).eval() | |
| _model_cache[model_name] = model | |
| return _model_cache[model_name] | |
| def separate_stems(audio_file: str, model_name: str = "htdemucs", progress=gr.Progress()) -> tuple[str, str, str, str]: | |
| """ | |
| Separate audio into stems using Demucs. | |
| Args: | |
| audio_file: Path to the input audio file (MP3, WAV, FLAC, OGG, M4A) | |
| model_name: Demucs model - htdemucs (fast) or htdemucs_ft (better quality) | |
| Returns: | |
| Tuple of paths to separated stems: (vocals, drums, bass, other) | |
| """ | |
| if audio_file is None: | |
| raise gr.Error("Please upload an audio file") | |
| progress(0.1, desc="Loading model...") | |
| model = load_model(model_name) | |
| progress(0.2, desc="Loading audio...") | |
| audio_data, sr = sf.read(audio_file) | |
| if audio_data.ndim == 1: | |
| waveform = torch.from_numpy(audio_data.astype(np.float32)).unsqueeze(0) | |
| else: | |
| waveform = torch.from_numpy(audio_data.T.astype(np.float32)) | |
| target_sr = model.samplerate | |
| if sr != target_sr: | |
| progress(0.25, desc="Resampling...") | |
| waveform_np = waveform.numpy() | |
| resampled = [librosa.resample(waveform_np[ch], orig_sr=sr, target_sr=target_sr) | |
| for ch in range(waveform_np.shape[0])] | |
| waveform = torch.from_numpy(np.stack(resampled)) | |
| if waveform.shape[0] == 1: | |
| waveform = waveform.repeat(2, 1) | |
| elif waveform.shape[0] > 2: | |
| waveform = waveform[:2, :] | |
| waveform = waveform.unsqueeze(0).to(device) | |
| progress(0.3, desc="Separating stems...") | |
| with torch.no_grad(): | |
| sources = apply_model(model, waveform, device=device, progress=True) | |
| progress(0.9, desc="Saving...") | |
| output_dir = tempfile.mkdtemp() | |
| output_files = {} | |
| for i, name in enumerate(model.sources): | |
| stem = sources[0, i].cpu().numpy() | |
| path = os.path.join(output_dir, f"{name}.wav") | |
| sf.write(path, stem.T, target_sr) | |
| output_files[name] = path | |
| return output_files.get("vocals"), output_files.get("drums"), output_files.get("bass"), output_files.get("other") | |
| with gr.Blocks(title="Demucs Stem Separator") as demo: | |
| gr.Markdown("# 🎵 Demucs Stem Separator\nSeparate audio → **Vocals** | **Drums** | **Bass** | **Other**") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| audio_input = gr.Audio(label="Audio File", type="filepath", sources=["upload"]) | |
| model_choice = gr.Dropdown(["htdemucs", "htdemucs_ft"], value="htdemucs", label="Model", | |
| info="htdemucs: Fast | htdemucs_ft: Better quality") | |
| separate_btn = gr.Button("Separate Stems", variant="primary", size="lg") | |
| with gr.Column(scale=2): | |
| with gr.Row(): | |
| vocals = gr.Audio(label="🎤 Vocals", type="filepath") | |
| drums = gr.Audio(label="🥁 Drums", type="filepath") | |
| with gr.Row(): | |
| bass = gr.Audio(label="🎸 Bass", type="filepath") | |
| other = gr.Audio(label="🎹 Other", type="filepath") | |
| separate_btn.click(separate_stems, [audio_input, model_choice], [vocals, drums, bass, other], api_name="separate") | |
| gr.Examples( | |
| examples=[["example.wav", "htdemucs"]], | |
| inputs=[audio_input, model_choice], | |
| outputs=[vocals, drums, bass, other], | |
| fn=separate_stems, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Example" | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(mcp_server=True, show_error=True) | |