Update app.py

app.py

@@ -5,6 +5,8 @@ import cv2
 import PIL 
 from controlnet_aux import OpenposeDetector
 from transformers import pipeline
+from rembg import remove
+from diffusers.models import AutoencoderKL
 
 
 #import spaces #[uncomment to use ZeroGPU]
@@ -79,6 +81,12 @@ def is_lora(model_name):
     return model_name == "sand74/changpu_lora"
 
 
+def remove_background(image):
+    image = remove(image)
+    return image
+
+
+
 #@spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
     model_id,
@@ -98,6 +106,7 @@ def infer(
     use_ip_adapter=False,
     ip_adapter_image=None,
     ip_adapter_scale=None,
+    rm_background=True,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -119,8 +128,10 @@ def infer(
     else:
         base_model_id = model_id
 
+    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch_dtype)
+    
     if not use_controlnet:
-        pipe = StableDiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype)
+        pipe = StableDiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype, vae=vae)
     else:
         controlnet_image = cv2.resize(controlnet_image, (width, height), interpolation=cv2.INTER_AREA)
         controlnet = ControlNetModel.from_pretrained(
@@ -158,6 +169,9 @@ def infer(
     
     image = pipe(**pipe_params).images[0]
 
+    if rm_background:
+        image = remove(image)
+        
     return image, seed
 
 
@@ -195,6 +209,7 @@ with gr.Blocks(css=css) as demo:
             run_button = gr.Button("Run", scale=0, variant="primary")
 
         result = gr.Image(label="Result", show_label=False)
+        rm_background = gr.Checkbox(label="Remove background?", scale=1, value=True)
 
         with gr.Group(visible=True) as lora_section:
             title = gr.Markdown(" ### LoRA section")
@@ -360,6 +375,7 @@ with gr.Blocks(css=css) as demo:
             use_ip_adapter,
             ip_adapter_image,
             ip_adapter_scale,
+            rm_background,
         ],
         outputs=[result, seed],
     )