Update app logic and deploy to Hugging Face

app/main.py

@@ -4,9 +4,11 @@ import tempfile
 import subprocess
 import cv2 # type: ignore
 import numpy as np
-from fastapi import FastAPI, UploadFile, File, Response, BackgroundTasks, Query
+import asyncio
+from fastapi import FastAPI, UploadFile, File, Response, BackgroundTasks, Query, HTTPException
 from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.concurrency import run_in_threadpool
 from rembg import new_session, remove
 from enum import Enum
 
@@ -29,6 +31,11 @@ class ModelName(str, Enum):
 # Cache sessions to avoid reloading models on every request
 sessions = {}
 
+# Global semaphore to limit concurrent processing.
+# Free tiers have limited CPU/RAM. We limit to 1 concurrent heavy task to prevent OOM/Crashes.
+MAX_CONCURRENT_PROCESSING = 1
+processing_semaphore = asyncio.Semaphore(MAX_CONCURRENT_PROCESSING)
+
 def get_session(model_name: str):
     if model_name not in sessions:
         print(f"Loading model: {model_name}...")
@@ -46,7 +53,7 @@ async def startup_event():
 
 @app.get("/")
 def read_root():
-    return {"message": "Background Removal API is running"}
+    return {"message": "Background Removal API is running", "concurrent_limit": MAX_CONCURRENT_PROCESSING}
 
 @app.post("/image-bg-removal")
 async def image_bg_removal(
@@ -61,26 +68,35 @@ async def image_bg_removal(
     Removes background from an image.
     Returns the image with transparent background (PNG).
     """
+    # Read file content first (IO bound, doesn't need semaphore)
     input_image = await file.read()
     
     session = get_session(model.value)
     
-    try:
-        output_image = remove(
-            input_image, 
-            session=session,
-            alpha_matting=alpha_matting,
-            alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,
-            alpha_matting_background_threshold=alpha_matting_background_threshold,
-            alpha_matting_erode_size=alpha_matting_erode_size
-        )
-    except Exception as e:
-        print(f"Error with alpha matting: {e}")
-        if alpha_matting:
-            print("Falling back to standard background removal (alpha_matting=False)...")
-            output_image = remove(input_image, session=session, alpha_matting=False)
-        else:
-            raise e
+    # Acquire semaphore before heavy processing
+    if processing_semaphore.locked():
+        print("Waiting for processing slot...")
+        
+    async with processing_semaphore:
+        try:
+            # Run blocking 'remove' function in a separate thread to avoid blocking the event loop
+            output_image = await run_in_threadpool(
+                remove,
+                input_image, 
+                session=session,
+                alpha_matting=alpha_matting,
+                alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,
+                alpha_matting_background_threshold=alpha_matting_background_threshold,
+                alpha_matting_erode_size=alpha_matting_erode_size
+            )
+        except Exception as e:
+            print(f"Error with alpha matting: {e}")
+            if alpha_matting:
+                print("Falling back to standard background removal (alpha_matting=False)...")
+                # Fallback also runs in thread pool
+                output_image = await run_in_threadpool(remove, input_image, session=session, alpha_matting=False)
+            else:
+                raise e
     
     return Response(content=output_image, media_type="image/png")
 
@@ -94,14 +110,20 @@ async def video_bg_removal(
     Removes background from a video.
     Returns WebM with Alpha.
     """
-    # Create temp file for input
+    # Create temp file for input (IO bound)
     with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_input:
         shutil.copyfileobj(file.file, tmp_input)
         tmp_input_path = tmp_input.name
 
     try:
-        # Pass model name to processing function
-        output_path = process_video(tmp_input_path, model.value)
+        # Acquire semaphore for the heavy video processing
+        if processing_semaphore.locked():
+             print("Waiting for video processing slot...")
+
+        async with processing_semaphore:
+            # Pass model name to processing function, run in thread pool
+            output_path = await run_in_threadpool(process_video, tmp_input_path, model.value)
+            
     except Exception as e:
         if os.path.exists(tmp_input_path):
             os.remove(tmp_input_path)