Update app.py

app.py

@@ -408,14 +408,14 @@ def main(device, segment_type):
         canny_image = Image.fromarray(image)
         return canny_image
 
-    def get_depth(image):
+    def get_depth(image, height, weight):
         image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
         with torch.no_grad(), torch.autocast("cuda"):
             depth_map = depth_estimator(image).predicted_depth
 
         depth_map = torch.nn.functional.interpolate(
             depth_map.unsqueeze(1),
-            size=(1024, 1024),
+            size=(height, weight),
             mode="bicubic",
             align_corners=False,
         )
@@ -472,7 +472,7 @@ def main(device, segment_type):
             kwargs['height'] = height
             kwargs['width'] = width
             condition_img = resize_and_center_crop(Image.fromarray(condition_img), (width, height))
-            spatial_condition = get_depth(condition_img)
+            spatial_condition = get_depth(condition_img, height, width)
         else:
             spatial_condition = None