Update app.py

app.py

@@ -2,6 +2,7 @@ import torch
 torch.jit.script = lambda f: f
 from zoedepth.utils.misc import colorize, save_raw_16bit
 from zoedepth.utils.geometry import depth_to_points, create_triangles
+from diffusers import DiffusionPipeline
 import gradio as gr
 import spaces
 
@@ -29,6 +30,9 @@ css = """
 DEVICE = 'cuda'
 model = torch.hub.load('isl-org/ZoeDepth', "ZoeD_N", pretrained=True).to("cpu").eval()
 
+CHECKPOINT = "prs-eth/marigold-v1-0"
+pipe = DiffusionPipeline.from_pretrained(CHECKPOINT)
+
 # ----------- Depth functions
 @spaces.GPU(enable_queue=True)
 def save_raw_16bit(depth, fpath="raw.png"):
@@ -46,10 +50,18 @@ def process_image(image: Image.Image):
     global model
     image = image.convert("RGB")
 
-    model.to(DEVICE)
-    out = model.infer_pil(image)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # model.to(device)
+    # depth = model.infer_pil(image)
+
+    # processed_array = save_raw_16bit(colorize(depth)[:, :, 0])
+    # return Image.fromarray(processed_array)
+
+    model.to(device)
+    
+    # # inference
+    processed_array = pipe(image)["depth"]
 
-    processed_array = save_raw_16bit(colorize(out)[:, :, 0])
     return Image.fromarray(processed_array)
 
 # ----------- Depth functions