Sebastian Krüger
6b2068e803
- Add artifact_civitai_download.sh with beautiful purple/magenta CLI - Rename artifact_comfyui_download.sh to artifact_hugginface_download.sh - Remove comfyui_models.example.yaml (moved to runpod repo) Features: - Dedicated downloader for CivitAI models - Beautiful CLI with progress bars and retry logic - Same architecture as HuggingFace downloader 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
710 lines
22 KiB
Bash
Executable File
710 lines
22 KiB
Bash
Executable File
#!/bin/bash
|
||
#
|
||
# CivitAI Model Downloader - A Beautiful CLI Tool
|
||
# Downloads NSFW AI models from CivitAI and creates symlinks to ComfyUI directories
|
||
#
|
||
# Usage: ./artifact_civitai_download.sh [COMMAND] [options]
|
||
#
|
||
# Commands: download, link, both (default)
|
||
#
|
||
|
||
set -euo pipefail
|
||
|
||
# ============================================================================
|
||
# COLOR PALETTE - Beautiful Terminal Colors (Purple/Magenta Theme)
|
||
# ============================================================================
|
||
|
||
# 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="P"
|
||
WARNING=" "
|
||
INFO="9"
|
||
ARROW_RIGHT="’"
|
||
DOUBLE_ARROW="»"
|
||
BOX_LIGHT=" "
|
||
BOX_HEAVY=""
|
||
BOX_DOUBLE="P"
|
||
ART="<¨"
|
||
PAINT="=¼"
|
||
|
||
# ============================================================================
|
||
# 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/models_civitai.yaml" ]]; then
|
||
CONFIG_FILE="${HOME}/Projects/runpod/models_civitai.yaml"
|
||
elif [[ -f "${PROJECT_ROOT}/models_civitai.yaml" ]]; then
|
||
CONFIG_FILE="${PROJECT_ROOT}/models_civitai.yaml"
|
||
elif [[ -f "${SCRIPT_DIR}/models_civitai.yaml" ]]; then
|
||
CONFIG_FILE="${SCRIPT_DIR}/models_civitai.yaml"
|
||
else
|
||
CONFIG_FILE="" # No config file by default
|
||
fi
|
||
|
||
# Default cache directory - detect RunPod or use local
|
||
if [[ -d "/workspace" ]]; then
|
||
# RunPod environment
|
||
CACHE_DIR="${CACHE_DIR:-/workspace/models/civitai}"
|
||
OUTPUT_DIR="${OUTPUT_DIR:-/workspace/ComfyUI/models}"
|
||
else
|
||
# Local environment
|
||
CACHE_DIR="${CACHE_DIR:-${HOME}/.cache/civitai}"
|
||
OUTPUT_DIR="${OUTPUT_DIR:-${HOME}/ComfyUI/models}"
|
||
fi
|
||
|
||
# Default command
|
||
COMMAND="both"
|
||
|
||
# CivitAI API key from environment or .env file
|
||
# Initialize CIVITAI_API_KEY if not set
|
||
CIVITAI_API_KEY="${CIVITAI_API_KEY:-}"
|
||
|
||
# Try multiple locations for .env file
|
||
if [[ -z "${CIVITAI_API_KEY}" ]] && [[ -f "${PROJECT_ROOT}/.env" ]]; then
|
||
CIVITAI_API_KEY=$(grep ^CIVITAI_API_KEY "${PROJECT_ROOT}/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
|
||
fi
|
||
if [[ -z "${CIVITAI_API_KEY}" ]] && [[ -f "${HOME}/Projects/runpod/.env" ]]; then
|
||
CIVITAI_API_KEY=$(grep ^CIVITAI_API_KEY "${HOME}/Projects/runpod/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
|
||
fi
|
||
if [[ -z "${CIVITAI_API_KEY}" ]] && [[ -f "/workspace/.env" ]]; then
|
||
CIVITAI_API_KEY=$(grep ^CIVITAI_API_KEY "/workspace/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
|
||
fi
|
||
if [[ -z "${CIVITAI_API_KEY}" ]] && [[ -f "/workspace/ai/.env" ]]; then
|
||
CIVITAI_API_KEY=$(grep ^CIVITAI_API_KEY "/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_MAGENTA}${BOX_DOUBLE}$(printf '%.0s'"${BOX_DOUBLE}" $(seq 1 $width))${BOX_DOUBLE}${RESET}"
|
||
echo -e "${BOLD_MAGENTA}${BOX_DOUBLE}$(printf '%.0s ' $(seq 1 $padding))${BOLD_CYAN}${text}$(printf '%.0s ' $(seq 1 $padding))${BOLD_MAGENTA}${BOX_DOUBLE}${RESET}"
|
||
echo -e "${BOLD_MAGENTA}${BOX_DOUBLE}$(printf '%.0s'"${BOX_DOUBLE}" $(seq 1 $width))${BOX_DOUBLE}${RESET}"
|
||
echo -e ""
|
||
}
|
||
|
||
print_section() {
|
||
local text="$1"
|
||
echo -e "\n${BOLD_MAGENTA}${DOUBLE_ARROW} ${text}${RESET}"
|
||
echo -e "${MAGENTA}$(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_MAGENTA}${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}${MAGENTA}${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_MAGENTA}Progress: ${RESET}["
|
||
printf "${BG_MAGENTA}${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/models/civitai')))
|
||
print("OUTPUT_DIR={0}".format(settings.get('output_dir', '/workspace/ComfyUI/models')))
|
||
print("PARALLEL_DOWNLOADS={0}".format(settings.get('parallel_downloads', 1)))
|
||
print("RATE_LIMIT_DELAY={0}".format(settings.get('rate_limit_delay', 5)))
|
||
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:
|
||
name = model.get('name', '')
|
||
version_id = model.get('version_id', '')
|
||
model_id = model.get('model_id', '')
|
||
description = model.get('description', '')
|
||
size_gb = model.get('size_gb', 0)
|
||
essential = model.get('essential', False)
|
||
model_type = model.get('type', 'checkpoints')
|
||
print('{0}|{1}|{2}|{3}|{4}|{5}|{6}'.format(name, version_id, model_id, description, size_gb, essential, model_type))
|
||
else:
|
||
sys.exit(1)
|
||
except Exception as e:
|
||
print("ERROR: {0}".format(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 requests" 2>/dev/null; then
|
||
print_warning "requests not installed, installing..."
|
||
pip3 install requests -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_MAGENTA}${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: ${MAGENTA}${CONFIG_FILE}${RESET}"
|
||
else
|
||
print_warning "No configuration file specified"
|
||
fi
|
||
|
||
# CIVITAI_API_KEY only required for download and both commands
|
||
if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
|
||
if [[ -z "$CIVITAI_API_KEY" ]]; then
|
||
print_error "CIVITAI_API_KEY not set. Please set it in .env file or environment."
|
||
print_info "Get your API key from: https://civitai.com/user/account"
|
||
exit 1
|
||
fi
|
||
print_success "CivitAI API key configured: ${DIM}${CIVITAI_API_KEY:0:10}...${RESET}"
|
||
fi
|
||
|
||
# Cache directory
|
||
if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
|
||
if [[ ! -d "$CACHE_DIR" ]]; then
|
||
print_info "Creating cache directory: ${MAGENTA}${CACHE_DIR}${RESET}"
|
||
mkdir -p "$CACHE_DIR"
|
||
fi
|
||
print_success "Cache directory ready: ${MAGENTA}${CACHE_DIR}${RESET}"
|
||
else
|
||
print_info "Cache directory: ${MAGENTA}${CACHE_DIR}${RESET}"
|
||
fi
|
||
|
||
# Output directory
|
||
if [[ "$COMMAND" == "link" ]] || [[ "$COMMAND" == "both" ]]; then
|
||
print_info "Output directory: ${MAGENTA}${OUTPUT_DIR}${RESET}"
|
||
fi
|
||
}
|
||
|
||
# Download a single model from CivitAI
|
||
download_civitai_model() {
|
||
local name="$1"
|
||
local version_id="$2"
|
||
local description="$3"
|
||
local size_gb="$4"
|
||
|
||
print_detail "Name: ${BOLD_WHITE}${name}${RESET}"
|
||
print_detail "Version ID: ${version_id}"
|
||
print_detail "Description: ${description}"
|
||
print_detail "Size: ${BOLD_YELLOW}${size_gb}GB${RESET}"
|
||
|
||
# Download using Python
|
||
python3 - <<EOPYDOWNLOAD
|
||
import os
|
||
import sys
|
||
import time
|
||
import requests
|
||
from pathlib import Path
|
||
|
||
cache_dir = '${CACHE_DIR}'
|
||
api_key = '${CIVITAI_API_KEY}'
|
||
version_id = '${version_id}'
|
||
name = '${name}'
|
||
|
||
# Create cache directory
|
||
Path(cache_dir).mkdir(parents=True, exist_ok=True)
|
||
|
||
# CivitAI API endpoint
|
||
url = f"https://civitai.com/api/download/models/{version_id}"
|
||
headers = {
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
max_retries = 3
|
||
for attempt in range(max_retries):
|
||
try:
|
||
# Make request
|
||
response = requests.get(url, headers=headers, stream=True, allow_redirects=True)
|
||
|
||
# Handle rate limiting
|
||
if response.status_code == 429:
|
||
retry_after = int(response.headers.get('Retry-After', 60))
|
||
print(f"Rate limited. Waiting {retry_after} seconds...", file=sys.stderr)
|
||
time.sleep(retry_after)
|
||
continue
|
||
|
||
response.raise_for_status()
|
||
|
||
# Extract filename from Content-Disposition header
|
||
content_disposition = response.headers.get('Content-Disposition', '')
|
||
if 'filename=' in content_disposition:
|
||
filename = content_disposition.split('filename=')[1].strip('"').strip("'")
|
||
else:
|
||
filename = f"{name}.safetensors"
|
||
|
||
output_path = Path(cache_dir) / filename
|
||
|
||
# Check if file already exists
|
||
if output_path.exists():
|
||
print(f"File already exists: {output_path}", file=sys.stderr)
|
||
print("SUCCESS")
|
||
sys.exit(0)
|
||
|
||
# Download with progress
|
||
total_size = int(response.headers.get('content-length', 0))
|
||
downloaded = 0
|
||
block_size = 8192
|
||
|
||
with open(output_path, 'wb') as f:
|
||
for chunk in response.iter_content(chunk_size=block_size):
|
||
if chunk:
|
||
f.write(chunk)
|
||
downloaded += len(chunk)
|
||
if total_size:
|
||
percent = (downloaded / total_size) * 100
|
||
print(f"\rDownloading: {percent:.1f}%", end='', file=sys.stderr)
|
||
|
||
print("", file=sys.stderr) # New line after progress
|
||
print("SUCCESS")
|
||
sys.exit(0)
|
||
|
||
except requests.exceptions.RequestException as e:
|
||
print(f"Error: {e}", file=sys.stderr)
|
||
if attempt < max_retries - 1:
|
||
wait_time = 2 ** attempt # Exponential backoff
|
||
print(f"Retrying in {wait_time} seconds...", file=sys.stderr)
|
||
time.sleep(wait_time)
|
||
else:
|
||
print("ERROR: Max retries reached", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
print("ERROR: Failed to download", file=sys.stderr)
|
||
sys.exit(1)
|
||
EOPYDOWNLOAD
|
||
|
||
if [[ $? -eq 0 ]]; then
|
||
print_success "Downloaded ${BOLD_WHITE}${name}${RESET}"
|
||
return 0
|
||
else
|
||
print_error "Failed to download ${name}"
|
||
return 1
|
||
fi
|
||
}
|
||
|
||
# Create symlink for a model
|
||
link_civitai_model() {
|
||
local name="$1"
|
||
local version_id="$2"
|
||
local model_type="$3"
|
||
|
||
print_detail "Linking to: ${MAGENTA}${OUTPUT_DIR}/${model_type}/${RESET}"
|
||
|
||
# Create output subdirectory if it doesn't exist
|
||
local target_dir="${OUTPUT_DIR}/${model_type}"
|
||
if [[ ! -d "$target_dir" ]]; then
|
||
print_info "Creating directory: ${MAGENTA}${target_dir}${RESET}"
|
||
mkdir -p "$target_dir"
|
||
fi
|
||
|
||
# Find model file in cache (look for .safetensors files)
|
||
local model_files
|
||
model_files=$(find "$CACHE_DIR" -type f -name "*.safetensors" 2>/dev/null || true)
|
||
|
||
if [[ -z "$model_files" ]]; then
|
||
print_warning "No model files found in cache for ${name}"
|
||
return 1
|
||
fi
|
||
|
||
local linked_count=0
|
||
|
||
# Link all found model files
|
||
while IFS= read -r source_file; do
|
||
if [[ -f "$source_file" ]]; then
|
||
local filename=$(basename "$source_file")
|
||
local link_path="${target_dir}/${filename}"
|
||
|
||
# Remove existing symlink 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): ${filename}"
|
||
continue
|
||
fi
|
||
|
||
# Create symlink
|
||
ln -s "$source_file" "$link_path"
|
||
print_detail "${LINK} Linked: ${DIM}${filename}${RESET}"
|
||
linked_count=$((linked_count+1))
|
||
fi
|
||
done <<< "$model_files"
|
||
|
||
if [[ $linked_count -gt 0 ]]; then
|
||
print_success "Linked ${linked_count} file(s) for ${BOLD_WHITE}${name}${RESET}"
|
||
return 0
|
||
else
|
||
print_error "Failed to link files for ${name}"
|
||
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
|
||
|
||
# Get rate limit delay from settings
|
||
local rate_limit_delay=5
|
||
if settings=$(parse_yaml "$CONFIG_FILE" "settings" 2>/dev/null); then
|
||
rate_limit_delay=$(echo "$settings" | grep ^RATE_LIMIT_DELAY= | cut -d'=' -f2-)
|
||
fi
|
||
|
||
while IFS='|' read -r name version_id model_id description size_gb essential model_type; 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_civitai_model "$name" "$version_id" "$description" "$size_gb"; then
|
||
success=false
|
||
fi
|
||
|
||
# Rate limit protection
|
||
if [[ $current -lt $total_models ]] && [[ $success == true ]]; then
|
||
print_detail "Waiting ${rate_limit_delay}s for rate limit protection..."
|
||
sleep "$rate_limit_delay"
|
||
fi
|
||
fi
|
||
|
||
# Link if command is 'link' or 'both'
|
||
if [[ "$COMMAND" == "link" ]] || [[ "$COMMAND" == "both" ]]; then
|
||
if $success; then
|
||
if ! link_civitai_model "$name" "$version_id" "$model_type"; 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_MAGENTA}${STAR} Summary${RESET}"
|
||
echo -e "${MAGENTA}$(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} ${MAGENTA}${CACHE_DIR}${RESET}"
|
||
echo -e " ${BOLD_WHITE}Output Directory:${RESET} ${MAGENTA}${OUTPUT_DIR}${RESET}"
|
||
echo -e " ${BOLD_WHITE}Duration:${RESET} ${BOLD_YELLOW}${minutes}m ${seconds}s${RESET}"
|
||
echo -e "${MAGENTA}$(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 "${ART} CivitAI Model Downloader ${PAINT}"
|
||
|
||
echo -e "${BOLD_MAGENTA}A Beautiful CLI Tool for Downloading NSFW AI Models${RESET}"
|
||
echo -e "${DIM}Powered by CivitAI ${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
|
||
;;
|
||
--output-dir)
|
||
OUTPUT_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: auto-detect)"
|
||
echo " --cache-dir DIR Cache directory (default: auto-detect)"
|
||
echo " RunPod: /workspace/models/civitai"
|
||
echo " Local: ~/.cache/civitai"
|
||
echo " --output-dir DIR Output directory (default: auto-detect)"
|
||
echo " RunPod: /workspace/ComfyUI/models"
|
||
echo " Local: ~/ComfyUI/models"
|
||
echo " -h, --help Show this help message"
|
||
echo ""
|
||
echo "Environment Variables:"
|
||
echo " CIVITAI_API_KEY CivitAI API key (required)"
|
||
echo " Get yours at: https://civitai.com/user/account"
|
||
echo ""
|
||
echo "Examples:"
|
||
echo " $0 download -c models_civitai.yaml"
|
||
echo " $0 link --output-dir /opt/ComfyUI/models"
|
||
echo " $0 both -c models_civitai.yaml --cache-dir /data/civitai-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
|