# Kyutai TTS Gradio App

import gradio as gr
import numpy as np
import torch
from moshi.models.loaders import CheckpointInfo
from moshi.models.tts import DEFAULT_DSM_TTS_REPO, DEFAULT_DSM_TTS_VOICE_REPO, TTSModel

# Load model at startup
checkpoint_info = CheckpointInfo.from_hf_repo(DEFAULT_DSM_TTS_REPO)
tts_model = TTSModel.from_checkpoint_info(
    checkpoint_info, n_q=32, temp=0.6, device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
    dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

# Default voice
default_voice = "expresso/ex03-ex01_happy_001_channel1_334s.wav"
voice_path = tts_model.get_voice_path(default_voice)

def tts_function(text, cfg_coef):
    entries = tts_model.prepare_script([text], padding_between=1)
    condition_attributes = tts_model.make_condition_attributes(
        [voice_path], cfg_coef=cfg_coef
    )

    pcms = []

    def _on_frame(frame):
        if (frame != -1).all():
            pcm = tts_model.mimi.decode(frame[:, 1:, :]).cpu().numpy()
            pcms.append(np.clip(pcm[0, 0], -1, 1))

    tts_model.stream(entries, condition_attributes, on_frame=_on_frame)

    if pcms:
        audio = np.concatenate(pcms)
        return (44100, audio)
    else:
        return "No audio generated."

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🎙️ Kyutai TTS - Real-Time Streaming Voice AI")
    with gr.Row():
        text_input = gr.Textbox(label="Enter Text", lines=5, placeholder="Type something...")
        cfg_slider = gr.Slider(minimum=0.0, maximum=5.0, value=2.0, label="CFG Coefficient")
    generate_btn = gr.Button("Generate Audio")
    audio_output = gr.Audio(label="Generated Audio", type="numpy")

    generate_btn.click(fn=tts_function, inputs=[text_input, cfg_slider], outputs=audio_output)

if __name__ == "__main__":
    demo.launch()