Update app.py

app.py

@@ -20,6 +20,8 @@ from zim import zim_model_registry, ZimPredictor, ZimAutomaticMaskGenerator
 from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
 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
@@ -233,12 +235,18 @@ def get_examples():
     images = os.listdir(assets_dir)
     return [os.path.join(assets_dir, img) for img in images]
 
+def download_onnx_weights(repo_id="naver-iv/zim-anything-vitb", file_dir="zim_vit_b_2043"):
+    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)
+
+    
 if __name__ == "__main__":
     backbone = "vit_b"
     
     # load ZIM
-    ckpt_mat = "ckpts/zim_vit_b_2043"
-    zim = zim_model_registry[backbone](checkpoint=ckpt_mat)
+    zim = zim_model_registry[backbone](checkpoint=download_onnx_weights())
     if torch.cuda.is_available():
         zim.cuda()
     zim_predictor = ZimPredictor(zim)