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

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

def speechbrain(aud):
  est_sources = model.separate_file(path=aud)
  torchaudio.save("clean_audio_file.wav", est_sources[:, :, 0].detach().cpu(), 8000)
  return "clean_audio_file.wav"


title = "Speech Enhancement"
description = "Gradio demo for Speech Enhancement by SpeechBrain. 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/2010.13154' target='_blank'>Attention is All You Need in Speech Separation</a> | <a href='https://github.com/speechbrain/speechbrain/tree/develop/templates/enhancement' '_blank'>Github Repo</a></p>"
examples = [
    ['samples_audio_samples_test_mixture.wav']
]

demo = gr.Interface(
  fn=speechbrain,
  inputs=gr.Audio(type="filepath"),
  outputs=[
    gr.Audio(label="Output Audio", type="filepath")
  ],
  title=title,
  description=description,
  article=article,
  examples=examples
)
demo.launch()