Update marigold_depth_estimation_lcm.py

marigold_depth_estimation_lcm.py

@@ -42,7 +42,29 @@ from diffusers.utils import BaseOutput, check_min_version
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
 check_min_version("0.27.0.dev0")
 
-
+def get_save_width_height(input_image_size, depth_pred):
+  orig_width, orig_height = input_image_size
+  depth_height, depth_width = depth_pred.shape[:2]
+
+  print("Inside pipe!")
+  print("Depth size: {}x{}".format(depth_width, depth_height))
+
+  # Use height
+  ratio = orig_height / depth_height
+  new_width = depth_width * ratio
+  if new_width <= orig_width:
+      save_width = int(new_width)
+      save_height = int(orig_height)
+  else:
+      ratio = orig_width / depth_width
+      new_height = depth_height * ratio
+      if new_height <= orig_height:
+          save_width = int(orig_width)
+          save_height = int(new_height)
+      else:
+          print("WTF IS THIS??")
+          return None, None
+          
 class MarigoldDepthConsistencyOutput(BaseOutput):
     """
     Output class for Marigold monocular depth prediction pipeline.
@@ -255,8 +277,9 @@ class MarigoldDepthConsistencyPipeline(DiffusionPipeline):
 
         # Resize back to original resolution
         if match_input_res:
+            save_width, save_height = get_save_width_height(input_size, depth_pred)
             pred_img = Image.fromarray(depth_pred)
-            pred_img = pred_img.resize(input_size)
+            pred_img = pred_img.resize((save_width, save_height))
             depth_pred = np.asarray(pred_img)
 
         # Clip output range