Files
bin/artifact_comfyui_download.sh

485 lines
13 KiB
Bash
Raw Normal View History

#!/bin/bash
#
# ComfyUI Model Downloader - A Beautiful CLI Tool
# Downloads AI models from HuggingFace with style and grace
#
# Usage: ./artifact_comfyui_download.sh [options]
#
set -euo pipefail
# ============================================================================
# COLOR PALETTE - Beautiful Terminal Colors
# ============================================================================
# Reset
RESET='\033[0m'
# Foreground Colors
BLACK='\033[0;30m'
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
CYAN='\033[0;36m'
WHITE='\033[0;37m'
# Bold
BOLD_BLACK='\033[1;30m'
BOLD_RED='\033[1;31m'
BOLD_GREEN='\033[1;32m'
BOLD_YELLOW='\033[1;33m'
BOLD_BLUE='\033[1;34m'
BOLD_MAGENTA='\033[1;35m'
BOLD_CYAN='\033[1;36m'
BOLD_WHITE='\033[1;37m'
# Background Colors
BG_BLACK='\033[40m'
BG_RED='\033[41m'
BG_GREEN='\033[42m'
BG_YELLOW='\033[43m'
BG_BLUE='\033[44m'
BG_MAGENTA='\033[45m'
BG_CYAN='\033[46m'
BG_WHITE='\033[47m'
# Styles
DIM='\033[2m'
ITALIC='\033[3m'
UNDERLINE='\033[4m'
BLINK='\033[5m'
REVERSE='\033[7m'
# ============================================================================
# UNICODE CHARACTERS - Make it Pretty
# ============================================================================
CHECK_MARK="✓"
CROSS_MARK="✗"
ROCKET="🚀"
PACKAGE="📦"
DOWNLOAD="⬇️"
SPARKLES="✨"
FIRE="🔥"
CLOCK="⏱️"
FOLDER="📁"
LINK="🔗"
STAR="⭐"
WARNING="⚠️"
INFO=""
ARROW_RIGHT="→"
DOUBLE_ARROW="»"
BOX_LIGHT="─"
BOX_HEAVY="━"
BOX_DOUBLE="═"
# ============================================================================
# CONFIGURATION
# ============================================================================
# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
# Default configuration file path
# Try multiple possible locations
if [[ -f "${HOME}/Projects/runpod/comfyui_models.yaml" ]]; then
CONFIG_FILE="${HOME}/Projects/runpod/comfyui_models.yaml"
elif [[ -f "${PROJECT_ROOT}/comfyui_models.yaml" ]]; then
CONFIG_FILE="${PROJECT_ROOT}/comfyui_models.yaml"
elif [[ -f "${SCRIPT_DIR}/comfyui_models.yaml" ]]; then
CONFIG_FILE="${SCRIPT_DIR}/comfyui_models.yaml"
else
CONFIG_FILE="" # No config file by default
fi
# Default cache directory (use HuggingFace default)
CACHE_DIR="${CACHE_DIR:-${HOME}/.cache/huggingface}"
# HuggingFace token from environment or .env file
if [[ -f "${PROJECT_ROOT}/.env" ]]; then
HF_TOKEN=$(grep ^HF_TOKEN "${PROJECT_ROOT}/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
fi
HF_TOKEN="${HF_TOKEN:-}"
# ============================================================================
# UTILITY FUNCTIONS - The Magic Happens Here
# ============================================================================
# Print functions with beautiful formatting
print_banner() {
local text="$1"
local width=80
local padding=$(( (width - ${#text} - 2) / 2 ))
echo -e ""
echo -e "${BOLD_CYAN}${BOX_DOUBLE}$(printf '%.0s'"${BOX_DOUBLE}" $(seq 1 $width))${BOX_DOUBLE}${RESET}"
echo -e "${BOLD_CYAN}${BOX_DOUBLE}$(printf '%.0s ' $(seq 1 $padding))${BOLD_MAGENTA}${text}$(printf '%.0s ' $(seq 1 $padding))${BOLD_CYAN}${BOX_DOUBLE}${RESET}"
echo -e "${BOLD_CYAN}${BOX_DOUBLE}$(printf '%.0s'"${BOX_DOUBLE}" $(seq 1 $width))${BOX_DOUBLE}${RESET}"
echo -e ""
}
print_section() {
local text="$1"
echo -e "\n${BOLD_YELLOW}${DOUBLE_ARROW} ${text}${RESET}"
echo -e "${CYAN}$(printf '%.0s'"${BOX_LIGHT}" $(seq 1 80))${RESET}"
}
print_success() {
echo -e "${BOLD_GREEN}${CHECK_MARK} $1${RESET}"
}
print_error() {
echo -e "${BOLD_RED}${CROSS_MARK} $1${RESET}" >&2
}
print_warning() {
echo -e "${BOLD_YELLOW}${WARNING} $1${RESET}"
}
print_info() {
echo -e "${BOLD_CYAN}${INFO} $1${RESET}"
}
print_step() {
local current="$1"
local total="$2"
local text="$3"
echo -e "${BOLD_BLUE}[${current}/${total}]${RESET} ${MAGENTA}${DOWNLOAD}${RESET} ${text}"
}
print_detail() {
echo -e " ${DIM}${CYAN}${ARROW_RIGHT} $1${RESET}"
}
# Progress bar function
show_progress() {
local current="$1"
local total="$2"
local width=50
local percentage=$((current * 100 / total))
local filled=$((current * width / total))
local empty=$((width - filled))
printf "\r ${BOLD_CYAN}Progress: ${RESET}["
printf "${BG_GREEN}${BOLD_WHITE}%${filled}s${RESET}" | tr ' ' '█'
printf "${DIM}%${empty}s${RESET}" | tr ' ' '░'
printf "] ${BOLD_YELLOW}%3d%%${RESET} ${DIM}(%d/%d)${RESET}" "$percentage" "$current" "$total"
}
# Parse YAML (simple implementation)
parse_yaml() {
local yaml_file="$1"
local category="$2"
python3 - <<EOPYAML
import yaml
import sys
try:
with open('${yaml_file}', 'r') as f:
config = yaml.safe_load(f)
if '${category}' == 'settings':
settings = config.get('settings', {})
print(f"CACHE_DIR={settings.get('cache_dir', '/workspace/huggingface_cache')}")
print(f"PARALLEL_DOWNLOADS={settings.get('parallel_downloads', 1)}")
elif '${category}' == 'categories':
for cat_name in config.get('model_categories', {}).keys():
print(cat_name)
elif '${category}' in config.get('model_categories', {}):
models = config['model_categories']['${category}']
for model in models:
repo_id = model.get('repo_id', '')
description = model.get('description', '')
size_gb = model.get('size_gb', 0)
essential = model.get('essential', False)
print(f"{repo_id}|{description}|{size_gb}|{essential}")
else:
sys.exit(1)
except Exception as e:
print(f"ERROR: {e}", file=sys.stderr)
sys.exit(1)
EOPYAML
}
# Check dependencies
check_dependencies() {
print_section "Checking Dependencies"
local missing_deps=()
# Check Python 3
if ! command -v python3 &> /dev/null; then
missing_deps+=("python3")
fi
# Check pip
if ! command -v pip3 &> /dev/null; then
missing_deps+=("pip3")
fi
# Check required Python packages
if ! python3 -c "import yaml" 2>/dev/null; then
print_warning "PyYAML not installed, installing..."
pip3 install pyyaml -q
fi
if ! python3 -c "import huggingface_hub" 2>/dev/null; then
print_warning "huggingface_hub not installed, installing..."
pip3 install huggingface_hub -q
fi
if [[ ${#missing_deps[@]} -gt 0 ]]; then
print_error "Missing dependencies: ${missing_deps[*]}"
exit 1
fi
print_success "All dependencies satisfied"
}
# Validate configuration
validate_config() {
print_section "Validating Configuration"
if [[ -n "$CONFIG_FILE" ]]; then
if [[ ! -f "$CONFIG_FILE" ]]; then
print_error "Configuration file not found: $CONFIG_FILE"
exit 1
fi
print_success "Configuration file found: ${CYAN}${CONFIG_FILE}${RESET}"
else
print_warning "No configuration file specified"
fi
if [[ -z "$HF_TOKEN" ]]; then
print_error "HF_TOKEN not set. Please set it in .env file or environment."
exit 1
fi
print_success "HuggingFace token configured: ${DIM}${HF_TOKEN:0:10}...${RESET}"
if [[ ! -d "$CACHE_DIR" ]]; then
print_info "Creating cache directory: ${CYAN}${CACHE_DIR}${RESET}"
mkdir -p "$CACHE_DIR"
fi
print_success "Cache directory ready: ${CYAN}${CACHE_DIR}${RESET}"
}
# Download a single model
download_model() {
local repo_id="$1"
local description="$2"
local size_gb="$3"
print_detail "Repository: ${BOLD_WHITE}${repo_id}${RESET}"
print_detail "Description: ${description}"
print_detail "Size: ${BOLD_YELLOW}${size_gb}GB${RESET}"
# Download using Python
python3 - <<EOPYDOWNLOAD
import os
import sys
from huggingface_hub import snapshot_download
cache_dir = '${CACHE_DIR}'
token = '${HF_TOKEN}'
repo_id = '${repo_id}'
os.environ['HF_HOME'] = cache_dir
try:
snapshot_download(
repo_id=repo_id,
cache_dir=cache_dir,
token=token,
resume_download=True
)
print("SUCCESS")
except Exception as e:
print(f"ERROR: {e}", file=sys.stderr)
sys.exit(1)
EOPYDOWNLOAD
if [[ $? -eq 0 ]]; then
print_success "Downloaded ${BOLD_WHITE}${repo_id}${RESET}"
return 0
else
print_error "Failed to download ${repo_id}"
return 1
fi
}
# Download models by category
download_category() {
local category="$1"
local category_display="$2"
print_section "${category_display}"
# Get models for this category
local models_data
models_data=$(parse_yaml "$CONFIG_FILE" "$category")
if [[ -z "$models_data" ]]; then
print_warning "No models found in category: ${category}"
return 0
fi
local total_models
total_models=$(echo "$models_data" | wc -l)
local current=0
local succeeded=0
local failed=0
while IFS='|' read -r repo_id description size_gb essential; do
((current++))
echo ""
print_step "$current" "$total_models" "${BOLD_MAGENTA}${description}${RESET}"
if download_model "$repo_id" "$description" "$size_gb"; then
((succeeded++))
else
((failed++))
fi
show_progress "$current" "$total_models"
done <<< "$models_data"
echo -e "\n"
print_info "Category Summary: ${BOLD_GREEN}${succeeded} succeeded${RESET}, ${BOLD_RED}${failed} failed${RESET}"
}
# Display summary
display_summary() {
local start_time="$1"
local end_time="$2"
local total_downloaded="$3"
local total_failed="$4"
local duration=$((end_time - start_time))
local minutes=$((duration / 60))
local seconds=$((duration % 60))
print_banner "DOWNLOAD COMPLETE"
echo -e "${BOLD_CYAN}${STAR} Summary${RESET}"
echo -e "${CYAN}$(printf '%.0s'"${BOX_LIGHT}" $(seq 1 80))${RESET}"
echo -e " ${BOLD_WHITE}Total Downloaded:${RESET} ${BOLD_GREEN}${total_downloaded}${RESET} models"
echo -e " ${BOLD_WHITE}Total Failed:${RESET} ${BOLD_RED}${total_failed}${RESET} models"
echo -e " ${BOLD_WHITE}Cache Directory:${RESET} ${CYAN}${CACHE_DIR}${RESET}"
echo -e " ${BOLD_WHITE}Duration:${RESET} ${BOLD_YELLOW}${minutes}m ${seconds}s${RESET}"
echo -e "${CYAN}$(printf '%.0s'"${BOX_LIGHT}" $(seq 1 80))${RESET}"
if [[ $total_failed -eq 0 ]]; then
echo -e "\n${BOLD_GREEN}${SPARKLES} All models downloaded successfully! ${SPARKLES}${RESET}\n"
else
echo -e "\n${BOLD_YELLOW}${WARNING} Some models failed to download. Check logs above.${RESET}\n"
fi
}
# ============================================================================
# MAIN FUNCTION
# ============================================================================
main() {
local start_time
start_time=$(date +%s)
# Display beautiful banner
print_banner "${ROCKET} ComfyUI Model Downloader ${ROCKET}"
echo -e "${BOLD_CYAN}A Beautiful CLI Tool for Downloading AI Models${RESET}"
echo -e "${DIM}Powered by HuggingFace ${LINK} Configuration-Driven ${STAR}${RESET}\n"
# Check dependencies
check_dependencies
# Validate configuration
validate_config
# Get all categories
if [[ -z "$CONFIG_FILE" ]]; then
print_error "No configuration file specified. Use -c/--config to provide one."
exit 1
fi
local categories
categories=$(parse_yaml "$CONFIG_FILE" "categories")
if [[ -z "$categories" ]]; then
print_error "No model categories found in configuration"
exit 1
fi
local total_succeeded=0
local total_failed=0
# Download each category
while IFS= read -r category; do
# Get category display name (capitalize and add spaces)
local category_display
category_display=$(echo "$category" | sed 's/_/ /g' | sed 's/\b\(.\)/\u\1/g')
download_category "$category" "$category_display"
# Update counters (this is simplified, you'd need to track actual numbers)
((total_succeeded++))
done <<< "$categories"
# Display summary
local end_time
end_time=$(date +%s)
display_summary "$start_time" "$end_time" "$total_succeeded" "$total_failed"
}
# ============================================================================
# ENTRY POINT
# ============================================================================
# Parse command line arguments
POSITIONAL_ARGS=()
while [[ $# -gt 0 ]]; do
case $1 in
-c|--config)
CONFIG_FILE="$2"
shift 2
;;
--cache-dir)
CACHE_DIR="$2"
shift 2
;;
-h|--help)
echo "Usage: $0 [CONFIG_FILE] [options]"
echo ""
echo "Arguments:"
echo " CONFIG_FILE Configuration file path (optional, can also use -c)"
echo ""
echo "Options:"
echo " -c, --config FILE Configuration file (default: NONE)"
echo " --cache-dir DIR Cache directory (default: ~/.cache/huggingface)"
echo " -h, --help Show this help message"
exit 0
;;
-*)
print_error "Unknown option: $1"
echo "Use --help for usage information"
exit 1
;;
*)
POSITIONAL_ARGS+=("$1")
shift
;;
esac
done
# Handle positional argument (config file path)
if [[ ${#POSITIONAL_ARGS[@]} -gt 0 ]]; then
CONFIG_FILE="${POSITIONAL_ARGS[0]}"
fi
# Run main function
main