fix: correct typo in filename (hugginface -> huggingface)
🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
767
artifact_huggingface_download.sh
Executable file
767
artifact_huggingface_download.sh
Executable file
@@ -0,0 +1,767 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# ComfyUI Model Downloader - A Beautiful CLI Tool
|
||||
# Downloads AI models from HuggingFace and creates symlinks to ComfyUI directories
|
||||
#
|
||||
# Usage: ./artifact_comfyui_download.sh [COMMAND] [options]
|
||||
#
|
||||
# Commands: download, link, both (default)
|
||||
#
|
||||
|
||||
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 - detect RunPod or use HuggingFace default
|
||||
if [[ -d "/workspace" ]]; then
|
||||
# RunPod environment
|
||||
CACHE_DIR="${CACHE_DIR:-/workspace/huggingface_cache}"
|
||||
COMFYUI_DIR="${COMFYUI_DIR:-/workspace/ComfyUI/models}"
|
||||
else
|
||||
# Local environment
|
||||
CACHE_DIR="${CACHE_DIR:-${HOME}/.cache/huggingface}"
|
||||
COMFYUI_DIR="${COMFYUI_DIR:-${HOME}/ComfyUI/models}"
|
||||
fi
|
||||
|
||||
# Default command
|
||||
COMMAND="both"
|
||||
|
||||
# HuggingFace token from environment or .env file
|
||||
# Initialize HF_TOKEN if not set
|
||||
HF_TOKEN="${HF_TOKEN:-}"
|
||||
|
||||
# Try multiple locations for .env file
|
||||
if [[ -z "${HF_TOKEN}" ]] && [[ -f "${PROJECT_ROOT}/ai/.env" ]]; then
|
||||
HF_TOKEN=$(grep ^HF_TOKEN "${PROJECT_ROOT}/ai/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
|
||||
fi
|
||||
if [[ -z "${HF_TOKEN}" ]] && [[ -f "${PROJECT_ROOT}/.env" ]]; then
|
||||
HF_TOKEN=$(grep ^HF_TOKEN "${PROJECT_ROOT}/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
|
||||
fi
|
||||
if [[ -z "${HF_TOKEN}" ]] && [[ -f "/workspace/ai/.env" ]]; then
|
||||
HF_TOKEN=$(grep ^HF_TOKEN "/workspace/ai/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
|
||||
fi
|
||||
|
||||
# ============================================================================
|
||||
# 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 - "$yaml_file" "$category" <<EOPYAML
|
||||
import yaml
|
||||
import sys
|
||||
|
||||
yaml_file = sys.argv[1]
|
||||
category = sys.argv[2]
|
||||
|
||||
try:
|
||||
with open(yaml_file, 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
if category == 'settings':
|
||||
settings = config.get('settings', {})
|
||||
print("CACHE_DIR={0}".format(settings.get('cache_dir', '/workspace/huggingface_cache')))
|
||||
print("PARALLEL_DOWNLOADS={0}".format(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)
|
||||
model_type = model.get('type', 'checkpoints')
|
||||
filename = model.get('filename', '')
|
||||
print('{0}|{1}|{2}|{3}|{4}|{5}'.format(repo_id, description, size_gb, essential, model_type, filename))
|
||||
else:
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print("ERROR: {0}".format(e), file=sys.stderr)
|
||||
sys.exit(1)
|
||||
EOPYAML
|
||||
}
|
||||
|
||||
# Parse file mappings for a specific model
|
||||
parse_file_mappings() {
|
||||
local yaml_file="$1"
|
||||
local category="$2"
|
||||
local repo_id="$3"
|
||||
|
||||
python3 - "$yaml_file" "$category" "$repo_id" <<EOPYTHON
|
||||
import yaml
|
||||
import sys
|
||||
|
||||
yaml_file = sys.argv[1]
|
||||
category = sys.argv[2]
|
||||
repo_id = sys.argv[3]
|
||||
|
||||
try:
|
||||
with open(yaml_file, 'r') as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
if category in config.get('model_categories', {}):
|
||||
models = config['model_categories'][category]
|
||||
for model in models:
|
||||
if model.get('repo_id', '') == repo_id:
|
||||
files = model.get('files', [])
|
||||
if files:
|
||||
for file_mapping in files:
|
||||
source = file_mapping.get('source', '')
|
||||
dest = file_mapping.get('dest', source)
|
||||
if source:
|
||||
print('{0}|{1}'.format(source, dest))
|
||||
sys.exit(0)
|
||||
|
||||
# No file mappings found
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print("ERROR: {0}".format(e), file=sys.stderr)
|
||||
sys.exit(1)
|
||||
EOPYTHON
|
||||
}
|
||||
|
||||
# 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"
|
||||
|
||||
# Show current command
|
||||
print_info "Command: ${BOLD_CYAN}${COMMAND}${RESET}"
|
||||
|
||||
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
|
||||
|
||||
# HF_TOKEN only required for download and both commands
|
||||
if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
|
||||
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}"
|
||||
fi
|
||||
|
||||
# Cache directory
|
||||
if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
|
||||
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}"
|
||||
else
|
||||
print_info "Cache directory: ${CYAN}${CACHE_DIR}${RESET}"
|
||||
fi
|
||||
|
||||
# ComfyUI directory
|
||||
if [[ "$COMMAND" == "link" ]] || [[ "$COMMAND" == "both" ]]; then
|
||||
print_info "ComfyUI directory: ${CYAN}${COMFYUI_DIR}${RESET}"
|
||||
fi
|
||||
}
|
||||
|
||||
# Find model files in HuggingFace cache
|
||||
find_model_files() {
|
||||
local repo_id="$1"
|
||||
local filename_filter="$2"
|
||||
|
||||
python3 - "$CACHE_DIR" "$repo_id" "$filename_filter" <<EOPYFINDFIND
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
cache_dir = sys.argv[1]
|
||||
repo_id = sys.argv[2]
|
||||
filename_filter = sys.argv[3]
|
||||
|
||||
# HuggingFace cache structure: cache_dir/models--org--name/snapshots/hash/
|
||||
# Try both with and without 'hub/' subdirectory for compatibility
|
||||
cache_path = Path(cache_dir)
|
||||
repo_path = repo_id.replace('/', '--')
|
||||
model_dir = cache_path / 'models--{0}'.format(repo_path)
|
||||
|
||||
# Fallback to hub/ subdirectory if direct path doesn't exist
|
||||
if not model_dir.exists():
|
||||
model_dir = cache_path / 'hub' / 'models--{0}'.format(repo_path)
|
||||
|
||||
if not model_dir.exists():
|
||||
sys.exit(1)
|
||||
|
||||
# Find the latest snapshot
|
||||
snapshots_dir = model_dir / 'snapshots'
|
||||
if not snapshots_dir.exists():
|
||||
sys.exit(1)
|
||||
|
||||
# Get all snapshot directories sorted by modification time
|
||||
snapshots = sorted(snapshots_dir.iterdir(), key=lambda x: x.stat().st_mtime, reverse=True)
|
||||
if not snapshots:
|
||||
sys.exit(1)
|
||||
|
||||
latest_snapshot = snapshots[0]
|
||||
|
||||
# Find model files
|
||||
for file_path in latest_snapshot.rglob('*'):
|
||||
if file_path.is_file():
|
||||
# If filename filter is specified, only match those files
|
||||
if filename_filter and filename_filter not in file_path.name:
|
||||
continue
|
||||
# Skip metadata files
|
||||
if file_path.name.endswith(('.json', '.txt', '.md', '.gitattributes')):
|
||||
continue
|
||||
print(str(file_path))
|
||||
EOPYFINDFIND
|
||||
}
|
||||
|
||||
# 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("ERROR: {0}".format(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
|
||||
}
|
||||
|
||||
# Create symlink for a model
|
||||
link_model() {
|
||||
local repo_id="$1"
|
||||
local model_type="$2"
|
||||
local filename_filter="$3"
|
||||
local file_mappings="$4" # Optional: explicit source|dest mappings
|
||||
|
||||
print_detail "Linking to: ${CYAN}${COMFYUI_DIR}/${model_type}/${RESET}"
|
||||
|
||||
# Create ComfyUI subdirectory if it doesn't exist
|
||||
local target_dir="${COMFYUI_DIR}/${model_type}"
|
||||
if [[ ! -d "$target_dir" ]]; then
|
||||
print_info "Creating directory: ${CYAN}${target_dir}${RESET}"
|
||||
mkdir -p "$target_dir"
|
||||
fi
|
||||
|
||||
# Clean existing symlinks for this repo in the target directory
|
||||
# This ensures we start fresh and don't have stale links
|
||||
find "$target_dir" -type l -lname "*${repo_id/\//-}*" -delete 2>/dev/null || true
|
||||
|
||||
# Find model files in cache
|
||||
local model_files
|
||||
model_files=$(find_model_files "$repo_id" "$filename_filter")
|
||||
|
||||
if [[ -z "$model_files" ]]; then
|
||||
print_warning "No model files found in cache for ${repo_id}"
|
||||
return 1
|
||||
fi
|
||||
|
||||
local linked_count=0
|
||||
|
||||
# If explicit file mappings are provided, use them
|
||||
if [[ -n "$file_mappings" ]]; then
|
||||
print_detail "Using explicit file mappings from YAML"
|
||||
|
||||
while IFS='|' read -r source_pattern dest_filename; do
|
||||
if [[ -z "$source_pattern" ]]; then
|
||||
continue
|
||||
fi
|
||||
|
||||
# Find the file matching the source pattern in model_files
|
||||
local source_file
|
||||
source_file=$(echo "$model_files" | grep -F "/$source_pattern" | head -n1)
|
||||
|
||||
if [[ -z "$source_file" ]] || [[ ! -f "$source_file" ]]; then
|
||||
print_warning "Source file not found: ${source_pattern}"
|
||||
continue
|
||||
fi
|
||||
|
||||
local link_path="${target_dir}/${dest_filename}"
|
||||
|
||||
# Remove existing symlink or file if it exists
|
||||
if [[ -L "$link_path" ]]; then
|
||||
rm -f "$link_path"
|
||||
elif [[ -e "$link_path" ]]; then
|
||||
print_warning "File already exists (not a symlink): ${dest_filename}"
|
||||
continue
|
||||
fi
|
||||
|
||||
# Create symlink
|
||||
ln -s "$source_file" "$link_path"
|
||||
print_detail "${LINK} Linked: ${DIM}${dest_filename}${RESET}"
|
||||
linked_count=$((linked_count+1))
|
||||
done <<< "$file_mappings"
|
||||
else
|
||||
# Fallback: use automatic prefixing for files without explicit mappings
|
||||
print_detail "No file mappings found, using automatic prefixing"
|
||||
|
||||
# Extract model name from repo_id for prefixing filenames
|
||||
# e.g., "facebook/musicgen-medium" -> "musicgen-medium"
|
||||
local model_name=$(echo "$repo_id" | sed 's/.*\///')
|
||||
|
||||
while IFS= read -r source_file; do
|
||||
if [[ -f "$source_file" ]]; then
|
||||
local filename=$(basename "$source_file")
|
||||
|
||||
# Add model name prefix to filename for better organization
|
||||
# e.g., "pytorch_model.bin" -> "musicgen-medium-pytorch_model.bin"
|
||||
local prefixed_filename="${model_name}-${filename}"
|
||||
local link_path="${target_dir}/${prefixed_filename}"
|
||||
|
||||
# Remove existing symlink or file if it exists
|
||||
if [[ -L "$link_path" ]]; then
|
||||
rm -f "$link_path"
|
||||
elif [[ -e "$link_path" ]]; then
|
||||
print_warning "File already exists (not a symlink): ${prefixed_filename}"
|
||||
continue
|
||||
fi
|
||||
|
||||
# Create symlink
|
||||
ln -s "$source_file" "$link_path"
|
||||
print_detail "${LINK} Linked: ${DIM}${prefixed_filename}${RESET}"
|
||||
linked_count=$((linked_count+1))
|
||||
fi
|
||||
done <<< "$model_files"
|
||||
fi
|
||||
|
||||
if [[ $linked_count -gt 0 ]]; then
|
||||
print_success "Linked ${linked_count} file(s) for ${BOLD_WHITE}${repo_id}${RESET}"
|
||||
return 0
|
||||
else
|
||||
print_error "Failed to link files for ${repo_id}"
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Process models by category
|
||||
process_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 model_type filename; do
|
||||
current=$((current+1))
|
||||
|
||||
echo ""
|
||||
print_step "$current" "$total_models" "${BOLD_MAGENTA}${description}${RESET}"
|
||||
|
||||
local success=true
|
||||
|
||||
# Download if command is 'download' or 'both'
|
||||
if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
|
||||
if ! download_model "$repo_id" "$description" "$size_gb"; then
|
||||
success=false
|
||||
fi
|
||||
fi
|
||||
|
||||
# Link if command is 'link' or 'both'
|
||||
if [[ "$COMMAND" == "link" ]] || [[ "$COMMAND" == "both" ]]; then
|
||||
if $success; then
|
||||
# Parse file mappings from YAML for this model
|
||||
local file_mappings
|
||||
file_mappings=$(parse_file_mappings "$CONFIG_FILE" "$category" "$repo_id")
|
||||
|
||||
# Pass file mappings to link_model (empty string if no mappings found)
|
||||
if ! link_model "$repo_id" "$model_type" "$filename" "$file_mappings"; then
|
||||
success=false
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
if $success; then
|
||||
succeeded=$((succeeded+1))
|
||||
else
|
||||
failed=$((failed+1))
|
||||
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
|
||||
|
||||
# Process 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')
|
||||
|
||||
process_category "$category" "$category_display"
|
||||
|
||||
# Update counters (this is simplified, you'd need to track actual numbers)
|
||||
total_succeeded=$((total_succeeded+1))
|
||||
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
|
||||
;;
|
||||
--comfyui-dir)
|
||||
COMFYUI_DIR="$2"
|
||||
shift 2
|
||||
;;
|
||||
download|link|both)
|
||||
COMMAND="$1"
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
echo "Usage: $0 [COMMAND] [options]"
|
||||
echo ""
|
||||
echo "Commands:"
|
||||
echo " download Download models only (default: both)"
|
||||
echo " link Create symlinks only (models must already be downloaded)"
|
||||
echo " both Download and create symlinks (default)"
|
||||
echo ""
|
||||
echo "Options:"
|
||||
echo " -c, --config FILE Configuration file (default: NONE)"
|
||||
echo " --cache-dir DIR Cache directory (default: auto-detect)"
|
||||
echo " RunPod: /workspace/huggingface_cache"
|
||||
echo " Local: ~/.cache/huggingface"
|
||||
echo " --comfyui-dir DIR ComfyUI models directory (default: auto-detect)"
|
||||
echo " RunPod: /workspace/ComfyUI/models"
|
||||
echo " Local: ~/ComfyUI/models"
|
||||
echo " -h, --help Show this help message"
|
||||
echo ""
|
||||
echo "Examples:"
|
||||
echo " $0 download -c models.yaml"
|
||||
echo " $0 link --comfyui-dir /opt/ComfyUI/models"
|
||||
echo " $0 both -c models.yaml --cache-dir /data/hf-cache"
|
||||
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
|
||||
Reference in New Issue
Block a user