Update app.py

app.py

@@ -9,7 +9,7 @@ from slowapi import Limiter, _rate_limit_exceeded_handler
 from slowapi.util import get_remote_address
 from slowapi.errors import RateLimitExceeded
 from slowapi.middleware import SlowAPIMiddleware
-
+import sys # Import sys module
 import config
 from services import generation
 from routers import ideas, images, videos
@@ -25,22 +25,40 @@ limiter = Limiter(key_func=get_remote_address, default_limits=[config.RATE_LIMIT
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Startup: Assign essential state FIRST
-    app.state.limiter = limiter # <<<--- MOVE ASSIGNMENT HERE
+    app.state.limiter = limiter
     logger.info("Rate limiter assigned to app state.")
 
-    # Now attempt to load models
-    logger.info("Application startup: Loading models...")
+    models_loaded_successfully = False
     try:
-        generation.load_models() # This might still fail
+        logger.info("Application startup: Loading models...")
+        generation.load_models() # This might raise an exception
+        models_loaded_successfully = True # Set flag ONLY if load_models completes
         logger.info("Models loaded successfully.")
     except Exception as e:
-        # Log the error, but the app continues. Endpoints needing models might fail.
-        logger.error(f"ERROR: Model loading failed during startup: {e}", exc_info=True)
-
-    yield # Application is now ready to serve requests
-
-    # Shutdown: Clean up (optional, less critical in stateless services)
-    logger.info("Application shutdown.")
+        logger.error(f"FATAL: Model loading failed during startup: {e}", exc_info=True)
+        # Option 1: Exit the application (cleaner for Docker environments)
+        # logger.critical("Exiting application due to model loading failure.")
+        # sys.exit(1)
+        # Option 2: Raise the exception again to make uvicorn aware of failure
+        # This might depend on how uvicorn handles lifespan exceptions
+        # raise # Re-raise the exception
+
+    # >>> Only yield if models loaded <<<
+    if models_loaded_successfully:
+        yield # Application is now ready to serve requests
+    else:
+        # If models didn't load, we don't yield, preventing Uvicorn
+        # from reporting "Application startup complete."
+        # You might need to manually stop the process if Option 1 above isn't used.
+        logger.error("Application startup failed due to model loading errors. Server will not serve requests effectively.")
+        # Keep the process running but indicate failure. Or use sys.exit(1) above.
+        # If you don't yield or exit, Uvicorn might hang or exit depending on version.
+        # Testing needed here - sys.exit(1) is often simplest in Docker.
+        # For now, just logging the failure and not yielding.
+
+    # --- Shutdown Logic ---
+    # This part might not be reached if sys.exit was called
+    logger.info("Application shutdown sequence starting.")
     if "generation" in globals() and hasattr(generation, 'model_cache'):
          generation.model_cache.clear()
     # Add any other cleanup here