lenML/ChatTTS-Forge

1import os2from typing import Union3 4import soundfile as sf5import torch6import torchaudio7 8from modules.Denoiser.AudioDenoiser import AudioDenoiser9from modules.devices import devices10from modules.utils.constants import MODELS_DIR11 12ad: Union[AudioDenoiser, None] = None13 14 15class TTSAudioDenoiser:16 17    def load_ad(self):18        global ad19        if ad is None:20            ad = AudioDenoiser(21                os.path.join(22                    MODELS_DIR,23                    "Denoise",24                    "audio-denoiser-512-32-v1",25                ),26                device=devices.device,27            )28            ad.model.to(devices.device)29        return ad30 31    def denoise(self, audio_data, sample_rate, auto_scale=False):32        ad = self.load_ad()33        sr = ad.model_sample_rate34        return sr, ad.process_waveform(audio_data, sample_rate, auto_scale)35 36 37if __name__ == "__main__":38    tts_deno = TTSAudioDenoiser()39    data, sr = sf.read("test.wav")40    audio_tensor = torch.from_numpy(data).unsqueeze(0).float()41    print(audio_tensor)42 43    # data, sr = torchaudio.load("test.wav")44    # print(data)45    # data = data.to(devices.device)46 47    sr, denoised = tts_deno.denoise(audio_data=audio_tensor, sample_rate=sr)48    denoised = denoised.cpu()49    torchaudio.save("denoised.wav", denoised, sample_rate=sr)50