strive33/Kokoro-TTS-FastAPI
0
1import os2import time3import asyncio4from pathlib import Path5from typing import Optional6import numpy as np7import soundfile as sf8import torch9from fastapi import FastAPI, HTTPException, BackgroundTasks10from fastapi.responses import FileResponse, JSONResponse, StreamingResponse11from fastapi.staticfiles import StaticFiles12from fastapi.middleware.cors import CORSMiddleware13from pydantic import BaseModel14import logging15import io16 17logging.basicConfig(level=logging.INFO)18logger = logging.getLogger(__name__)19 20# ── App Setup ──────────────────────────────────────────────────────────────────21app = FastAPI(22 title="Kokoro TTS API",23 description="Text-to-Speech API powered by Kokoro-82M",24 version="1.0.0",25)26 27app.add_middleware(28 CORSMiddleware,29 allow_origins=["*"],30 allow_methods=["*"],31 allow_headers=["*"],32)33 34# ── Directories ────────────────────────────────────────────────────────────────35STATIC_DIR = Path("static")36STATIC_DIR.mkdir(exist_ok=True)37 38app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")39 40# ── Model Loading ──────────────────────────────────────────────────────────────41CUDA_AVAILABLE = torch.cuda.is_available()42device = "cuda" if CUDA_AVAILABLE else "cpu"43logger.info(f"Device: {device}")44 45model = None46pipelines = {}47 48def load_model():49 global model, pipelines50 try:51 from kokoro import KPipeline, KModel52 logger.info("Loading Kokoro model...")53 model = KModel(repo_id="hexgrad/Kokoro-82M").to(device).eval()54 55 lang_codes = ["a", "b", "e", "f", "h", "i", "j", "p", "z"]56 for code in lang_codes:57 try:58 pipelines[code] = KPipeline(lang_code=code, model=False)59 except Exception as e:60 logger.warning(f"Pipeline '{code}' failed: {e}")61 62 if "a" in pipelines:63 pipelines["a"].g2p.lexicon.golds["kokoro"] = "kˈOkəɹO"64 if "b" in pipelines:65 pipelines["b"].g2p.lexicon.golds["kokoro"] = "kˈQkəɹQ"66 67 logger.info(f"Model loaded. Pipelines: {list(pipelines.keys())}")68 except Exception as e:69 logger.error(f"Model load failed: {e}")70 71# Load on startup72@app.on_event("startup")73async def startup_event():74 loop = asyncio.get_event_loop()75 await loop.run_in_executor(None, load_model)76 77# ── Voice Registry ─────────────────────────────────────────────────────────────78VOICES = {79 # American English80 "af_heart": {"label": "Heart", "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},81 "af_bella": {"label": "Bella", "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},82 "af_nicole": {"label": "Nicole", "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},83 "af_aoede": {"label": "Aoede", "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},84 "af_kore": {"label": "Kore", "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},85 "af_sarah": {"label": "Sarah", "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},86 "af_nova": {"label": "Nova", "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},87 "af_sky": {"label": "Sky", "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},88 "af_river": {"label": "River", "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},89 "am_michael": {"label": "Michael", "lang": "en-US", "gender": "male", "flag": "🇺🇸", "code": "a"},90 "am_fenrir": {"label": "Fenrir", "lang": "en-US", "gender": "male", "flag": "🇺🇸", "code": "a"},91 "am_puck": {"label": "Puck", "lang": "en-US", "gender": "male", "flag": "🇺🇸", "code": "a"},92 "am_echo": {"label": "Echo", "lang": "en-US", "gender": "male", "flag": "🇺🇸", "code": "a"},93 "am_eric": {"label": "Eric", "lang": "en-US", "gender": "male", "flag": "🇺🇸", "code": "a"},94 "am_liam": {"label": "Liam", "lang": "en-US", "gender": "male", "flag": "🇺🇸", "code": "a"},95 "am_adam": {"label": "Adam", "lang": "en-US", "gender": "male", "flag": "🇺🇸", "code": "a"},96 # British English97 "bf_emma": {"label": "Emma", "lang": "en-GB", "gender": "female", "flag": "🇬🇧", "code": "b"},98 "bf_isabella": {"label": "Isabella","lang": "en-GB", "gender": "female", "flag": "🇬🇧", "code": "b"},99 "bf_alice": {"label": "Alice", "lang": "en-GB", "gender": "female", "flag": "🇬🇧", "code": "b"},100 "bf_lily": {"label": "Lily", "lang": "en-GB", "gender": "female", "flag": "🇬🇧", "code": "b"},101 "bm_george": {"label": "George", "lang": "en-GB", "gender": "male", "flag": "🇬🇧", "code": "b"},102 "bm_fable": {"label": "Fable", "lang": "en-GB", "gender": "male", "flag": "🇬🇧", "code": "b"},103 "bm_lewis": {"label": "Lewis", "lang": "en-GB", "gender": "male", "flag": "🇬🇧", "code": "b"},104 "bm_daniel": {"label": "Daniel", "lang": "en-GB", "gender": "male", "flag": "🇬🇧", "code": "b"},105 # Spanish106 "ef_dora": {"label": "Dora", "lang": "es", "gender": "female", "flag": "🇪🇸", "code": "e"},107 "em_alex": {"label": "Alex", "lang": "es", "gender": "male", "flag": "🇪🇸", "code": "e"},108 # French109 "ff_siwis": {"label": "Siwis", "lang": "fr", "gender": "female", "flag": "🇫🇷", "code": "f"},110 # Hindi111 "hf_alpha": {"label": "Alpha", "lang": "hi", "gender": "female", "flag": "🇮🇳", "code": "h"},112 "hf_beta": {"label": "Beta", "lang": "hi", "gender": "female", "flag": "🇮🇳", "code": "h"},113 "hm_omega": {"label": "Omega", "lang": "hi", "gender": "male", "flag": "🇮🇳", "code": "h"},114 "hm_psi": {"label": "Psi", "lang": "hi", "gender": "male", "flag": "🇮🇳", "code": "h"},115 # Italian116 "if_sara": {"label": "Sara", "lang": "it", "gender": "female", "flag": "🇮🇹", "code": "i"},117 "im_nicola": {"label": "Nicola", "lang": "it", "gender": "male", "flag": "🇮🇹", "code": "i"},118 # Japanese119 "jf_alpha": {"label": "Alpha", "lang": "ja", "gender": "female", "flag": "🇯🇵", "code": "j"},120 "jf_gongitsune":{"label": "Gongitsune","lang": "ja", "gender": "female", "flag": "🇯🇵", "code": "j"},121 "jf_nezumi": {"label": "Nezumi", "lang": "ja", "gender": "female", "flag": "🇯🇵", "code": "j"},122 "jm_kumo": {"label": "Kumo", "lang": "ja", "gender": "male", "flag": "🇯🇵", "code": "j"},123 # Portuguese124 "pf_dora": {"label": "Dora", "lang": "pt", "gender": "female", "flag": "🇧🇷", "code": "p"},125 "pm_alex": {"label": "Alex", "lang": "pt", "gender": "male", "flag": "🇧🇷", "code": "p"},126 # Chinese127 "zf_xiaobei": {"label": "Xiaobei", "lang": "zh", "gender": "female", "flag": "🇨🇳", "code": "z"},128 "zf_xiaoxiao": {"label": "Xiaoxiao", "lang": "zh", "gender": "female", "flag": "🇨🇳", "code": "z"},129 "zm_yunjian": {"label": "Yunjian", "lang": "zh", "gender": "male", "flag": "🇨🇳", "code": "z"},130 "zm_yunxi": {"label": "Yunxi", "lang": "zh", "gender": "male", "flag": "🇨🇳", "code": "z"},131}132 133# ── Pydantic Models ────────────────────────────────────────────────────────────134class TTSRequest(BaseModel):135 text: str136 voice: str = "af_heart"137 speed: float = 1.0138 output_format: str = "wav" # "wav" or "mp3"139 140 141# ── Helper ─────────────────────────────────────────────────────────────────────142def _synthesize_to_bytes(text: str, voice: str, speed: float, output_format: str) -> tuple:143 if model is None:144 raise RuntimeError("Model not loaded yet")145 voice_info = VOICES.get(voice)146 if not voice_info:147 raise ValueError(f"Unknown voice: {voice}")148 pipeline = pipelines.get(voice_info["code"])149 if not pipeline:150 raise ValueError(f"No pipeline for lang code: {voice_info['code']}")151 voice_pack = pipeline.load_voice(voice)152 all_audio = []153 for _, ps, _ in pipeline(text, voice, speed, split_pattern=r"\n+"):154 ref_s = voice_pack[len(ps) - 1].to(device)155 all_audio.append(model(ps, ref_s, speed).cpu().numpy())156 if not all_audio:157 raise RuntimeError("No audio generated")158 final_audio = np.concatenate(all_audio)159 duration = len(final_audio) / 24000160 buf = io.BytesIO()161 sf.write(buf, final_audio, 24000, format="WAV")162 return buf.getvalue(), duration163 164# ── Routes ─────────────────────────────────────────────────────────────────────165 166@app.get("/")167async def root():168 return FileResponse("static/index.html")169 170@app.get("/health")171async def health():172 return {173 "status": "ok",174 "model_loaded": model is not None,175 "device": device,176 "cuda": CUDA_AVAILABLE,177 "pipelines": list(pipelines.keys()),178 }179 180@app.get("/voices")181async def list_voices():182 available = {183 k: v for k, v in VOICES.items() if v["code"] in pipelines184 }185 # Group by language186 grouped = {}187 for vid, info in available.items():188 lang = info["lang"]189 if lang not in grouped:190 grouped[lang] = []191 grouped[lang].append({"id": vid, **info})192 return {"voices": available, "grouped": grouped, "total": len(available)}193 194@app.post("/tts")195async def text_to_speech(request: TTSRequest):196 if not request.text.strip():197 raise HTTPException(400, "text cannot be empty")198 if request.voice not in VOICES:199 raise HTTPException(400, f"Unknown voice. GET /voices for list.")200 if not 0.5 <= request.speed <= 2.0:201 raise HTTPException(400, "speed must be between 0.5 and 2.0")202 if request.output_format not in ("wav", "mp3"):203 raise HTTPException(400, "output_format must be wav or mp3")204 if model is None:205 raise HTTPException(503, "Model is still loading, please retry in a moment")206 207 try:208 loop = asyncio.get_event_loop()209 audio_bytes, duration = await loop.run_in_executor(210 None,211 lambda: _synthesize_to_bytes(request.text, request.voice, request.speed, request.output_format),212 )213 except Exception as e:214 logger.error(f"TTS error: {e}")215 raise HTTPException(500, str(e))216 217 fmt = request.output_format218 return StreamingResponse(219 io.BytesIO(audio_bytes),220 media_type="audio/mpeg" if fmt == "mp3" else "audio/wav",221 headers={222 "Content-Disposition": f'attachment; filename="kokoro_{request.voice}.{fmt}"',223 "X-Duration-Seconds": str(round(duration, 2)),224 },225 )226 227