Update app.py

app.py

@@ -6,7 +6,14 @@ from PIL import Image, ImageOps
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 pipeline = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
 
+def squarify_image(img):
+    if(img.height > img.width): bg_size = img.height
+    else:  bg_size = img.width
+    bg = Image.new(mode="RGB", size=(bg_size,bg_size), color="white")
+    bg.paste(img, ( int((bg.width - bg.width)/2), 0) )
 
+    return bg
+    
 def divisible_by_8(image):
     width, height = image.size
     
@@ -24,8 +31,10 @@ def generate(image_editor, prompt, neg_prompt, strength, guidance):
     image = image_editor['background'].convert('RGB')
     image.thumbnail((1024, 1024))
     image = divisible_by_8(image)
+    image = squarify_image(image)
 
     layer = image_editor["layers"][0].resize(image.size)
+    layer = squarify_image(layer)
     
     mask = Image.new("RGBA", image.size, "WHITE") 
     mask.paste(layer, (0, 0), layer)