main.py

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
from routes import base, data
from helpers.db import init_chroma, chroma_client
from helpers.configs import get_settings
from pathlib import Path
import os
import urllib.request

ASSETS_DIR = "assets"
MODEL_URL = "https://github.com/YapaLab/yolo-face/releases/download/1.0.0/yolov11n-face.pt"
MODEL_PATH = os.path.join(ASSETS_DIR, "yolov11n-face.pt")

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("Initializing ChromaDB...")
    init_chroma()
    print("ChromaDB Ready.")
    if not Path(MODEL_PATH).exists():
        Path(ASSETS_DIR).mkdir(exist_ok=True)
        print("⬇ Downloading YOLO face model...")
        urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
        print("Model ready.")
    else:
        print("Model already exists.")
    yield
    if chroma_client is not None:
        try:
            chroma_client.persist()
            print("ChromaDB persisted & closed.")
        except:
            pass


app = FastAPI(lifespan=lifespan)

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

app.include_router(base.base_router)
app.include_router(data.data_router)


if __name__ == "__main__":
    import uvicorn

    settings = get_settings()
    print(f"Server running at http://127.0.0.1:{settings.port}")
    uvicorn.run(app, host=settings.host, port=settings.port)