from kokoro import KModel, KPipeline
import os
import random
import numpy as np
import kokoro
import torch


model = KModel().to('cpu').eval()
pipeline = KPipeline(lang_code='a', model=False)
    
# def generate_tts(text, voice='af_heart', speed=1):
#     pack = pipeline.load_voice(voice)
#     audio_chunks = []
#     for _, ps, _ in pipeline(text, voice, speed):
#         ref_s = pack[len(ps)-1]
#         try:
#             audio = model(ps, ref_s, speed)
#             audio_chunks.append(audio.numpy())
#         except:
#             print("lol there was an issue idk")
#         # yield 24000, audio.numpy()
#     if audio_chunks:
#         concatenated_audio = np.concatenate(audio_chunks)
#         print(concatenated_audio.shape)
#         return 24000, concatenated_audio
#     else:
#         return 24000, np.array([])


def generate_tts(text, voice='af_heart', speed=1):
    pack = pipeline.load_voice(voice)
    first = True
    for _, ps, _ in pipeline(text, voice, speed):
        ref_s = pack[len(ps)-1]
        try:
            audio = model(ps, ref_s, speed)
        except:
            print("lol there was an issue idk")
        yield 24000, audio.numpy()
        if first:
            first = False
            yield 24000, torch.zeros(1).numpy()