Supports images with transparent background.

app.py

@@ -48,16 +48,52 @@ def read_content(file_path: str) -> str:
 
     return content
 
-def predict(img, prompt="", seed=0):
-    img = img.convert("RGB")
+def predict(img, prompt="", seed=0, use_removal=True):
+    img = img.convert("RGBA")
     img = resize_with_padding(img, (512, 512))
-    mask = remover.process(img, type='map')
-    mask = ImageOps.invert(mask)
+
+    if use_removal:
+        # remover で背景を自動抽出
+        mask = remover.process(img.convert("RGB"), type='map')
+        mask = ImageOps.invert(mask)
+    else:
+        # 背景削除済みPNGのアルファチャンネルをマスクに利用
+        if "A" in img.getbands():  # RGBAかどうか
+            mask = img.getchannel("A")
+            mask = ImageOps.invert(mask)  # 前景を黒、背景を白に
+        else:
+            # JPEGなどアルファがない場合 → 背景編集しない（全部黒）
+            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.convert("RGB"),
+            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.convert("RGB"),
+            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
     
@@ -108,6 +144,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 +156,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 ])