Update app.py

app.py

@@ -9,6 +9,8 @@ from huggingface_hub import hf_hub_download
 from controlnet_union import ControlNetModel_Union
 from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 
+from PIL import Image, ImageDraw
+
 MODELS = {
  "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
 }
@@ -55,13 +57,31 @@ prompt = "high quality"
 
 @spaces.GPU
 def fill_image(image, model_selection):
- source = image["background"]
- mask = image["layers"][0]
 
- alpha_channel = mask.split()[3]
- binary_mask = alpha_channel.point(lambda p: p > 0 and 255)
- cnet_image = source.copy()
- cnet_image.paste(0, (0, 0), binary_mask)
+ margin = 100
+ # Open the original image
+ source = image["image"] # Changed from image["background"] to match new input format
+ 
+ # Calculate new output size
+ output_size = (source.width + 2*margin, source.height + 2*margin)
+ 
+ # Create a white background
+ background = Image.new('RGB', output_size, (255, 255, 255))
+ 
+ # Calculate position to paste the original image
+ position = (margin, margin)
+ 
+ # Paste the original image onto the white background
+ background.paste(source, position)
+ 
+ # Create the mask
+ mask = Image.new('L', output_size, 255) # Start with all white
+ mask_draw = ImageDraw.Draw(mask)
+ mask_draw.rectangle([position, (position[0] + source.width, position[1] + source.height)], fill=0)
+ 
+ # Prepare the image for ControlNet
+ cnet_image = background.copy()
+ cnet_image.paste(0, (0, 0), mask)
 
  for image in pipe(
  prompt_embeds=prompt_embeds,
@@ -73,9 +93,9 @@ def fill_image(image, model_selection):
  yield image, cnet_image
 
  image = image.convert("RGBA")
- cnet_image.paste(image, (0, 0), binary_mask)
+ cnet_image.paste(image, (0, 0), mask)
 
- yield source, cnet_image
+ yield background, cnet_image
 
 
 def clear_result():