Update app.py

app.py

@@ -1,99 +1,40 @@
 import gradio as gr
-import numpy as np
 from PIL import Image
-import torch
-import warnings
+from rembg import remove
+import io
 
-# Suppress warnings for cleaner output
-warnings.filterwarnings("ignore")
-
-# Initialize models dictionary to cache loaded models
-models = {}
-
-def load_bria_model():
-    from transformers import AutoModelForImageSegmentation, AutoImageProcessor
-    model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)
-    processor = AutoImageProcessor.from_pretrained("briaai/RMBG-1.4")
-    return model, processor
-
-def load_rembg_model(model_name):
-    from rembg import new_session
-    return new_session(model_name)
-
-def load_isnet_model(model_url):
-    # Placeholder - you would implement proper ISNet loading here
-    return None
-
-def apply_bria(image, model, processor):
-    inputs = processor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    mask = outputs.logits.squeeze().cpu().numpy()
-    mask = (mask - mask.min()) / (mask.max() - mask.min())
-    mask = (mask * 255).astype(np.uint8)
-    return Image.fromarray(mask)
-
-def apply_rembg(image, session):
-    from rembg import remove
-    return remove(image, session=session)
-
-def apply_isnet(image, model):
-    # Placeholder for ISNet implementation
-    return image
-
-def remove_background(image):
+# Lightweight background removal (works on CPU)
+def remove_background(input_image):
     try:
-        # Convert input to PIL Image
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-        
-        # Initialize models if not already loaded
-        if "bria" not in models:
-            bria_model, bria_processor = load_bria_model()
-            models["bria"] = (bria_model, bria_processor)
-        
-        if "u2net" not in models:
-            models["u2net"] = load_rembg_model("u2net")
-        
-        if "isnet" not in models:
-            models["isnet"] = load_isnet_model("https://example.com/isnet.pth")
-        
-        # Apply models in sequence
-        results = []
-        
-        # BRIA
-        bria_model, bria_processor = models["bria"]
-        bria_result = apply_bria(image, bria_model, bria_processor)
-        results.append(bria_result)
-        
-        # U2Net
-        u2net_result = apply_rembg(image, models["u2net"])
-        results.append(u2net_result)
-        
-        # Combine results (simple average for demonstration)
-        combined = np.zeros_like(np.array(results[0]), dtype=np.float32)
-        for res in results:
-            combined += np.array(res).astype(np.float32) / len(results)
-        combined = np.clip(combined, 0, 255).astype(np.uint8)
-        
-        # Apply mask to original image
-        final = image.copy()
-        final.putalpha(Image.fromarray(combined))
-        
-        return final
-    
+        # Convert Gradio input to PIL Image
+        if isinstance(input_image, dict):  # Handle paste/drop events
+            img = Image.open(io.BytesIO(input_image["bytes"]))
+        else:
+            img = Image.fromarray(input_image)
+        
+        # Process with optimized settings
+        result = remove(
+            img,
+            session=new_session("u2net"),  # Smallest working model
+            alpha_matting=False,  # Disable memory-heavy feature
+            only_mask=False
+        )
+        
+        return result
     except Exception as e:
-        print(f"Error: {e}")
-        return image  # Return original image on error
-
-# Create Gradio interface
-interface = gr.Interface(
-    fn=remove_background,
-    inputs=gr.Image(label="Input Image"),
-    outputs=gr.Image(label="Result with Transparent Background"),
-    title="Advanced Background Removal",
-    description="Combines multiple state-of-the-art models for high-quality background removal"
-)
-
-if __name__ == "__main__":
-    interface.launch(share=True)  # Set share=True for public link
+        print(f"Error: {str(e)}")
+        return input_image  # Return original if fails
+
+# Simple interface
+with gr.Blocks(title="Free BG Remover") as demo:
+    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
+    gr.Button("Remove Background").click(
+        remove_background,
+        inputs=input_img,
+        outputs=output_img
+    )
+
+demo.launch(show_error=True)