Fix model loading and file paths

backend/baselines/ViT/ViT_LRP.py

@@ -445,12 +445,13 @@ def _conv_filter(state_dict, patch_size=16):
 
 # ViT_LRP
 def vit_base_patch16_224(pretrained=False, **kwargs):
+    checkpoint_dir = kwargs.pop('checkpoint_dir', None)
     model = VisionTransformer(
         patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, **kwargs)
     model.default_cfg = default_cfgs['vit_base_patch16_224']
     if pretrained:
         load_pretrained(
-            model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3), filter_fn=_conv_filter)
+            model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3), filter_fn=_conv_filter, checkpoint_dir=checkpoint_dir)
     return model
 
 def vit_large_patch16_224(pretrained=False, **kwargs):

backend/flask_backend.py

@@ -96,7 +96,7 @@ logger.info(f"Using checkpoint directory: {checkpoint_dir}")
 
 import json 
 
-all_model_details_json_list= list(json.load(open("models_details.json")))
+all_model_details_json_list= list(json.load(open(os.path.join(_script_dir, "models_details.json"))))
 all_model_details_dict = dict()
 for i, item in enumerate(all_model_details_json_list):
     all_model_details_dict[item["model_id"]] = item
@@ -168,10 +168,10 @@ def load_model_by_id(model_id):
         if model_id == "vit_base_patch16_224.augreg2_in21k_ft_in1k":
             # Original default model logic
             # Removed incorrect checkpoint_dir argument
-            model = vit_base_patch16_224(pretrained=True).to(target_device)
+            model = vit_base_patch16_224(pretrained=True, checkpoint_dir=checkpoint_dir).to(target_device)
             model.eval()
             
-            inference_model = vit_orig(pretrained=True).to(target_device)
+            inference_model = vit_orig(pretrained=True, checkpoint_dir=checkpoint_dir).to(target_device)
             inference_model.eval()
             
         elif model_id == "vit_base_patch14_reg4_dinov2.lvd142m":