import gradio as gr
import torch
from transformers import AutoProcessor, AutoModel

# Load Kokoro TTS Model
device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "hexgrad/Kokoro-82M"
model = AutoModel.from_pretrained(model_name).to(device)
processor = AutoProcessor.from_pretrained(model_name)

def text_to_speech(text):
    """Convert input text to speech using Kokoro TTS"""
    inputs = processor(text, return_tensors="pt").to(device)
    with torch.no_grad():
        output = model.generate(**inputs)
    return output.cpu().numpy()

# Gradio Interface
description = "Enter text and listen to the Kokoro TTS model read it aloud."

demo = gr.Interface(
    fn=text_to_speech,
    inputs=gr.Textbox(placeholder="Type something here..."),
    outputs=gr.Audio(type="numpy"),
    title="Kokoro TTS - Text-to-Speech",
    description=description,
)

demo.launch()