Update app.py

app.py

@@ -7,7 +7,6 @@ import numpy as np
 from transformers import AutoModelForImageSegmentation
 from io import BytesIO
 from loadimg import load_img
-from contextlib import asynccontextmanager
 
 # -------------------------
 # Model Setup
@@ -16,13 +15,18 @@ MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-birefnet = None  # Will be initialized in lifespan
+birefnet = None  # will initialize on startup
 
 # -------------------------
-# Lifespan (startup/shutdown)
+# FastAPI App
 # -------------------------
-@asynccontextmanager
-async def lifespan(app: FastAPI):
+app = FastAPI(title="Background Removal API")
+
+# -------------------------
+# Startup Event
+# -------------------------
+@app.on_event("startup")
+async def load_model():
     global birefnet
     if birefnet is None:
         print("Loading BiRefNet model...")
@@ -34,10 +38,6 @@ async def lifespan(app: FastAPI):
         )
         birefnet.to(device).eval()
         print("Model loaded successfully.")
-    yield  # Hand control back to FastAPI
-    # No special shutdown actions needed
-
-app = FastAPI(title="Background Removal API", lifespan=lifespan)
 
 # -------------------------
 # Image Preprocessing
@@ -64,7 +64,7 @@ def process_image(image: Image.Image) -> Image.Image:
     return image
 
 # -------------------------
-# API Endpoint
+# Remove Background Endpoint
 # -------------------------
 @app.post("/remove-background")
 async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):