Update infer.py

infer.py

@@ -64,10 +64,11 @@ def infer_pipe(pipe, test_image, task_name, seed, device, video_depth=False):
             task_emb=task_emb,
             ).images[0]
 
-        # Post-process the prediction
+        # NEW CODE:
         if task_name == 'depth':
             output_npy = pred.mean(axis=-1)
-            output_color = colorize_depth_map(output_npy, reverse_color=True)
+            depth_normalized = ((output_npy - output_npy.min()) / (output_npy.max() - output_npy.min()) * 255).astype(np.uint8)
+            output_color = Image.fromarray(depth_normalized, mode='L')
         else:
             output_npy = pred
             output_color = Image.fromarray((output_npy * 255).astype(np.uint8))