| import torch |
| import timm |
|
|
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| def load_model(model_path, model_type="vit"): |
| if model_type == "clip": |
| model = timm.create_model("vit_base_patch16_clip_224.openai", pretrained=False, num_classes=1) |
| else: |
| model = timm.create_model("vit_base_patch16_224", pretrained=False, num_classes=1) |
|
|
| checkpoint = torch.load(model_path, map_location="cpu") |
|
|
| if isinstance(checkpoint, dict): |
| state_dict = checkpoint.get("model_state_dict", checkpoint.get("state_dict", checkpoint)) |
| else: |
| state_dict = checkpoint |
|
|
| model.load_state_dict(state_dict, strict=False) |
|
|
| model.to(DEVICE) |
| model.eval() |
|
|
| return model |
