Update app.py

app.py

@@ -7,42 +7,38 @@ import io
 
 def remove_background(input_image):
     try:
-        # Initialize the pipeline with trust_remote_code=True
+        # Initialize the pipeline with correct parameters
         segmentor = pipeline(
-            "image-segmentation", 
+            task="image-segmentation",
             model="briaai/RMBG-1.4",
-            trust_remote_code=True,
-            device="cpu"
+            trust_remote_code=True
         )
         
+        # Convert input to PIL Image if it's not already
+        if not isinstance(input_image, Image.Image):
+            input_image = Image.fromarray(input_image)
+            
         # Process the image
-        result = segmentor(input_image)
-        return result['output_image']
+        result = segmentor(input_image, return_mask=True)
+        
+        # Create transparent background image
+        output_image = Image.new('RGBA', input_image.size, (0, 0, 0, 0))
+        output_image.paste(input_image, mask=result)
+        
+        return output_image
+
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")
 
-# Custom theme with valid colors
-theme = gr.themes.Soft(
-    primary_hue="orange",    # Changed from "gold" to "orange"
-    secondary_hue="blue",
-    neutral_hue="gray"
-).set(
-    body_background_fill="linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 100%)",
-    body_text_color="#ffffff",
-    button_primary_background_fill="linear-gradient(45deg, #FFD700, #FFA500)",
-    button_primary_text_color="#000000",
-    border_color_primary="#FFD700"
-)
-
 # Create Gradio interface
-with gr.Blocks(theme=theme) as demo:
+with gr.Blocks(theme=gr.themes.Default()) as demo:
     gr.HTML(
         """
         <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
-            <h1 style="font-size: 2.5rem; margin-bottom: 1rem; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
+            <h1 style="font-size: 2.5rem; margin-bottom: 1rem;">
                 AI Background Remover
             </h1>
-            <p style="color: #cccccc; font-size: 1.2rem; margin-bottom: 2rem;">
+            <p style="color: #666; font-size: 1.1rem;">
                 Remove backgrounds instantly using RMBG V1.4 model
             </p>
         </div>