Update app.py

app.py

@@ -12,9 +12,27 @@ import torchvision.transforms.functional as F
 from decord import VideoReader
 from nncore.engine import load_checkpoint
 from nncore.nn import build_model
+from contextlib import asynccontextmanager
 
-app = FastAPI(title="R2-Tuning API")
+# Global variables for model and config
+model, cfg = None, None
+
+# Lifespan handler to manage startup and shutdown
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup: Load the model and config
+    global model, cfg
+    print("Loading model on startup...")
+    model, cfg = init_model(CONFIG, WEIGHT)
+    print("Model loaded successfully.")
+    yield  # Application runs here
+    # Shutdown: Clean up (if needed)
+    print("Shutting down...")
+
+# Initialize FastAPI app with lifespan
+app = FastAPI(title="R2-Tuning API", lifespan=lifespan)
 
+# Enable CORS for React app
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["http://localhost:3000"],
@@ -23,11 +41,10 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Configuration
 CONFIG = 'configs/qvhighlights/r2_tuning_qvhighlights.py'
 WEIGHT = 'r2_tuning_qvhighlights-ed516355.pth'
 
-model, cfg = None, None
-
 def convert_time(seconds):
     minutes, seconds = divmod(round(max(seconds, 0)), 60)
     return f'{minutes:02d}:{seconds:02d}'
@@ -73,12 +90,6 @@ def process_video(video_path: str, query: str, model, cfg) -> dict:
     hd = [{"x": i * 2, "y": y} for i, y in enumerate(hd)]
     return {"moment_retrieval": mr, "highlight_detection": hd}
 
-@app.on_event("startup")
-async def startup_event():
-    global model, cfg
-    model, cfg = init_model(CONFIG, WEIGHT)
-    print("Model loaded successfully.")
-
 @app.post("/predict")
 async def predict(video: UploadFile = File(...), query: str = Form(...)):
     try: