Add remove-bg

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import numpy as np
 import random
 from typing import Optional
+from rembg import remove
 
 # import spaces #[uncomment to use ZeroGPU]
 from diffusers import StableDiffusionPipeline, StableDiffusionControlNetPipeline, ControlNetModel
@@ -45,6 +46,8 @@ def infer(lora_model_id: Optional[str] = "osmr/stable-diffusion-v1-4-lora-iv-ghi
           ipadapter_scale: float = DEFAULT_IPA_SCALE,
           ipadapter_image: object = None,
 
+          do_remove_bg: bool = False,
+
           progress = gr.Progress(track_tqdm=True)):
 
     use_lora = (lora_model_id in [
@@ -133,6 +136,9 @@ def infer(lora_model_id: Optional[str] = "osmr/stable-diffusion-v1-4-lora-iv-ghi
         ip_adapter_image=ipadapter_image
     ).images[0]
 
+    if do_remove_bg:
+        image = remove(image)
+
     return image, seed
 
 
@@ -244,6 +250,8 @@ with gr.Blocks(css=css) as demo:
                     value=DEFAULT_LS,
                 )
 
+            do_remove_bg = gr.Checkbox(label="Remove background", value=False)
+
         with gr.Accordion("ControlNet Settings", open=False):
             controlnet_type = gr.Dropdown(
                 choices=[
@@ -300,6 +308,7 @@ with gr.Blocks(css=css) as demo:
             controlnet_image,
             ipadapter_scale,
             ipadapter_image,
+            do_remove_bg,
         ],
         outputs=[result, seed],
     )

requirements.txt

@@ -5,4 +5,5 @@ invisible_watermark
 torch
 transformers
 xformers
-gradio
+gradio
+rembg