Fix labels

app_caption.py

@@ -11,11 +11,11 @@ from prismer_model import Model
 
 def create_demo():
     model = Model()
-
+    model.mode = 'caption'
     with gr.Row():
         with gr.Column():
-            model_name = gr.Dropdown(label='Model', choices=['Prismer-Base'], value='Prismer-Base')
             image = gr.Image(label='Input', type='filepath')
+            model_name = gr.Dropdown(label='Model', choices=['Prismer-Base, Prismer-Large'], value='Prismer-Base')
             run_button = gr.Button('Run')
         with gr.Column(scale=1.5):
             caption = gr.Text(label='Caption')

prismer_model.py

@@ -58,7 +58,7 @@ def run_experts(image_path: str) -> tuple[str | None, ...]:
 
     keys = ['depth', 'edge', 'normal', 'seg_coco', 'obj_detection', 'ocr_detection']
     results = [pathlib.Path('prismer/helpers/labels') / key / 'helpers/images/image.png' for key in keys]
-    return tuple(path.as_posix() if path.exists() else None for path in results)
+    return tuple(path.as_posix() for path in results)
 
 
 class Model:
@@ -67,24 +67,42 @@ class Model:
         self.model = None
         self.tokenizer = None
         self.exp_name = ''
+        self.mode = ''
 
     def set_model(self, exp_name: str) -> None:
         if exp_name == self.exp_name:
             return
 
-        config = {
-            'dataset': 'demo',
-            'data_path': 'prismer/helpers',
-            'label_path': 'prismer/helpers/labels',
-            'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
-            'image_resolution': 480,
-            'prismer_model': 'prismer_base' if self.exp_name == 'Prismer-Base' else 'prismer_large',
-            'freeze': 'freeze_vision',
-            'prefix': 'A picture of',
-        }
-
-        model = PrismerCaption(config)
-        state_dict = torch.load(f'prismer/logging/caption_{exp_name}/pytorch_model.bin', map_location='cuda:0')
+        if self.mode == 'caption':
+            config = {
+                'dataset': 'demo',
+                'data_path': 'prismer/helpers',
+                'label_path': 'prismer/helpers/labels',
+                'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
+                'image_resolution': 480,
+                'prismer_model': 'prismer_base' if self.exp_name == 'Prismer-Base' else 'prismer_large',
+                'freeze': 'freeze_vision',
+                'prefix': 'A picture of',
+            }
+
+            model = PrismerCaption(config)
+            state_dict = torch.load(f'prismer/logging/caption_{exp_name}/pytorch_model.bin', map_location='cuda:0')
+
+        elif self.mode == 'vqa':
+            config = {
+                'dataset': 'demo',
+                'data_path': 'prismer/helpers',
+                'label_path': 'prismer/helpers/labels',
+                'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
+                'image_resolution': 480,
+                'prismer_model': 'prismer_base' if self.exp_name == 'Prismer-Base' else 'prismer_large',
+                'freeze': 'freeze_vision',
+                'prefix': 'A picture of',
+            }
+
+            model = PrismerCaption(config)
+            state_dict = torch.load(f'prismer/logging/caption_{exp_name}/pytorch_model.bin', map_location='cuda:0')
+
         model.load_state_dict(state_dict)
         model.eval()