Update app.py

app.py

@@ -1,36 +1,43 @@
 import gradio as gr
 from PIL import Image
-from rembg import remove
+from rembg import remove, new_session
 import io
 
-# Lightweight background removal (works on CPU)
+# Initialize session once (saves memory)
+session = new_session("u2net")  # Smallest working model
+
 def remove_background(input_image):
     try:
         # Convert Gradio input to PIL Image
-        if isinstance(input_image, dict):  # Handle paste/drop events
+        if isinstance(input_image, dict):  # Handle paste/drop
             img = Image.open(io.BytesIO(input_image["bytes"]))
         else:
             img = Image.fromarray(input_image)
         
-        # Process with optimized settings
+        # Auto-resize to prevent OOM errors
+        if max(img.size) > 1024:
+            img = img.resize((1024, 1024))
+        
+        # Process with memory-efficient settings
         result = remove(
             img,
-            session=new_session("u2net"),  # Smallest working model
+            session=session,
             alpha_matting=False,  # Disable memory-heavy feature
             only_mask=False
         )
-        
         return result
     except Exception as e:
         print(f"Error: {str(e)}")
         return input_image  # Return original if fails
 
-# Simple interface
+# Updated Gradio interface (compatible with latest version)
 with gr.Blocks(title="Free BG Remover") as demo:
-    gr.Markdown("""### 🆓 Free Background Remover (Works on CPU)""")
+    gr.Markdown("## 🆓 Free Background Remover (Works on CPU)")
+    
     with gr.Row():
-        gr.Image(type="pil", label="Upload").style(height=400) >> input_img
-        gr.Image(type="pil", label="Result").style(height=400) << output_img
+        input_img = gr.Image(label="Upload", type="pil", height=400)
+        output_img = gr.Image(label="Result", type="pil", height=400)
+    
     gr.Button("Remove Background").click(
         remove_background,
         inputs=input_img,