Spaces:
Sleeping
Sleeping
File size: 4,190 Bytes
ea520ad |
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 132 133 134 135 136 137 138 |
import os
import re
import io
import json
import time
import base64
import asyncio
from typing import AsyncGenerator, List
import numpy as np
from fastapi import FastAPI, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from kokoro_onnx import Kokoro # pip install kokoro-onnx
APP_NAME = "kokoro-onnx-fastapi-sse"
SAMPLE_RATE = 24000 # Kokoro output
CHANNELS = 1
CHUNK_SAMPLES = 2400 # 100 ms
MODEL_PATH = os.getenv("KOKORO_MODEL", "models/kokoro-v1.0.int8.onnx")
VOICES_PATH = os.getenv("KOKORO_VOICES", "models/voices-v1.0.bin")
app = FastAPI(title=APP_NAME)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]
)
kokoro = None
_model_lock = asyncio.Lock()
def split_text(text: str, max_len: int = 220) -> List[str]:
# Split on sentences, then fold long pieces
parts = re.split(r"(?<=[\.\!\?।])\s+", text.strip())
chunks: List[str] = []
for p in parts:
p = p.strip()
while len(p) > max_len:
cut = p.rfind(" ", 0, max_len)
if cut == -1:
cut = max_len
chunks.append(p[:cut].strip())
p = p[cut:].strip()
if p:
chunks.append(p)
return [c for c in chunks if c]
def floats_to_s16le(samples: np.ndarray) -> np.ndarray:
x = np.clip(samples, -1.0, 1.0)
return (x * 32767.0).astype(np.int16)
async def sse_generator(text: str, voice: str, speed: float, lang: str) -> AsyncGenerator[bytes, None]:
# Yields SSE messages with base64 PCM16 frames (100ms per chunk).
seq = 0
total_samples = 0
try:
# Warmup ping
yield b": keep-alive\n\n"
for sentence in split_text(text):
async with _model_lock:
samples, sr = kokoro.create(sentence, voice=voice, speed=speed, lang=lang)
assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE}, got {sr}"
pcm16 = floats_to_s16le(np.asarray(samples))
for i in range(0, len(pcm16), CHUNK_SAMPLES):
frame = pcm16[i:i+CHUNK_SAMPLES]
total_samples += len(frame)
payload = {
"seq": seq,
"sr": SAMPLE_RATE,
"ch": CHANNELS,
"format": "s16le",
"pcm16": base64.b64encode(frame.tobytes()).decode("ascii"),
}
msg = f"data: {json.dumps(payload, separators=(',',':'))}\n\n"
yield msg.encode("utf-8")
seq += 1
await asyncio.sleep(0) # give control back to loop
# Done event
done = {"total_chunks": seq, "total_samples": total_samples}
yield f"event: done\ndata: {json.dumps(done)}\n\n".encode("utf-8")
except asyncio.CancelledError:
# Client disconnected
return
@app.on_event("startup")
async def _load_model():
global kokoro
kokoro = Kokoro(MODEL_PATH, VOICES_PATH)
@app.get("/healthz")
async def healthz():
return {"status": "ok", "model": os.path.basename(MODEL_PATH)}
@app.get("/v1/voices")
async def list_voices():
try:
import numpy as _np
with _np.load(VOICES_PATH) as z:
names = sorted(list(z.files))
return {"voices": names}
except Exception:
fallback = [
"af", "af_bella", "af_nicole", "af_sarah", "af_sky",
"am_adam", "am_michael",
"bf_emma", "bf_isabella",
"bm_george", "bm_lewis",
]
return {"voices": fallback, "note": "fallback list; voices file not parsed"}
@app.get("/v1/tts.sse")
async def tts_sse(
text: str = Query(..., description="Text to synthesize"),
voice: str = Query("af_sarah"),
speed: float = Query(1.0, ge=0.5, le=1.5),
lang: str = Query("en-us"),
):
headers = {
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no", # for nginx
}
return StreamingResponse(
sse_generator(text, voice, speed, lang),
media_type="text/event-stream",
headers=headers,
)
|