Update app.py

app.py

@@ -1,6 +1,5 @@
 import spaces
 import gradio as gr
-from gradio_imageslider import ImageSlider
 import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
@@ -11,99 +10,59 @@ print("Loading Fibo-Edit-RMBG model...")
 pipe = DiffusionPipeline.from_pretrained(
     "briaai/Fibo-Edit-RMBG",
     torch_dtype=torch.bfloat16,
-    trust_remote_code=True
+    trust_remote_code=True,
 )
 
-# Check if CUDA is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe.to(device)
+# pipe.to(device)
 print(f"Model loaded on {device}")
 
-@spaces.GPU
-def process_image(image, num_steps=10, guidance_scale=1.0):
-    """
-    Process an image with Fibo-Edit-RMBG
-    """
-    if image is None:
-        return None, None, "Please upload an image"
-
-    # Fixed instruction for background removal
-    instruction = "\"{'edit_instruction':'Generate a detailed grayscale alpha matte. Map the opaque foreground to white and the background to black. Produce soft, anti-aliased grayscale gradients at the edges of the subject to represent fine details and transparency.'}\""
-
-    try:
-        # Run the pipeline to get the grayscale alpha matte
-        mask = pipe(
-            image=image,
-            prompt=instruction,
-            num_inference_steps=num_steps,
-            guidance_scale=guidance_scale
-        ).images[0]
-
-        # Convert original image to RGBA
-        original_rgba = image.convert("RGBA")
-
-        # Convert mask to grayscale (in case it's not already)
-        mask_gray = mask.convert("L")
-
-        # Resize mask to match original image dimensions
-        if mask_gray.size != original_rgba.size:
-            mask_gray = mask_gray.resize(original_rgba.size, Image.LANCZOS)
+INSTRUCTION = "{'edit_instruction':'Generate a detailed grayscale alpha matte. Map the opaque foreground to white and the background to black. Produce soft, anti-aliased grayscale gradients at the edges of the subject to represent fine details and transparency.'}"
 
-        # Use the mask as the alpha channel
-        original_rgba.putalpha(mask_gray)
 
-        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
-        original_rgba.save(tmp.name, format="PNG")
-
-        return (original_rgba, image), tmp.name, "✅ Image processed successfully!"
+@spaces.GPU
+def process(image, num_steps=10, guidance_scale=1.0):
+    if image is None:
+        return None, None
 
-    except Exception as e:
-        return None, None, f"❌ Error: {str(e)}"
+    mask = pipe(
+        image=image,
+        prompt=INSTRUCTION,
+        num_inference_steps=int(num_steps),
+        guidance_scale=guidance_scale,
+    ).images[0]
 
-# css = """
-# .image-container img {
-#     object-fit: contain !important;
-# }
-# """
+    original_rgba = image.convert("RGBA")
+    mask_gray = mask.convert("L")
+    if mask_gray.size != original_rgba.size:
+        mask_gray = mask_gray.resize(original_rgba.size, Image.Resampling.LANCZOS)
+    original_rgba.putalpha(mask_gray)
+    result = original_rgba
 
-# css = """
-# .image-container img,
-# .image-container .fixed img,
-# .image-container .fixed {
-#     object-fit: contain !important;
-#     max-height: 400px !important;
-# }
-# """
+    # Save full-res result for download
+    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+    result.save(tmp.name, format="PNG")
 
-css = """
-#output-slider img {
-    object-fit: contain !important;
-}
-"""
+    return (image.convert("RGBA"), result), tmp.name
 
 
-# Create Gradio interface
-with gr.Blocks(title="Fibo-Edit-RMBG - Background Removal", css=css) as demo:
+with gr.Blocks(title="Fibo-Edit-RMBG - Background Removal", theme=gr.themes.Soft(primary_hue="sky")) as demo:
     gr.Markdown("""
-    # Fibo-Edit-RMBG - Background Removal
+# Fibo-Edit-RMBG - Background Removal
 
-    Powered by [Bria AI's Fibo-Edit-RMBG model](https://huggingface.co/briaai/Fibo-Edit-RMBG)
+Powered by [Bria AI's Fibo-Edit-RMBG model](https://huggingface.co/briaai/Fibo-Edit-RMBG)
 
-    This model performs background removal on any image
+This model performs background removal on any image
 
-    ### How to use:
-    1. Upload your image
-    2. Adjust settings if needed (optional)
-    3. Click Process!
+### How to use:
+1. Upload your image
+2. Adjust settings if needed (optional)
+3. Click Process!
     """)
 
     with gr.Row():
         with gr.Column(scale=1):
