from speechbrain.pretrained import SpectralMaskEnhancement
import torchaudio
import os
import torch

model = SpectralMaskEnhancement.from_hparams(
    source="speechbrain/metricgan-plus-voicebank", savedir="tmp/denoiser"
)
def remove_noise(input_path, output_path):
    noisy_audio, sr = torchaudio.load(input_path)
    lengths = torch.tensor([noisy_audio.shape[1]]) / noisy_audio.shape[1]
    denoised = model.enhance_batch(noisy_audio, lengths)
    torchaudio.save(output_path, denoised.cpu(), sr)