Update app.py

app.py

@@ -25,24 +25,22 @@ executor = ThreadPoolExecutor(max_workers=4)
 # GPU model setup
 birefnet = None
 transform_image = None
+model_loaded = threading.Event()
 
 def load_model():
     global birefnet, transform_image
-    birefnet = AutoModelForImageSegmentation.from_pretrained(
-        "ZhengPeng7/BiRefNet", trust_remote_code=True
-    )
-    birefnet.to("cuda")
-    birefnet.eval()
-    transform_image = transforms.Compose([
-        transforms.Resize((1024, 1024)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ])
-
-# Lazy load the model on the first request
-@app.before_first_request
-def initialize():
-    threading.Thread(target=load_model).start()
+    if not model_loaded.is_set():
+        birefnet = AutoModelForImageSegmentation.from_pretrained(
+            "ZhengPeng7/BiRefNet", trust_remote_code=True
+        )
+        birefnet.to("cuda")
+        birefnet.eval()
+        transform_image = transforms.Compose([
+            transforms.Resize((1024, 1024)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
+        model_loaded.set()
 
 # Helper functions
 def decode_image_from_base64(image_data):
@@ -68,6 +66,9 @@ async def process_image(image):
     """Process the image asynchronously, including background removal."""
     global birefnet, transform_image
 
+    # Ensure the model is loaded
+    load_model()
+
     # Convert image to tensor
     input_images = transform_image(image).unsqueeze(0).to("cuda")