-            input_image = gr.Image(
-                label="Input Image",
-                type="pil",
-                height=400
-            )
+            input_image = gr.Image(label="Input Image", type="pil", height=400)
 
             with gr.Accordion("Advanced Settings", open=False):
                 num_steps = gr.Slider(
@@ -111,57 +70,55 @@ with gr.Blocks(title="Fibo-Edit-RMBG - Background Removal", css=css) as demo:
                     maximum=50,
                     value=10,
                     step=1,
-                    label="Number of Steps"
+                    label="Number of Steps",
                 )
-
                 guidance_scale = gr.Slider(
                     minimum=1.0,
                     maximum=15.0,
                     value=1.0,
                     step=0.5,
-                    label="Guidance Scale"
+                    label="Guidance Scale",
                 )
 
-            process_btn = gr.Button("🚀 Process Image", variant="primary")
+            process_btn = gr.Button("Remove Background", variant="primary", size="lg")
 
         with gr.Column(scale=1):
-            output_image = ImageSlider(
-                label="Result vs Original",
+            slider = gr.ImageSlider(
+                label="Original vs Background Removed",
                 type="pil",
-                height=400,
-                elem_id="output-slider"
+                image_mode="RGBA",
+                format="png",
+                slider_position=50,
+                max_height=500,
+                container=True,
             )
+            download_file = gr.File(label="Download Result (PNG with transparency)")
 
-            download_file = gr.File(label="Download PNG (with transparency)")
+    gr.Markdown("""
+---
+### About Fibo-Edit-RMBG
 
-            status_text = gr.Textbox(
-                label="Status",
-                interactive=False,
-                lines=2
-            )
+Fibo-Edit-RMBG is built on [Fibo-Edit](https://huggingface.co/briaai/Fibo-Edit), an 8B parameter image editing model
+that uses structured prompts for precise, deterministic editing.
+
+- **Model**: [briaai/Fibo-Edit-RMBG](https://huggingface.co/briaai/Fibo-Edit-RMBG)
+- **Architecture**: Based on FIBO with 8B parameters
+- **License**: Non-commercial use ([CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/))
+
+For commercial use, please [contact Bria AI](https://bria.ai/contact-us).
+    """)
 
-    # Connect the button
     process_btn.click(
-        fn=process_image,
+        fn=process,
         inputs=[input_image, num_steps, guidance_scale],
-        outputs=[output_image, download_file, status_text]
+        outputs=[slider, download_file],
     )
 
-    gr.Markdown("""
-    ---
-    ### About Fibo-Edit-RMBG
-
-    Fibo-Edit-RMBG is built on [Fibo-Edit](https://huggingface.co/briaai/Fibo-Edit), an 8B parameter image editing model
-    that uses structured prompts for precise, deterministic editing.
-
-    - **Model**: [briaai/Fibo-Edit-RMBG](https://huggingface.co/briaai/Fibo-Edit-RMBG)
-    - **Architecture**: Based on FIBO with 8B parameters
-    - **License**: Non-commercial use ([CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/))
-
-    For commercial use, please [contact Bria AI](https://bria.ai/contact-us).
-    """)
+    input_image.change(
+        fn=process,
+        inputs=[input_image],
+        outputs=[slider, download_file],
+    )
 
-# Launch the app
 if __name__ == "__main__":
-    demo.launch(show_api=False)
-
+    demo.launch()