#!/bin/bash
# NEIA VIDEO TRANSFER - Download de Modelos
# Uso: curl -sL https://neia.systems/baixar-modelos.sh | bash

set -e

echo ""
echo "========================================"
echo " NEIA VIDEO TRANSFER - DOWNLOAD ~46 GB"
echo " https://neia.systems"
echo "========================================"
echo ""

# Detectar pasta do ComfyUI
MODELS=""
for loc in /workspace/ComfyUI/models /workspace/comfyui/models /root/ComfyUI/models /opt/ComfyUI/models $HOME/ComfyUI/models /runpod-volume/ComfyUI/models; do
    [ -d "$loc" ] && MODELS="$loc" && break
done

[ -z "$MODELS" ] && { echo "[ERRO] ComfyUI nao encontrado. Defina MODELS= e rode novamente."; exit 1; }
echo "[OK] ComfyUI: $MODELS"

# Instalar aria2 se nao tiver (muito mais rapido)
if ! command -v aria2c &>/dev/null; then
    echo "[+] Instalando aria2 para downloads rapidos..."
    apt-get update -qq && apt-get install -y -qq aria2
fi

# Criar pastas
mkdir -p "$MODELS"/{diffusion_models/Wan22Animate,clip,vae,clip_vision,checkpoints,loras,vitpose,yolo}

# Funcao de download
dl() {
    local url="$1" out="$2" name="$3"
    [ -f "$out" ] && echo "[$name] OK - ja existe" && return
    echo "[$name] Baixando..."
    aria2c -x 16 -s 16 -k 1M -c -q --file-allocation=none -d "$(dirname "$out")" -o "$(basename "$out")" "$url"
    echo "[$name] OK"
}

echo ""
echo "Baixando modelos..."
echo ""

dl "https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/resolve/main/Wan22Animate/Wan2_2-Animate-14B_fp8_scaled_e4m3fn_KJ_v2.safetensors" \
   "$MODELS/diffusion_models/Wan22Animate/Wan2_2-Animate-14B_fp8_scaled_e4m3fn_KJ_v2.safetensors" \
   "1/13 Animate-14B (17.3GB)"

dl "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/umt5-xxl-enc-bf16.safetensors" \
   "$MODELS/clip/umt5-xxl-enc-bf16.safetensors" \
   "2/13 TextEncoder (11.4GB)"

dl "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Wan2_1_VAE_bf16.safetensors" \
   "$MODELS/vae/Wan2_1_VAE_bf16.safetensors" \
   "3/13 VAE (335MB)"

dl "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/clip_vision/clip_vision_h.safetensors" \
   "$MODELS/clip_vision/clip_vision_h.safetensors" \
   "4/13 CLIP Vision (1.26GB)"

dl "https://huggingface.co/VeryAladeen/Sec-4B/resolve/main/SeC-4B-fp8.safetensors" \
   "$MODELS/checkpoints/SeC-4B-fp8.safetensors" \
   "5/13 SeC-4B (4.26GB)"

dl "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/LoRAs/Wan22_relight/WanAnimate_relight_lora_fp16.safetensors" \
   "$MODELS/loras/WanAnimate_relight_lora_fp16.safetensors" \
   "6/13 LoRA Relight (1.44GB)"

dl "https://huggingface.co/lightx2v/Wan2.2-Lightning/resolve/main/Wan2.2-I2V-A14B-4steps-lora-rank64-Seko-V1/low_noise_model.safetensors" \
   "$MODELS/loras/Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors" \
   "7/13 LoRA Lightning (1.23GB)"

dl "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/FastWan/FastWan_T2V_14B_480p_lora_rank_128_bf16.safetensors" \
   "$MODELS/loras/FastWan_T2V_14B_480p_lora_rank_128_bf16.safetensors" \
   "8/13 LoRA FastWan (1.25GB)"

dl "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/Pusa/Wan21_PusaV1_LoRA_14B_rank512_bf16.safetensors" \
   "$MODELS/loras/Wan21_PusaV1_LoRA_14B_rank512_bf16.safetensors" \
   "9/13 LoRA PusaV1 (4.91GB)"

dl "https://huggingface.co/Kijai/WanVideo_comfy/resolve/main/LoRAs/Wan22_FunReward/Wan2.2-Fun-A14B-InP-LOW-HPS2.1_resized_dynamic_avg_rank_15_bf16.safetensors" \
   "$MODELS/loras/Wan2.2-Fun-A14B-InP-LOW-HPS2.1_resized_dynamic_avg_rank_15_bf16.safetensors" \
   "10/13 LoRA FunReward (102MB)"

dl "https://huggingface.co/Kijai/vitpose_comfy/resolve/main/onnx/vitpose_h_wholebody_model.onnx" \
   "$MODELS/vitpose/vitpose_h_wholebody_model.onnx" \
   "11/13 VitPose Model (420KB)"

dl "https://huggingface.co/Kijai/vitpose_comfy/resolve/main/onnx/vitpose_h_wholebody_data.bin" \
   "$MODELS/vitpose/vitpose_h_wholebody_data.bin" \
   "12/13 VitPose Data (2.55GB)"

dl "https://huggingface.co/Wan-AI/Wan2.2-Animate-14B/resolve/main/process_checkpoint/det/yolov10m.onnx" \
   "$MODELS/yolo/yolov10m.onnx" \
   "13/13 YOLO (61.7MB)"

echo ""
echo "========================================"
echo " DOWNLOAD COMPLETO!"
echo "========================================"
echo ""
echo "Instale os Custom Nodes no ComfyUI Manager:"
echo "  - ComfyUI-WanVideoWrapper"
echo "  - ComfyUI-WanAnimatePreprocess"
echo "  - ComfyUI-VideoHelperSuite"
echo "  - ComfyUI-KJNodes"
echo ""
echo "Pronto! Carregue o workflow e seja feliz!"
echo "https://neia.systems"
echo ""