#!/bin/bash
# Flux Model Download Script with Defensive Programming
# Handles model downloads with robust error handling and recovery
set -euo pipefail # Exit on error, undefined variables, and pipe failures
IFS=$'\n\t' # Set Internal Field Separator
# Script configuration
readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
readonly PROJECT_ROOT="$(dirname "${SCRIPT_DIR}")"
readonly MODELS_DIR="${PROJECT_ROOT}/models"
readonly LOG_FILE="${PROJECT_ROOT}/download-models.log"
readonly LOCK_FILE="${PROJECT_ROOT}/.download.lock"
# Model URLs and configurations
readonly T5_XXL_URL="https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp16.safetensors"
readonly CLIP_L_URL="https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors"
readonly VAE_URL="https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors"
readonly SCHNELL_URL="https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors"
# Colors
readonly YELLOW='\033[1;33m'
readonly GREEN='\033[0;32m'
readonly RED='\033[0;31m'
readonly BLUE='\033[0;34m'
readonly NC='\033[0m'
# Trap for cleanup
cleanup() {
local exit_code=$?
if [[ -f "${LOCK_FILE}" ]]; then
rm -f "${LOCK_FILE}"
fi
if [[ ${exit_code} -ne 0 ]]; then
log ERROR "Download script failed with exit code ${exit_code}"
fi
exit ${exit_code}
}
trap cleanup EXIT INT TERM
# Logging function
log() {
local level="$1"
shift
local message="$*"
local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
# Log to file
echo "[${timestamp}] [${level}] ${message}" >> "${LOG_FILE}"
# Display to console
case "${level}" in
ERROR)
echo -e "${RED}[ERROR] ${message}${NC}" >&2
;;
WARN)
echo -e "${YELLOW}[WARN] ${message}${NC}"
;;
INFO)
echo -e "${GREEN}[INFO] ${message}${NC}"
;;
DEBUG)
if [[ "${DEBUG:-0}" == "1" ]]; then
echo -e "${BLUE}[DEBUG] ${message}${NC}"
fi
;;
esac
}
# Prevent concurrent downloads
acquire_lock() {
if [[ -f "${LOCK_FILE}" ]]; then
log ERROR "Another download is already in progress"
log ERROR "If this is incorrect, remove ${LOCK_FILE}"
exit 1
fi
touch "${LOCK_FILE}"
}
# Validate network connectivity
check_network() {
log DEBUG "Checking network connectivity..."
local test_urls=(
"https://huggingface.co"
"https://github.com"
)
for url in "${test_urls[@]}"; do
if curl -s --head --fail "${url}" > /dev/null 2>&1; then
log DEBUG "Network check passed: ${url}"
return 0
fi
done
log ERROR "No network connectivity to required services"
return 1
}
# Check available disk space
check_disk_space() {
local required_gb="${1:-30}"
local target_dir="${2:-${MODELS_DIR}}"
# Create directory if it doesn't exist
mkdir -p "${target_dir}"
local available_gb
available_gb=$(df -BG "${target_dir}" | awk 'NR==2 {print int($4)}')
if [[ ${available_gb} -lt ${required_gb} ]]; then
log ERROR "Insufficient disk space: ${available_gb}GB available, ${required_gb}GB required"
return 1
fi
log DEBUG "Disk space check passed: ${available_gb}GB available"
return 0
}
# Verify file integrity
verify_file() {
local file="$1"
local min_size="${2:-1000000}" # Default 1MB minimum
if [[ ! -f "${file}" ]]; then
log ERROR "File does not exist: ${file}"
return 1
fi
local file_size
file_size=$(stat -f%z "${file}" 2>/dev/null || stat -c%s "${file}" 2>/dev/null || echo "0")
if [[ ${file_size} -lt ${min_size} ]]; then
log ERROR "File appears corrupted (too small): ${file}"
rm -f "${file}"
return 1
fi
log DEBUG "File verification passed: ${file} (${file_size} bytes)"
return 0
}
# Download with resume support and progress
download_file() {
local url="$1"
local output="$2"
local description="${3:-File}"
local max_retries="${4:-3}"
log INFO "Downloading ${description}..."
log DEBUG "URL: ${url}"
log DEBUG "Output: ${output}"
# Create output directory
local output_dir
output_dir=$(dirname "${output}")
mkdir -p "${output_dir}"
# Check if file already exists and is valid
if [[ -f "${output}" ]]; then
if verify_file "${output}" 1000000; then
log INFO "${description} already exists and is valid, skipping..."
return 0
fi
log WARN "Existing file appears corrupted, re-downloading..."
fi
# Download with retries
local attempt=1
while [[ ${attempt} -le ${max_retries} ]]; do
log DEBUG "Download attempt ${attempt}/${max_retries}"
# Use appropriate downloader
if command -v wget &> /dev/null; then
if wget -c "${url}" -O "${output}" --progress=bar:force 2>&1 | \
grep --line-buffered "%" | \
sed -u -e "s/\.//g" | \
awk '{printf("\r[PROGRESS] %s", $0)}'; then
echo # New line after progress
if verify_file "${output}" 1000000; then
log INFO "${description} downloaded successfully"
return 0
fi
fi
elif command -v curl &> /dev/null; then
if curl -L -C - "${url}" -o "${output}" --progress-bar; then
if verify_file "${output}" 1000000; then
log INFO "${description} downloaded successfully"
return 0
fi
fi
else
log ERROR "Neither wget nor curl is installed"
return 1
fi
log WARN "Download attempt ${attempt} failed"
((attempt++))
if [[ ${attempt} -le ${max_retries} ]]; then
log INFO "Retrying in 5 seconds..."
sleep 5
fi
done
log ERROR "Failed to download ${description} after ${max_retries} attempts"
return 1
}
# Check and install huggingface-cli
check_hf_cli() {
if command -v huggingface-cli &> /dev/null; then
log DEBUG "huggingface-cli is already installed"
return 0
fi
log INFO "Installing huggingface-cli..."
if command -v pip &> /dev/null; then
if pip install -q huggingface-hub; then
log INFO "huggingface-cli installed successfully"
return 0
fi
elif command -v pip3 &> /dev/null; then
if pip3 install -q huggingface-hub; then
log INFO "huggingface-cli installed successfully"
return 0
fi
fi
log WARN "Could not install huggingface-cli automatically"
log INFO "Please install it manually: pip install huggingface-hub"
return 1
}
# Validate Hugging Face authentication
check_hf_auth() {
log INFO "Checking Hugging Face authentication..."
# Check environment variable
if [[ -z "${HF_TOKEN:-}" ]]; then
log WARN "HF_TOKEN environment variable not set"
echo ""
echo "Flux.1-dev requires Hugging Face authentication:"
echo "1. Create account at: https://huggingface.co/join"
echo "2. Accept license at: https://huggingface.co/black-forest-labs/FLUX.1-dev"
echo "3. Create token at: https://huggingface.co/settings/tokens"
echo ""
read -p "Enter your Hugging Face token (or press Enter to skip): " -s hf_token
echo
if [[ -z "${hf_token}" ]]; then
log INFO "Skipping Flux.1-dev download"
return 1
fi
export HF_TOKEN="${hf_token}"
fi
# Validate token
if command -v huggingface-cli &> /dev/null; then
if huggingface-cli whoami &> /dev/null; then
log INFO "Hugging Face authentication successful"
return 0
else
log ERROR "Invalid Hugging Face token"
return 1
fi
else
log WARN "Cannot validate token without huggingface-cli"
return 0 # Assume token is valid
fi
}
# Download Flux.1-dev using huggingface-cli
download_flux_dev() {
log INFO "Preparing to download Flux.1-dev..."
if ! check_hf_cli; then
log WARN "huggingface-cli not available, skipping Flux.1-dev"
return 1
fi
if ! check_hf_auth; then
return 1
fi
local output_dir="${MODELS_DIR}/unet"
mkdir -p "${output_dir}"
log INFO "Downloading Flux.1-dev (this may take a while, ~23GB)..."
# Download with huggingface-cli
if huggingface-cli download \
black-forest-labs/FLUX.1-dev \
flux1-dev.safetensors \
--local-dir "${output_dir}" \
--local-dir-use-symlinks False \
--token "${HF_TOKEN}" 2>&1 | tee -a "${LOG_FILE}"; then
# Verify download
local model_file="${output_dir}/flux1-dev.safetensors"
if verify_file "${model_file}" 1000000000; then # 1GB minimum
log INFO "Flux.1-dev downloaded successfully"
return 0
fi
fi
log ERROR "Failed to download Flux.1-dev"
return 1
}
# Download Flux.1-schnell (non-gated)
download_flux_schnell() {
log INFO "Downloading Flux.1-schnell..."
local output_file="${MODELS_DIR}/unet/flux1-schnell.safetensors"
if ! check_disk_space 15; then
log ERROR "Insufficient disk space for Flux.1-schnell"
return 1
fi
download_file \
"${SCHNELL_URL}" \
"${output_file}" \
"Flux.1-schnell model (~12GB)" \
3
}
# Download CLIP encoders
download_clip_models() {
log INFO "Downloading CLIP text encoders..."
local clip_dir="${MODELS_DIR}/clip"
mkdir -p "${clip_dir}"
# T5-XXL encoder
download_file \
"${T5_XXL_URL}" \
"${clip_dir}/t5xxl_fp16.safetensors" \
"T5-XXL text encoder" \
3
# CLIP-L encoder
download_file \
"${CLIP_L_URL}" \
"${clip_dir}/clip_l.safetensors" \
"CLIP-L text encoder" \
3
}
# Download VAE
download_vae() {
log INFO "Downloading Flux VAE..."
local vae_dir="${MODELS_DIR}/vae"
mkdir -p "${vae_dir}"
download_file \
"${VAE_URL}" \
"${vae_dir}/ae.safetensors" \
"Flux VAE" \
3
}
# Download RMBG Background Removal Model
download_rmbg() {
log INFO "Downloading RMBG-1.4 background removal model..."
local rmbg_dir="${MODELS_DIR}/rmbg"
mkdir -p "${rmbg_dir}"
local rmbg_url="https://huggingface.co/briaai/RMBG-1.4/resolve/main/model.pth"
download_file \
"${rmbg_url}" \
"${rmbg_dir}/RMBG-1.4.pth" \
"RMBG-1.4 model" \
3
}
# Interactive menu
show_menu() {
echo ""
echo -e "${BLUE}╔══════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║ Flux Model Download Menu ║${NC}"
echo -e "${BLUE}╚══════════════════════════════════════════════════════════╝${NC}"
echo ""
echo "Select models to download:"
echo ""
echo "1) Flux.1-dev (best quality, requires HF auth, ~23GB)"
echo "2) Flux.1-schnell (good quality, no auth, ~12GB)"
echo "3) Both models"
echo "4) Text encoders only (CLIP + T5)"
echo "5) VAE only"
echo "6) Complete setup (recommended)"
echo "7) Skip download"
echo ""
read -p "Enter your choice (1-7): " -n 1 -r choice
echo
case "${choice}" in
1)
download_flux_dev || log WARN "Flux.1-dev download failed"
download_clip_models
download_vae
;;
2)
download_flux_schnell
download_clip_models
download_vae
;;
3)
download_flux_dev || log WARN "Flux.1-dev download failed"
download_flux_schnell
download_clip_models
download_vae
;;
4)
download_clip_models
;;
5)
download_vae
;;
6)
log INFO "Starting complete setup..."
download_flux_schnell
download_flux_dev || log WARN "Flux.1-dev skipped"
download_clip_models
download_vae
;;
7)
log INFO "Skipping model download"
return 0
;;
*)
log ERROR "Invalid choice: ${choice}"
return 1
;;
esac
}
# Verify installation
verify_installation() {
log INFO "Verifying model installation..."
local models_found=0
local issues=()
# Check for main models
if [[ -f "${MODELS_DIR}/unet/flux1-dev.safetensors" ]]; then
log INFO "✓ Flux.1-dev found"
((models_found++))
elif [[ -f "${MODELS_DIR}/unet/flux1-schnell.safetensors" ]]; then
log INFO "✓ Flux.1-schnell found"
((models_found++))
else
issues+=("No Flux model found in ${MODELS_DIR}/unet/")
fi
# Check CLIP encoders
if [[ -f "${MODELS_DIR}/clip/t5xxl_fp16.safetensors" ]]; then
log INFO "✓ T5-XXL encoder found"
else
issues+=("T5-XXL encoder not found")
fi
if [[ -f "${MODELS_DIR}/clip/clip_l.safetensors" ]]; then
log INFO "✓ CLIP-L encoder found"
else
issues+=("CLIP-L encoder not found")
fi
# Check VAE
if [[ -f "${MODELS_DIR}/vae/ae.safetensors" ]]; then
log INFO "✓ VAE found"
else
issues+=("VAE not found")
fi
# Report results
if [[ ${#issues[@]} -eq 0 ]]; then
log INFO "All required models are installed!"
return 0
else
log WARN "Missing components:"
for issue in "${issues[@]}"; do
log WARN " - ${issue}"
done
if [[ ${models_found} -eq 0 ]]; then
log ERROR "No Flux models installed. ComfyUI will not work properly."
return 1
fi
return 0
fi
}
# Main execution
main() {
# Initialize
echo "Model download started at $(date)" > "${LOG_FILE}"
log INFO "Starting model download script"
# Acquire lock
acquire_lock
# Run checks
check_network || exit 1
check_disk_space 50 || {
log WARN "Limited disk space available"
read -p "Continue anyway? (y/n): " -n 1 -r
echo
[[ ${REPLY} =~ ^[Yy]$ ]] || exit 1
}
# Create model directories
log INFO "Creating model directories..."
mkdir -p "${MODELS_DIR}"/{unet,clip,vae,checkpoints,loras,embeddings,rmbg}
# Show menu and download
show_menu
# Verify installation
verify_installation
log INFO "Model setup completed"
echo ""
echo -e "${GREEN}Model locations:${NC}"
echo " • Main models: ${MODELS_DIR}/unet/"
echo " • CLIP encoders: ${MODELS_DIR}/clip/"
echo " • VAE: ${MODELS_DIR}/vae/"
echo ""
echo -e "${GREEN}Log file: ${LOG_FILE}${NC}"
}
# Parse arguments
while [[ $# -gt 0 ]]; do
case "$1" in
--debug)
DEBUG=1
;;
--auto)
AUTO_MODE=1
;;
--help|-h)
cat << EOF
Usage: $0 [OPTIONS]
Options:
--debug Enable debug output
--auto Run in automatic mode (download Schnell + encoders)
--help Show this help message
This script downloads Flux models for ComfyUI.
EOF
exit 0
;;
*)
log ERROR "Unknown option: $1"
exit 1
;;
esac
shift
done
# Run main
main
