pad non square image before resize

app.py

@@ -28,8 +28,22 @@ pipe = DiffusionPipeline.from_pretrained(
 device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = pipe.to(device)
 
+def pad_image(image):
+    w, h = image.size
+    if w == h:
+        return image
+    elif w > h:
+        new_image = Image.new(image.mode, (w, w), (0, 0, 0))
+        new_image.paste(image, (0, (w - h) // 2))
+        return new_image
+    else:
+        new_image = Image.new(image.mode, (h, h), (0, 0, 0))
+        new_image.paste(image, ((h - w) // 2, 0))
+        return new_image
+
 
 def process_image(image, text, prompt):
+  image = pad_image(image)
   image = image.resize((512, 512))
   with autocast("cuda"):
       inpainted_image = pipe(image=image, text=text, prompt=prompt).images[0]