Update app.py

app.py

@@ -31,18 +31,18 @@ def process(image: Image.Image) -> Image.Image:
     input_tensor = transform_image(image).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        # RMBG v1.4 returns a tuple, we need the first element
         preds = model(input_tensor)
         
-        # 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
+        # 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
             
-        # Apply sigmoid and move to CPU
+        # Now apply sigmoid to the tensor
         mask = pred.sigmoid().cpu()
 
     # Process the mask
@@ -91,7 +91,6 @@ 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()