From 563b8beccb0f00678b778e630bfd5d6282855c62 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9D=80=E7=81=AB=E7=9A=84=E5=86=B0=E5=9D=97nya?= Date: Sat, 27 Jun 2026 13:27:26 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=B8=8E=E4=BC=98=E5=8C=96?= =?UTF-8?q?=20prepare=5Fmodel=5Fscript.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- prepare_model_script.py | 54 +++++++++++++++++++++++------------------ 1 file changed, 31 insertions(+), 23 deletions(-) diff --git a/prepare_model_script.py b/prepare_model_script.py index 16b906e..5b897e0 100644 --- a/prepare_model_script.py +++ b/prepare_model_script.py @@ -11,18 +11,19 @@ "name": "OFA-Sys/chinese-clip-vit-large-patch14", "desc": "chinese-clip-vit-large-patch14 (高性能，3GB)" }, - "3": { - "name": "OFA-Sys/chinese-clip-rn50", - "desc": "chinese-clip-rn50 (快速，体积小)" - } + # "3": { + # "name": "OFA-Sys/chinese-clip-rn50", + # "desc": "chinese-clip-rn50 (快速，体积小)" + # } } MIRRORS = { - "1": None, # 官方站点 - "2": "https://hf-mirror.com/", - "3": "https://hf-cdn.sufy.com/" + "1": None, # 官方站点 + "2": "https://hf-mirror.com", + "3": "https://hf-cdn.sufy.com" } + # --- 辅助函数 --- def select_model(): @@ -30,23 +31,24 @@ def select_model(): print("\n请选择要使用的模型：") for key, value in MODELS.items(): print(f"{key}) {value['desc']}") - - choice = input("请输入选择([1],2,3): ") + + choice = input("请输入选择([1],2): ") if not choice.strip() or choice not in MODELS: - choice = "1" # 默认选项 - + choice = "1" # 默认选项 + selected_model = MODELS[choice]["name"] print(f"选择模型: {selected_model}") return selected_model + def select_download_path(model_name): """交互式地提示用户选择下载位置。""" model_folder_name = model_name.split('/')[-1] - + print("\n请选择模型下载方式：") print("1) 将模型下载到当前目录") print("2) 自定义下载目录") - + choice = input("请输入选择([1],2): ") if not choice.strip() or choice == "1": # 下载到当前目录下的一个子文件夹 @@ -54,10 +56,11 @@ def select_download_path(model_name): else: custom_path = input("请输入自定义目录路径: ") local_dir = os.path.join(custom_path, model_folder_name) - + print(f"模型将下载到: {local_dir}") return local_dir + def select_mirror(): """交互式地提示用户选择 Hugging Face 镜像。""" print("\n请选择 HuggingFace 镜像站点：") @@ -81,14 +84,14 @@ def select_mirror(): if endpoint: print(f"镜像站设置为：{endpoint}") - # 设置环境变量，huggingface_hub 会自动使用 - os.environ['HF_ENDPOINT'] = endpoint + return endpoint else: print("将使用官方站点进行下载。") + def set_proxy(): """交互式地提示用户设置代理。""" - print() # 增加一个换行以获得更好的格式 + print() # 增加一个换行以获得更好的格式 proxy_url = input("设置代理（留空不设置）：").strip() if proxy_url: # 为 http 和 https 同时设置环境变量 @@ -101,23 +104,25 @@ def set_proxy(): os.environ.pop('HTTPS_PROXY', None) print("不使用代理。") + def check_model_existence(local_dir): """检查模型是否似乎已经下载。""" # 简单地检查一个通用文件是否存在，例如 config.json return os.path.exists(os.path.join(local_dir, 'config.json')) + # --- 主执行流程 --- if __name__ == "__main__": print("正在检测模型...") - + # 此处假设模型不存在并总是询问。 # 更复杂的检查可以记住上次选择的模型。 print("未检测到模型。") # 1. 选择模型 model_repo_id = select_model() - + # 2. 选择下载路径 download_dir = select_download_path(model_repo_id) @@ -130,8 +135,8 @@ def check_model_existence(local_dir): exit() # 4. 选择镜像 - select_mirror() - + endpoint = select_mirror() + # 5. 设置代理 set_proxy() @@ -140,7 +145,9 @@ def check_model_existence(local_dir): try: snapshot_download( repo_id=model_repo_id, - local_dir=download_dir + local_dir=download_dir, + endpoint=endpoint, + ignore_patterns='*.pt' ) print("下载完成！") @@ -148,8 +155,9 @@ def check_model_existence(local_dir): config_filename = "config.json" with open(config_filename, "w", encoding="utf-8") as f: import json + json.dump(config_data, f, ensure_ascii=False, indent=4) print(f"模型路径配置已保存到 {config_filename} 文件中。") except Exception as e: print(f"\n下载过程中发生错误: {e}") - print("请检查您的网络连接、镜像站点、代理设置或目录权限。") \ No newline at end of file + print("请检查您的网络连接、镜像站点、代理设置或目录权限。")