Revert to gradio_imageslider version

app.py

@@ -1,48 +1,10 @@
 import spaces
 import gradio as gr
+from gradio_imageslider import ImageSlider
 import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
-import base64
-from io import BytesIO
-
-
-def _img_to_b64(img: Image.Image) -> str:
-    buf = BytesIO()
-    img.save(buf, format="PNG")
-    return base64.b64encode(buf.getvalue()).decode()
-
-
-def make_slider_html(result: Image.Image, original: Image.Image) -> str:
-    r, o = _img_to_b64(result), _img_to_b64(original)
-    return f"""
-<div id="cmp" style="position:relative;width:100%;cursor:col-resize;user-select:none;line-height:0">
-  <img src="data:image/png;base64,{o}" style="width:100%;display:block">
-  <div id="cmp-clip" style="position:absolute;top:0;left:0;width:50%;overflow:hidden">
-    <img src="data:image/png;base64,{r}" style="width:200%;max-width:none;display:block">
-  </div>
-  <div id="cmp-line" style="position:absolute;top:0;left:50%;width:3px;height:100%;background:#fff;box-shadow:0 0 4px rgba(0,0,0,.8)"></div>
-  <div id="cmp-handle" style="position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);width:36px;height:36px;border-radius:50%;background:#fff;box-shadow:0 0 6px rgba(0,0,0,.6);display:flex;align-items:center;justify-content:center;font-size:14px;pointer-events:none">&#8646;</div>
-</div>
-<script>
-(function(){{
-  const el=document.getElementById('cmp'),
-        clip=document.getElementById('cmp-clip'),
-        line=document.getElementById('cmp-line'),
-        handle=document.getElementById('cmp-handle');
-  function setPos(x){{
-    const r=el.getBoundingClientRect(),p=Math.max(0,Math.min((x-r.left)/r.width,1))*100+'%';
-    clip.style.width=p; line.style.left=p; handle.style.left=p;
-  }}
-  let drag=false;
-  el.addEventListener('mousedown',e=>{{drag=true;setPos(e.clientX);}});
-  document.addEventListener('mouseup',()=>drag=false);
-  document.addEventListener('mousemove',e=>{{if(drag)setPos(e.clientX);}});
-  el.addEventListener('touchstart',e=>setPos(e.touches[0].clientX),{{passive:true}});
-  el.addEventListener('touchmove',e=>setPos(e.touches[0].clientX),{{passive:true}});
-}})();
-</script>"""
-
+import os
 
 # Load the model
 print("Loading Fibo-Edit-RMBG model...")
@@ -52,19 +14,24 @@ pipe = DiffusionPipeline.from_pretrained(
     trust_remote_code=True
 )
 
+# Check if CUDA is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 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 "", "Please upload an image"
+        return 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,
@@ -72,18 +39,25 @@ def process_image(image, num_steps=10, guidance_scale=1.0):
             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)
+
+        # Use the mask as the alpha channel
         original_rgba.putalpha(mask_gray)
 
-        return make_slider_html(original_rgba, image), "✅ Image processed successfully!"
+        return (original_rgba, image), "✅ Image processed successfully!"
 
     except Exception as e:
-        return "", f"❌ Error: {str(e)}"
-
+        return None, f"❌ Error: {str(e)}"
 
+# Create Gradio interface
 with gr.Blocks(title="Fibo-Edit-RMBG - Background Removal") as demo:
     gr.Markdown("""
     # Fibo-Edit-RMBG - Background Removal
@@ -114,6 +88,7 @@ with gr.Blocks(title="Fibo-Edit-RMBG - Background Removal") as demo:
                     step=1,
                     label="Number of Steps"
                 )
+
                 guidance_scale = gr.Slider(
                     minimum=1.0,
                     maximum=15.0,
@@ -125,17 +100,23 @@ with gr.Blocks(title="Fibo-Edit-RMBG - Background Removal") as demo:
             process_btn = gr.Button("🚀 Process Image", variant="primary")
 
         with gr.Column(scale=1):
-            output_slider = gr.HTML(label="Result vs Original")
+            output_image = ImageSlider(
+                label="Result vs Original",
+                type="pil",
+                height=400
+            )
+
             status_text = gr.Textbox(
                 label="Status",
                 interactive=False,
                 lines=2
             )
 
+    # Connect the button
     process_btn.click(
         fn=process_image,
         inputs=[input_image, num_steps, guidance_scale],
-        outputs=[output_slider, status_text]
+        outputs=[output_image, status_text]
     )
 
     gr.Markdown("""
@@ -152,6 +133,7 @@ with gr.Blocks(title="Fibo-Edit-RMBG - Background Removal") as demo:
     For commercial use, please [contact Bria AI](https://bria.ai/contact-us).
     """)
 
+# Launch the app
 if __name__ == "__main__":
     demo.launch(show_api=False)