#!/bin/bash

# Wan2.2 I2V A14B 模型下载脚本

# 从 HuggingFace 下载 Wan-AI/Wan2.2-I2V-A14B 模型

set -e # 遇到错误时退出

echo "=========================================="

echo "Wan2.2 I2V A14B 模型下载脚本"

echo "=========================================="

# 检查 huggingface-cli 是否安装

if ! command -v huggingface-cli &> /dev/null; then

echo "错误: huggingface-cli 未安装"

echo "请运行: pip install huggingface_hub[cli]"

exit 1

fi

# 创建目录结构

echo "创建模型目录..."

mkdir -p models

mkdir -p output

# 下载函数

download_model() {

local repo_id=$1

local local_dir=$2

local description=$3

echo "----------------------------------------"

echo "下载 $description..."

echo "仓库: $repo_id"

echo "目标目录: $local_dir"

echo "----------------------------------------"

if [ -d "$local_dir" ] && [ "$(ls -A $local_dir)" ]; then

echo "目录 $local_dir 已存在且非空，跳过下载"

return 0

fi

huggingface-cli download "$repo_id" --local-dir "$local_dir" --local-dir-use-symlinks False

if [ $? -eq 0 ]; then

echo "✅ $description 下载完成"

else

echo "❌ $description 下载失败"

return 1

fi

}

# 验证下载的文件

verify_files() {

echo "=========================================="

echo "验证下载的模型文件..."

echo "=========================================="

local required_files=(

"Wan2.1_VAE.pth"

"models_t5_umt5-xxl-enc-bf16.pth"

"low_noise_model"

"high_noise_model"

"google"

)

local missing_files=()

for file in "${required_files[@]}"; do

if [ ! -e "./models/$file" ]; then

missing_files+=("$file")

fi

done

if [ ${#missing_files[@]} -eq 0 ]; then

echo "✅ 所有必需文件都已下载"

return 0

else

echo "❌ 缺少以下文件:"

for file in "${missing_files[@]}"; do

echo " - $file"

done

return 1

fi

}

# 检查磁盘空间

check_disk_space() {

local required_space=130 # GB (I2V A14B model is ~126GB)

local available_space=$(df . | tail -1 | awk '{print $4}')

local available_gb=$((available_space / 1024 / 1024))

echo "检查磁盘空间..."

echo "可用空间: ${available_gb}GB"

echo "所需空间: ${required_space}GB (I2V A14B 模型约 126GB)"

if [ $available_gb -lt $required_space ]; then

echo "⚠️ 警告: 磁盘空间不足，需要至少 ${required_space}GB 可用空间"

read -p "是否继续? (y/N): " -n 1 -r

echo

if [[ ! $REPLY =~ ^[Yy]$ ]]; then

exit 1

fi

fi

}

# 主下载流程

main() {

check_disk_space

echo "开始下载 I2V A14B 模型..."

echo "注意: 模型约 126GB，下载时间取决于您的网络速度"

echo ""

# 下载 I2V A14B 模型

download_model "Wan-AI/Wan2.2-I2V-A14B" "./models" "I2V A14B 模型"

# 验证下载的文件

verify_files

# 显示下载结果

echo ""

echo "=========================================="

echo "下载完成！"

echo "=========================================="

echo "模型文件位置: ./models/"

echo " ├── Wan2.1_VAE.pth"

echo " ├── models_t5_umt5-xxl-enc-bf16.pth"

echo " ├── low_noise_model/"

echo " ├── high_noise_model/"

echo " └── google/"

echo ""

echo "输出目录: ./output/"

echo ""

echo "现在可以运行以下命令构建并启动 Docker 容器:"

echo " docker-compose up -d --build"

echo ""

}

# 清理函数

cleanup() {

echo ""

echo "清理临时文件..."

# 可以在这里添加清理逻辑

}

# 设置信号处理

trap cleanup EXIT

# 运行主函数

main "$@"