Update app.py

app.py

@@ -1,5 +1,7 @@
 import gradio as gr
 import torch
+import numpy as np
+from PIL import Image
 
 from diffusers import AutoPipelineForInpainting, UNet2DConditionModel
 import diffusers
@@ -16,6 +18,14 @@ def read_content(file_path: str) -> str:
 
     return content
 
+def resize_image(img, target_shape):
+    """
+    Resize the image to the target shape while preserving aspect ratio.
+    """
+    img = Image.fromarray(img)
+    img = img.resize(target_shape[::-1], Image.ANTIALIAS)
+    return np.array(img)
+
 def predict(dict, prompt="", negative_prompt="", guidance_scale=7.5, steps=30, strength=1.0, scheduler="DPMSolverMultistepScheduler-Karras"):
     if negative_prompt == "":
         negative_prompt = None
@@ -34,7 +44,9 @@ def predict(dict, prompt="", negative_prompt="", guidance_scale=7.5, steps=30, s
     mask = dict["mask"]
     
     output = pipe(prompt = prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength, clip_skip=1)
-    return output.images[0], gr.update(visible=True)
+    input_shape = init_image.shape[:2]
+    output_image = resize_image(output.images[0], input_shape)
+    return output_image, gr.update(visible=True)
 
 
 css = '''