app/main.py

from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles

from app.api.health import router as health_router
from app.api.jobs import router as jobs_router
from app.api.models import router as models_router
from app.api.speech import router as speech_router
from app.api.ui import router as ui_router
from app.config import get_settings
from app.db import setup_database
from app.errors import register_exception_handlers
from app.services.engine_registry import EngineRegistry, set_engine_registry
from app.services.job_worker import start_worker, stop_worker
from app.services.kitten_engine import KittenEngine
from app.services.kokoro_engine import KokoroEngine
from app.services.melo_engine import MeloEngine
from app.services.model_assets import validate_kokoro_assets, validate_model_assets
from app.services.supertonic_engine import SupertonicEngine


@asynccontextmanager
async def lifespan(app: FastAPI):
    settings = get_settings()
    validate_model_assets(settings.model_dir, allow_missing_voice_files=True)
    validate_kokoro_assets(settings.kokoro_model_dir, settings.kokoro_model_file)
    supertonic_engine = SupertonicEngine.from_settings(settings)
    kitten_engine = KittenEngine.from_settings(settings)
    kokoro_engine = KokoroEngine.from_settings(settings)
    melo_model_to_engine: dict[str, object] = {}
    if settings.enable_melo_tts:
        melo_engine = MeloEngine.from_settings(settings)
        if not melo_engine.is_available():
            if settings.melo_fail_fast:
                raise RuntimeError(
                    "ENABLE_MELO_TTS is true but the MeloTTS dependency is unavailable. "
                    "Install the optional Melo stack before enabling it."
                )
        else:
            if settings.melo_fail_fast:
                melo_engine.warmup()
            melo_model_to_engine = {
                model: melo_engine for model in melo_engine.public_model_ids
            }
    registry = EngineRegistry(
        model_to_engine={
            **{
                model: supertonic_engine for model in supertonic_engine.public_model_ids
            },
            **{model: kitten_engine for model in kitten_engine.public_model_ids},
            **{model: kokoro_engine for model in kokoro_engine.public_model_ids},
            **melo_model_to_engine,
        }
    )
    set_engine_registry(registry)

    # Web UI infrastructure
    setup_database()
    start_worker()

    # Mount static files for the web UI
    static_dir = Path(__file__).resolve().parent / "static"
    if static_dir.exists():
        app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")

    yield

    stop_worker()


app = FastAPI(title="CPU OpenAI-Compatible TTS", version="0.2.0", lifespan=lifespan)

register_exception_handlers(app)
app.include_router(health_router)
app.include_router(models_router)
app.include_router(speech_router)
app.include_router(jobs_router)
app.include_router(ui_router)