Update app.py

app.py

@@ -6,6 +6,7 @@ import torch
 from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline
 from diffusers.utils import load_image
 from peft import PeftModel, LoraConfig
+from rembg import remove
 
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -38,6 +39,7 @@ def infer(
     ip_adapter_checkbox=False,
     ip_adapter_scale=0.0,
     ip_adapter_image=None,
+    del_background=False,
     progress=gr.Progress(track_tqdm=True),    
 ):  
     ckpt_dir='./model_output'
@@ -116,7 +118,11 @@ def infer(
 
     pipe.to(device)
 
-    return pipe(**params).images[0]
+    if del_background:
+        return remove(pipe(**params).images[0])
+    else:
+        return pipe(**params).images[0]
+
 
 css = """
 #col-container {
@@ -184,6 +190,11 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                 step=1,
                 value=20,  # Replace with defaults that work for your model
             )
+        with gr.Row():
+            del_background = gr.Checkbox(
+                label="Delete background?",
+                value=False
+            )
         with gr.Row():
             controlnet_checkbox = gr.Checkbox(
                 label="ControlNet",
@@ -283,6 +294,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
             ip_adapter_checkbox,
             ip_adapter_scale,
             ip_adapter_image,  
+            del_background,
         ],
         outputs=[result],
     )