Update app.py

app.py

@@ -1,17 +1,16 @@
 import os
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Query
 from fastapi.responses import StreamingResponse, HTMLResponse
-from PIL import Image, ImageSequence
+from PIL import Image
 import torch
 import numpy as np
 from transformers import AutoModelForImageSegmentation
 from io import BytesIO
 import requests
 import uvicorn
-from concurrent.futures import ThreadPoolExecutor
 
 # -------------------------
-# Optional HEIC Support
+# Optional HEIC/HEIF Support
 # -------------------------
 try:
     import pillow_heif
@@ -26,6 +25,7 @@ except ImportError:
 MODEL_DIR = "models/BiRefNet"
 os.makedirs(MODEL_DIR, exist_ok=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 print("Loading BiRefNet model...")
 birefnet = AutoModelForImageSegmentation.from_pretrained(
@@ -34,17 +34,13 @@ birefnet = AutoModelForImageSegmentation.from_pretrained(
     trust_remote_code=True,
     revision="main"
 )
-birefnet.to(device)
-if device == "cuda":
-    birefnet = birefnet.half()  # FP16 for faster GPU inference
-birefnet.eval()
+birefnet.to(device, dtype=dtype).eval()
 print("Model loaded successfully.")
 
 # -------------------------
 # FastAPI App
 # -------------------------
 app = FastAPI(title="Background Remover API")
-executor = ThreadPoolExecutor(max_workers=4)
 
 # -------------------------
 # Utility Functions
@@ -53,27 +49,27 @@ def load_image_from_url(url: str) -> Image.Image:
     try:
         response = requests.get(url, timeout=10)
         response.raise_for_status()
-        return Image.open(BytesIO(response.content))
+        return Image.open(BytesIO(response.content)).convert("RGB")
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Error loading image from URL: {str(e)}")
 
-def transform_image(image: Image.Image, target_size=(512, 512)) -> torch.Tensor:
-    image = image.resize(target_size)
+def transform_image(image: Image.Image, resolution: int = 512) -> torch.Tensor:
+    image = image.resize((resolution, resolution))
     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)
     arr = (arr - mean) / std
     arr = np.transpose(arr, (2, 0, 1))  # HWC -> CHW
-    tensor = torch.from_numpy(arr).unsqueeze(0).to(torch.float16 if device=="cuda" else torch.float32).to(device)
+    tensor = torch.from_numpy(arr).unsqueeze(0).to(dtype).to(device)
     return tensor
 
-def process_image(image: Image.Image) -> Image.Image:
-    image_size = image.size
-    input_tensor = transform_image(image)
+def process_image(image: Image.Image, resolution: int = 512) -> Image.Image:
+    orig_size = image.size
+    input_tensor = transform_image(image, resolution)
     with torch.no_grad():
         preds = birefnet(input_tensor)[-1].sigmoid().cpu()
     pred = preds[0, 0]
-    mask = Image.fromarray((pred.numpy() * 255).astype(np.uint8)).resize(image_size)
+    mask = Image.fromarray((pred.numpy() * 255).astype(np.uint8)).resize(orig_size)
     image = image.convert("RGBA")
     image.putalpha(mask)
     return image
@@ -82,30 +78,28 @@ def process_image(image: Image.Image) -> Image.Image:
 # /remove-background Endpoint
 # -------------------------
 @app.post("/remove-background")
-async def remove_background(file: UploadFile = File(None), image_url: str = Form(None)):
+async def remove_background(
+    file: UploadFile = File(None),
+    image_url: str = Form(None),
+    resolution: int = Form(512)
+):
     """
     Remove background from an image.
-    Accepts file upload or image URL.
+    Accepts a file upload or image URL.
+    Optional resolution (default 512) for faster inference.
     Returns PNG with transparent background.
     """
     try:
         if file:
-            image = Image.open(BytesIO(await file.read()))
+            image = Image.open(BytesIO(await file.read())).convert("RGB")
         elif image_url:
             image = load_image_from_url(image_url)
         else:
-            raise HTTPException(status_code=400, detail="Provide 'file' or 'image_url'.")
-        
-        # Handle multi-frame images (GIF, PDF)
-        if getattr(image, "is_animated", False):
-            frames = [process_image(frame.convert("RGBA")) for frame in ImageSequence.Iterator(image)]
-            buf = BytesIO()
-            frames[0].save(buf, format="PNG", save_all=True, append_images=frames[1:])
-        else:
-            result = process_image(image.convert("RGBA"))
-            buf = BytesIO()
-            result.save(buf, format="PNG")
+            raise HTTPException(status_code=400, detail="Provide either 'file' or 'image_url'.")
 
+        result = process_image(image, resolution)
+        buf = BytesIO()
+        result.save(buf, format="PNG")
         buf.seek(0)
         return StreamingResponse(buf, media_type="image/png")
     except Exception as e:
@@ -135,55 +129,61 @@ async def index():
             <h2 class="mb-4">Background Remover API Tester</h2>
             <form id="uploadForm" class="mb-4" enctype="multipart/form-data">
                 <div class="mb-3">
-                    <label class="form-label">Upload Image (any format):</label>
+                    <label for="fileInput" class="form-label">Upload Image (any format, e.g. JPG, PNG, HEIC):</label>
                     <input class="form-control" type="file" id="fileInput" name="file" accept="image/*">
                 </div>
+                <div class="mb-3">
+                    <label for="resInput" class="form-label">Resolution (default 512):</label>
+                    <input class="form-control" type="number" id="resInput" name="resolution" value="512" min="64" max="2048">
+                </div>
                 <button class="btn btn-primary" type="submit">Remove Background</button>
             </form>
             <div class="mb-4">OR</div>
             <form id="urlForm" class="mb-4">
                 <div class="mb-3">
-                    <label class="form-label">Enter Image URL:</label>
+                    <label for="urlInput" class="form-label">Enter Image URL:</label>
                     <input class="form-control" type="text" id="urlInput" placeholder="https://example.com/image.jpg">
                 </div>
+                <div class="mb-3">
+                    <label for="urlResInput" class="form-label">Resolution (default 512):</label>
+                    <input class="form-control" type="number" id="urlResInput" name="resolution" value="512" min="64" max="2048">
+                </div>
                 <button class="btn btn-success" type="submit">Remove Background</button>
             </form>
             <div id="resultContainer" class="mt-4">
                 <h5>Result:</h5>
                 <img id="resultImg" src="" alt="">
-                <a id="downloadLink" class="btn btn-info mt-2" download="result.png" style="display:none;">Download PNG</a>
             </div>
         </div>
         <script>
             const uploadForm = document.getElementById("uploadForm");
             const urlForm = document.getElementById("urlForm");
             const resultImg = document.getElementById("resultImg");
-            const downloadLink = document.getElementById("downloadLink");
 
             uploadForm.addEventListener("submit", async e => {
                 e.preventDefault();
                 const fileInput = document.getElementById("fileInput");
+                const res = document.getElementById("resInput").value || 512;
                 if (!fileInput.files.length) return alert("Please select a file!");
                 const formData = new FormData();
                 formData.append("file", fileInput.files[0]);
-                const res = await fetch("/remove-background", { method: "POST", body: formData });
-                const blob = await res.blob();
+                formData.append("resolution", res);
+                const response = await fetch("/remove-background", { method: "POST", body: formData });
+                const blob = await response.blob();
                 resultImg.src = URL.createObjectURL(blob);
-                downloadLink.href = resultImg.src;
-                downloadLink.style.display = "inline-block";
             });
 
             urlForm.addEventListener("submit", async e => {
                 e.preventDefault();
                 const url = document.getElementById("urlInput").value.trim();
+                const res = document.getElementById("urlResInput").value || 512;
                 if (!url) return alert("Please enter an image URL!");
                 const formData = new FormData();
                 formData.append("image_url", url);
-                const res = await fetch("/remove-background", { method: "POST", body: formData });
-                const blob = await res.blob();
+                formData.append("resolution", res);
+                const response = await fetch("/remove-background", { method: "POST", body: formData });
+                const blob = await response.blob();
                 resultImg.src = URL.createObjectURL(blob);
-                downloadLink.href = resultImg.src;
-                downloadLink.style.display = "inline-block";
             });
         </script>
     </body>
@@ -192,7 +192,7 @@ async def index():
     return HTMLResponse(html)
 
 # -------------------------
-# Run App on Spaces
+# Run App
 # -------------------------
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)