#!/bin/bash
# 分块并行下载 + sha256 校验 + 原子落位
set -u
BASE=https://hf-mirror.com
TMP=/home/zyw/yue2-par
mkdir -p "$TMP"

get() {  # repo rev rel destdir chunks [sha256]
  local repo="$1" rev="$2" rel="$3" dest="$4" chunks="${5:-8}" want="${6:-}"
  local url="$BASE/$repo/resolve/$rev/$rel" out="$dest/$rel"
  mkdir -p "$dest/$(dirname "$rel")"
  local total; total=$(curl -sIL -m 40 "$url" | grep -i '^content-length' | tail -1 | tr -d '\r' | awk '{print $2}')
  [ -z "$total" ] && { echo "FAIL 取不到大小 $rel"; return 1; }
  if [ -s "$out" ] && [ "$(stat -c%s "$out")" = "$total" ]; then
    if [ -z "$want" ] || [ "$(sha256sum "$out" | cut -d' ' -f1)" = "$want" ]; then
      echo "SKIP 已完整且校验通过 $rel"; return 0
    fi
    echo "WARN 大小对但 sha256 不符，重下 $rel"
  fi
  local work="$TMP/$(echo "$repo-$rel" | tr '/' '_')"
  rm -rf "$work"; mkdir -p "$work"
  local part=$(( total / chunks + 1 )) pids=()
  for i in $(seq 0 $((chunks-1))); do
    local s=$(( i * part )) e=$(( i * part + part - 1 ))
    [ $e -ge "$total" ] && e=$(( total - 1 )); [ $s -gt $e ] && continue
    ( curl -L --retry 8 --retry-delay 3 --retry-all-errors -sS -r "$s-$e" -o "$work/p$i" "$url" ) &
    pids+=($!)
  done
  for p in "${pids[@]}"; do wait "$p"; done
  cat "$work"/p* > "$work/full" 2>/dev/null
  local got; got=$(stat -c%s "$work/full" 2>/dev/null || echo 0)
  if [ "$got" != "$total" ]; then echo "FAIL $rel 大小 $got/$total"; return 1; fi
  if [ -n "$want" ]; then
    local have; have=$(sha256sum "$work/full" | cut -d' ' -f1)
    if [ "$have" != "$want" ]; then echo "FAIL $rel sha256 不符: $have != $want"; return 1; fi
    echo "sha256 OK $rel"
  fi
  mv "$work/full" "$out"; rm -rf "$work"
  echo "OK $rel $(stat -c%s "$out")"
}

verify() {  # file want
  local f="$1" want="$2"
  [ -s "$f" ] || { echo "缺失 $f"; return 1; }
  local have; have=$(sha256sum "$f" | cut -d' ' -f1)
  if [ "$have" = "$want" ]; then echo "校验通过 $(basename "$(dirname "$f")")/$(basename "$f")"; else echo "校验失败 $f: $have != $want"; fi
}

M=/home/zyw/ComfyUI/models
echo "=== 校验已有文件 $(date +%T) ==="
verify "$M/MERT2/MERT-v2-FullSong/model.safetensors" e6dd2ab187d6dd62b6521cd7d8f932e237acf0c5757745a7232082e28391350d
python3 - <<'PY'
import json,os,hashlib,glob
M="/home/zyw/ComfyUI/models"
for repo,dest in (("YuE2/YuE2-Vae",f"{M}/YuE2/YuE2-Vae"),("SheetSage2",f"{M}/SheetSage2")):
    mp=os.path.join(dest,"weights_manifest.json")
    if not os.path.exists(mp): print(f"  {repo}: 无 manifest"); continue
    d=json.load(open(mp))
    files=d.get("files") or {d.get("filename","model.safetensors"):d}
    for name,info in files.items():
        p=os.path.join(dest,name); want=info.get("sha256")
        if not os.path.exists(p): print(f"  {repo}/{name}: 缺失"); continue
        h=hashlib.sha256(open(p,'rb').read()).hexdigest() if os.path.getsize(p)<2**30*2 else "过大跳过"
        print(f"  {repo}/{name}: {'OK' if h==want else '不符' if h!='过大跳过' else '跳过(自行校验)'}")
PY

echo "=== 下载 YuE2-3B 主权重 $(date +%T) ==="
get m-a-p/YuE2-3B main model.safetensors "$M/YuE2/YuE2-3B" 8 1d55c42c1a9875c34f5d736e15078449992b044e807ce2a138e6cf289a1e59e9
echo "=== 完成 $(date +%T) ==="
ls -l "$M/YuE2/YuE2-3B/model.safetensors"
df -h / | tail -1