Update app.py

app.py

@@ -23,11 +23,13 @@ pillow_heif.register_heif_opener()
 executor = ThreadPoolExecutor(max_workers=os.cpu_count() or 4)
 
 # -------------------------
-# Model Setup (Loaded Once)
+# Model Setup (Load Once)
 # -------------------------
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 
 print("Loading BiRefNet model (first run may take a while)...")
 birefnet = AutoModelForImageSegmentation.from_pretrained(
@@ -37,13 +39,15 @@ birefnet = AutoModelForImageSegmentation.from_pretrained(
 )
 birefnet.to(device)
 birefnet.eval()
-print("Model loaded successfully.")
+print(f"Model loaded successfully on {device}.")
 
 # -------------------------
 # Image Preprocessing
 # -------------------------
+TARGET_SIZE = (512, 512)  # Lower resolution for faster inference
+
 def transform_image(image: Image.Image) -> torch.Tensor:
-    image = image.resize((1024, 1024))
+    image = image.resize(TARGET_SIZE)
     arr = np.array(image).astype(np.float32) / 255.0
     mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
     std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
@@ -53,16 +57,22 @@ def transform_image(image: Image.Image) -> torch.Tensor:
     return tensor
 
 def process_image_sync(image: Image.Image) -> BytesIO:
-    """Process image synchronously and return PNG bytes (no files written)."""
+    """Process image synchronously and return PNG bytes (in-memory)."""
     image_size = image.size
     input_tensor = transform_image(image)
+
     with torch.no_grad():
-        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+        if device == "cuda":
+            # Mixed precision for GPU
+            with torch.cuda.amp.autocast():
+                preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+        else:
+            # CPU fallback
+            preds = birefnet(input_tensor)[-1].sigmoid().cpu()
 
     pred = preds[0, 0].numpy()
     mask = Image.fromarray((pred * 255).astype(np.uint8)).resize(image_size)
 
-    # Apply alpha mask and keep only in-memory result
     image = image.copy()
     image.putalpha(mask)
 
@@ -85,24 +95,24 @@ def open_image_safely(file_bytes: bytes) -> Image.Image:
         img = Image.open(BytesIO(file_bytes))
         fmt = (img.format or "").lower()
 
-        # Handle PDF: first page only
+        # Handle PDF: first page
         if fmt == "pdf":
             from pdf2image import convert_from_bytes
             pdf_images = convert_from_bytes(file_bytes, first_page=1, last_page=1)
             return pdf_images[0].convert("RGB")
 
-        # Handle animated GIF: first frame only
+        # Handle GIF: first frame
         if fmt == "gif" and getattr(img, "is_animated", False):
             img.seek(0)
             return img.convert("RGB")
 
-        # Handle SVG: rasterize to PNG
+        # Handle SVG
         if fmt == "svg":
             import cairosvg
             png_bytes = cairosvg.svg2png(bytestring=file_bytes)
             return Image.open(BytesIO(png_bytes)).convert("RGB")
 
-        # Other formats (HEIC, HEIF, JPG, PNG, etc.)
+        # Other formats (HEIC, HEIF, JPG, PNG)
         return img.convert("RGB")
 
     except Exception as e:
@@ -118,15 +128,11 @@ app = FastAPI(title="Background Removal API", description="Removes image backgro
 # -------------------------
 @app.post("/remove_bg_file")
 async def remove_bg_file(file: UploadFile = File(...)):
-    """Upload an image and get transparent PNG."""
     try:
         contents = await file.read()
         image = open_image_safely(contents)
         output_buffer = await process_image_async(image)
-
-        # Return directly from memory
         return StreamingResponse(output_buffer, media_type="image/png")
-
     except HTTPException as e:
         raise e
     except Exception as e:
@@ -134,7 +140,6 @@ async def remove_bg_file(file: UploadFile = File(...)):
 
 @app.post("/remove_bg_url")
 async def remove_bg_url(image_url: str = Form(...)):
-    """Provide image URL and get transparent PNG."""
     try:
         image = load_img(image_url, output_type="pil").convert("RGB")
         output_buffer = await process_image_async(image)
@@ -143,7 +148,7 @@ async def remove_bg_url(image_url: str = Form(...)):
         raise HTTPException(status_code=500, detail=f"Error processing URL: {e}")
 
 # -------------------------
-# Web Interface (no saving)
+# Web Interface
 # -------------------------
 @app.get("/", response_class=HTMLResponse)
 async def index():
@@ -211,7 +216,7 @@ async def index():
             fileForm.addEventListener('submit', async (e) => {
                 e.preventDefault();
                 const fileInput = document.getElementById('fileInput');
-                if (fileInput.files.length === 0) return alert("Select a file!");
+                if (!fileInput.files.length) return alert("Select a file!");
                 const file = fileInput.files[0];
                 beforeImg.src = URL.createObjectURL(file);
                 const formData = new FormData();
@@ -249,16 +254,8 @@ async def index():
     return HTMLResponse(content=html)
 
 # -------------------------
-# Run Server
+# Run Server (Auto-detect filename)
 # -------------------------
 if __name__ == "__main__":
-    import sys
-    import os
-    import uvicorn
-
-    # Get current filename without .py
     module_name = os.path.splitext(os.path.basename(__file__))[0]
-
-    # Run uvicorn using the detected module name
     uvicorn.run(f"{module_name}:app", host="0.0.0.0", port=7860, workers=2)
-