#!/bin/bash
# ComfyUI Model Setup Script
# Edit the downloads section to add/remove models

set -e  # Exit on error

# --- Credentials (set these as vast.ai environment variables, not hardcoded) ---
HF_TOKEN="${HF_TOKEN:-}"
CIVITAI_TOKEN="${CIVITAI_TOKEN:-}"

# --- Base paths (standard for most vast.ai ComfyUI images) ---
COMFY_ROOT="/workspace/ComfyUI"
MODELS_DIR="$COMFY_ROOT/models"

# --- Helper functions ---
download_hf() {
  local dest_dir="$1"
  local repo="$2"       # e.g. "black-forest-labs/FLUX.1-dev"
  local filepath="$3"
  local filename="$4"   # e.g. "flux1-dev.safetensors"
  local url="https://huggingface.co/$repo/resolve/main/$filepath/$filename"

  if [ -f "$dest_dir/$filename" ]; then
    echo "Skipping $filename (already exists)"
    return 0
  fi

  mkdir -p "$dest_dir"
  echo "Downloading $filename from HuggingFace..."
  wget -q --show-progress \
    ${HF_TOKEN:+--header="Authorization: Bearer $HF_TOKEN"} \
    -O "$dest_dir/$filename" \
    "$url"
}

download_civitai() {
  local dest_dir="$1"
  local model_id="$2"   # numeric model version ID from the CivitAI URL
  local filename="$3"

  if [ -f "$dest_dir/$filename" ]; then
    echo "Skipping $filename (already exists)"
    return 0
  fi

  mkdir -p "$dest_dir"
  echo "Downloading $filename from CivitAI..."
  wget -q --show-progress \
    --content-disposition \
    -O "$dest_dir/$filename" \
    "https://civitai.com/api/download/models/$model_id?token=$CIVITAI_TOKEN"
}

# --- Downloads ---
# diffusion_models
download_hf "$MODELS_DIR/diffusion_models" "Kijai/LTX2.3_comfy" "diffusion_models" "ltx-2.3-22b-distilled_transformer_only_fp8_scaled.safetensors"

# latent_upscale_models
download_hf "$MODELS_DIR/latent_upscale_models" "Lightricks/LTX-2.3" "ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors"

# Checkpoints

# VAEs
download_hf "$MODELS_DIR/vae" "Comfy-Org/Wan_2.2_ComfyUI_Repackaged" "split_files/vae" "wan_2.1_vae.safetensors"
download_hf "$MODELS_DIR/vae" "Kijai/LTX2.3_comfy" "vae" "LTX23_video_vae_bf16.safetensors"
download_hf "$MODELS_DIR/vae" "Kijai/LTX2.3_comfy" "vae" "LTX23_audio_vae_bf16.safetensors"
download_hf "$MODELS_DIR/vae" "Kijai/LTX2.3_comfy" "vae" "taeltx2_3.safetensors"


# LoRAs

# CLIP / Text encoders
download_hf "$MODELS_DIR/test_encoders" "Comfy-Org/Wan_2.2_ComfyUI_Repackaged" "split_files/text_encoders" "umt5_xxl_fp16.safetensors"
download_hf "$MODELS_DIR/text_encoders" "Comfy-Org/ltx-2" "split_files/text_encoders" "gemma_3_12B_it_fp4_mixed.safetensors"
download_hf "$MODELS_DIR/text_encoders" "Kijai/LTX2.3_comfy" "text_encoders" "ltx-2.3_text_projection_bf16.safetensors"

echo "All models downloaded"