Update src/predict.py

src/predict.py

@@ -11,6 +11,15 @@ from src.models.loupe.modeling_loupe import LoupeModel
 from src.models.loupe.image_precessing_loupe import LoupeImageProcessor
 from src.lit_model import LitModel
 
+from huggingface_hub import hf_hub_download
+
+ckpt_path = hf_hub_download(
+    repo_id="xxwyyds/Loupe",
+    filename="loupe_model/pretrained_weights/pe/PE-Core-L14-336.pt"
+)
+
+
+seg_ckpt = hf_hub_download(repo_id="xxwyyds/Loupe", filename="loupe_model/model_weigths/seg/model.safetensors")
 
 # Initialize hydra
 hydra.initialize(config_path="../configs", version_base=None)
@@ -19,16 +28,19 @@ hydra.initialize(config_path="../configs", version_base=None)
 # Load model configuration
 
 cfg = hydra.compose(config_name="infer")
+cfg.model.backbone_path = ckpt_path
 # seg:/home/xxw/Loupe/model_weigths/seg/model.safetensors
-cfg.ckpt.checkpoint_paths = ["model_weigths/seg/model.safetensors"]
+cfg.ckpt.checkpoint_paths = [seg_ckpt]
 loupe_config = LoupeConfig(stage=cfg.stage.name, **cfg.model)
 loupe = LoupeModel(loupe_config)
 model = LitModel(cfg, loupe)
 processor = LoupeImageProcessor(loupe_config)
 
 # cls:/home/xxw/Loupe/model_weigths/cls/model.safetensors
+cls_ckpt = hf_hub_download(repo_id="xxwyyds/Loupe", filename="loupe_model/model_weigths/cls/model.safetensors")
 cfc = hydra.compose(config_name="infer")
-cfc.ckpt.checkpoint_paths = ["model_weigths/cls/model.safetensors"]
+cfc.ckpt.checkpoint_paths = [cls_ckpt]
+cfc.model.backbone_path = ckpt_path
 cls_loupe_config = LoupeConfig(stage=cfc.stage.name, **cfc.model)
 cls_loupe = LoupeModel(cls_loupe_config)
 cls_model = LitModel(cfc, cls_loupe)