Update app.py

app.py

@@ -162,10 +162,18 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     margin_x = max(0, min(margin_x + x_offset, target_size[0] - new_width))
     margin_y = max(0, min(margin_y + y_offset, target_size[1] - new_height))
     
+    if image.mode == "RGBA":
+        background = Image.new("RGBA", target_size, (255, 255, 255, 255)
+    else:
+        background = Image.new("RGB", target_size, (255, 255, 255)
+
     background = Image.new('RGB', target_size, (255, 255, 255))
-    background.paste(source, (margin_x, margin_y))
+    # Paste the source (resized image) onto the background
+    background.paste(source, (margin_x, margin_y)
+
+    # Create the generated mask (L mode)
+    mask = Image.new("L", target_size, 255)
     
-    mask = Image.new('L', target_size, 255)
     mask_draw = ImageDraw.Draw(mask)
     white_gaps_patch = 2
 
@@ -228,8 +236,18 @@ def outpaint(image, width, height, overlap_percentage, num_inference_steps, resi
     background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom, x_offset, y_offset)
     if not can_expand(background.width, background.height, width, height, alignment):
         alignment = "Middle"
+
+    # Extract original alpha from the input image
+    original_alpha = background.split()[3] if background.mode == "RGBA" else Image.new("L", background.size, 255)
+    original_alpha_mask = original_alpha.point(lambda p: 0 if p < 255 else 255)  # 0 where transparent
+
+    # Combine original and generated mask
+    combined_mask = ImageChops.logical_or(original_alpha_mask, mask)  # 0 where to fill
+
+    # Use the combined_mask in the pipeline
     cnet_image = background.copy()
-    cnet_image.paste(0, (0, 0), mask)
+    cnet_image.paste(0, (0, 0), combined_mask)  # Overlay black on combined_mask area
+
     final_prompt = f"score_9, score_8_up, score_7_up, {prompt_input} , high quality, 4k"
     print(f"Outpainting using SDXL model: {pipe.config.model_name}")
     (
@@ -247,8 +265,11 @@ def outpaint(image, width, height, overlap_percentage, num_inference_steps, resi
         num_inference_steps=num_inference_steps
     ):
         yield cnet_image, image
+    # Invert the combined_mask and paste the generated image back
+    filled_mask = combined_mask.point(lambda p: 255 - p)
     image = image.convert("RGBA")
-    cnet_image.paste(image, (0, 0), mask)
+    cnet_image.paste(image, (0, 0), filled_mask)
+
     yield background, cnet_image
     
 @spaces.GPU(duration=7)