import os
import sys
from huggingface_hub import snapshot_download

# 设置保存路径
local_dir = os.path.join(os.getcwd(), "FLUX_Model")

print(f"正在下载 FLUX.1-Kontext-dev 模型到: {local_dir}")
print("注意：此模型非常大 (约 20GB+)，请确保磁盘空间充足。")
print("这可能需要很长时间...")

try:
    snapshot_download(
        repo_id="black-forest-labs/FLUX.1-Kontext-dev",
        local_dir=local_dir,
        local_dir_use_symlinks=False,
        resume_download=True,
        max_workers=8,
        # 可以排除一些不必要的文件以节省时间和空间
        # ignore_patterns=["*.msgpack", "*.bin"] 
    )
    print("FLUX 模型下载成功！")
except Exception as e:
    print(f"下载出错: {e}")
    sys.exit(1)