from speechbrain.inference.separation import SepformerSeparation as separator
import torchaudio
import gradio as gr

model = separator.from_hparams(source="speechbrain/sepformer-wham16k-enhancement", savedir='pretrained_models/sepformer-wham16k-enhancement')

def speechbrain(aud):
  

    # for custom file, change path
    est_sources = model.separate_file(path=aud) 
    torchaudio.save("enhanced_wham16k.wav", est_sources[:, :, 0].detach().cpu(), 16000)
    return "enhanced_wham16k.wav"

inputs = gr.inputs.Audio(label="Input Audio", type="filepath")
outputs =  [
  gr.outputs.Audio(label="Output Audio", type="filepath"),
]

title = "Speech Enhancer"
description = "Gradio demo for Speech Enhance by SpeechBrain. To use it, simply upload your audio, or click one of the examples to load them"
examples = [
    ['example_wham16k.wav']
]
gr.Interface(speechbrain, inputs, outputs, title=title, description=description, article=article, examples=examples).launch()