import gradio as gr
import torch
import torchaudio
from denoisers import WaveUNetModel

MODEL = WaveUNetModel.from_pretrained("wrice/waveunet-vctk-24khz")

def denoise(inputs):
    sr, audio = inputs
    audio = torch.from_numpy(audio)[None]
    audio = audio / 32768.0

    if sr != MODEL.config.sample_rate:
        audio = torchaudio.functional.resample(audio, sr, MODEL.config.sample_rate)

    chunk_size = MODEL.config.max_length

    padding = abs(audio.size(-1) % chunk_size - chunk_size)
    padded = torch.nn.functional.pad(audio, (0, padding))

    clean = []
    for i in range(0, padded.shape[-1], chunk_size):
        audio_chunk = padded[:, i:i + chunk_size]
        with torch.no_grad():
            clean_chunk = MODEL(audio_chunk[None]).logits
        clean.append(clean_chunk.squeeze(0))

    denoised = torch.concat(clean)[:, :audio.shape[-1]].squeeze().clamp(-1.0, 1.0)
    return MODEL.config.sample_rate, denoised.numpy()

iface = gr.Interface(fn=denoise, inputs="audio", outputs="audio")
iface.launch()