main.py

import os
import time
import asyncio
from pathlib import Path
from typing import Optional
import numpy as np
import soundfile as sf
import torch
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import logging
import io

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ── App Setup ──────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Kokoro TTS API",
    description="Text-to-Speech API powered by Kokoro-82M",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Directories ────────────────────────────────────────────────────────────────
STATIC_DIR = Path("static")
STATIC_DIR.mkdir(exist_ok=True)

app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

# ── Model Loading ──────────────────────────────────────────────────────────────
CUDA_AVAILABLE = torch.cuda.is_available()
device = "cuda" if CUDA_AVAILABLE else "cpu"
logger.info(f"Device: {device}")

model = None
pipelines = {}

def load_model():
    global model, pipelines
    try:
        from kokoro import KPipeline, KModel
        logger.info("Loading Kokoro model...")
        model = KModel(repo_id="hexgrad/Kokoro-82M").to(device).eval()

        lang_codes = ["a", "b", "e", "f", "h", "i", "j", "p", "z"]
        for code in lang_codes:
            try:
                pipelines[code] = KPipeline(lang_code=code, model=False)
            except Exception as e:
                logger.warning(f"Pipeline '{code}' failed: {e}")

        if "a" in pipelines:
            pipelines["a"].g2p.lexicon.golds["kokoro"] = "kˈOkəɹO"
        if "b" in pipelines:
            pipelines["b"].g2p.lexicon.golds["kokoro"] = "kˈQkəɹQ"

        logger.info(f"Model loaded. Pipelines: {list(pipelines.keys())}")
    except Exception as e:
        logger.error(f"Model load failed: {e}")

# Load on startup
@app.on_event("startup")
async def startup_event():
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, load_model)

