app.py

import spaces 
import torch
import torchaudio
import gradio as gr
from demucs import pretrained
from demucs.apply import apply_model
from audiotools import AudioSignal
from typing import Dict
from pyharp import *


DEMUX_MODELS = ["mdx_extra_q", "mdx_extra", "htdemucs", "mdx_q"]

STEM_CHOICES = {
    "Vocals": 3,
    "Drums": 0,
    "Bass": 1,
    "Other": 2,
    "Instrumental (No Vocals)": "instrumental"
}

@spaces.GPU(duration = 180)
def separate_stem(audio_file_path: str, model_name: str, stem_choice: str) -> AudioSignal:
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = pretrained.get_model(model_name)
    model.to(device)
    model.eval()

    waveform, sr = torchaudio.load(audio_file_path)
    is_mono = waveform.shape[0] == 1
    if is_mono:
        waveform = waveform.repeat(2, 1)

    waveform = waveform.to(device)

    with torch.no_grad():
        stems_batch = apply_model(
            model,
            waveform.unsqueeze(0),
            overlap=0.2,
            shifts=1,
            split=True
        )

    stems = stems_batch[0]

    if stem_choice == "Instrumental (No Vocals)":
        stem = stems[0] + stems[1] + stems[2]
    else:
        stem_index = STEM_CHOICES[stem_choice]
        stem = stems[stem_index]

    if is_mono:
        stem = stem.mean(dim=0, keepdim=True)

    return AudioSignal(stem.cpu().numpy().astype('float32'), sample_rate=sr)

# Gradio Callback Function

def process_fn_stem(audio_file_path: str, demucs_model: str, stem_choice: str):
    """
    PyHARP process function:
      - Separates the chosen stem using Demucs.
      - Saves the stem as a .wav file.
    """
    stem_signal = separate_stem(audio_file_path, model_name=demucs_model, stem_choice=stem_choice)
    stem_path = save_audio(stem_signal, f"{stem_choice.lower().replace(' ', '_')}.wav")
    return stem_path, LabelList(labels=[])


# Model Card 
model_card = ModelCard(
    name="Demucs Stem Separator",
    description="Uses Demucs to separate a music track into a selected stem.",
    author="Alexandre Défossez, Nicolas Usunier, Léon Bottou, Francis Bach",
    tags=["demucs", "source-separation", "pyharp", "stems"]
)

# Gradio UI
with gr.Blocks() as demo:

    dropdown_model = gr.Dropdown(
        label="Select Demucs Model",
        choices=DEMUX_MODELS,
        value="mdx_extra_q"
    )

    dropdown_stem = gr.Dropdown(
        label="Select Stem to Separate",
        choices=list(STEM_CHOICES.keys()),
        value="Vocals"
    )

    app = build_endpoint(
        model_card=model_card,
        components=[dropdown_model, dropdown_stem],
        process_fn=process_fn_stem
    )

demo.queue()
demo.launch(show_error=True)