Update app.py

app.py

@@ -1,13 +1,28 @@
 import time
 import torch
-from transformers import AutoModelForImageSegmentation
+import gc
+from transformers import AutoConfig, AutoModelForImageSegmentation
 from PIL import Image
 from torchvision import transforms
 import gradio as gr
-import gc
 
 def load_model():
-    model = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet_lite', trust_remote_code=True)
+    # Fetch the config first (with trust_remote_code=True)
+    config = AutoConfig.from_pretrained("zhengpeng7/BiRefNet_lite", trust_remote_code=True)
+    
+    # Ensure it's not treated as a seq2seq model
+    config.is_encoder_decoder = False
+
+    # Optionally, block calls to get_text_config if needed:
+    # config.get_text_config = lambda decoder=True: None
+
+    # Now load the model with our tweaked config
+    model = AutoModelForImageSegmentation.from_pretrained(
+        "zhengpeng7/BiRefNet_lite",
+        config=config,
+        trust_remote_code=True
+    )
+
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     model.to(device)
     model.eval()
@@ -30,13 +45,17 @@ def run_inference(images, model, device):
         original_sizes.append(img.size)
         inputs.append(transform_image(img))
     input_tensor = torch.stack(inputs).to(device)
+
     try:
         with torch.no_grad():
+            # If the last layer is returned as [-1], 
+            # adjust accordingly or see how your model outputs are structured
             preds = model(input_tensor)[-1].sigmoid().cpu()
     except torch.OutOfMemoryError:
         del input_tensor
         torch.cuda.empty_cache()
         raise
+
     # Post-process
     results = []
     for i, img in enumerate(images):
@@ -46,6 +65,7 @@ def run_inference(images, model, device):
         result = Image.new("RGBA", original_sizes[i], (0, 0, 0, 0))
         result.paste(img, mask=mask)
         results.append(result)
+
     # Cleanup
     del input_tensor, preds
     gc.collect()
@@ -61,9 +81,8 @@ def binary_search_max(images):
         mid = (low + high) // 2
         batch = images[:mid]
         try:
-            # Re-load model to avoid leftover memory fragmentation
             global birefnet, device
-            birefnet, device = load_model()
+            birefnet, device = load_model()  # re-init to reduce memory fragmentation
             res = run_inference(batch, birefnet, device)
             best = res
             best_count = mid
@@ -84,7 +103,7 @@ def extract_objects(filepaths):
         summary = f"Total request time: {total_time:.2f}s\nProcessed {len(images)} images successfully."
         return results, summary
     except torch.OutOfMemoryError:
-        # OOM occurred, try to find feasible batch size now
+        # OOM occurred, try fallback
         oom_time = time.time()
         initial_attempt_time = oom_time - start_time
         
@@ -114,7 +133,8 @@ iface = gr.Interface(
     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 On-Demand Fallback",
-    description="Upload as many images as you want. If OOM occurs, a fallback will find the max feasible number. Extra cleanup steps and reinitialization for more consistent results."
+    description="Upload as many images as you want. If OOM occurs, fallback logic will find the max feasible number."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()