from models.tacotron2 import Tacotron2Model
from models.hifigan import HiFiGANModel


class SpeechPipeline:
    def __init__(self):
        self.tacotron = Tacotron2Model()
        self.hifigan = HiFiGANModel()

    def text_to_mel(self, text):
        mel, _, _ = self.tacotron.text_to_mel(text)
        return mel

    def text_to_audio(self, text):
        mel = self.text_to_mel(text)
        audio = self.hifigan.mel_to_audio(mel)
        return audio