Update app.py

app.py

@@ -6,23 +6,22 @@ from PIL import Image
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
 from typing import Union, List
-from loadimg import load_img
+from loadimg import load_img  # Your helper to load from URL or file
 
 torch.set_float32_matmul_precision("high")
 
-# Load RMBG v1.4 model
-model = AutoModelForImageSegmentation.from_pretrained(
-    "briaai/RMBG-1.4", 
-    trust_remote_code=True
+# Load BiRefNet model
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+birefnet.to(device)
 
-# Transform for RMBG v1.4
+# Image transformation
 transform_image = transforms.Compose([
     transforms.Resize((1024, 1024)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
 @spaces.GPU
@@ -31,28 +30,12 @@ def process(image: Image.Image) -> Image.Image:
     input_tensor = transform_image(image).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        preds = model(input_tensor)
-        
-        # Handle list output - extract the tensor from the list
-        if isinstance(preds, list):
-            # Usually the mask is the last or first element
-            pred = preds[-1] if len(preds) > 0 else preds[0]
-        elif isinstance(preds, tuple):
-            pred = preds[0]
-        else:
-            pred = preds
-            
-        # Now apply sigmoid to the tensor
-        mask = pred.sigmoid().cpu()
+        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
 
-    # Process the mask
-    mask_tensor = mask[0].squeeze()
-    mask_pil = transforms.ToPILImage()(mask_tensor).resize(image_size).convert("L")
-    
-    # Create binary mask with threshold
-    binary_mask = mask_pil.point(lambda p: 255 if p > 127 else 0)
+    pred = preds[0].squeeze()
+    mask = transforms.ToPILImage()(pred).resize(image_size).convert("L")
+    binary_mask = mask.point(lambda p: 255 if p > 127 else 0)
 
-    # Apply mask with white background
     white_bg = Image.new("RGB", image_size, (255, 255, 255))
     result = Image.composite(image, white_bg, binary_mask)
     return result
@@ -62,6 +45,7 @@ def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str]
     results = []
 
     try:
+        # Single image upload
         if image is not None:
             image = image.convert("RGB")
             processed = process(image)
@@ -69,6 +53,7 @@ def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str]
             processed.save(filename)
             return filename
 
+        # Single image from URL
         if image_url:
             im = load_img(image_url, output_type="pil").convert("RGB")
             processed = process(im)
@@ -76,6 +61,7 @@ def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str]
             processed.save(filename)
             return filename
 
+        # Batch of URLs
         if batch_urls:
             urls = [u.strip() for u in batch_urls.split(",") if u.strip()]
             for url in urls:
@@ -91,11 +77,10 @@ def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str]
 
     except Exception as e:
         print("General error:", e)
-        import traceback
-        traceback.print_exc()
 
     return None
 
+# Interface
 demo = gr.Interface(
     fn=handler,
     inputs=[
@@ -104,9 +89,9 @@ demo = gr.Interface(
         gr.Textbox(label="Comma-separated Image URLs (Batch)"),
     ],
     outputs=gr.File(label="Output File(s)", file_count="multiple"),
-    title="Background Remover (RMBG v1.4)",
+    title="Background Remover (White Fill)",
     description="Upload an image, paste a URL, or send a batch of URLs to remove the background and replace it with white.",
 )
 
 if __name__ == "__main__":
-    demo.launch(show_error=True, mcp_server=True)
+    demo.launch(show_error=True, mcp_server=True)