File size: 5,488 Bytes
1425afc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import os
from pathlib import Path
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) ---
BASE_DIR = Path(__file__).resolve().parent
tokenizer_cache = None
kokoro_cache = {}
model_status = "not_loaded"
model_error = None
def get_style_vector_choices(directory="voices"):
directory_path = BASE_DIR / directory
return [file.name for file in directory_path.iterdir() if file.suffix == ".pt"]
def get_onnx_models(directory="weights"):
directory_path = BASE_DIR / directory
return [file.name for file in directory_path.iterdir() if file.suffix == ".onnx"]
def local_tts(
text: str,
model_path: str,
style_vector: str,
output_file_format: str = "wav",
speed: float = 1.0
):
global tokenizer_cache, model_status, model_error
if len(text) > 0:
try:
style_vector_path = str(BASE_DIR / "voices" / style_vector)
model_path_full = str(BASE_DIR / "weights" / model_path)
cache_key = (model_path_full, style_vector_path)
if tokenizer_cache is None:
tokenizer_cache = Tokenizer()
if cache_key not in kokoro_cache:
model_status = "loading"
model_error = None
kokoro_cache[cache_key] = Kokoro(model_path_full, style_vector_path, tokenizer=tokenizer_cache, lang='en-us')
model_status = "ready"
inference = kokoro_cache[cache_key]
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:
model_status = "failed"
model_error = str(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.", str(BASE_DIR / "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!", str(BASE_DIR / "assets" / "fun_fact.wav")),
("Thanks", "Thank you for listening to this audio. It was generated by the Kokoro TTS model.", str(BASE_DIR / "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)
|