Source code for minestudio.models.utils.download

'''
Date: 2024-12-14 01:46:36
LastEditors: muzhancun muzhancun@126.com
LastEditTime: 2024-12-14 02:00:17
FilePath: /MineStudio/minestudio/models/utils/download.py
'''
import huggingface_hub
import os
import pathlib

[docs] def download_model(model_name: str, local_dir: str = "downloads") -> str: """Downloads a specified model from Hugging Face Hub if it doesn't exist locally. Prompts the user for confirmation before downloading. The model is saved to a subdirectory within `local_dir` named after `model_name`. :param model_name: The name of the model to download. Valid names are "ROCKET-1", "VPT", "GROOT", "STEVE-1". :type model_name: str :param local_dir: The base directory to save downloaded models. Defaults to "downloads". This will be relative to the parent directory of this script's location. :type local_dir: str :returns: The local path to the downloaded model directory if successful or if download was skipped by user choice but the directory was expected to exist (though it might not in that case). Returns None if download is skipped and directory doesn't exist, or if an error occurs. :rtype: str :raises AssertionError: if `model_name` is not one of the recognized names. """ assert model_name in ["ROCKET-1", "VPT", "GROOT", "STEVE-1"], f"Unknown model: {model_name}" local_dir = os.path.join(pathlib.Path(__file__).parent.parent.absolute(), local_dir, model_name) download = False response = input(f"Detecting missing {model_name}, do you want to download it from huggingface (Y/N)?\n") if not os.path.exists(local_dir): while True: if response == 'Y' or response == 'y': download = True break elif response == 'N' or response == 'n': break else: response = input("Please input Y or N:\n") if not download: return None print(f"Downloading {model_name} to {local_dir}") huggingface_hub.hf_hub_download(repo_id=f'CraftJarvis/{model_name}', filename='.', local_dir=local_dir, repo_type='model') return local_dir