Update app.py

app.py

@@ -48,19 +48,48 @@ def read_content(file_path: str) -> str:
 
     return content
 
-def predict(img, prompt="", seed=0):
+def predict(img, prompt="", seed=0, use_removal=True):
     img = img.convert("RGB")
     img = resize_with_padding(img, (512, 512))
-    mask = remover.process(img, type='map')
-    mask = ImageOps.invert(mask)
+    
+    if use_removal:
+        mask = remover.process(img, type='map')
+        mask = ImageOps.invert(mask)
+    else:
+        mask = Image.new("L", img.size, 0)
+
     with torch.autocast("cuda"):
         generator = torch.Generator(device='cuda').manual_seed(seed)
-        output_controlnet = pipe(generator=generator, prompt=prompt, image=img, mask_image=mask, control_image=mask, num_images_per_prompt=1, num_inference_steps=20, guess_mode=False, controlnet_conditioning_scale=1.0, guidance_scale=7.5).images[0]
+        output_controlnet = pipe(
+            generator=generator,
+            prompt=prompt,
+            image=img,
+            mask_image=mask,
+            control_image=mask,
+            num_images_per_prompt=1,
+            num_inference_steps=20,
+            guess_mode=False,
+            controlnet_conditioning_scale=1.0,
+            guidance_scale=7.5
+        ).images[0]
+
         generator = torch.Generator(device='cuda').manual_seed(seed)
-        output_sd2 = pipe(generator=generator, prompt=prompt, image=img, mask_image=mask, control_image=mask, num_images_per_prompt=1, num_inference_steps=20, guess_mode=False, controlnet_conditioning_scale=0.0, guidance_scale=7.5).images[0]
+        output_sd2 = pipe(
+            generator=generator,
+            prompt=prompt,
+            image=img,
+            mask_image=mask,
+            control_image=mask,
+            num_images_per_prompt=1,
+            num_inference_steps=20,
+            guess_mode=False,
+            controlnet_conditioning_scale=0.0,
+            guidance_scale=7.5
+        ).images[0]
+
     torch.cuda.empty_cache()
     return output_controlnet, output_sd2, mask
-    
+
 css = '''
 .container {max-width: 1150px;margin: auto;padding-top: 1.5rem}
 #image_upload{min-height:400px}
@@ -108,6 +137,7 @@ with image_blocks as demo:
             with gr.Column(variant='compact', ):
                 image = gr.Image(value=bird_image, sources=['upload'], elem_id="image_upload", type="pil", label="Upload an image", width=512, height=512)
                 with gr.Row(variant='compact', elem_id="prompt-container", equal_height=True):
+                    use_removal = gr.Checkbox(label="Use BG Remover", value=True)
                     prompt = gr.Textbox(label='prompt', placeholder = 'What you want in the background?', show_label=True, elem_id="input-text")
                     seed = gr.Number(label="seed", value=13)
                     btn = gr.Button("Generate Background!")
@@ -119,7 +149,7 @@ with image_blocks as demo:
                 mask_out = gr.Image(value=bird_mask, label="Background Mask", elem_id="output-mask", width=512, height=512)
             with gr.Column(variant='compact', ):
                 sd2_out = gr.Image(value=bird_sd2, label="SD2 Output", elem_id="output-sd2", width=512, height=512)
-        btn.click(fn=predict, inputs=[image, prompt, seed], outputs=[controlnet_out, sd2_out, mask_out ])
+        btn.click(fn=predict, inputs=[image, prompt, seed, use_removal], outputs=[controlnet_out, sd2_out, mask_out ])