Update app.py

app.py

@@ -19,8 +19,8 @@ transform_image = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-def try_inference(images):
-    # Convert images to tensors
+def run_inference(images):
+    # Convert all images into a batch tensor
     inputs = []
     original_sizes = []
     for img in images:
@@ -28,14 +28,9 @@ def try_inference(images):
         inputs.append(transform_image(img))
     input_tensor = torch.stack(inputs).to(device)
 
-    # Attempt inference
-    try:
-        with torch.no_grad():
-            preds = birefnet(input_tensor)[-1].sigmoid().cpu()
-    except torch.OutOfMemoryError:
-        # Clear CUDA cache and return None to indicate OOM
-        torch.cuda.empty_cache()
-        return None
+    # Run inference
+    with torch.no_grad():
+        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
 
     # Post-process
     results = []
@@ -49,55 +44,66 @@ def try_inference(images):
     return results
 
 def extract_objects(filepaths):
-    # Open all images
     images = [Image.open(p).convert("RGB") for p in filepaths]
+    start_time = time.time()
 
-    total_start = time.time()
-    # If you have N images, start by trying them all
-    low = 1
-    high = len(images)
-    best = None
-    best_batch_size = 0
-
-    # Binary search to find max batch size that doesn't OOM
-    while low <= high:
-        mid = (low + high) // 2
-        batch_test = images[:mid]
+    # Attempt to process all at once
+    try:
+        results = run_inference(images)
+        end_time = time.time()
+        total_time = end_time - start_time
+        summary = f"Total request time: {total_time:.2f}s\nProcessed {len(images)} images successfully."
+        return results, summary
+    except torch.OutOfMemoryError:
+        # Only if we fail, do we attempt to find a feasible batch size
+        torch.cuda.empty_cache()
 
-        start = time.time()
-        results = try_inference(batch_test)
-        end = time.time()
+        fail_time = time.time()
+        initial_attempt_time = fail_time - start_time
 
-        if results is not None:
-            # Succeeded with 'mid' batch size
-            best = results
-            best_batch_size = mid
-            low = mid + 1  # try a bigger batch
-        else:
-            # OOM, try smaller batch
-            high = mid - 1
+        # Binary search to find max feasible batch size
+        low, high = 1, len(images)
+        best = None
+        best_count = 0
+        
+        while low <= high:
+            mid = (low + high) // 2
+            batch = images[:mid]
+            try:
+                res = run_inference(batch)
+                best = res
+                best_count = mid
+                low = mid + 1  # try bigger
+            except torch.OutOfMemoryError:
+                torch.cuda.empty_cache()
+                high = mid - 1  # try smaller
 
-    total_end = time.time()
+        end_time = time.time()
+        total_time = end_time - start_time
 
-    if best is None:
-        # Even a single image caused OOM
-        summary = "Could not process even a single image without OOM."
-        return [], summary
-    else:
-        # Process the final chosen batch size fully (we already have results)
-        summary = (
-            f"Total request time: {total_end - total_start:.2f} s\n"
-            f"Successfully processed {best_batch_size} images in a single batch.\n"
-            f"Could not handle more than {best_batch_size} images without OOM."
-        )
-        return best, summary
+        if best is None:
+            # Not even 1 image works
+            summary = (
+                f"Initial attempt OOM after {initial_attempt_time:.2f}s.\n"
+                f"Could not process even a single image.\n"
+                f"Total time with fallback attempts: {total_time:.2f}s."
+            )
+            return [], summary
+        else:
+            summary = (
+                f"Initial attempt OOM after {initial_attempt_time:.2f}s. "
+                f"After fallback tests, found that {best_count} images can be processed.\n"
+                f"Total time including fallback: {total_time:.2f}s.\n"
+                f"Next time, try using up to {best_count} images."
+            )
+            return best, summary
 
 iface = gr.Interface(
     fn=extract_objects,
     inputs=gr.Files(label="Upload Multiple Images", type="filepath", file_count="multiple"),
     outputs=[gr.Gallery(label="Processed Images"), gr.Textbox(label="Timing Info")],
-    title="BiRefNet Dynamic Batch OOM Test",
-    description="Upload images. The system will try to process all at once, and if OOM occurs, it will try smaller batches automatically, quickly finding the largest feasible batch size."
+    title="BiRefNet Bulk Background Removal with On-Demand Fallback",
+    description="Upload as many images as you want. If OOM occurs, a quick fallback will find the max feasible number of images without adding overhead unless needed."
 )
 
 iface.launch()