#!/bin/bash
# 用 hf-mirror 的 resolve 直链下载 YuE2 全套模型（绕过 hf CLI 的 Xet 401 问题）
set -u
M=/home/zyw/ComfyUI/models
API=https://hf-mirror.com/api/models
BASE=https://hf-mirror.com

dl_repo() {
  local repo="$1" rev="$2" dest="$3"
  mkdir -p "$dest"
  echo "=== repo $repo @ $rev -> $dest ==="
  local files
  files=$(curl -s -m 60 "$API/$repo/revision/$rev" | python3 -c "
import sys,json
try: d=json.load(sys.stdin)
except Exception: sys.exit(1)
for f in d.get('siblings',[]): print(f['rfilename'])
")
  if [ -z "$files" ]; then echo "!! 无法获取文件列表 $repo"; return 1; fi
  while IFS= read -r f; do
    [ -z "$f" ] && continue
    case "$f" in */.gitattributes|.gitattributes) continue;; esac
    local out="$dest/$f"
    mkdir -p "$(dirname "$out")"
    local exp
    exp=$(curl -sIL -m 30 "$BASE/$repo/resolve/$rev/$f" | grep -i '^content-length' | tail -1 | tr -d '\r' | awk '{print $2}')
    if [ -s "$out" ] && [ -n "$exp" ] && [ "$(stat -c%s "$out")" = "$exp" ]; then
      echo "SKIP $f"; continue
    fi
    echo "GET  $f (期望 ${exp:-?} 字节)"
    curl -L -C - --retry 8 --retry-delay 3 --retry-all-errors -sS -o "$out" \
      "$BASE/$repo/resolve/$rev/$f" \
      && echo "OK   $f $(stat -c%s "$out")" || echo "FAIL $f"
  done <<< "$files"
}

dl_repo m-a-p/YuE2-3B          main                                     $M/YuE2/YuE2-3B
dl_repo m-a-p/YuE2-Vae         main                                     $M/YuE2/YuE2-Vae
dl_repo m-a-p/SheetSage2       main                                     $M/SheetSage2
dl_repo m-a-p/MERT-v2-FullSong d8ba1c745e733b3908ce6ad16ebeb17ac7600a42 $M/MERT2/MERT-v2-FullSong

echo "=== 汇总 ==="
du -sh $M/YuE2/YuE2-3B $M/YuE2/YuE2-Vae $M/SheetSage2 $M/MERT2/MERT-v2-FullSong
find $M/YuE2 $M/SheetSage2 $M/MERT2 -name "*.safetensors" -printf "%10s  %p\n"
echo "ALL_DONE"