Update app.py

app.py

@@ -45,10 +45,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 = {}
 
@@ -67,7 +67,8 @@ class get_pretrained_model(nn.Module):
         print(f"Loading pretrained [{model_name}] model")
 
         self.backbone = AutoModel.from_pretrained(
-            MODEL_REGISTRY[model_name],
+            # MODEL_REGISTRY[model_name],
+            model_name,
             trust_remote_code=True)
 
         hidden_size = self.backbone.config.hidden_size
@@ -252,7 +253,8 @@ def run_diagnosis(
     x = preprocess_fn(input_image).unsqueeze(0)
 
     # Resolve backbone
-    ckpt_path = os.path.join(CKPT_ROOT, MODEL_REGISTRY[backbone_name])
+    # ckpt_path = os.path.join(CKPT_ROOT, MODEL_REGISTRY[backbone_name])
+    kpt_path = os.path.join(CKPT_ROOT, backbone_name)
 
     if not os.path.exists(ckpt_path):
         raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
@@ -260,7 +262,7 @@ def run_diagnosis(
     # Load model (cache for speed)
     if backbone_name not in MODEL_CACHE:
       MODEL_CACHE[backbone_name] = modelModule.load_from_checkpoint(
-        ckpt_path, backbone_model_name=MODEL_REGISTRY[backbone_name], num_layers_to_unfreeze = 2)
+        ckpt_path, backbone_model_name=backbone_name, num_layers_to_unfreeze = 2)
     model = MODEL_CACHE[backbone_name]
 
     model.eval()