Update app.py (get weights using hf_hub_download)

app.py

@@ -18,6 +18,8 @@ import cv2
 from zim import zim_model_registry, ZimPredictor, ZimAutomaticMaskGenerator
 from zim.utils import show_mat_anns
 
+from huggingface_hub import hf_hub_download
+
 def get_shortest_axis(image):
     h, w, _ = image.shape
     return h if h < w else w
@@ -213,12 +215,17 @@ def get_examples():
     images = os.listdir(assets_dir)
     return [os.path.join(assets_dir, img) for img in images]
 
-if __name__ == "__main__":
+def download_onnx_weights(repo_id="naver-iv/zim-anything-vitl", file_dir="zim_vit_l_2092"):
+    hf_hub_download(repo_id=repo_id, filename=f"{file_dir}/encoder.onnx")
+    filepath = hf_hub_download(repo_id=repo_id, filename=f"{file_dir}/decoder.onnx")
+
+    return os.path.dirname(filepath)
 
-    backbone = "vit_l"
-    ckpt_p = "ckpts/zim_vit_l_2092"
 
-    model = zim_model_registry[backbone](checkpoint=ckpt_p)
+if __name__ == "__main__":
+    backbone = "vit_l"
+    model = zim_model_registry[backbone](checkpoint=download_onnx_weights())
+    
     if torch.cuda.is_available():
         model.cuda()