import os
from scipy.io.wavfile import write
import gradio as gr

def separator(audio):
    os.makedirs("output", exist_ok=True)
    write('test.wav', audio[0], audio[1])
    result = os.system("audio-separator test.wav --model_filename model_bs_roformer_ep_317_sdr_12.9755.ckpt --output_dir=./output --output_format=wav")

    files = [
        "./output/test_(Instrumental)_model_bs_roformer_ep_317_sdr_12.wav",
        "./output/test_(Vocals)_model_bs_roformer_ep_317_sdr_12.wav"
    ]

    for file in files:
        if not os.path.isfile(file):
            print(f"File not found: {file}")
        else:
            print(f"File exists: {file}")

    return files

with gr.Blocks() as demo:
    audio_input = gr.Audio(type="numpy", label="Input")
    
    output1 = gr.Audio(type="filepath", label="Stem 1", interactive=False)
    output2 = gr.Audio(type="filepath", label="Stem 2", interactive=False)
    submit_button = gr.Button("Separate")
    submit_button.click(separator, inputs=audio_input, outputs=[output1, output2])

demo.launch()