22

app.py

@@ -95,13 +95,14 @@ print("pipe ready.")
 
 def prepare(edit_dict):
     # return edit_dict['background']
-    return edit_dict['composite']
+    return image_utils.convert_transparent_to_white(edit_dict['layers'][0])
 
 @spaces.GPU
 def inference(
     prompt,
     negative_prompt,
     edit_dict,
+    mask_image,
     # image_scale=1.0,
     # control_mode='Canny',
     control_context_scale = 0.75,
@@ -148,7 +149,7 @@ def inference(
     # control_image = control_image.resize((width, height))
     # control_image_torch = get_image_latent(control_image, sample_size=sample_size)[:, :, 0]
 
-    mask_image = edit_dict['layers'][0]
+    # mask_image = edit_dict['layers'][0]
     # mask_image = edit_dict['composite']
     if mask_image is not None:
         mask_image = get_image_latent(mask_image, sample_size=sample_size)[:, :1, 0]
@@ -208,7 +209,7 @@ with gr.Blocks() as demo:
                     height=600, 
                     sources=['upload', 'clipboard'], 
                     type="pil",
-                    label="Mask Image"
+                    label="Edit Image"
                 )
                 
                 prompt = gr.Textbox(
@@ -270,7 +271,7 @@ with gr.Blocks() as demo:
                 # polished_prompt = gr.Textbox(label="Polished prompt", interactive=False)
 
                 with gr.Accordion("Preprocessor data", open=True):
-                    mask_image = gr.Image(label="Mask image", show_label=False)
+                    mask_image = gr.Image(label="Mask image")
                     
         gr.Examples(examples=examples, inputs=[edit_dict, prompt])
         gr.Markdown(read_file("static/footer.md"))
@@ -286,6 +287,7 @@ with gr.Blocks() as demo:
             prompt,
             negative_prompt,
             edit_dict,
+            mask_image,
             # image_scale,
             # control_mode,
             control_context_scale,

utils/image_utils.py

@@ -19,6 +19,20 @@ def rescale_image(img, scale, nearest=32, max_size=1280):
 
     return img.resize((new_w, new_h), Image.LANCZOS), new_w, new_h
 
+def convert_transparent_to_white(input_image):
+    """
+    Converts transparent pixels in a PIL Image to white.
+    """
+    if input_image.mode in ('RGBA', 'LA') or \
+       (input_image.mode == 'P' and 'transparency' in input_image.info):
+        alpha = input_image.convert('RGBA').split()[-1]
+        bg = Image.new("RGB", input_image.size, (255, 255, 255))
+        converted_image = Image.composite(input_image.convert("RGB"), bg, alpha)
+        return converted_image
+    else:
+        # If the image doesn't have an alpha channel, return it as is
+        return input_image
+    
 def padding_image(images, new_width, new_height):
     new_image = Image.new('RGB', (new_width, new_height), (255, 255, 255))