satoko-chatbot/tts-service/app.py

"""Service local Kokoro CPU : API WAV compatible avec le client vocal."""
import asyncio
import logging
import os
import time
from contextlib import asynccontextmanager
from io import BytesIO

import numpy as np
import soundfile as sf
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from kokoro import KPipeline
from pydantic import BaseModel, Field

logger = logging.getLogger('uvicorn.error')
pipeline = None
generation_lock = asyncio.Lock()


def generate(text, voice):
    with torch.inference_mode():
        chunks = [result.audio.cpu().numpy() for result in pipeline(text, voice=voice)]
    if not chunks:
        raise ValueError('Aucun audio généré')
    audio = np.concatenate(chunks)
    output = BytesIO()
    sf.write(output, audio, 24000, format='WAV', subtype='PCM_16')
    return output.getvalue(), len(audio) / 24000


@asynccontextmanager
async def lifespan(app):
    global pipeline
    torch.set_num_threads(int(os.getenv('TTS_THREADS', '4')))
    torch.set_num_interop_threads(1)
    pipeline = KPipeline(lang_code='f', repo_id='hexgrad/Kokoro-82M', device='cpu')
    await asyncio.to_thread(generate, 'Bonjour, je suis prête.', 'ff_siwis')
    yield


app = FastAPI(lifespan=lifespan)


class SpeechRequest(BaseModel):
    input: str = Field(min_length=1, max_length=300)
    voice: str = 'ff_siwis'


@app.get('/healthz')
def health():
    return {'ready': pipeline is not None, 'model': 'Kokoro-82M', 'voice': 'ff_siwis', 'device': 'cpu'}


@app.post('/v1/audio/speech')
async def speech(request: SpeechRequest):
    if request.voice != 'ff_siwis' or not request.input.strip():
        raise HTTPException(422, 'Voix inconnue ou texte vide')
    if generation_lock.locked():
        raise HTTPException(429, 'Synthèse occupée')
    async with generation_lock:
        started = time.monotonic()
        # Shield : ne pas libérer le verrou pendant une génération annulée par le client.
        task = asyncio.create_task(asyncio.to_thread(generate, request.input, request.voice))
        try:
            payload, duration = await asyncio.shield(task)
        except asyncio.CancelledError:
            await task
            raise
        elapsed = time.monotonic() - started
        logger.info('Synthèse : caractères=%d génération=%.3fs audio=%.3fs', len(request.input), elapsed, duration)
        return Response(payload, media_type='audio/wav', headers={
            'X-Generation-Seconds': f'{elapsed:.3f}', 'X-Audio-Seconds': f'{duration:.3f}'})