import gradio as gr
import os
import random
import subprocess

def separate_audio(audio_path):
    stem_count = 2
    print(f"{audio_path=}")
    head, tail = os.path.split(audio_path)

    gradio_temp_path = head
    audio_filename = tail.split('.')[0]
    print(f"{gradio_temp_path=}")
    print(f"{audio_filename=}")
    print(f"{stem_count=}")
    os.makedirs(gradio_temp_path + f"/separated_audio/{audio_filename}/", exist_ok=True)
    command = f"spleeter separate -p spleeter:{stem_count}stems {audio_path} -o {gradio_temp_path}/separated_audio/"
    command = command.split()
    print(f"{command=}")

    result = subprocess.run(command)
    print(result)

    paths = []
    paths.append(gradio_temp_path + f"/separated_audio/{audio_filename}/accompaniment.wav")
    paths.append(gradio_temp_path + f"/separated_audio/{audio_filename}/vocals.wav")
    return [gr.Audio(path) for path in paths]

iface = gr.Interface(
    fn=separate_audio,
    inputs=[gr.Audio(type="filepath", label="Audio File")],
    outputs=[gr.Audio(label='Accompaniment'), gr.Audio(label='Vocals')]
)

iface.launch()