Update app.py

app.py

@@ -7,6 +7,7 @@ import numpy as np
 from transformers import AutoModelForImageSegmentation
 from io import BytesIO
 from loadimg import load_img
+from contextlib import asynccontextmanager
 
 # -------------------------
 # Model Setup
@@ -14,19 +15,13 @@ from loadimg import load_img
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
 birefnet = None  # will initialize on startup
 
 # -------------------------
-# FastAPI App
-# -------------------------
-app = FastAPI(title="Background Removal API")
-
-# -------------------------
-# Startup Event
+# Lifespan (Startup/Shutdown)
 # -------------------------
-@app.on_event("startup")
-async def load_model():
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     global birefnet
     if birefnet is None:
         print("Loading BiRefNet model...")
@@ -34,10 +29,17 @@ async def load_model():
             "ZhengPeng7/BiRefNet",
             cache_dir=MODEL_DIR,
             trust_remote_code=True,
-            revision="main"  # pin revision to avoid unexpected updates
+            revision="main"
         )
         birefnet.to(device).eval()
         print("Model loaded successfully.")
+    yield
+    # Optional shutdown logic here (nothing needed for this model)
+
+# -------------------------
+# FastAPI App
+# -------------------------
+app = FastAPI(title="Background Removal API", lifespan=lifespan)
 
 # -------------------------
 # Image Preprocessing