import gradio as gr
import os
from audio_separator.separator import Separator

# Initialize the Separator

separator = Separator()

def separate_audio(input_file, output_dir, model_name, output_format):
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)


    # Separate the audio
    separator.separate_audio_file(
        input_file,
        output_dir,
        model_name=model_name,
        denoise=True,
        output_format=output_format,
        normalization_threshold=0.9,
        mdx_segment_size=256,
        mdx_overlap=8,
        primary_stem_only=False
    )

    # Rename the output files to match the requested format
    os.rename(os.path.join(output_dir, 'Vocals.wav'), os.path.join(output_dir, '1_main_vocal.wav'))
    os.rename(os.path.join(output_dir, 'Other.wav'), os.path.join(output_dir, '2_backing_vocal.wav'))
    os.rename(os.path.join(output_dir, 'Instrumental.wav'), os.path.join(output_dir, '3_instrumental.wav'))

    return [
        os.path.join(output_dir, '1_main_vocal.wav'),
        os.path.join(output_dir, '2_backing_vocal.wav'),
        os.path.join(output_dir, '3_instrumental.wav')
    ]

def process_audio(audio_file, model_name):
    output_dir = "output"
    return separate_audio(audio_file.name, output_dir, model_name)

# Define the Gradio interface
iface = gr.Blocks(theme='hev832/Applio')

with iface:
    gr.Markdown("# Hex Separator")
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(type="filepath", label="Input Audio")
            model_name = gr.Dropdown(
                choices=["UVR-MDX-NET-Inst_HQ_3", "UVR_MDXNET_KARA_2", "UVR-MDX-NET-Inst_HQ_4"],
                label="Model",
                value="UVR-MDX-NET-Inst_HQ_3"
            )
        with gr.Column():
            audio_format = gr.Dropdown(
                choices=["wav", "mp3", "flac"], 
                label="Audio Format",
                value="wav"
            )
            
            submit_btn = gr.Button("Separate Audio")
        with gr.Column():
            with gr.Row():
                vocal_output = gr.Audio(label="Main Vocal")
                backing_vocal_output = gr.Audio(label="Backing Vocal")
            instrumental_output = gr.Audio(label="Instrumental")

    submit_btn.click(
        process_audio,
        inputs=[audio_input, model_name, audio_format],
        outputs=[vocal_output, backing_vocal_output, instrumental_output]
    )

iface.launch()