Update app.py

app.py

@@ -7,6 +7,8 @@ from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionC
 from diffusers.utils import load_image
 from peft import PeftModel, LoraConfig
 
+from rembg import remove 
+
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_id_default = "stable-diffusion-v1-5/stable-diffusion-v1-5"
@@ -31,13 +33,13 @@ def infer(
     guidance_scale=7.0,
     lora_scale=1.0,
     num_inference_steps=20,
-    controlnet_checkbox=False,
-    controlnet_strength=0.0,
-    controlnet_mode="edge_detection",
-    controlnet_image=None,
-    ip_adapter_checkbox=False,
-    ip_adapter_scale=0.0,
-    ip_adapter_image=None,
+    controlnet_checkbox=False,        
+    controlnet_strength=0.0,          
+    controlnet_mode="edge_detection", 
+    controlnet_image=None,            
+    ip_adapter_checkbox=False,      
+    ip_adapter_scale=0.0,   
+    remove_bg=None,          
     progress=gr.Progress(track_tqdm=True),    
 ):  
     ckpt_dir='./lora_pussinboots_logos'
@@ -118,7 +120,13 @@ def infer(
 
     pipe.to(device)
 
-    return pipe(**params).images[0]
+    image = pipe(**params).images[0]
+
+    # Если выбрано удаление фона
+    if remove_bg:
+        image = remove(image)       
+
+    return image
 
 examples = [
     "Puss in Boots wearing a sombrero crosses the Grand Canyon on a tightrope with a guitar.",
@@ -272,6 +280,17 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                     step=32,
                     value=512,  # Replace with defaults that work for your model
                 )
+
+                # Удаление фона------------------------------------------------------------------------------------------------
+                # Checkbox для удаления фона
+                with gr.Blocks():
+                    remove_bg = gr.Checkbox(
+                        label="Remove Background",
+                        value=False,
+                        interactive=True
+                    )
+                # -------------------------------------------------------------------------------------------------------------
+
         
         gr.Examples(examples=examples, inputs=[prompt])
 
@@ -298,7 +317,8 @@ with gr.Blocks(css=css, fill_height=True) as demo:
             controlnet_image,
             ip_adapter_checkbox,
             ip_adapter_scale,
-            ip_adapter_image,  
+            ip_adapter_image, 
+            remove_bg, 
         ],
         outputs=[result],
     )