hw6

app.py

@@ -60,6 +60,8 @@ with gr.Blocks(css=css) as demo:
                 visible=True,
             )
 
+            delete_background = gr.Checkbox(label="Delete background?")
+
             use_controlnet = gr.Checkbox(label="Use ControlNet")
             control_strength = gr.Slider(
                 label="ControlNet strength",

infer.py

@@ -6,6 +6,7 @@ import os
 import spaces
 import gradio as gr
 
+from rembg import remove
 from PIL import Image
 from transformers import pipeline
 from controlnet_aux import MLSDdetector, HEDdetector, NormalBaeDetector, LineartDetector
@@ -188,6 +189,7 @@ def infer(
     use_ip_adapter,
     ip_adapter_scale,
     ip_adapter_image,
+    delete_background,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -269,4 +271,7 @@ def infer(
             generator=generator
         ).images[0]
 
+    if delete_background:
+        image = remove(image)
+
     return image, seed

requirements.txt

@@ -3,4 +3,5 @@ peft
 torch
 opencv-python
 mediapipe
-controlnet_aux
+controlnet_aux
+rembg