import os
import gdown

# 检查checkpoints文件夹是否存在，如果不存在则创建
checkpoints_dir = './checkpoints'
if not os.path.exists(checkpoints_dir):
    os.makedirs(checkpoints_dir)
    print(f"已创建文件夹: {checkpoints_dir}")

# 设置checkpoint.pt的完整路径
checkpoint_file = os.path.join(checkpoints_dir, 'checkpoint.pt')

# 检查checkpoint.pt是否已存在
if not os.path.exists(checkpoint_file):
    # 从你的链接中提取的文件ID
    file_id = '1cuvdiu3bXyni71A2hoeZSWT1NOsNfeD_'
    
    # 构建下载URL
    download_url = f'https://drive.google.com/uc?id={file_id}'
    
    # 下载文件
    print(f"正在下载文件到: {checkpoint_file}")
    gdown.download(download_url, checkpoint_file, quiet=False)
    print("下载完成！")
else:
    print(f"文件已存在: {checkpoint_file}")