import os
import shutil
import gradio as gr
from demucs.separate import main

def separate_stems(audio_file):
    input_path = "input.mp3"
    shutil.copy(audio_file, input_path)

    output_dir = "output"
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir, exist_ok=True)

    # Run Demucs
    main(["-n", "htdemucs", "-o", output_dir, input_path])

    # Build list of stems to return
    base = os.path.splitext(os.path.basename(input_path))[0]
    stem_path = os.path.join(output_dir, "htdemucs", base)
    stems = [os.path.join(stem_path, f"{stem}.mp3") for stem in ["vocals", "drums", "bass", "other"]]
    return stems

demo = gr.Interface(
    fn=separate_stems,
    inputs=gr.Audio(type="filepath", label="Upload Song"),
    outputs=[
        gr.Audio(label="Vocals"),
        gr.Audio(label="Drums"),
        gr.Audio(label="Bass"),
        gr.Audio(label="Other"),
    ],
    title="Demucs v4 Stem Separator",
    description="Upload a song to separate vocals, drums, bass, and other using Facebook's Demucs model.",
)

demo.launch()