Download model checkpoint from HF hub

app.py

@@ -1,9 +1,11 @@
 import glob
 import gradio as gr
 import numpy as np
+import spaces
 import torch
 import tempfile
 import uuid
+from huggingface_hub import hf_hub_download
 from PIL import Image, ImageOps, ImageEnhance
 from pathlib import Path
 from zipfile import ZipFile, is_zipfile
@@ -81,6 +83,12 @@ model_configs = {
     'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
     'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
 }
+encoder2name = {
+    'vits': 'Small',
+    'vitb': 'Base',
+    'vitl': 'Large',
+    'vitg': 'Giant', # we are undergoing company review procedures to release our giant model checkpoint
+}
 
 title = "# Depth Anything V2"
 description = """Looking Glass demo for **Depth Anything V2**.
@@ -113,9 +121,11 @@ def upscale_image(image, model, background, discard_alpha):
         image = model.infer(image)
     return image.convert("RGB") if discard_alpha else image
 
-def on_submit(image, batch_images, book, config, upscale_model, upscale_method, denoise_level, discard_alpha, progress=gr.Progress()):
-    model = DepthAnythingV2(**model_configs[config])
-    state_dict = torch.load(f'checkpoints/depth_anything_v2_{config}.pth', map_location="cpu")
+def on_submit(image, batch_images, book, encoder, upscale_model, upscale_method, denoise_level, discard_alpha, progress=gr.Progress()):
+    model_name = encoder2name[encoder]
+    model = DepthAnythingV2(**model_configs[encoder])
+    filepath = hf_hub_download(repo_id=f"depth-anything/Depth-Anything-V2-{model_name}", filename=f"depth_anything_v2_{encoder}.pth", repo_type="model")
+    state_dict = torch.load(filepath, map_location="cpu")
     model.load_state_dict(state_dict)
     model = model.to(DEVICE).eval()