Update app.py

app.py

@@ -46,10 +46,10 @@ configs = {
     "THRESHOLD": 0.5
     }
 
-# MODEL_REGISTRY = {
-#     "CheXFormer-small": "m42-health/CXformer-small",
-#     "ViT-base-16": "google/vit-base-patch16-224",
-# }
+MODEL_REGISTRY = {
+    "CheXFormer-small": "m42-health/CXformer-small",
+    "ViT-base-16": "google/vit-base-patch16-224",
+}
 
 MODEL_CACHE = {}
 
@@ -68,8 +68,8 @@ class get_pretrained_model(nn.Module):
         print(f"Loading pretrained [{model_name}] model")
 
         self.backbone = AutoModel.from_pretrained(
-            # MODEL_REGISTRY[model_name],
-            model_name,
+            MODEL_REGISTRY[model_name],
+            # model_name,
             trust_remote_code=True)
 
         hidden_size = self.backbone.config.hidden_size