Update app.py

app.py

@@ -5,7 +5,7 @@ from PIL import Image
 from torchvision import transforms
 import gradio as gr
 
-# Load model
+# Load the model
 birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet_lite', trust_remote_code=True)
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 birefnet.to(device)
@@ -19,72 +19,85 @@ transform_image = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-def process_batch(img_batch):
+def try_inference(images):
+    # Convert images to tensors
     inputs = []
     original_sizes = []
-    for img in img_batch:
+    for img in images:
         original_sizes.append(img.size)
         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
 
+    # Post-process
     results = []
-    for i, img in enumerate(img_batch):
+    for i, img in enumerate(images):
         pred = preds[i].squeeze()
         pred_pil = transforms.ToPILImage()(pred)
         mask = pred_pil.resize(original_sizes[i])
         result = Image.new("RGBA", original_sizes[i], (0, 0, 0, 0))
         result.paste(img, mask=mask)
         results.append(result)
-
     return results
 
 def extract_objects(filepaths):
-    # Open all images from the uploaded file paths
-    images = [Image.open(path).convert("RGB") for path in filepaths]
-
-    # You can define a batch size here (e.g., batch_size = 5)
-    # This prevents trying to process all images at once if too large
-    batch_size = 5
-    batches = [images[i:i+batch_size] for i in range(0, len(images), batch_size)]
+    # Open all images
+    images = [Image.open(p).convert("RGB") for p in filepaths]
 
     total_start = time.time()
-    all_results = []
-    batch_times = []
-    for b_idx, batch in enumerate(batches):
-        b_start = time.time()
-        res = process_batch(batch)
-        if res is None:
-            # Handle OOM gracefully
-            all_results.extend([Image.new("RGBA", (100, 100), (255,0,0,255)) for _ in batch])
-            batch_times.append(f"Batch {b_idx+1}: OOM Error")
+    # 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]
+
+        start = time.time()
+        results = try_inference(batch_test)
+        end = time.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:
-            all_results.extend(res)
-            b_end = time.time()
-            batch_times.append(f"Batch {b_idx+1}: {(b_end - b_start):.2f} s")
-    total_end = time.time()
+            # OOM, try smaller batch
+            high = mid - 1
 
-    summary = (
-        f"Total request time: {total_end - total_start:.2f} s\n"
-        "Batch times:\n" + "\n".join(batch_times)
-    )
+    total_end = time.time()
 
-    return all_results, summary
+    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
 
 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 Bulk Background Removal with Queue & Batch",
-    description="Upload multiple images. The request is queued and processed in batches to avoid OOM errors."
+    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."
 )
 
-# Enable the queue with defined concurrency to prevent multiple large requests at once
-# You can adjust concurrency_count and max_size as needed.
-iface.queue(concurrency_count=1, max_size=10).launch()
+iface.launch()