#!/usr/bin/env bash
set -e

# ============================================================
# CONFIGURATION
# ============================================================
# By default, it uses the current directory. 
# Change this if your ComfyUI is located elsewhere.
COMFYUI_ROOT=$(pwd)
COMFYUI_MODELS="${COMFYUI_ROOT}/models"

# Optional: Add your HF Token here if you encounter permission issues
# export HF_TOKEN="your_token_here"

# ============================================================
# HELPER FUNCTION
# ============================================================
download_model() {
    local url="$1"
    local dest="$2"
    
    if [ -f "$dest" ]; then
        echo -e "\e[36m>>> Skipping: $dest (Already exists)\e[0m"
    else
        echo -e "\e[33m>>> Downloading: $dest\e[0m"
        mkdir -p "$(dirname "$dest")"
        
        local auth_header=""
        if [ -n "$HF_TOKEN" ]; then
            auth_header="Authorization: Bearer $HF_TOKEN"
        fi

        if [ -n "$auth_header" ]; then
            wget -q --show-progress --header "$auth_header" -O "$dest" "$url"
        else
            wget -q --show-progress -O "$dest" "$url"
        fi
    fi
}

echo -e "\e[32mStarting downloads to: ${COMFYUI_MODELS}\e[0m"
echo "--------------------------------------------------"

# ============================================================
# 1. FLUX.2 KLEIN 9B
# ============================================================
echo "--- FLUX.2 KLEIN ---"

# Transformer
download_model \
  "https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-fp8/resolve/main/flux-2-klein-9b-fp8.safetensors" \
  "${COMFYUI_MODELS}/diffusion_models/flux-2-klein-9b-fp8.safetensors"

# Text Encoder
download_model \
  "https://huggingface.co/Comfy-Org/flux2-klein-9B/resolve/main/split_files/text_encoders/qwen_3_8b_fp8mixed.safetensors" \
  "${COMFYUI_MODELS}/text_encoders/qwen_3_8b_fp8mixed.safetensors"

# VAE
download_model \
  "https://huggingface.co/Comfy-Org/flux2-dev/resolve/main/split_files/vae/flux2-vae.safetensors" \
  "${COMFYUI_MODELS}/vae/flux2-vae.safetensors"

# LoRAs
download_model \
  "https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap/resolve/main/bfs_head_v1_flux-klein_9b_step3500_rank128.safetensors" \
  "${COMFYUI_MODELS}/loras/bfs_head_v1_flux-klein_9b_step3500_rank128.safetensors"

download_model \
  "https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap/resolve/main/bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors" \
  "${COMFYUI_MODELS}/loras/bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors"


# ============================================================
# 2. LTX-VIDEO 2 (LTXV-2)
# ============================================================
echo "--- LTX-VIDEO 2 ---"

# Transformer (DEV is Recommended)
download_model \
  "https://huggingface.co/Kijai/LTXV2_comfy/resolve/main/diffusion_models/ltx-2-19b-dev-fp8_transformer_only.safetensors" \
  "${COMFYUI_MODELS}/diffusion_models/ltx-2-19b-dev-fp8_transformer_only.safetensors"

# Text Encoders & Connector
download_model \
  "https://huggingface.co/Comfy-Org/ltx-2/resolve/main/split_files/text_encoders/gemma_3_12B_it_fp8_scaled.safetensors" \
  "${COMFYUI_MODELS}/text_encoders/gemma_3_12B_it_fp8_scaled.safetensors"

download_model \
  "https://huggingface.co/Kijai/LTXV2_comfy/resolve/main/text_encoders/ltx-2-19b-embeddings_connector_dev_bf16.safetensors" \
  "${COMFYUI_MODELS}/text_encoders/ltx-2-19b-embeddings_connector_dev_bf16.safetensors"

# VAEs
download_model \
  "https://huggingface.co/Kijai/LTXV2_comfy/resolve/main/VAE/LTX2_video_vae_bf16.safetensors" \
  "${COMFYUI_MODELS}/vae/LTX2_video_vae_bf16.safetensors"

download_model \
  "https://huggingface.co/Kijai/LTXV2_comfy/resolve/main/VAE/LTX2_audio_vae_bf16.safetensors" \
  "${COMFYUI_MODELS}/vae/LTX2_audio_vae_bf16.safetensors"

# Video LoRAs
download_model \
  "https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap-Video/resolve/main/ltx-2/head_swap_v1_13500_first_frame.safetensors" \
  "${COMFYUI_MODELS}/loras/head_swap_v1_13500_first_frame.safetensors"

download_model \
  "https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap-Video/resolve/main/ltx-2/head_swap_v1_8750_first_and_last_frame.safetensors" \
  "${COMFYUI_MODELS}/loras/head_swap_v1_8750_first_and_last_frame.safetensors"


# ============================================================
# 3. UPSCALER & EXTRAS
# ============================================================
echo "--- UPSCALER & EXTRAS ---"

download_model \
  "https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-spatial-upscaler-x2-1.0.safetensors" \
  "${COMFYUI_MODELS}/latent_upscale_models/ltx-2-spatial-upscaler-x2-1.0.safetensors"

download_model \
  "https://huggingface.co/Kijai/MelBandRoFormer_comfy/resolve/main/MelBandRoformer_fp32.safetensors" \
  "${COMFYUI_MODELS}/diffusion_models/MelBandRoformer_fp32.safetensors"

# ============================================================
# DONE
# ============================================================
echo "--------------------------------------------------"
echo -e "\e[32mAll models downloaded successfully!\e[0m"