from speechbrain.pretrained import SpectralMaskEnhancement
import torchaudio

model = SpectralMaskEnhancement.from_hparams(source="speechbrain/metricgan-plus-voicebank", )

from speechbrain.pretrained import SpectralMaskEnhancement
import torchaudio
import torch

model = SpectralMaskEnhancement.from_hparams(source="speechbrain/metricgan-plus-voicebank")

def remove_noise(input_path, output_path):
    enhanced = model.enhance_file(input_path)
    waveform, sample_rate = enhanced

    if waveform.dim() == 0:
        raise ValueError(f"Enhanced waveform is empty for file: {input_path}")
    elif waveform.dim() == 1:
        waveform = waveform.unsqueeze(0)

    torchaudio.save(output_path, waveform, sample_rate)