# ── Voice Registry ─────────────────────────────────────────────────────────────
VOICES = {
    # American English
    "af_heart":   {"label": "Heart",    "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},
    "af_bella":   {"label": "Bella",    "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},
    "af_nicole":  {"label": "Nicole",   "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},
    "af_aoede":   {"label": "Aoede",    "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},
    "af_kore":    {"label": "Kore",     "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},
    "af_sarah":   {"label": "Sarah",    "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},
    "af_nova":    {"label": "Nova",     "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},
    "af_sky":     {"label": "Sky",      "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},
    "af_river":   {"label": "River",    "lang": "en-US", "gender": "female", "flag": "🇺🇸", "code": "a"},
    "am_michael": {"label": "Michael",  "lang": "en-US", "gender": "male",   "flag": "🇺🇸", "code": "a"},
    "am_fenrir":  {"label": "Fenrir",   "lang": "en-US", "gender": "male",   "flag": "🇺🇸", "code": "a"},
    "am_puck":    {"label": "Puck",     "lang": "en-US", "gender": "male",   "flag": "🇺🇸", "code": "a"},
    "am_echo":    {"label": "Echo",     "lang": "en-US", "gender": "male",   "flag": "🇺🇸", "code": "a"},
    "am_eric":    {"label": "Eric",     "lang": "en-US", "gender": "male",   "flag": "🇺🇸", "code": "a"},
    "am_liam":    {"label": "Liam",     "lang": "en-US", "gender": "male",   "flag": "🇺🇸", "code": "a"},
    "am_adam":    {"label": "Adam",     "lang": "en-US", "gender": "male",   "flag": "🇺🇸", "code": "a"},
    # British English
    "bf_emma":     {"label": "Emma",    "lang": "en-GB", "gender": "female", "flag": "🇬🇧", "code": "b"},
    "bf_isabella": {"label": "Isabella","lang": "en-GB", "gender": "female", "flag": "🇬🇧", "code": "b"},
    "bf_alice":    {"label": "Alice",   "lang": "en-GB", "gender": "female", "flag": "🇬🇧", "code": "b"},
    "bf_lily":     {"label": "Lily",    "lang": "en-GB", "gender": "female", "flag": "🇬🇧", "code": "b"},
    "bm_george":   {"label": "George",  "lang": "en-GB", "gender": "male",   "flag": "🇬🇧", "code": "b"},
    "bm_fable":    {"label": "Fable",   "lang": "en-GB", "gender": "male",   "flag": "🇬🇧", "code": "b"},
    "bm_lewis":    {"label": "Lewis",   "lang": "en-GB", "gender": "male",   "flag": "🇬🇧", "code": "b"},
    "bm_daniel":   {"label": "Daniel",  "lang": "en-GB", "gender": "male",   "flag": "🇬🇧", "code": "b"},
    # Spanish
    "ef_dora":  {"label": "Dora",  "lang": "es", "gender": "female", "flag": "🇪🇸", "code": "e"},
    "em_alex":  {"label": "Alex",  "lang": "es", "gender": "male",   "flag": "🇪🇸", "code": "e"},
    # French
    "ff_siwis": {"label": "Siwis", "lang": "fr", "gender": "female", "flag": "🇫🇷", "code": "f"},
    # Hindi
    "hf_alpha": {"label": "Alpha", "lang": "hi", "gender": "female", "flag": "🇮🇳", "code": "h"},
    "hf_beta":  {"label": "Beta",  "lang": "hi", "gender": "female", "flag": "🇮🇳", "code": "h"},
    "hm_omega": {"label": "Omega", "lang": "hi", "gender": "male",   "flag": "🇮🇳", "code": "h"},
    "hm_psi":   {"label": "Psi",   "lang": "hi", "gender": "male",   "flag": "🇮🇳", "code": "h"},
    # Italian
    "if_sara":    {"label": "Sara",   "lang": "it", "gender": "female", "flag": "🇮🇹", "code": "i"},
    "im_nicola":  {"label": "Nicola", "lang": "it", "gender": "male",   "flag": "🇮🇹", "code": "i"},
    # Japanese
    "jf_alpha":     {"label": "Alpha",     "lang": "ja", "gender": "female", "flag": "🇯🇵", "code": "j"},
    "jf_gongitsune":{"label": "Gongitsune","lang": "ja", "gender": "female", "flag": "🇯🇵", "code": "j"},
    "jf_nezumi":    {"label": "Nezumi",    "lang": "ja", "gender": "female", "flag": "🇯🇵", "code": "j"},
    "jm_kumo":      {"label": "Kumo",      "lang": "ja", "gender": "male",   "flag": "🇯🇵", "code": "j"},
    # Portuguese
    "pf_dora": {"label": "Dora", "lang": "pt", "gender": "female", "flag": "🇧🇷", "code": "p"},
    "pm_alex": {"label": "Alex", "lang": "pt", "gender": "male",   "flag": "🇧🇷", "code": "p"},
    # Chinese
    "zf_xiaobei":  {"label": "Xiaobei",  "lang": "zh", "gender": "female", "flag": "🇨🇳", "code": "z"},
    "zf_xiaoxiao": {"label": "Xiaoxiao", "lang": "zh", "gender": "female", "flag": "🇨🇳", "code": "z"},
    "zm_yunjian":  {"label": "Yunjian",  "lang": "zh", "gender": "male",   "flag": "🇨🇳", "code": "z"},
    "zm_yunxi":    {"label": "Yunxi",    "lang": "zh", "gender": "male",   "flag": "🇨🇳", "code": "z"},
}

# ── Pydantic Models ────────────────────────────────────────────────────────────
class TTSRequest(BaseModel):
    text: str
    voice: str = "af_heart"
    speed: float = 1.0
    output_format: str = "wav"  # "wav" or "mp3"


# ── Helper ─────────────────────────────────────────────────────────────────────
def _synthesize_to_bytes(text: str, voice: str, speed: float, output_format: str) -> tuple:
    if model is None:
        raise RuntimeError("Model not loaded yet")
    voice_info = VOICES.get(voice)
    if not voice_info:
        raise ValueError(f"Unknown voice: {voice}")
    pipeline = pipelines.get(voice_info["code"])
    if not pipeline:
        raise ValueError(f"No pipeline for lang code: {voice_info['code']}")
    voice_pack = pipeline.load_voice(voice)
    all_audio = []
    with torch.no_grad():
        for _, ps, _ in pipeline(text, voice, speed, split_pattern=r"[.!?\n]+"):
            ref_s = voice_pack[len(ps) - 1].to(device)
            all_audio.append(model(ps, ref_s, speed).cpu().numpy())
    if not all_audio:
        raise RuntimeError("No audio generated")
    final_audio = np.concatenate(all_audio)
    duration = len(final_audio) / 24000
    
    buf = io.BytesIO()
    sf.write(buf, final_audio, 24000, format="WAV")
    
    if output_format == "mp3":
        buf.seek(0)
        from pydub import AudioSegment
        audio_segment = AudioSegment.from_wav(buf)
        mp3_buf = io.BytesIO()
        audio_segment.export(mp3_buf, format="mp3", parameters=["-q:a", "2"])
        return mp3_buf.getvalue(), duration
        
    return buf.getvalue(), duration

# ── Routes ─────────────────────────────────────────────────────────────────────

@app.get("/")
async def root():
    return FileResponse("static/index.html")

@app.get("/health")
async def health():
    return {
        "status": "ok",
        "model_loaded": model is not None,
        "device": device,
        "cuda": CUDA_AVAILABLE,
        "pipelines": list(pipelines.keys()),
    }

@app.get("/voices")
async def list_voices():
    available = {
        k: v for k, v in VOICES.items() if v["code"] in pipelines
    }
    # Group by language
    grouped = {}
    for vid, info in available.items():
        lang = info["lang"]
        if lang not in grouped:
            grouped[lang] = []
        grouped[lang].append({"id": vid, **info})
    return {"voices": available, "grouped": grouped, "total": len(available)}

@app.post("/tts")
async def text_to_speech(request: TTSRequest):
    if not request.text.strip():
        raise HTTPException(400, "text cannot be empty")
    if request.voice not in VOICES:
        raise HTTPException(400, f"Unknown voice. GET /voices for list.")
    if not 0.5 <= request.speed <= 2.0:
        raise HTTPException(400, "speed must be between 0.5 and 2.0")
    if request.output_format not in ("wav", "mp3"):
        raise HTTPException(400, "output_format must be wav or mp3")
    if model is None:
        raise HTTPException(503, "Model is still loading, please retry in a moment")

    try:
        loop = asyncio.get_event_loop()
        audio_bytes, duration = await loop.run_in_executor(
            None,
            lambda: _synthesize_to_bytes(request.text, request.voice, request.speed, request.output_format),
        )
    except Exception as e:
        logger.error(f"TTS error: {e}")
        raise HTTPException(500, str(e))

    fmt = request.output_format
    return StreamingResponse(
        io.BytesIO(audio_bytes),
        media_type="audio/mpeg" if fmt == "mp3" else "audio/wav",
        headers={
            "Content-Disposition": f'attachment; filename="kokoro_{request.voice}.{fmt}"',
            "X-Duration-Seconds": str(round(duration, 2)),
        },
    )