import gradio
import numpy as np
import torch
from hifi_gan_bwe import BandwidthExtender

model = BandwidthExtender.from_pretrained("hifi-gan-bwe-05-d3abf04-vctk-48kHz")


def extend(audio):
    fs, x = audio
    x = x.astype(np.float32) / 32767.0

    with torch.no_grad():
        y = model(torch.from_numpy(x), fs).numpy()
        fs = int(model.sample_rate)
    y = (y * 32767.0).astype(np.int16)

    return fs, y


gradio.Interface(
  fn=extend,
  inputs="audio",
  outputs="audio",
).launch()