import gradio as gr
from voicefixer import VoiceFixer

voicefixer = VoiceFixer()

def inference(audio, mode):
    voicefixer.restore(input=audio.name,
                       output="output.wav",
                       cuda=False,
                       mode=int(mode))
    return "output.wav"

input_audio = gr.inputs.Audio(label="Input Audio")
mode = gr.inputs.Radio(choices=['0', '1', '2'], label="Mode", default='0')
output_audio = gr.outputs.Audio(label="Output Audio")

title = "Voice Fixer"
description = "Gradio demo for VoiceFixer: Toward General Speech Restoration With Neural Vocoder. To use it, simply add 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/2109.13731' target='_blank'>VoiceFixer: Toward General Speech Restoration With Neural Vocoder</a> | <a href='https://github.com/haoheliu/voicefixer_main' target='_blank'>Github Repo</a></p>"

examples = [["bruce.wav", "2"]]

gr.Interface(inference, 
             inputs=[input_audio, mode], 
             outputs=output_audio,
             title=title, 
             description=description, 
             article=article, 
             examples=examples, 
             enable_queue=True).launch()