Update app.py

app.py

@@ -28,6 +28,7 @@ def depth_estimation(image, model_path, encoder='vits'):
         config_path = 'configs/config_v2_vits14.json'
         
         # Ensure model path exists or download if needed
+        model_path="checkpoint/latest.pth"
         if not os.path.exists(model_path):
             return "Model checkpoint not found. Please upload a valid model path."
         
@@ -36,7 +37,7 @@ def depth_estimation(image, model_path, encoder='vits'):
         # Preprocess image
         rgb = torch.from_numpy(np.array(image)).permute(2, 0, 1).to(device)  # C, H, W
         predictions = model.infer(rgb)
-        depth = predictions["depth"].squeeze().to('cpu').numpy()
+        depth = predictions["depth"].squeeze().to(device).numpy()
 
         min_depth = depth.min()
         max_depth = depth.max()
@@ -87,7 +88,6 @@ def main():
         fn=depth_estimation,
         inputs=[
             gr.Image(type="numpy", label="Input Image"),
-            gr.Textbox(value='checkpoint/latest.pth', label='Model Path'),
             gr.Dropdown(choices=['vits', 'vitb', 'vitl', 'vitg'], value='vits', label='Encoder'),
         ],
         outputs=[