Fix model initialization with lifespan handler

pocket_tts/main.py

@@ -3,6 +3,7 @@ import logging
 import os
 import tempfile
 import threading
+from contextlib import asynccontextmanager
 from pathlib import Path
 from queue import Queue
 
@@ -42,8 +43,24 @@ cli_app = typer.Typer(
 tts_model = None
 global_model_state = None
 
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global tts_model, global_model_state
+    logger.info("Loading TTS Model in lifespan...")
+    tts_model = TTSModel.load_model(DEFAULT_VARIANT)
+    global_model_state = tts_model.get_state_for_audio_prompt(DEFAULT_AUDIO_PROMPT)
+    logger.info(f"The size of the model state is {size_of_dict(global_model_state) // 1e6} MB")
+    yield
+    # Clean up if needed (optional)
+    logger.info("Shutting down TTS Model...")
+
+
 web_app = FastAPI(
-    title="Kyutai Pocket TTS API", description="Text-to-Speech generation API", version="1.0.0"
+    title="Kyutai Pocket TTS API",
+    description="Text-to-Speech generation API",
+    version="1.0.0",
+    lifespan=lifespan,
 )
 web_app.add_middleware(
     CORSMiddleware,
@@ -232,14 +249,7 @@ def serve(
     reload: Annotated[bool, typer.Option(help="Enable auto-reload")] = False,
 ):
     """Start the FastAPI server."""
-
-    global tts_model, global_model_state
-    tts_model = TTSModel.load_model(DEFAULT_VARIANT)
-
-    # Pre-load the voice prompt
-    global_model_state = tts_model.get_state_for_audio_prompt(voice)
-    logger.info(f"The size of the model state is {size_of_dict(global_model_state) // 1e6} MB")
-
+    # Model loading is now handled by the lifespan context manager
     uvicorn.run("pocket_tts.main:web_app", host=host, port=port, reload=reload)