import os
from huggingface_hub import snapshot_download

def download_rectified_noise():
    # 本地保存路径（可自行修改）
    local_dir = "/gemini/space/zhaozy/guzhenyu/UAVFlow/UAV_Flow_base/exps/jsflow-experiment/samples/VIRTUAL_imagenet256_labeled"

    print(f"开始下载模型仓库到 {local_dir} ...")

    try:
        # 使用 HF 官方推荐的 snapshot_download
        snapshot_download(
            repo_id="xiangzai/reference",
            repo_type="model",                 # 这是 model 仓库
            local_dir=local_dir,
            local_dir_use_symlinks=False,      # 下载真实文件，而不是软链接
            resume_download=True               # 支持断点续传
        )
        print("下载完成！")

    except Exception as e:
        print(f"下载出错: {e}")

if __name__ == "__main__":
    download_rectified_noise()