Update app.py

app.py

@@ -92,8 +92,13 @@ def load_model(select_skpt):
         global_pool=args.global_pool,
     ).to(device)
 
-    args.resume = os.path.join(CKPT_SAVE_PATH, ckpt)
-    args.resume = CKPT_PATH[ckpt]
+    args.resume = os.path.join(CKPT_SAVE_PATH, CKPT_PATH[ckpt])
+    if os.path.isfile(args.resume) == False:
+        hf_hub_download(local_dir=CKPT_SAVE_PATH,
+                        local_dir_use_symlinks=False,
+                        repo_id='Wolowolo/fsfm-3c',
+                        filename=CKPT_PATH[ckpt])
+    args.resume = os.path.join(CKPT_SAVE_PATH, CKPT_PATH[ckpt])
     checkpoint = torch.load(args.resume, map_location=device)
     model.load_state_dict(checkpoint['model'], strict=False)
     model.eval()
@@ -246,15 +251,10 @@ CKPT_NAME = [
     'DfD-Checkpoint_Fine-tuned_on_FF++',
     'FAS-Checkpoint_Fine-tuned_on_MCIO',
 ]
-# CKPT_PATH = {
-#     '✨Unified-detector_v1_Fine-tuned_on_4_classes': 'finetuned_models/Unified-detector/v1_Fine-tuned_on_4_classes/checkpoint-min_val_loss.pth',
-#     'DfD-Checkpoint_Fine-tuned_on_FF++': 'finetuned_models/FF++_c23_32frames/checkpoint-min_val_loss.pth',
-#     'FAS-Checkpoint_Fine-tuned_on_MCIO': 'finetuned_models/MCIO_protocol/Both_MCIO/checkpoint-min_val_loss.pth',
-# }
 CKPT_PATH = {
-    '✨Unified-detector_v1_Fine-tuned_on_4_classes': './checkpoints/checkpoint-min_train_loss.pth',
-    'DfD-Checkpoint_Fine-tuned_on_FF++': '/mnt/localDisk2/wgj/FSFM/released/FSFM-main/fsfm-3c/finuetune/cross_dataset_DfD/checkpoint/finetuned_models/ft_on_FF++_c23_32frames/pt_from_VF2_ViT-B_epoch600/checkpoint-min_val_loss.pth',
-    'FAS-Checkpoint_Fine-tuned_on_MCIO': '/mnt/localDisk2/wgj/FSFM/FSFM-3C/codespace/fsfm-3c/finuetune/cross_dataset_DfD/finetuned_models/FAS_MCIO/checkpoint-199.pth',
+    '✨Unified-detector_v1_Fine-tuned_on_4_classes': 'finetuned_models/Unified-detector/v1_Fine-tuned_on_4_classes/checkpoint-min_train_loss.pth',
+    'DfD-Checkpoint_Fine-tuned_on_FF++': 'finetuned_models/FF++_c23_32frames/checkpoint-min_val_loss.pth',
+    'FAS-Checkpoint_Fine-tuned_on_MCIO': 'finetuned_models/MCIO_protocol/Both_MCIO/checkpoint-min_val_loss.pth',
 }
 CKPT_CLASS = {
     '✨Unified-detector_v1_Fine-tuned_on_4_classes': 4,
@@ -321,7 +321,7 @@ with gr.Blocks(css=".custom-label { font-weight: bold !important; font-size: 16p
 if __name__ == "__main__":
     args = get_args_parser()
     args = args.parse_args()
-    ckpt = 'DfD-Checkpoint_Fine-tuned_on_FF++'
+    ckpt = '✨Unified-detector_v1_Fine-tuned_on_4_classes'
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     args.nb_classes = CKPT_CLASS[ckpt]
     model = models_vit.__dict__[CKPT_MODEL[ckpt]](
@@ -329,8 +329,13 @@ if __name__ == "__main__":
         drop_path_rate=args.drop_path,
         global_pool=args.global_pool,
     ).to(device)
-    args.resume = os.path.join(CKPT_SAVE_PATH, ckpt)
-    args.resume = CKPT_PATH[ckpt]
+    args.resume = os.path.join(CKPT_SAVE_PATH, CKPT_PATH[ckpt])
+    if os.path.isfile(args.resume) == False:
+        hf_hub_download(local_dir=CKPT_SAVE_PATH,
+                        local_dir_use_symlinks=False,
+                        repo_id='Wolowolo/fsfm-3c',
+                        filename=CKPT_PATH[ckpt])
+    args.resume = os.path.join(CKPT_SAVE_PATH, CKPT_PATH[ckpt])
     checkpoint = torch.load(args.resume, map_location=device)
     model.load_state_dict(checkpoint['model'], strict=False)
     model.eval()