Add hf_hub models

gradio_app.py

@@ -25,6 +25,7 @@ from demo.defaults import DefaultPredictor
 from demo.visualizer import Visualizer, ColorMode
 
 import gradio as gr
+from huggingface_hub import hf_hub_download
 
 KEY_DICT = {"Cityscapes (19 classes)": "cityscapes",
             "COCO (133 classes)": "coco",
@@ -34,17 +35,25 @@ SWIN_CFG_DICT = {"cityscapes": "configs/cityscapes/oneformer_swin_large_IN21k_38
             "coco": "configs/coco/oneformer_swin_large_IN21k_384_bs16_100ep.yaml",
             "ade20k": "configs/ade20k/oneformer_swin_large_IN21k_384_bs16_160k.yaml",}
 
-SWIN_MODEL_DICT = {"cityscapes": "models/250_16_swin_l_oneformer_cityscapes_90k.pth",
-              "coco": "models/150_16_swin_l_oneformer_coco_100ep.pth",
-              "ade20k": "models/250_16_swin_l_oneformer_ade20k_160k.pth"}
+SWIN_MODEL_DICT = {"cityscapes": hf_hub_download(repo_id="shi-labs/swin_l_oneformer_cityscapes", 
+                                            filename="250_16_swin_l_oneformer_cityscapes_90k.pth"),
+              "coco": hf_hub_download(repo_id="shi-labs/swin_l_oneformer_coco", 
+                                            filename="150_16_swin_l_oneformer_coco_100ep.pth"),
+              "ade20k": hf_hub_download(repo_id="shi-labs/swin_l_oneformer_ade20k", 
+                                            filename="250_16_swin_l_oneformer_ade20k_160k.pth")
+            }
 
 DINAT_CFG_DICT = {"cityscapes": "configs/cityscapes/oneformer_dinat_large_bs16_90k.yaml",
             "coco": "configs/coco/oneformer_dinat_large_bs16_100ep.yaml",
             "ade20k": "configs/ade20k/oneformer_dinat_large_IN21k_384_bs16_160k.yaml",}
 
-DINAT_MODEL_DICT = {"cityscapes": "models/250_16_dinat_l_oneformer_cityscapes_90k.pth",
-              "coco": "models/150_16_dinat_l_oneformer_coco_100ep.pth",
-              "ade20k": "models/250_16_dinat_l_oneformer_ade20k_160k.pth"}
+DINAT_MODEL_DICT = {"cityscapes": hf_hub_download(repo_id="shi-labs/dinat_l_oneformer_cityscapes", 
+                                            filename="250_16_dinat_l_oneformer_cityscapes_90k.pth"),
+              "coco": hf_hub_download(repo_id="shi-labs/dinat_l_oneformer_coco", 
+                                            filename="150_16_dinat_l_oneformer_coco_100ep.pth"),
+              "ade20k": hf_hub_download(repo_id="shi-labs/dinat_l_oneformer_ade20k", 
+                                            filename="250_16_dinat_l_oneformer_ade20k_160k.pth")
+            }
 
 MODEL_DICT = {"DiNAT-L": DINAT_MODEL_DICT,
         "Swin-L": SWIN_MODEL_DICT }

requirements.txt

@@ -26,6 +26,8 @@ ftfy
 regex
 inflect
 diffdist
+gdown
+huggingface_hub
 pytorch_lightning
 tqdm
 gradio==3.1.4