Update app.py

app.py

@@ -91,6 +91,8 @@ def predict_images(dict, depth, prompt="", negative_prompt="", guidance_scale=7.
     if negative_prompt == "":
         negative_prompt = None
 
+    og_size = dict["image"].shape[:2]
+
     init_image = cv2.resize(dict["image"], (512, 512))
 
     mask = Image.fromarray(cv2.resize(dict["mask"], (512, 512))[:,:,0])
@@ -116,7 +118,7 @@ def predict_images(dict, depth, prompt="", negative_prompt="", guidance_scale=7.
     output_depth_vis = output_depth_vis.astype("uint8")
     output_depth = Image.fromarray(output_depth_vis)
 
-    return output.rgb[0], output_depth, gr.update(visible=True)
+    return output.rgb[0].resize(og_size), output_depth.resize(og_size), gr.update(visible=True)
 
 
 css = '''