Spaces:
Runtime error
Runtime error
| import os | |
| import logging | |
| from huggingface_hub import hf_hub_download | |
| # Configure logger | |
| logger = logging.getLogger(__name__) | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") | |
| # Supported Hugging Face models (key → repo + file) | |
| HF_MODELS = { | |
| # Depth Estimation | |
| "dpt_hybrid_384": ("isl-org/MiDaS", "dpt_hybrid_384.pt"), | |
| "midas_v21_small_256": ("isl-org/MiDaS", "midas_v21_small_256.pt"), | |
| "midas_v21_384": ("isl-org/MiDaS", "midas_v21_384.pt"), | |
| "dpt_swin2_large_384": ("isl-org/MiDaS", "dpt_swin2_large_384.pt"), | |
| "dpt_beit_large_512": ("isl-org/MiDaS", "dpt_beit_large_512.pt"), | |
| # Object Detection | |
| "yolov5n": ("ultralytics/yolov5", "yolov5n.pt"), | |
| "yolov5s": ("ultralytics/yolov5", "yolov5s.pt"), | |
| "yolov8s": ("ultralytics/yolov8", "yolov8s.pt"), | |
| "yolov8l": ("ultralytics/yolov8", "yolov8l.pt"), | |
| "yolov11b": ("Ultralytics/YOLO11", "yolov11b.pt"), | |
| "rtdetr": ("IDEA-Research/RT-DETR", "rtdetr_r50vd_detr.pth"), | |
| # Semantic Segmentation | |
| "segformer_b0": ("nvidia/segformer-b0-finetuned-ade-512-512", "model.safetensors"), | |
| "segformer_b5": ("nvidia/segformer-b5-finetuned-ade-512-512", "model.safetensors"), | |
| "deeplabv3_resnet50": ("facebook/deeplabv3-resnet50", "pytorch_model.bin"), | |
| } | |
| def download_model_if_needed(model_key: str, save_path: str): | |
| """ | |
| Downloads the model from Hugging Face Hub if it's not already present. | |
| Args: | |
| model_key (str): Key from HF_MODELS dict. | |
| save_path (str): Local path to store the downloaded model. | |
| Raises: | |
| ValueError: If the model_key is not supported. | |
| """ | |
| if model_key not in HF_MODELS: | |
| logger.error(f"❌ Model key '{model_key}' not found in registry.") | |
| raise ValueError(f"Unsupported model key: {model_key}") | |
| repo_id, filename = HF_MODELS[model_key] | |
| if os.path.exists(save_path): | |
| logger.info(f"✅ Model '{model_key}' already exists at '{save_path}'. Skipping download.") | |
| return | |
| os.makedirs(os.path.dirname(save_path), exist_ok=True) | |
| logger.info(f"⬇️ Downloading '{model_key}' from Hugging Face Hub...") | |
| try: | |
| hf_hub_download( | |
| repo_id=repo_id, | |
| filename=filename, | |
| cache_dir=os.path.dirname(save_path), | |
| force_download=True # Set to False later if you want to cache | |
| ) | |
| logger.info(f"✅ Successfully downloaded '{model_key}' to '{save_path}'") | |
| except Exception as e: | |
| logger.error(f"❌ Download failed for '{model_key}': {e}") | |
| raise | |