Update app.py

app.py

@@ -18,7 +18,7 @@ model = AutoModelForImageSegmentation.from_pretrained(
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# RMBG v1.4 uses different preprocessing
+# Transform for RMBG v1.4
 transform_image = transforms.Compose([
     transforms.Resize((1024, 1024)),
     transforms.ToTensor(),
@@ -31,20 +31,25 @@ def process(image: Image.Image) -> Image.Image:
     input_tensor = transform_image(image).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        # RMBG v1.4 returns the mask directly
+        # RMBG v1.4 returns a tuple, we need the first element
         preds = model(input_tensor)
-        # Get the mask - RMBG returns different structure than BiRefNet
-        if isinstance(preds, list):
-            pred = preds[-1]
+        
+        # Handle different return types
+        if isinstance(preds, tuple):
+            pred = preds[0]  # Take first element if tuple
+        elif isinstance(preds, list):
+            pred = preds[-1]  # Take last element if list
         else:
             pred = preds
-        
-        pred = pred.sigmoid().cpu()
+            
+        # Apply sigmoid and move to CPU
+        mask = pred.sigmoid().cpu()
 
-    mask = pred[0].squeeze()
-    mask_pil = transforms.ToPILImage()(mask).resize(image_size).convert("L")
+    # Process the mask
+    mask_tensor = mask[0].squeeze()
+    mask_pil = transforms.ToPILImage()(mask_tensor).resize(image_size).convert("L")
     
-    # Create binary mask
+    # Create binary mask with threshold
     binary_mask = mask_pil.point(lambda p: 255 if p > 127 else 0)
 
     # Apply mask with white background
@@ -52,7 +57,6 @@ def process(image: Image.Image) -> Image.Image:
     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 = []
@@ -87,6 +91,9 @@ def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str]
 
     except Exception as e:
         print("General error:", e)
+        # Add debug info
+        import traceback
+        traceback.print_exc()
 
     return None