import gradio as gr
from spleeter.separator import Separator
from spleeter.audio.adapter import get_default_audio_adapter

# Initialize the separator with the 2stems model (vocals and accompaniment)
separator = Separator('spleeter:2stems')
audio_adapter = get_default_audio_adapter()

def remove_vocals(input_audio):
    # Load the audio file
    waveform, sample_rate = audio_adapter.load(input_audio)
    
    # Separate the audio into vocals and accompaniment
    prediction = separator.separate(waveform)
    
    # Extract the accompaniment (non-vocal part)
    accompaniment = prediction['accompaniment']
    
    # Save the accompaniment to a temporary file
    output_audio = "accompaniment.wav"
    audio_adapter.save(output_audio, accompaniment, sample_rate)
    
    return output_audio

# Create the Gradio interface
iface = gr.Interface(
    fn=remove_vocals,
    inputs="audio",
    outputs="audio",
    title="AI Vocal Remover",
    description="Upload an audio file and the AI will remove the vocals, leaving the accompaniment.",
    examples=[["example.mp3"]]
)

# Launch the interface
iface.launch()