#!/bin/bash
set -e
REPO="https://huggingface.co/LuckyOda/comfyui-full-pack/resolve/main"
COMFY="${COMFYUI_DIR:-/workspace/ComfyUI}"
CD="$COMFY/models"
PIP="${VENV_PIP:-/venv/main/bin/pip}"
[ ! -f "$PIP" ] && PIP="pip"

dl() {
    local f="$1" d="$2"
    mkdir -p "$d"
    local target="$d/$(basename "$f")"
    [ -f "$target" ] && echo "  [SKIP] $f" && return
    echo "  [DL] $f"
    wget -q --show-progress -O "$target" "$REPO/$(echo $f | sed 's/ /%20/g')" || true
}

echo "=== Wan Video Ref Pack ==="
echo ">>> Workflow: Ген Видео Фон берется с Видео Рефа"

# ============================================
# CUSTOM NODES
# ============================================
echo ">>> Installing custom nodes"
cd "$COMFY/custom_nodes"
repos=(
    "https://github.com/kijai/ComfyUI-WanVideoWrapper.git"
    "https://github.com/kijai/ComfyUI-WanAnimatePreprocess.git"
    "https://github.com/kijai/ComfyUI-KJNodes.git"
    "https://github.com/rgthree/rgthree-comfy.git"
    "https://github.com/pythongosssss/ComfyUI-Custom-Scripts.git"
    "https://github.com/yolain/ComfyUI-Easy-Use.git"
    "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git"
    "https://github.com/kijai/ComfyUI-segment-anything-2.git"
    "https://github.com/ltdrdata/ComfyUI-Impact-Pack.git"
    "https://github.com/chflame163/ComfyUI_LayerStyle.git"
    "https://github.com/cubiq/ComfyUI_essentials.git"
    "https://github.com/Fannovel16/comfyui_controlnet_aux.git"
    "https://github.com/Comfy-Org/ComfyUI-Manager.git"
)
for r in "${repos[@]}"; do
    dir=$(basename "$r" .git)
    [ -d "$dir" ] && echo "  [SKIP] $dir" && continue
    echo "  [CLONE] $dir"
    git clone "$r" --recursive 2>/dev/null || true
done

echo ">>> Installing pip dependencies (minimal — preserving base packages)"
$PIP install -q piexif PyWavelets 2>/dev/null || true
# NOTE: We do NOT bulk-install requirements.txt from every node
# as it can overwrite working package versions and break generation.
# If a specific node fails to load, install its requirements individually.

# ============================================
# MODELS
# ============================================
echo ">>> diffusion_models"
dl "wan2.2_animate_14B_bf16.safetensors" "$CD/diffusion_models"

echo ">>> text_encoders"
dl "EchoVault_T9.safetensors" "$CD/text_encoders"

echo ">>> vae"
dl "GlassRoot_D2.safetensors" "$CD/vae"

echo ">>> clip_vision"
dl "IronSight_V7.safetensors" "$CD/clip_vision"

echo ">>> loras"
dl "SolarFlint_L2.safetensors" "$CD/loras"
dl "VelvetRush_Q4.safetensors" "$CD/loras"
dl "FrostByte_K7.safetensors" "$CD/loras"
dl "PhantomWeave_R5.safetensors" "$CD/loras"
dl "NovaMind_X1.safetensors" "$CD/loras"

echo ">>> detection"
dl "vitpose_h_wholebody_data.bin" "$CD/detection"
dl "vitpose_h_wholebody_model.onnx" "$CD/detection"
dl "yolov10m.onnx" "$CD/detection"

echo ">>> sam2"
dl "sam2.1_hiera_base_plus-fp16.safetensors" "$CD/sam2"

echo ">>> rife (frame interpolation)"
mkdir -p "$CD/../custom_nodes/ComfyUI-WanVideoWrapper/rife"
RIFE_DIR="$CD/../custom_nodes/ComfyUI-WanVideoWrapper/rife"
if [ ! -f "$RIFE_DIR/rife49.pth" ]; then
    echo "  [DL] rife49.pth"
    wget -q --show-progress -O "$RIFE_DIR/rife49.pth" "https://huggingface.co/Kijai/rife_onnx/resolve/main/rife49.pth" || true
fi

echo ">>> Restarting ComfyUI"
supervisorctl restart comfyui 2>/dev/null || true

echo "=== Wan Video Ref Pack Complete (~55GB) ==="