Update app.py

app.py

@@ -1,114 +1,22 @@
-import os
-import torch
 import gradio as gr
-from gradio_imageslider import ImageSlider
-# Ensure loadimg.py is in your directory. 
-# If not, we use a fallback to PIL.
-try:
-    from loadimg import load_img
-except ImportError:
-    def load_img(path, output_type="pil"):
-        from PIL import Image
-        import requests
-        from io import BytesIO
-        if path.startswith('http'):
-            response = requests.get(path)
-            img = Image.open(BytesIO(response.content))
-        else:
-            img = Image.open(path)
-        return img
-
-from transformers import AutoModelForImageSegmentation
-from torchvision import transforms
-
-# --- Hardware Setup ---
-# Force CPU if CUDA is not available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Current Hardware: {device.upper()}")
-
-# --- Model Initialization ---
-# RMBG-2.0 is heavy; we use trust_remote_code=True for the BiRefNet architecture
-print("Loading model... this may take a minute on CPU.")
-model = AutoModelForImageSegmentation.from_pretrained(
-    "briaai/RMBG-2.0", 
-    trust_remote_code=True
+from rembg import remove
+from PIL import Image
+import io
+
+def fast_process(input_image):
+    if input_image is None:
+        return None
+    # Rembg handles the ONNX optimization internally
+    output = remove(input_image)
+    return output
+
+demo = gr.Interface(
+    fn=fast_process,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Ultra-Fast CPU Background Remover",
+    description="Using rembg + ONNX for 1-2 second processing on most CPUs."
 )
-model.to(device)
-model.eval()
-
-# Standard ImageNet normalization used by BiRefNet
-preprocess = transforms.Compose([
-    transforms.Resize((1024, 1024)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-])
-
-output_folder = 'output_images'
-os.makedirs(output_folder, exist_ok=True)
-
-def remove_background(image_input):
-    if image_input is None:
-        return None, None
-    
-    # Handle both filepath (str) and PIL Image inputs
-    if isinstance(image_input, str):
-        orig_img = load_img(image_input, output_type="pil").convert("RGB")
-    else:
-        orig_img = image_input.convert("RGB")
-    
-    w, h = orig_img.size
-    
-    # Preprocess
-    input_tensor = preprocess(orig_img).unsqueeze(0).to(device)
-    
-    # Inference
-    with torch.no_grad():
-        # BiRefNet returns a list of preds; we take the last one
-        result = model(input_tensor)[-1].sigmoid().cpu()
-    
-    # Post-process Mask
-    mask = transforms.ToPILImage()(result[0].squeeze())
-    mask = mask.resize((w, h))
-    
-    # Create Final Transparent Image
-    no_bg_img = orig_img.copy()
-    no_bg_img.putalpha(mask)
-    
-    save_path = os.path.join(output_folder, "result.png")
-    no_bg_img.save(save_path)
-    
-    return (no_bg_img, orig_img), save_path
-
-# --- UI Setup ---
-with gr.Blocks(title="RMBG 2.0 CPU") as demo:
-    gr.Markdown("# RMBG-2.0 Background Remover")
-    gr.Markdown("Optimized for CPU/GPU deployment.")
-    
-    with gr.Tab("Image Upload"):
-        with gr.Row():
-            in_img = gr.Image(label="Upload Image", type="pil")
-            out_slider = ImageSlider(label="Comparison", type="pil")
-        
-        out_file = gr.File(label="Download PNG")
-        submit_btn = gr.Button("Remove Background", variant="primary")
-        
-        submit_btn.click(
-            fn=remove_background, 
-            inputs=in_img, 
-            outputs=[out_slider, out_file]
-        )
-
-    with gr.Tab("URL / Batch"):
-        url_input = gr.Textbox(label="Paste Image URL")
-        url_slider = ImageSlider(label="Comparison", type="pil")
-        url_file = gr.File(label="Download PNG")
-        url_btn = gr.Button("Process URL")
-        
-        url_btn.click(
-            fn=remove_background,
-            inputs=url_input,
-            outputs=[url_slider, url_file]
-        )
 
 if __name__ == "__main__":
     demo.launch()