import gradio as gr
import os

def inference(input_file_path,model_name):
    if model_name == "" or model_name == None:
        model_name = "htdemucs"
    # Make sure the output directory exists
    temp_dir = os.path.dirname(input_file_path)
    temp_out_dir = os.path.join(temp_dir, 'out')
    os.makedirs(f'{temp_out_dir}', exist_ok=True)


    print(input_file_path)
    # Define the output file paths
    filename = os.path.basename(input_file_path).split('.')[0]
    if not filename:
        filename="test"

    vocals_out = os.path.join(temp_out_dir, model_name,filename,"vocals.mp3")
    no_vocals_out = os.path.join(temp_out_dir, model_name,filename,"no_vocals.mp3")

    cmd = f'python -m demucs.separate -n {model_name} --mp3 --two-stems vocals -o {temp_out_dir} "{input_file_path}"'
    print(cmd)
    os.system(cmd)

    print("process finish:",vocals_out,no_vocals_out)
    # Return the paths of the output files
    return vocals_out, no_vocals_out

title = "Facebook UVR Test"
description = "Seperate vocals and music using UVR"
article = "Made by line"
examples=[['test.mp3']]
# Create a Gradio interface

demo = gr.Interface(
    fn=inference,
    inputs=[
        gr.Audio(type="filepath", label="Input Audio"),
        gr.Dropdown(["htdemucs", "htdemucs_ft", "htdemucs_6s", "piano"], label="Model Name")
    ],
    outputs=[
        gr.Audio(type="filepath", label="Vocals out"),
        gr.Audio(type="filepath", label="Music out"),
    ],
    title=title,
    description=description,
    article=article,
    examples=examples
)

if __name__ == "__main__":
    demo.launch()