import torch
import torchaudio
import gradio as gr
from clearvoice import ClearVoice

myClearVoice = ClearVoice(task='speech_enhancement', model_names=['FRCRN_SE_16K'])

def fn_clearvoice(aud):
    # Load and add fake batch dimension
    """
    noisy = enhance_model.load_audio(
    aud
    ).unsqueeze(0)
    enhanced = enhance_model.enhance_batch(noisy, lengths=torch.tensor([1.]))
    """
    output_wav_dict = myClearVoice(input_path='input.wav', online_write=False)
    key = next(iter(output_wav_dict))
    output_wav = output_wav_dict[key]
    torchaudio.save('enhanced.wav', output_wav.cpu(), 16000)
    return 'enhanced.wav'
  
inputs = gr.Audio(sources=["upload"], label="Input Audio", type="filepath")
outputs = gr.Audio(label="Output Audio", type="filepath")
title = "ClearVoice"
description = "Gradio demo for Speech enhancement with ClearVoice. 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/2206.07293' target='_blank'>FRCRN: Boosting Feature Representation Using Frequency Recurrence for Monaural Speech Enhancement</a> | <a href='https://github.com/speechbrain/speechbrain' target='_blank'>Github Repo</a></p>"
examples = [
    ['input.wav']
]
gr.Interface(fn_clearvoice, inputs, outputs, title=title, description=description, article=article, examples=examples).launch()