Update app.py

app.py

@@ -8,121 +8,149 @@ import io
 def remove_background(input_image):
     try:
         # Initialize the pipeline with trust_remote_code=True
-        segmentor = pipeline("image-segmentation", 
-                           model="briaai/RMBG-1.4",
-                           device=-1,
-                           trust_remote_code=True)
+        segmentor = pipeline(
+            "image-segmentation", 
+            model="briaai/RMBG-1.4",
+            trust_remote_code=True,
+            device="cpu"
+        )
         
-        # Process the image
-        result = segmentor(input_image)
-        return result['output_image']
+        # Process the image and get result
+        result = segmentor(input_image, return_mask=True)
+        
+        # Convert result to RGBA
+        if isinstance(result, Image.Image):
+            # Create transparent background
+            output = Image.new('RGBA', result.size, (0, 0, 0, 0))
+            output.paste(input_image, mask=result)
+        else:
+            output = result['output_image']
+        
+        return output
+
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")
 
-# Custom CSS for mobile-friendly design
+# Custom theme and styling
+theme = gr.themes.Soft(
+    primary_hue="gold",
+    secondary_hue="orange",
+).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"
+)
+
 css = """
 .gradio-container {
-    font-family: 'Segoe UI', sans-serif;
-    max-width: 100% !important;
-    padding: 10px !important;
-}
-.container {
-    display: flex;
-    flex-direction: column;
-    gap: 20px;
+    max-width: 1200px !important;
+    margin: 0 auto !important;
+    padding: 20px !important;
 }
 .image-container {
-    width: 100%;
-    max-width: 500px;
-    margin: 0 auto;
+    border-radius: 15px !important;
+    border: 2px solid rgba(255, 215, 0, 0.3) !important;
+    padding: 10px !important;
+    background: rgba(255, 255, 255, 0.1) !important;
+    transition: transform 0.3s ease !important;
+}
+.image-container:hover {
+    transform: translateY(-5px) !important;
 }
 .gr-button {
-    background: linear-gradient(45deg, #FFD700, #FFA500);
-    border: none !important;
-    color: black !important;
-    padding: 12px 20px !important;
-    border-radius: 8px !important;
-    font-weight: bold !important;
-    margin: 10px 0 !important;
-    width: 100% !important;
-    max-width: 300px !important;
+    min-width: 200px !important;
+    height: 45px !important;
+    font-size: 16px !important;
+    margin: 10px !important;
+    transition: all 0.3s ease !important;
+}
+.gr-button:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 5px 15px rgba(255, 215, 0, 0.3) !important;
 }
-@media (max-width: 768px) {
-    .gradio-container {
-        padding: 5px !important;
-    }
-    .gr-button {
-        padding: 10px 15px !important;
-    }
+.footer {
+    text-align: center;
+    margin-top: 20px;
+    color: #666;
 }
 """
 
 # Create Gradio interface
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(theme=theme, css=css) as demo:
     gr.HTML(
         """
-        <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
-            <h1 style="font-size: 2rem; margin-bottom: 1rem; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
-                AI Background Remover
+        <div style="text-align: center; margin-bottom: 2rem;">
+            <h1 style="font-size: 3rem; margin-bottom: 1rem; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
+                AI Background Remover Pro
             </h1>
-            <p style="color: #666; font-size: 1rem; margin-bottom: 2rem;">
-                Powered by RMBG V1.4 model
+            <p style="color: #cccccc; font-size: 1.2rem;">
+                Remove backgrounds instantly using advanced AI technology
             </p>
         </div>
         """
     )
     
-    with gr.Row(equal_height=True):
+    with gr.Row():
         with gr.Column():
             input_image = gr.Image(
-                label="Upload Image",
+                label="Upload Your Image",
                 type="pil",
-                sources=["upload", "clipboard"],
-                interactive=True
+                elem_classes="image-container"
             )
             
+            with gr.Row():
+                clear_btn = gr.Button("Clear", variant="secondary")
+                process_btn = gr.Button("Remove Background", variant="primary")
+                download_btn = gr.Button("Download", variant="primary", visible=False)
+
         with gr.Column():
             output_image = gr.Image(
                 label="Result",
                 type="pil",
-                interactive=False
+                elem_classes="image-container"
             )
-    
-    with gr.Row():
-        clear_btn = gr.Button("Clear", variant="secondary")
-        process_btn = gr.Button("Remove Background", variant="primary")
-
-    # Status message for feedback
-    status_msg = gr.Textbox(label="Status", interactive=False, visible=False)
+            
+    # Status message
+    status_msg = gr.Textbox(
+        label="Status",
+        placeholder="Ready to process your image...",
+        interactive=False
+    )
 
-    # Event handlers with concurrency limit
+    # Event handlers
     def process_and_update(image):
         if image is None:
-            return None, "Please upload an image first"
+            return None, "Please upload an image first", gr.Button.update(visible=False)
         try:
             result = remove_background(image)
-            return result, "Background removed successfully!"
+            return (
+                result, 
+                "✨ Background removed successfully!", 
+                gr.Button.update(visible=True)
+            )
         except Exception as e:
-            return None, f"Error: {str(e)}"
+            return None, f"❌ Error: {str(e)}", gr.Button.update(visible=False)
 
     process_btn.click(
         fn=process_and_update,
         inputs=[input_image],
-        outputs=[output_image, status_msg],
-        concurrency_limit=1,  # Limit concurrent processing
-        show_progress=True    # Show processing progress
+        outputs=[output_image, status_msg, download_btn],
     )
     
     clear_btn.click(
-        fn=lambda: (None, None, ""),
-        outputs=[input_image, output_image, status_msg],
-        queue=False  # No need for queue in clear operation
+        fn=lambda: (None, None, "Ready to process your image...", gr.Button.update(visible=False)),
+        outputs=[input_image, output_image, status_msg, download_btn],
     )
 
-# Launch the app with updated parameters
-demo.launch(
-    share=True,
-    max_threads=1,     # Control maximum worker threads
-    show_error=True,
-    server_name="0.0.0.0"
-)
+    gr.HTML(
+        """
+        <div class="footer">
+            <p>Powered by BRIA AI's RMBG V1.4 Model</p>
+        </div>
+        """
+    )
+
+# Launch the app
+demo.launch()