models/hifigan.py

import torch
from speechbrain.pretrained import HIFIGAN


class HiFiGANModel:
    def __init__(self, device=None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        self.model = HIFIGAN.from_hparams(
            source="speechbrain/tts-hifigan-ljspeech",
            savedir="checkpoints/hifigan",
            run_opts={"device": self.device},
        )

    @torch.inference_mode()
    def mel_to_audio(self, mel):
        wav = self.model.decode_batch(mel)
        if wav.dim() == 3:
            wav = wav.squeeze(1)
        return wav[0]

    def architecture(self):
        return self.model