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

MAX_LENGTH = 30.0

model = BandwidthExtender.from_pretrained("hifi-gan-bwe-10-42890e3-vctk-48kHz")


def extend(audio):
    fs, x = audio
    x = x[:int(MAX_LENGTH * fs)]
    x = x.astype(np.float32) / 32767.0
    if len(x.shape) == 1:
        x = x[:, np.newaxis]

    with torch.no_grad():
        y = np.stack([model(torch.from_numpy(x), fs) for x in x.T]).T
        y = (y * 32767.0).astype(np.int16)
        fs = int(model.sample_rate)

    return fs, y


gradio.Interface(
  fn=extend,
  examples=[["examples/example.wav"]],
  title="HiFi-GAN+",
  description="""
An audio bandwidth extension model from the paper
[Bandwidth Extension is All You Need](https://doi.org/10.1109/ICASSP39728.2021.9413575)
by Su, et al.
For more information, see this
[blog post](https://brentspell.com/2022/hifi-gan-bwe).
""",
  thumbnail="https://brentspell.com/favicon-32x32.png",
  inputs="audio",
  outputs="audio",
).launch()