Update app.py

app.py

@@ -1,11 +1,10 @@
 import torch
 torch.jit.script = lambda f: f
-from zoedepth.utils.config import get_config
-from zoedepth.models.builder import build_model
 from zoedepth.utils.misc import colorize, save_raw_16bit
 from zoedepth.utils.geometry import depth_to_points, create_triangles
 import gradio as gr
 import spaces
+
 from PIL import Image
 import numpy as np
 import trimesh
@@ -31,6 +30,7 @@ DEVICE = 'cuda'
 model = torch.hub.load('isl-org/ZoeDepth', "ZoeD_N", pretrained=True).to("cpu").eval()
 
 # ----------- Depth functions
+@spaces.GPU(enable_queue=True)
 def save_raw_16bit(depth, fpath="raw.png"):
     if isinstance(depth, torch.Tensor):
         depth = depth.squeeze().cpu().numpy()
@@ -42,7 +42,8 @@ def save_raw_16bit(depth, fpath="raw.png"):
     return depth
 
 @spaces.GPU(enable_queue=True)
-def process_image(model, image: Image.Image):
+def process_image(image: Image.Image):
+    global model
     image = image.convert("RGB")
 
     model.to(DEVICE)
@@ -54,8 +55,9 @@ def process_image(model, image: Image.Image):
 # ----------- Depth functions
 
 # ----------- Mesh functions
-
+@spaces.GPU(enable_queue=True)
 def depth_edges_mask(depth):
+    global model
     """Returns a mask of edges in the depth map.
     Args:
     depth: 2D numpy array of shape (H, W) with dtype float32.
@@ -72,12 +74,14 @@ def depth_edges_mask(depth):
 
 @spaces.GPU(enable_queue=True)
 def predict_depth(model, image):
+    global model
     model.to(DEVICE)
     depth = model.infer_pil(image)
     return depth
 
 @spaces.GPU(enable_queue=True)
-def get_mesh(model, image: Image.Image, keep_edges=True):
+def get_mesh(image: Image.Image, keep_edges=True):
+    global model
     image.thumbnail((1024,1024))  # limit the size of the input image
 
     depth = predict_depth(model, image)
@@ -117,7 +121,8 @@ with gr.Blocks(css=css) as API:
             inputs=gr.Image(label="Input Image", type='pil', height=500)  # Input is an image
             outputs=gr.Image(label="Depth Map", type='pil', height=500)  # Output is also an image
         generate_btn = gr.Button(value="Generate")
-        generate_btn.click(partial(process_image, model), inputs=inputs, outputs=outputs, api_name="generate_depth")
+        # generate_btn.click(partial(process_image, model), inputs=inputs, outputs=outputs, api_name="generate_depth")
+        generate_btn.click(process_image, inputs=inputs, outputs=outputs, api_name="generate_depth")
         
     with gr.Tab("Image to 3D"):
         with gr.Row():
@@ -125,7 +130,8 @@ with gr.Blocks(css=css) as API:
                 inputs=[gr.Image(label="Input Image", type='pil', height=500), gr.Checkbox(label="Keep occlusion edges", value=True)]
             outputs=gr.Model3D(label="3D Mesh", clear_color=[1.0, 1.0, 1.0, 1.0], height=500)
         generate_btn = gr.Button(value="Generate")
-        generate_btn.click(partial(get_mesh, model), inputs=inputs, outputs=outputs, api_name="generate_mesh")
+        # generate_btn.click(partial(get_mesh, model), inputs=inputs, outputs=outputs, api_name="generate_mesh")
+        generate_btn.click(get_mesh, inputs=inputs, outputs=outputs, api_name="generate_mesh")
 
 if __name__ == '__main__':
     API.launch()