Update app.py

app.py

@@ -6,6 +6,7 @@ import random
 from diffusers import DiffusionPipeline, StableDiffusionPipeline
 from peft import PeftModel, LoraConfig
 import torch
+from rembg import remove
 
 model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
 
@@ -60,6 +61,9 @@ def infer(
         generator=generator,
     ).images[0]
 
+    if remove_background:
+        image = remove(image)
+
     return image, seed
 
 css = """
@@ -113,7 +117,11 @@ with gr.Blocks(css=css) as demo:
             )
 
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
+            with gr.Row():
+                remove_background = gr.Checkbox(
+                    label="Delete background?", value=True
+                )
+                
             with gr.Row():
                 width = gr.Slider(
                     label="Width",
@@ -170,6 +178,7 @@ with gr.Blocks(css=css) as demo:
             guidance_scale,
             lora_scale,
             num_inference_steps,
+            remove_backgroung
         ],
         outputs=[result, seed],
     )