Spaces:
Runtime error
Runtime error
| """Model loader for HuggingFace Hub and local checkpoint cache.""" | |
| import os | |
| from typing import Optional | |
| import torch | |
| from src.models.model_factory import DRSeverityModel, build_model | |
| from src.utils.checkpoint import load_checkpoint | |
| def load_app_model( | |
| repo_id: Optional[str] = None, | |
| filename: str = "coral_best.pth", | |
| local_checkpoint_path: Optional[str] = "checkpoints/coral_best.pth", | |
| device: torch.device = torch.device("cpu"), | |
| ) -> DRSeverityModel: | |
| """Load model from Hugging Face Hub or local path with CPU optimization. | |
| Args: | |
| repo_id: HuggingFace repository ID (e.g. 'username/dr-severity-grading'). | |
| filename: Checkpoint filename in HF repository. | |
| local_checkpoint_path: Fallback local path if repo_id is None. | |
| device: Torch device (defaults to CPU for HF Spaces free tier). | |
| Returns: | |
| Loaded DRSeverityModel in eval mode. | |
| """ | |
| checkpoint_file = None | |
| if repo_id: | |
| try: | |
| from huggingface_hub import hf_hub_download | |
| checkpoint_file = hf_hub_download(repo_id=repo_id, filename=filename) | |
| except Exception as e: | |
| print(f"HF Hub download failed: {e}. Checking local cache...") | |
| if not checkpoint_file or not os.path.exists(checkpoint_file): | |
| if local_checkpoint_path and os.path.exists(local_checkpoint_path): | |
| checkpoint_file = local_checkpoint_path | |
| class Config: | |
| backbone = "efficientnet_b0" | |
| pretrained = False | |
| num_classes = 5 | |
| drop_rate = 0.2 | |
| drop_path_rate = 0.1 | |
| model = build_model(Config(), loss_name="coral") | |
| if checkpoint_file and os.path.exists(checkpoint_file): | |
| ckpt = load_checkpoint(checkpoint_file, device=device) | |
| model.load_state_dict(ckpt["state_dict"]) | |
| print(f"Loaded weights from {checkpoint_file}") | |
| else: | |
| print("Warning: No checkpoint loaded. Model initialized with random weights.") | |
| model.to(device) | |
| model.eval() | |
| return model | |