Update with md5sum and half precision inference

app.py

@@ -20,6 +20,7 @@ description = """
 # Prismer 
 The official demo for **Prismer: A Vision-Language Model with An Ensemble of Experts**.
 Please refer to our [project page](https://shikun.io/projects/prismer) or [github](https://github.com/NVlabs/prismer) for more details.
+Expert labels will be only computed once for the same image checked with md5sum.
 """
 
 if (SPACE_ID := os.getenv('SPACE_ID')) is not None:

label_prettify.py

@@ -8,6 +8,7 @@ import numpy as np
 import shutil
 
 from prismer.utils import create_ade20k_label_colormap
+matplotlib.use('agg')
 
 obj_label_map = torch.load('prismer/dataset/detection_features.pt')['labels']
 coco_label_map = torch.load('prismer/dataset/coco_features.pt')['labels']

prismer/configs/experts.yaml

@@ -1,3 +1,3 @@
 data_path: helpers
-im_name: 4265e527d8897c032baabea78e590c18
+im_name: ca2a8e1305af24483124b85c53bd24b3
 save_path: helpers/labels

prismer_model.py

@@ -34,10 +34,10 @@ def download_models() -> None:
         subprocess.run(shlex.split('python download_checkpoints.py --download_experts=True'), cwd='prismer')
 
     model_names = [
-        'vqa_prismer_base',
-        'vqa_prismer_large',
+        # 'vqa_prismer_base',
+        # 'vqa_prismer_large',
         'pretrain_prismer_base',
-        'pretrain_prismer_large',
+        # 'pretrain_prismer_large',
     ]
     for model_name in model_names:
         if pathlib.Path(f'prismer/logging/{model_name}').exists():
@@ -78,6 +78,7 @@ def run_experts(image_path: str) -> Tuple[str, Tuple[str, ...]]:
 
     config = yaml.load(open('prismer/configs/experts.yaml', 'r'), Loader=yaml.Loader)
     config['im_name'] = im_name
+    print(im_name)
     with open('prismer/configs/experts.yaml', 'w') as yaml_file:
         yaml.dump(config, yaml_file, default_flow_style=False)
 
@@ -89,7 +90,7 @@ def run_experts(image_path: str) -> Tuple[str, Tuple[str, ...]]:
         run_expert('depth')
         with concurrent.futures.ProcessPoolExecutor() as executor:
             executor.map(run_expert, expert_names)
-        executor.shutdown(wait=True)
+            executor.shutdown(wait=True)
 
         # no parallelization just to be safe
         # expert_names = ['depth', 'edge', 'normal', 'objdet', 'ocrdet', 'segmentation']