Update app.py

app.py

@@ -6,22 +6,23 @@ from PIL import Image
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
 from typing import Union, List
-from loadimg import load_img  # Your helper to load from URL or file
+from loadimg import load_img
 
 torch.set_float32_matmul_precision("high")
 
-# Load BiRefNet model
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
+# Load RMBG v1.4 model
+model = AutoModelForImageSegmentation.from_pretrained(
+    "briaai/RMBG-1.4", 
+    trust_remote_code=True
 )
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-birefnet.to(device)
+model.to(device)
 
-# Image transformation
+# RMBG v1.4 uses different preprocessing
 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
@@ -30,22 +31,33 @@ def process(image: Image.Image) -> Image.Image:
     input_tensor = transform_image(image).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+        # RMBG v1.4 returns the mask directly
+        preds = model(input_tensor)
+        # Get the mask - RMBG returns different structure than BiRefNet
+        if isinstance(preds, list):
+            pred = preds[-1]
+        else:
+            pred = preds
+        
+        pred = pred.sigmoid().cpu()
 
-    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)
+    mask = pred[0].squeeze()
+    mask_pil = transforms.ToPILImage()(mask).resize(image_size).convert("L")
+    
+    # Create binary mask
+    binary_mask = mask_pil.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
 
+# Rest of your code remains the same...
 @spaces.GPU
 def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str], None]:
     results = []
 
     try:
-        # Single image upload
         if image is not None:
             image = image.convert("RGB")
             processed = process(image)
@@ -53,7 +65,6 @@ 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)
@@ -61,7 +72,6 @@ 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:
@@ -80,7 +90,6 @@ def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str]
 
     return None
 
-# Interface
 demo = gr.Interface(
     fn=handler,
     inputs=[
@@ -89,9 +98,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 (White Fill)",
+    title="Background Remover (RMBG v1.4)",
     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)