import gradio as gr
import torch
import demucs.separate
import shlex
import os
import spaces

# Check if CUDA is available
device = "cuda" if torch.cuda.is_available() else "cpu"

# Define the inference function
@spaces.GPU
def inference(audio_file, model_name, two_stems, mp3, mp3_bitrate):
    """
    Performs inference using Demucs.

    Args:
        audio_file: The audio file to separate.
        model_name: The name of the Demucs model to use.
        two_stems: The name of the stem to separate (for two-stems mode).
        mp3: Whether to save the output as MP3.
        mp3_bitrate: The bitrate of the output MP3 file.

    Returns:
        A list of separated audio files.
    """

    # Construct the command line arguments
    cmd = f"demucs -n {model_name} --clip-mode clamp --shifts=1"
    if two_stems:
        cmd += f" --two-stems={two_stems}"
    if mp3:
        cmd += f" --mp3 --mp3-bitrate={mp3_bitrate}"
    cmd += f" {audio_file.name}"

    # Run Demucs
    demucs.separate.main(shlex.split(cmd))

    # Get the output file paths
    output_dir = os.path.join("separated", model_name, os.path.splitext(os.path.basename(audio_file.name))[0])
    output_files = [os.path.join(output_dir, f) for f in os.listdir(output_dir) if os.path.isfile(os.path.join(output_dir, f))]

    return output_files

# Define the Gradio interface
iface = gr.Interface(
    fn=inference,
    inputs=[
        gr.Audio(source="upload", type="filepath"),
        gr.Dropdown(["htdemucs", "htdemucs_ft", "htdemucs_6s", "hdemucs_mmi", "mdx", "mdx_extra", "mdx_q", "mdx_extra_q"], label="Model Name"),
        gr.Dropdown(["vocals", "drums", "bass", "other"], label="Two Stems (Optional)", optional=True),
        gr.Checkbox(label="Save as MP3"),
        gr.Slider(128, 320, step=32, label="MP3 Bitrate", visible=False),
    ],
    outputs=gr.Files(),
    title="Demucs Music Source Separation",
    description="Separate vocals, drums, bass, and other instruments from your music using Demucs.",
)

# Launch the Gradio interface
iface.launch()