app.py

import os
import gradio as gr
import tempfile
import soundfile as sf
from models import Tokenizer, Kokoro
from fastapi import FastAPI, Request
from fastapi.responses import FileResponse
import uvicorn

# --- EXISTING LOGIC (UNCHANGED) ---

def get_style_vector_choices(directory="voices"):
    return [file for file in os.listdir(directory) if file.endswith(".pt")]

def get_onnx_models(directory="weights"):
    return [file for file in os.listdir(directory) if file.endswith(".onnx")]

def local_tts(
        text: str,
        model_path: str,
        style_vector: str,
        output_file_format: str = "wav",
        speed: float = 1.0
):
    if len(text) > 0:
        try:
            tokenizer = Tokenizer()
            style_vector_path = os.path.join("voices", style_vector)
            model_path_full = os.path.join("weights", model_path)

            inference = Kokoro(model_path_full, style_vector_path, tokenizer=tokenizer, lang='en-us')

            audio, sample_rate = inference.generate_audio(text, speed=speed)

            with tempfile.NamedTemporaryFile(suffix=f".{output_file_format}", delete=False) as temp_file:
                sf.write(temp_file.name, audio, sample_rate)
                temp_file_path = temp_file.name

            return temp_file_path

        except Exception as e:
            raise gr.Error(f"An error occurred during TTS inference: {str(e)}")
    else:
        raise gr.Error("Input text cannot be empty.")

style_vector_choices = get_style_vector_choices()
onnx_models_choices = get_onnx_models()

sample_outputs = [
    ("Educational Note", "Machine learning models rely on large datasets and complex algorithms to identify patterns and make predictions.", "assets/edu_note.wav"),
    ("Fun Fact", "Did you know that honey never spoils? Archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still edible!", "assets/fun_fact.wav"),
    ("Thanks", "Thank you for listening to this audio. It was generated by the Kokoro TTS model.", "assets/thanks.wav")
]

example_texts = [
    ["Machine learning models rely on large datasets and complex algorithms to identify patterns and make predictions."],
    ["Did you know that honey never spoils? Archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still edible!"],
    ["Thank you for listening to this audio. It was generated by the Kokoro TTS model."]
]

# --- GRADIO INTERFACE (UNCHANGED) ---

with gr.Blocks() as demo:
    gr.Markdown("## <center> Kokoro TTS ONNX Inference | [GitHub Link](https://github.com/yakhyo/kokoro-onnx) </center>")
    with gr.Row(variant="panel"):
        model_path = gr.Dropdown(choices=onnx_models_choices, label="ONNX Model Path", value=onnx_models_choices[0])
        style_vector = gr.Dropdown(choices=style_vector_choices, label="Style Vector", value=style_vector_choices[0])
        output_file_format = gr.Dropdown(choices=["wav", "mp3"], label="Output Format", value="wav")
        speed = gr.Slider(minimum=0.5, maximum=2.0, value=1.0, step=0.1, label="Speed")

    text = gr.Textbox(label="Input Text", placeholder="Enter text to convert to speech.")
    btn = gr.Button("Generate Speech")
    output_audio = gr.Audio(label="Generated Audio", type="filepath")

    btn.click(fn=local_tts, inputs=[text, model_path, style_vector, output_file_format, speed], outputs=output_audio)

    gr.Examples(examples=example_texts, inputs=[text], label="Click an example to populate the input text")
    gr.Markdown("### Sample Texts and Audio")
    for topic, sample_text, sample_audio in sample_outputs:
        with gr.Row():
            gr.Textbox(value=sample_text, label=topic, interactive=False)
            gr.Audio(value=sample_audio, label="Example Audio", type="filepath", interactive=False)

# --- FASTAPI WRAPPER & STARTUP ---

app = FastAPI()

@app.post("/v1/audio/speech")
async def api_speech(request: Request):
    """
    OpenAI-compatible /v1/audio/speech endpoint for automation.
    Expects JSON: {"input": "text", "voice": "voice_file.pt", "model": "model_file.onnx"}
    """
    data = await request.json()
    input_text = data.get("input", "")
    m_path = data.get("model", onnx_models_choices[0])
    s_vec = data.get("voice", style_vector_choices[0])
    spd = float(data.get("speed", 1.0))
    
    file_path = local_tts(input_text, m_path, s_vec, speed=spd)
    return FileResponse(file_path, media_type="audio/wav")

# Mount Gradio into the FastAPI app
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    # Runs on port 7860 as expected by the Dockerfile/Hugging Face
    uvicorn.run(app, host="0.0.0.0", port=7860)