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

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

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

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

    def generate_atoks(self, text, speaker="default", show_progress_bar=True):
        if speaker == "default":
            speaker = "3645" if "3645" in self.s2a.speaker_map else next(iter(self.s2a.speaker_map.keys()))
        text = text.replace("\n", " ")
        stoks = self.t2s.generate(text, T=.5, top_k=3, show_progress_bar=show_progress_bar)
        atoks = self.s2a.generate(stoks, [speaker], T=2, top_k=8, show_progress_bar=show_progress_bar)
        return atoks
        
    def generate(self, text, speaker="default", show_progress_bar=True):
        return self.vocoder.decode(self.generate_atoks(text, speaker, show_progress_bar=show_progress_bar))
    
    def generate_to_file(self, fname, text, speaker="default", show_progress_bar=True):
        self.vocoder.decode_to_file(fname, self.generate_atoks(text, speaker, show_progress_bar=show_progress_bar))
        
    def generate_to_notebook(self, text, speaker="default", show_progress_bar=True):
        self.vocoder.decode_to_notebook(self.generate_atoks(text, speaker, show_progress_bar=show_progress_bar))
