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


def inference(audio, is_fast):
    os.makedirs("out", exist_ok=True)
    write('test.wav', audio[0], audio[1])
    os.environ['DANNA_CHECKPOINTS'] = './checkpoints'
    cmd = "danna_sep --outdir out test.wav"
    if is_fast:
        cmd += " --fast"
    os.system(cmd)
    return "./out/test_vocals.wav", "./out/test_bass.wav",\
        "./out/test_drums.wav", "./out/test_other.wav"


title = "Danna-Sep"
description = "Gradio demo for Danna-Sep: Unite to separate them all. To use it, simply upload your audio, or click one of the examples to load them. Read more at the links below."
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.03752' target='_blank'>Danna-Sep: Unite to separate them all</a> | <a href='https://github.com/yoyololicon/danna-sep' target='_blank'>Github Repo</a></p>"

examples = []
gr.Interface(
    inference,
    [
        gr.inputs.Audio(type="numpy", label="Input"),
        gr.inputs.Checkbox(label="Faster inference without X-UMX")
    ],
    [gr.outputs.Audio(type="file", label="Vocals"), gr.outputs.Audio(type="file", label="Bass"), gr.outputs.Audio(
        type="file", label="Drums"), gr.outputs.Audio(type="file", label="Other")],
    title=title,
    description=description,
    article=article,
    examples=examples
).launch(enable_queue=True)