Update Gradio app with multiple files

app.py

@@ -2,25 +2,25 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 import torch
-from transformers import AutoImageProcessor, U2NetForImageSegmentation
-from rembg import remove, new_session
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
 import spaces
 import os
 
-# Load models
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Load RMBG-2.0 model
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True).eval().to(device)
 
-# Load U2Net model for background removal
-u2net_processor = AutoImageProcessor.from_pretrained("briaai/RMBG-1.4")
-u2net_model = U2NetForImageSegmentation.from_pretrained("briaai/RMBG-1.4")
-if device == "cuda":
-    u2net_model = u2net_model.to(device)
-
-# Initialize rembg session
-rembg_session = new_session('u2net')
+# Data settings
+image_size = (1024, 1024)
+transform_image = transforms.Compose([
+    transforms.Resize(image_size),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
 
 @spaces.GPU(duration=30)
-def remove_background_u2net(image):
+def remove_background(image):
     """
     Remove background from image using RMBG-2.0 model.
     
@@ -49,53 +49,6 @@ def remove_background_u2net(image):
     
     return image_rgba
 
-@spaces.GPU(duration=20)
-def remove_background_rembg(image):
-    """
-    Remove background from image using rembg library.
-    
-    Args:
-        image (PIL.Image): Input image to process
-    
-    Returns:
-        PIL.Image: Image with background removed
-    """
-    if image is None:
-        return None
-    
-    # Convert PIL to bytes for rembg
-    import io
-    img_byte_arr = io.BytesIO()
-    image.save(img_byte_arr, format='PNG')
-    img_bytes = img_byte_arr.getvalue()
-    
-    # Remove background
-    output_bytes = remove(img_bytes, session=rembg_session)
-    
-    # Convert back to PIL
-    output_image = Image.open(io.BytesIO(output_bytes))
-    
-    return output_rgba
-
-def process_image(image, method):
-    """
-    Process image background removal based on selected method.
-    
-    Args:
-        image (PIL.Image): Input image
-        method (str): Background removal method ('u2net' or 'rembg')
-    
-    Returns:
-        PIL.Image: Processed image
-    """
-    if image is None:
-        return None
-    
-    if method == "u2net":
-        return remove_background_u2net(image)
-    else:
-        return remove_background_rembg(image)
-
 def create_collage(original, processed):
     """
     Create a side-by-side comparison of original and processed images.
@@ -152,7 +105,7 @@ with gr.Blocks(title="Background Removal App", theme=gr.themes.Soft()) as demo:
         # Background Removal App
         Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)
         
-        Upload an image to remove its background using AI models. Choose between two advanced background removal methods.
+        Upload an image to remove its background using the advanced RMBG-2.0 AI model from BRIA AI.
         """
     )
     
@@ -165,16 +118,6 @@ with gr.Blocks(title="Background Removal App", theme=gr.themes.Soft()) as demo:
                 sources=["upload", "webcam", "clipboard"]
             )
             
-            method = gr.Radio(
-                choices=[
-                    ("RMBG-2.0 (BRIA AI)", "u2net"),
-                    ("RMBG (Rembg)", "rembg")
-                ],
-                value="u2net",
-                label="Background Removal Method",
-                info="Choose the AI model for background removal"
-            )
-            
             process_btn = gr.Button("Remove Background", variant="primary", size="lg")
             
             with gr.Accordion("Advanced Options", open=False):
@@ -205,19 +148,19 @@ with gr.Blocks(title="Background Removal App", theme=gr.themes.Soft()) as demo:
     # Example images
     gr.Examples(
         examples=[
-            ["https://gradio-builds.s3.amazonaws.com/assets/cheetah-003.jpg", "u2net"],
-            ["https://gradio-builds.s3.amazonaws.com/assets/TheCheethcat.jpg", "rembg"],
+            ["https://gradio-builds.s3.amazonaws.com/assets/cheetah-003.jpg"],
+            ["https://gradio-builds.s3.amazonaws.com/assets/TheCheethcat.jpg"],
         ],
-        inputs=[input_image, method],
+        inputs=input_image,
         outputs=output_image,
-        fn=process_image,
+        fn=remove_background,
         cache_examples=True
     )
     
     # Event handlers
     process_btn.click(
-        fn=process_image,
-        inputs=[input_image, method],
+        fn=remove_background,
+        inputs=input_image,
         outputs=output_image,
         show_progress=True
     ).then(
@@ -238,13 +181,12 @@ with gr.Blocks(title="Background Removal App", theme=gr.themes.Soft()) as demo:
     )
 
 # MCP Server Functions
-def remove_background_mcp(image_path: str, method: str = "u2net") -> str:
+def remove_background_mcp(image_path: str) -> str:
     """
     Remove background from an image file and save the result.
     
     Args:
         image_path (str): Path to the input image file
-        method (str): Background removal method ('u2net' or 'rembg')
     
     Returns:
         str: Path to the output image file with background removed
@@ -253,11 +195,8 @@ def remove_background_mcp(image_path: str, method: str = "u2net") -> str:
         # Load image
         image = Image.open(image_path)
         
-        # Process based on method
-        if method == "u2net":
-            result = remove_background_u2net(image)
-        else:
-            result = remove_background_rembg(image)
+        # Process image
+        result = remove_background(image)
         
         # Save result
         output_path = image_path.replace('.', '_no_bg.')
@@ -267,13 +206,12 @@ def remove_background_mcp(image_path: str, method: str = "u2net") -> str:
     except Exception as e:
         raise Exception(f"Error processing image: {str(e)}")
 
-def remove_background_base64(image_data: str, method: str = "u2net") -> str:
+def remove_background_base64(image_data: str) -> str:
     """
     Remove background from base64 encoded image data.
     
     Args:
         image_data (str): Base64 encoded image data
-        method (str): Background removal method ('u2net' or 'rembg')
     
     Returns:
         str: Base64 encoded image with background removed
@@ -286,11 +224,8 @@ def remove_background_base64(image_data: str, method: str = "u2net") -> str:
         image_bytes = base64.b64decode(image_data)
         image = Image.open(io.BytesIO(image_bytes))
         
-        # Process based on method
-        if method == "u2net":
-            result = remove_background_u2net(image)
-        else:
-            result = remove_background_rembg(image)
+        # Process image
+        result = remove_background(image)
         
         # Encode result back to base64
         output_buffer = io.BytesIO()

requirements.txt

@@ -1,7 +1,7 @@
 gradio[mcp]
 torch
 transformers
-rembg
+torchvision
 Pillow
 numpy
 spaces