import torch
from scipy.io.wavfile import write

from main_pipeline import CleaningPipeline

import gradio as gr

title = "Audio denoising and speaker diarization "

example_list = [
    ["dialog.mp3"]
]


def app_pipeline(audio):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    cleaning_pipeline = CleaningPipeline(device)

    audio_path = 'test.wav'
    write(audio_path, audio[0], audio[1])
    result = cleaning_pipeline(audio_path)
    if result != []:
        return result


app = gr.Interface(
    app_pipeline,
    gr.Audio(type="numpy", label="Input_audio"),
    [gr.Audio(visible=True, label='denoised_audio' if i == 0 else f'speaker{i}') for i in range(20)],
    title=title,
    examples=example_list,
    cache_examples=False,

)

app.launch(debug=True, enable_queue=True,
           )