from kokoro import KPipeline
from IPython.display import display, Audio
import soundfile as sf
import torch

class TTS:

    pipeline = KPipeline(lang_code='a')

    def __init__(self):
        pass

    # Should be saved as .wav
    def create_audio(self, text: str, save_path: str) -> None:
        generator = self.pipeline(text, voice='af_heart')
        for i, (gs, ps, audio) in enumerate(generator):
            # print(i, gs, ps)
            display(Audio(data=audio, rate=24000, autoplay=i==0))
            sf.write(save_path, audio, 24000)