Update app.py

app.py

@@ -1,28 +1,43 @@
+import os
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.responses import StreamingResponse, HTMLResponse
 from PIL import Image
-import torch, numpy as np
+import torch
+import numpy as np
 from transformers import AutoModelForImageSegmentation
 from io import BytesIO
 from loadimg import load_img
-import os
 
 # -------------------------
 # Model Setup
 # -------------------------
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print("Loading BiRefNet model...")
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet",
-    cache_dir=MODEL_DIR,
-    trust_remote_code=True
-)
-birefnet.to(device).eval()
-print("Model loaded.")
+birefnet = None  # will initialize on startup
+
+# -------------------------
+# FastAPI App
+# -------------------------
+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...")
+        birefnet = AutoModelForImageSegmentation.from_pretrained(
+            "ZhengPeng7/BiRefNet",
+            cache_dir=MODEL_DIR,
+            trust_remote_code=True,
+            revision="main"  # pin revision to avoid unexpected updates
+        )
+        birefnet.to(device).eval()
+        print("Model loaded successfully.")
 
 # -------------------------
 # Image Preprocessing
@@ -49,12 +64,7 @@ def process_image(image: Image.Image) -> Image.Image:
     return image
 
 # -------------------------
-# FastAPI App
-# -------------------------
-app = FastAPI(title="Background Removal API")
-
-# -------------------------
-# Single /remove-background endpoint
+# Remove Background Endpoint
 # -------------------------
 @app.post("/remove-background")
 async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
@@ -79,7 +89,7 @@ async def remove_background(file: UploadFile = File(None), image_url: str = Form
         raise HTTPException(status_code=500, detail=str(e))
 
 # -------------------------
-# Simple HTML frontend
+# Web Interface
 # -------------------------
 @app.get("/", response_class=HTMLResponse)
 async def index():