File size: 4,151 Bytes
ac29381 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | #!/bin/bash
# DreamZero 模型权重一键下载脚本
#
# 用法:
# bash scripts/data/download_checkpoints.sh
# bash scripts/data/download_checkpoints.sh --mirror # 使用 HF 镜像（中国用户）
# bash scripts/data/download_checkpoints.sh --offline /path/to/cache # 从本地缓存加载
#
# 下载内容:
# - Wan2.2-TI2V-5B: Wan2.2 视频 DiT 骨干 (~32G)
# - umt5-xxl: T5 tokenizer (~49G)
# - CLIP encoder: 图像编码器 (~2.4G)
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
CHECKPOINT_DIR="${CHECKPOINT_DIR:-$REPO_ROOT/checkpoints}"
MIRROR=false
OFFLINE_CACHE=""
# Parse arguments
while [[ $# -gt 0 ]]; do
case $1 in
--mirror)
MIRROR=true
shift
;;
--offline)
OFFLINE_CACHE="$2"
shift 2
;;
*)
echo "Unknown option: $1"
echo "Usage: $0 [--mirror] [--offline /path/to/cache]"
exit 1
;;
esac
done
# ============ 网络配置 ============
if [ "$MIRROR" = true ]; then
export HF_ENDPOINT=https://hf-mirror.com
echo ">>> 使用 HF 镜像: hf-mirror.com"
fi
# ============ 离线模式 ============
if [ -n "$OFFLINE_CACHE" ]; then
if [ ! -d "$OFFLINE_CACHE" ]; then
echo "ERROR: 离线缓存目录不存在: $OFFLINE_CACHE"
exit 1
fi
echo ">>> 离线模式，从 $OFFLINE_CACHE 复制..."
cp -r "$OFFLINE_CACHE"/* "$CHECKPOINT_DIR/"
echo ">>> 复制完成！"
exit 0
fi
# ============ 下载 Wan2.2-TI2V-5B ============
WAN22_DIR="$CHECKPOINT_DIR/Wan2.2-TI2V-5B"
if [ -d "$WAN22_DIR" ] && [ -n "$(ls -A "$WAN22_DIR" 2>/dev/null)" ]; then
echo "[SKIP] Wan2.2-TI2V-5B 已存在: $WAN22_DIR"
else
echo "[DOWNLOAD] Wan2.2-TI2V-5B → $WAN22_DIR"
mkdir -p "$WAN22_DIR"
huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_DIR"
echo "[OK] Wan2.2-TI2V-5B 下载完成"
fi
# ============ 下载 umt5-xxl ============
TOKENIZER_DIR="$CHECKPOINT_DIR/umt5-xxl"
if [ -d "$TOKENIZER_DIR" ] && [ -n "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then
echo "[SKIP] umt5-xxl 已存在: $TOKENIZER_DIR"
else
echo "[DOWNLOAD] umt5-xxl → $TOKENIZER_DIR"
mkdir -p "$TOKENIZER_DIR"
huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR"
echo "[OK] umt5-xxl 下载完成"
fi
# ============ 下载 CLIP encoder ============
CLIP_DIR="$CHECKPOINT_DIR/clip-encoder"
CLIP_FILE="$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"
if [ -f "$CLIP_FILE" ]; then
echo "[SKIP] CLIP encoder 已存在: $CLIP_FILE"
else
echo "[DOWNLOAD] CLIP encoder → $CLIP_DIR"
mkdir -p "$CLIP_DIR"
# 尝试从 dreamzero 独立仓库下载
if huggingface-cli download poet70/clip-encoder \
--local-dir "$CLIP_DIR" 2>/dev/null; then
echo "[OK] CLIP encoder 下载完成"
else
echo "[WARN] 独立 CLIP 仓库暂不可用，尝试从 Wan2.1 提取..."
# 后备：下载 Wan2.1-I2V-14B-480P（包含 CLIP 权重）
WAN21_DIR="$CHECKPOINT_DIR/Wan2.1-I2V-14B-480P"
if [ ! -d "$WAN21_DIR" ]; then
huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN21_DIR"
fi
# 复制 CLIP 权重
if [ -f "$WAN21_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then
cp "$WAN21_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" "$CLIP_DIR/"
echo "[OK] CLIP encoder 从 Wan2.1 提取完成"
else
echo "[ERROR] 无法获取 CLIP encoder"
echo "请手动从 DreamZero-DROID 检查点提取:"
echo " cp /path/to/DreamZero-DROID/models_clip_*.pth $CLIP_DIR/"
exit 1
fi
fi
fi
echo ""
echo "============================================"
echo " 所有权重下载完成！"
echo "============================================"
echo "Wan2.2: $WAN22_DIR"
echo "umt5-xxl: $TOKENIZER_DIR"
echo "CLIP: $CLIP_DIR"
echo "============================================"
|