# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/7. Pipeline.ipynb.

# %% auto 0
__all__ = ['Pipeline']

# %% ../nbs/7. Pipeline.ipynb 1
import torch
from spear_tts_pytorch.t2s_up import TSARTransformer
from spear_tts_pytorch.s2a_delar_mup import SAARTransformer
from spear_tts_pytorch.a2wav import Vocoder

# %% ../nbs/7. Pipeline.ipynb 2
class Pipeline:
    def __init__(self):
        self.t2s = TSARTransformer.load_model(local_filename='nbs/t2s_up.model').cuda()
        self.s2a = SADelARTransformer.load_model(local_filename='nbs/s2a_up.model').cuda()
        self.vocoder = Vocoder()

    def generate(self, text):
        stoks = self.t2s.generate(text, T=0.6)
        atoks = self.s2a.generate(stoks, T=0.6)
        audio = self.vocoder.decode(atoks)
        return audio
    
    def generate_to_file(self, fname, text):
        stoks = self.t2s.generate(text)
        atoks = self.s2a.generate(stoks)
        self.vocoder.decode_to_file(fname, atoks)
