#!/usr/bin/env bash
# Triton 小模型「原始可下载来源」-> /srv/mirror/models/triton/<name>/
# 注: torch/paddle2onnx/tf2onnx 等导出步骤在构建机执行(见各 1/README.md)；
#     本脚本只镜像公网可直接下载的源权重/产物。P2(movinet/deepfake/syncnet/yolov10)默认不抓。
set -uo pipefail
OUT=/srv/mirror/models/triton
mkdir -p "$OUT"

dl(){ # url out
  mkdir -p "$(dirname "$2")"; echo "[dl] $2"
  curl -fL --retry 5 --retry-delay 2 -C - -o "$2" "$1" || {
    case "$1" in https://github.com/*|https://raw.githubusercontent.com/*)
      echo "   direct failed -> ghproxy"; curl -fL --retry 5 -C - -o "$2" "https://ghproxy.net/$1" || echo "   FAILED: $2";;
      *) echo "   FAILED: $2";; esac; }
}

# nudenet —— 官方已发布 ONNX, 直接可用
dl https://github.com/notAI-tech/NudeNet/releases/download/v3.4-weights/640m.onnx "$OUT/nudenet/detector_640.onnx"

# paddleocr —— PP-OCRv4 det+rec 原始 inference tar (构建机 paddle2onnx 导出 det.onnx/rec.onnx)
dl https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_det_infer.tar "$OUT/paddleocr/ch_PP-OCRv4_det_infer.tar"
dl https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_rec_infer.tar "$OUT/paddleocr/ch_PP-OCRv4_rec_infer.tar"

# rtdetr —— RT-DETRv2 r18vd 原始权重 (构建机 export_onnx)
dl https://github.com/lyuwenyu/storage/releases/download/v0.2/rtdetrv2_r18vd_120e_coco_rerun_48.1.pth "$OUT/rtdetr/rtdetrv2_r18vd_120e_coco_rerun_48.1.pth"

# sed (PANNs CNN14 16k) —— 原始 checkpoint + AudioSet 类表 (构建机 torch->onnx)
dl "https://zenodo.org/record/3987831/files/Cnn14_16k_mAP%3D0.438.pth?download=1" "$OUT/sed/Cnn14_16k_mAP=0.438.pth"
dl "http://storage.googleapis.com/us_audioset/youtube_corpus/v1/csv/class_labels_indices.csv" "$OUT/sed/class_labels_indices.csv"

# ecapa_tdnn —— speechbrain 原始模型目录 (构建机 torch->onnx)
HF=/home/ecs-user/dlvenv/bin/hf
HF_ENDPOINT=https://hf-mirror.com HF_HUB_ENABLE_HF_TRANSFER=1 "$HF" download speechbrain/spkrec-ecapa-voxceleb --local-dir "$OUT/ecapa_tdnn/spkrec-ecapa-voxceleb" 2>&1 | tail -3 || echo "   FAILED: ecapa"

# 生成各子目录 .files 清单
for d in "$OUT"/*/; do
  ( cd "$d" && find . -type f ! -name .files -printf '%P\n' | sort > .files )
done
echo "=== triton small models done ==="
find "$OUT" -type f ! -name .files -printf '%s\t%p\n' | sort -k2
