Load models from Hugging Face

app.py

@@ -9,7 +9,7 @@ from torchvision.transforms import Compose
 import tempfile
 from gradio_imageslider import ImageSlider
 
-from depth_anything.dpt import DPT_DINOv2
+from depth_anything.dpt import DepthAnything
 from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
 
 css = """
@@ -24,8 +24,7 @@ css = """
     }
 """
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-model = DPT_DINOv2(encoder='vitl', features=256, out_channels=[256, 512, 1024, 1024]).to(DEVICE).eval()
-model.load_state_dict(torch.load('checkpoints/depth_anything_vitl14.pth'))
+model = DepthAnything.from_pretrained('LiheYoung/depth_anything_vitl14').to(DEVICE).eval()
 
 title = "# Depth Anything"
 description = """Official demo for **Depth Anything: Unleashing the Power of Large-Scale Unlabeled Data**.
@@ -49,7 +48,6 @@ transform = Compose([
 def predict_depth(model, image):
     return model(image)
 
-
 with gr.Blocks(css=css) as demo:
     gr.Markdown(title)
     gr.Markdown(description)
@@ -93,4 +91,4 @@ with gr.Blocks(css=css) as demo:
 
 
 if __name__ == '__main__':
-    demo.queue().launch()
+    demo.queue().launch()