Update app.py

app.py

@@ -5,13 +5,13 @@ from PIL import Image
 from torchvision import transforms
 import gradio as gr
 
-# Load the model from Hugging Face
+# Load model
 birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet_lite', trust_remote_code=True)
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 birefnet.to(device)
 birefnet.eval()
 
-# Define the transform to preprocess the input images
+# Preprocessing
 image_size = (1024, 1024)
 transform_image = transforms.Compose([
     transforms.Resize(image_size),
@@ -19,63 +19,72 @@ transform_image = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-def extract_objects(filepaths):
-    # Open all images from the uploaded file paths
-    images = [Image.open(path).convert("RGB") for path in filepaths]
-
-    start_time = time.time()
+def process_batch(img_batch):
     inputs = []
     original_sizes = []
-    for img in images:
+    for img in img_batch:
         original_sizes.append(img.size)
         inputs.append(transform_image(img))
     input_tensor = torch.stack(inputs).to(device)
 
-    # Inference
-    inf_start = time.time()
-    with torch.no_grad():
-        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
-    inf_end = time.time()
+    try:
+        with torch.no_grad():
+            preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+    except torch.OutOfMemoryError:
+        torch.cuda.empty_cache()
+        return None
 
-    # Post-process results
     results = []
-    image_times = []
-    for i, img in enumerate(images):
-        t_start = time.time()
+    for i, img in enumerate(img_batch):
         pred = preds[i].squeeze()
         pred_pil = transforms.ToPILImage()(pred)
         mask = pred_pil.resize(original_sizes[i])
-        
-        # Create a transparent background image
         result = Image.new("RGBA", original_sizes[i], (0, 0, 0, 0))
         result.paste(img, mask=mask)
         results.append(result)
-        t_end = time.time()
-        image_times.append(t_end - t_start)
 
-    end_time = time.time()
-    total_time = end_time - start_time
-    inference_time = inf_end - inf_start
-    prep_post_time = total_time - inference_time
+    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)]
+
+    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")
+        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()
 
-    # Create a summary of timings
     summary = (
-        f"Total request time: {total_time:.2f} s\n"
-        f"Inference time (batch): {inference_time:.2f} s\n"
-        f"Pre/Post-processing time: {prep_post_time:.2f} s\n"
-        "Per-image post-processing times:\n" +
-        "\n".join([f"  Image {i+1}: {t:.2f} s" for i, t in enumerate(image_times)])
+        f"Total request time: {total_end - total_start:.2f} s\n"
+        "Batch times:\n" + "\n".join(batch_times)
     )
 
-    return results, summary
+    return all_results, 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",
-    description="Upload multiple images and process them in one request. Timing information for the full request and per-image processing is provided."
+    title="BiRefNet Bulk Background Removal with Queue & Batch",
+    description="Upload multiple images. The request is queued and processed in batches to avoid OOM errors."
 )
 
-if __name__ == "__main__":
-    iface.launch()
+# 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()