Update app.py

added delete background

app.py

@@ -5,6 +5,9 @@ import random
 # import spaces #[uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline
 from peft import PeftModel, PeftConfig
+from rembg import remove
+from PIL import Image
+import io
 import torch
 from typing import Optional
 
@@ -94,6 +97,7 @@ def infer(
     num_inference_steps: int = 20,
     scheduler_name: Optional[str] = None,
     lora_scale: float = 1.0,
+    remove_background: bool = False,
     progress=gr.Progress(track_tqdm=True),
 ):
     # получаем/загружаем нужный pipe
@@ -132,6 +136,18 @@ def infer(
         generator=generator,
     ).images[0]
 
+    if remove_background:
+        # Конвертируем PIL Image в bytes
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format='PNG')
+        img_byte_arr = img_byte_arr.getvalue()
+        
+        # Удаляем фон
+        output_image = remove(img_byte_arr)
+        
+        # Конвертируем обратно в PIL Image
+        image = Image.open(io.BytesIO(output_image))
+
     return image, seed
 
 
@@ -174,7 +190,7 @@ with gr.Blocks(css=css) as demo:
             minimum=0.0,
             maximum=3.0,
             step=0.1,
-            value=0.7,
+            value=1.8,
             visible=False,  # Initially hidden
         )
 
@@ -192,6 +208,12 @@ with gr.Blocks(css=css) as demo:
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
+            remove_background = gr.Checkbox(
+                label="Remove background from generated image",
+                value=False,
+                info="Use rembg to remove background from the generated image"
+            )
+            
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
@@ -272,7 +294,8 @@ with gr.Blocks(css=css) as demo:
             guidance_scale,
             num_inference_steps,
             scheduler_select,
-            lora_scale_slider
+            lora_scale_slider,
+            remove_background
         ],
         outputs=[result, seed],
     )