#!/bin/bash
set -euo pipefail

# =============================================
# ComfyUI models on persistent /workspace + symlink /ComfyUI/models
# Adds: FLUX ControlNet Union (InstantX)
# =============================================

# Персистентный путь под модели
WORK_PATH="/workspace/runpod-slim/ComfyUI"
STORAGE_MODELS="$WORK_PATH/models"

# Подпапки моделей
CLIP_DIR="$STORAGE_MODELS/clip"
DIFF_DIR="$STORAGE_MODELS/diffusion_models"   # Flux checkpoint
VAE_DIR="$STORAGE_MODELS/vae"
LORA_DIR="$STORAGE_MODELS/loras"
CN_DIR="$STORAGE_MODELS/controlnet"           # <-- ControlNet сюда

# Создаём структуру
mkdir -p "$CLIP_DIR" "$DIFF_DIR" "$VAE_DIR" "$LORA_DIR" "$CN_DIR"

echo ">>> Using persistent storage: $STORAGE_MODELS"

# Симлинк /ComfyUI/models -> /workspace/comfyui/models
# if [ ! -L /ComfyUI/models ]; then
#   echo ">>> Linking /ComfyUI/models -> $STORAGE_MODELS"
#   rm -rf /ComfyUI/models
#   ln -s "$STORAGE_MODELS" /ComfyUI/models
# else
#   echo ">>> Symlink /ComfyUI/models already exists"
# fi

# --- TE-only encoders (FLAN-T5-XXL FP16 & ViT-L TE-only) ---
echo ">>> Downloading TE-only encoders -> $CLIP_DIR"
curl -fL -C - --retry 5 --retry-all-errors \
  "https://huggingface.co/easygoing0114/flan-t5-xxl-fused/resolve/main/flan_t5_xxl_TE-only_FP16.safetensors" \
  -o "$CLIP_DIR/flan_t5_xxl_TE-only_FP16.safetensors"

curl -fL -C - --retry 5 --retry-all-errors \
  "https://huggingface.co/zer0int/CLIP-GmP-ViT-L-14/resolve/main/ViT-L-14-TEXT-detail-improved-hiT-GmP-TE-only-HF.safetensors" \
  -o "$CLIP_DIR/vit_l_te_only.safetensors"

# --- VAE (официальный через HF_TOKEN или публичный фолбэк) ---
echo ">>> Downloading VAE -> $VAE_DIR"
if [[ -n "${HF_TOKEN:-}" ]]; then
  echo ">>> HF_TOKEN present; trying official BFL VAE"
  if curl -fL -C - --retry 5 --retry-all-errors \
       -H "Authorization: Bearer $HF_TOKEN" \
       "https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/ae.safetensors" \
       -o "$VAE_DIR/ae.safetensors"; then
    echo ">>> Official VAE downloaded"
  else
    echo ">>> Official VAE failed; using public mirror"
    curl -fL -C - --retry 5 --retry-all-errors \
      "https://huggingface.co/ffxvs/vae-flux/resolve/main/ae.safetensors" \
      -o "$VAE_DIR/ae.safetensors"
  fi
else
  echo ">>> No HF_TOKEN; using public mirror VAE"
  curl -fL -C - --retry 5 --retry-all-errors \
    "https://huggingface.co/ffxvs/vae-flux/resolve/main/ae.safetensors" \
    -o "$VAE_DIR/ae.safetensors"
fi

# --- Flux1-dev checkpoint (Comfy-Org, public) ---
echo ">>> Downloading flux1-dev checkpoint -> $DIFF_DIR"
curl -fL -C - --retry 5 --retry-all-errors \
  "https://huggingface.co/Comfy-Org/flux1-dev/resolve/main/flux1-dev.safetensors" \
  -o "$DIFF_DIR/flux1-dev.safetensors"

# --- ControlNet Union for FLUX.1-dev (InstantX) ---
# Файлы: diffusion_pytorch_model.safetensors (~6.6 GB) + config.json
echo ">>> Downloading FLUX ControlNet Union -> $CN_DIR"
curl -fL -C - --retry 5 --retry-all-errors \
  "https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union/resolve/main/diffusion_pytorch_model.safetensors" \
  -o "$CN_DIR/flux1-dev-controlnet-union.safetensors"

curl -fL -C - --retry 5 --retry-all-errors \
  "https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union/resolve/main/config.json" \
  -o "$CN_DIR/flux1-dev-controlnet-union.json"

# --- LoRAs (public) ---
echo ">>> Downloading LoRAs -> $LORA_DIR"
for f in Ernie_LoRA_FLUX Ostin_LoRA_FLUX Sheep_LoRA_FLUX; do
  curl -fL -C - --retry 5 --retry-all-errors \
    "https://huggingface.co/playrix/loras/resolve/main/${f}.safetensors" \
    -o "$LORA_DIR/${f}.safetensors"
done

echo ">>> DONE. Models stored under $STORAGE_MODELS and visible via /ComfyUI/models"
echo ">>> ComfyUI: Reload Models."
echo ">>> Load Diffusion:   models/diffusion_models/flux1-dev.safetensors"
echo ">>> Load TE-only:     models/clip/flan_t5_xxl_TE-only_FP16.safetensors + models/clip/vit_l_te_only.safetensors"
echo ">>> Load VAE:         models/vae/ae.safetensors"
echo ">>> Load ControlNet:  models/controlnet/flux1-dev-controlnet-union.safetensors"