Update app.py

app.py

@@ -3,11 +3,14 @@ import numpy as np
 import random
 from diffusers import DiffusionPipeline
 import torch
-from rembg import remove
-from PIL import Image
+
+###########################################
+from rembg import remove 
+from PIL import Image 
+###############################################
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+model_repo_id = "stabilityai/sdxl-turbo"  
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
@@ -21,7 +24,6 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
     model,
     prompt,
@@ -31,9 +33,9 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
+    remove_bg, #################################################################
     progress=gr.Progress(track_tqdm=True),
 ):
-
     global model_repo_id
     if model != model_repo_id:
         print(model, model_repo_id)
@@ -52,6 +54,12 @@ def infer(
         generator=generator,
     ).images[0]
 
+    #############################################################
+    # Если выбрано удаление фона
+    if remove_bg:
+        image = remove(image)
+    ##############################################################
+    
     return image, seed
 
 
@@ -113,14 +121,14 @@ with gr.Blocks(css=css) as demo:
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=7.5,  # Replace with defaults that work for your model
+                    value=7.5,  
         )
         num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
                     maximum=100,
                     step=1,
-                    value=30,  # Replace with defaults that work for your model
+                    value=30,  
         )
 
         with gr.Accordion("Advanced Settings", open=False):
@@ -130,7 +138,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,  # Replace with defaults that work for your model
+                    value=512,  
                 )
 
                 height = gr.Slider(
@@ -138,9 +146,18 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,  # Replace with defaults that work for your model
+                    value=512,  
                 )
 
+        #########################################################
+        # Добавляем Checkbox для удаления фона
+        remove_bg = gr.Checkbox(
+            label="Remove Background",
+            value=False,
+            interactive=True
+        )
+        ##########################################################
+
         gr.Examples(examples=examples, inputs=[prompt])
         gr.Examples(examples=examples_negative, inputs=[negative_prompt])
         
@@ -159,10 +176,11 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            remove_bg,  # Добавляем remove_bg в inputs ###############################################
         ],
         outputs=[result, seed],
     )
 
 if __name__ == "__main__":
     demo.launch()
-    
+