Sebastian Krüger
8291a3b662
Complete rewrite of both model download scripts with: - Beautiful colorful CLI output with progress indicators - Pure bash/curl downloads (no Python dependencies for downloading) - yq-based YAML parsing (consistent with arty.sh) - Three commands: download, link, verify - Filtering by --category and --repo-id (comma-separated) - --dry-run mode for previewing operations - Respects format field for file extensions (.safetensors, .pt, etc.) - Uses type field for output subdirectories (checkpoints, embeddings, loras) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
621 lines
17 KiB
Bash
Executable File
621 lines
17 KiB
Bash
Executable File
#!/bin/bash
|
||
#
|
||
# CivitAI Model Downloader - A Beautiful CLI Tool
|
||
# Downloads AI models from CivitAI and creates symlinks to output directories
|
||
#
|
||
# Usage: ./artifact_civitai_download.sh [COMMAND] [OPTIONS]
|
||
#
|
||
# Commands:
|
||
# download Download models to cache directory (default)
|
||
# link Create symlinks from cache to output directory
|
||
# verify Verify symlinks in output directory
|
||
#
|
||
# Options:
|
||
# -c, --config FILE Configuration YAML file (required)
|
||
# --cache-dir DIR Cache directory
|
||
# --output-dir DIR Output/installation directory
|
||
# --category CAT1,CAT2 Filter by category (comma-separated)
|
||
# --repo-id ID1,ID2 Filter by repo_id (comma-separated)
|
||
# --auth-token TOKEN CivitAI API token
|
||
# -n, --dry-run Show what would be done
|
||
# -h, --help Show help
|
||
#
|
||
|
||
set -euo pipefail
|
||
|
||
# ============================================================================
|
||
# COLOR PALETTE - Beautiful Terminal Colors
|
||
# ============================================================================
|
||
|
||
RESET='\033[0m'
|
||
|
||
# Foreground Colors
|
||
RED='\033[0;31m'
|
||
GREEN='\033[0;32m'
|
||
YELLOW='\033[0;33m'
|
||
BLUE='\033[0;34m'
|
||
MAGENTA='\033[0;35m'
|
||
CYAN='\033[0;36m'
|
||
|
||
# Bold
|
||
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
|
||
BG_MAGENTA='\033[45m'
|
||
|
||
# Styles
|
||
DIM='\033[2m'
|
||
|
||
# ============================================================================
|
||
# UNICODE CHARACTERS
|
||
# ============================================================================
|
||
|
||
CHECK_MARK="✓"
|
||
CROSS_MARK="✗"
|
||
ROCKET="🚀"
|
||
DOWNLOAD="⬇️"
|
||
LINK_ICON="🔗"
|
||
WARNING="⚠️"
|
||
INFO="ℹ️"
|
||
SPARKLES="✨"
|
||
ARROW_RIGHT="→"
|
||
BOX_LIGHT="─"
|
||
BOX_DOUBLE="═"
|
||
|
||
# ============================================================================
|
||
# CONFIGURATION
|
||
# ============================================================================
|
||
|
||
CONFIG_FILE=""
|
||
COMMAND="download"
|
||
DRY_RUN=false
|
||
CATEGORY_FILTER=""
|
||
REPO_ID_FILTER=""
|
||
|
||
# Default directories - detect RunPod or local
|
||
if [[ -d "/workspace" ]]; then
|
||
CACHE_DIR="${CACHE_DIR:-/workspace/models/civitai}"
|
||
OUTPUT_DIR="${OUTPUT_DIR:-/workspace/ComfyUI/models}"
|
||
else
|
||
CACHE_DIR="${CACHE_DIR:-${HOME}/.cache/civitai}"
|
||
OUTPUT_DIR="${OUTPUT_DIR:-${HOME}/ComfyUI/models}"
|
||
fi
|
||
|
||
# CivitAI API key from environment
|
||
CIVITAI_API_KEY="${CIVITAI_API_KEY:-}"
|
||
|
||
# Try to load from .env files
|
||
load_env_token() {
|
||
if [[ -n "$CIVITAI_API_KEY" ]]; then
|
||
return 0
|
||
fi
|
||
|
||
local env_files=(
|
||
"${HOME}/.env"
|
||
"${HOME}/Projects/runpod/.env"
|
||
"/workspace/.env"
|
||
"/workspace/ai/.env"
|
||
)
|
||
|
||
for env_file in "${env_files[@]}"; do
|
||
if [[ -f "$env_file" ]]; then
|
||
local token
|
||
token=$(grep "^CIVITAI_API_KEY=" "$env_file" 2>/dev/null | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
|
||
if [[ -n "$token" ]]; then
|
||
CIVITAI_API_KEY="$token"
|
||
return 0
|
||
fi
|
||
fi
|
||
done
|
||
}
|
||
|
||
# ============================================================================
|
||
# LOGGING FUNCTIONS
|
||
# ============================================================================
|
||
|
||
print_banner() {
|
||
local text="$1"
|
||
local width=70
|
||
local text_len=${#text}
|
||
local padding=$(( (width - text_len) / 2 ))
|
||
|
||
echo ""
|
||
echo -e "${BOLD_MAGENTA}${BOX_DOUBLE}$(printf '%0.s═' $(seq 1 $width))${BOX_DOUBLE}${RESET}"
|
||
echo -e "${BOLD_MAGENTA}║$(printf '%*s' $padding '')${BOLD_CYAN}${text}$(printf '%*s' $((width - padding - text_len)) '')${BOLD_MAGENTA}║${RESET}"
|
||
echo -e "${BOLD_MAGENTA}${BOX_DOUBLE}$(printf '%0.s═' $(seq 1 $width))${BOX_DOUBLE}${RESET}"
|
||
echo ""
|
||
}
|
||
|
||
print_section() {
|
||
local text="$1"
|
||
echo -e "\n${BOLD_MAGENTA}» ${text}${RESET}"
|
||
echo -e "${MAGENTA}$(printf '%0.s─' $(seq 1 70))${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}${MAGENTA}${ARROW_RIGHT} $1${RESET}"
|
||
}
|
||
|
||
show_progress() {
|
||
local current="$1"
|
||
local total="$2"
|
||
local width=40
|
||
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"
|
||
}
|
||
|
||
# ============================================================================
|
||
# YAML PARSING (using yq)
|
||
# ============================================================================
|
||
|
||
check_yq() {
|
||
if ! command -v yq &>/dev/null; then
|
||
print_error "yq is not installed. Please install yq first."
|
||
print_info "Install: https://github.com/mikefarah/yq"
|
||
exit 1
|
||
fi
|
||
}
|
||
|
||
# Get total count of models
|
||
get_model_count() {
|
||
local config="$1"
|
||
yq eval '. | length' "$config" 2>/dev/null || echo "0"
|
||
}
|
||
|
||
# Get model at index
|
||
get_model_field() {
|
||
local config="$1"
|
||
local index="$2"
|
||
local field="$3"
|
||
local value
|
||
value=$(yq eval ".[$index].$field // \"\"" "$config" 2>/dev/null)
|
||
# Remove quotes if present
|
||
echo "$value" | sed 's/^"//;s/"$//'
|
||
}
|
||
|
||
# Check if model matches filters
|
||
matches_filters() {
|
||
local repo_id="$1"
|
||
local category="$2"
|
||
|
||
# Check category filter
|
||
if [[ -n "$CATEGORY_FILTER" ]]; then
|
||
local match=false
|
||
IFS=',' read -ra cats <<< "$CATEGORY_FILTER"
|
||
for cat in "${cats[@]}"; do
|
||
cat=$(echo "$cat" | xargs) # trim whitespace
|
||
if [[ "$category" == "$cat" ]]; then
|
||
match=true
|
||
break
|
||
fi
|
||
done
|
||
if [[ "$match" == false ]]; then
|
||
return 1
|
||
fi
|
||
fi
|
||
|
||
# Check repo_id filter
|
||
if [[ -n "$REPO_ID_FILTER" ]]; then
|
||
local match=false
|
||
IFS=',' read -ra repos <<< "$REPO_ID_FILTER"
|
||
for repo in "${repos[@]}"; do
|
||
repo=$(echo "$repo" | xargs) # trim whitespace
|
||
if [[ "$repo_id" == "$repo" ]]; then
|
||
match=true
|
||
break
|
||
fi
|
||
done
|
||
if [[ "$match" == false ]]; then
|
||
return 1
|
||
fi
|
||
fi
|
||
|
||
return 0
|
||
}
|
||
|
||
# ============================================================================
|
||
# DOWNLOAD FUNCTIONS
|
||
# ============================================================================
|
||
|
||
download_model() {
|
||
local repo_id="$1"
|
||
local version_id="$2"
|
||
local format="$3"
|
||
local description="$4"
|
||
|
||
local filename="${repo_id}.${format}"
|
||
local output_path="${CACHE_DIR}/${filename}"
|
||
|
||
print_detail "Repository: ${BOLD_WHITE}${repo_id}${RESET}"
|
||
print_detail "Version ID: ${version_id}"
|
||
print_detail "Format: ${format}"
|
||
[[ -n "$description" ]] && print_detail "Description: ${description}"
|
||
print_detail "Output: ${CYAN}${output_path}${RESET}"
|
||
|
||
# Check if already exists
|
||
if [[ -f "$output_path" ]]; then
|
||
print_success "Already downloaded: ${filename}"
|
||
return 0
|
||
fi
|
||
|
||
# Dry-run mode
|
||
if [[ "$DRY_RUN" == true ]]; then
|
||
print_info "DRY-RUN: Would download ${BOLD_WHITE}${filename}${RESET}"
|
||
return 0
|
||
fi
|
||
|
||
# Create cache directory
|
||
mkdir -p "$CACHE_DIR"
|
||
|
||
# Download with curl
|
||
local url="https://civitai.com/api/download/models/${version_id}"
|
||
print_detail "Downloading from CivitAI..."
|
||
|
||
local http_code
|
||
http_code=$(curl -L -w "%{http_code}" \
|
||
-H "Content-Type: application/json" \
|
||
-H "Authorization: Bearer ${CIVITAI_API_KEY}" \
|
||
--progress-bar \
|
||
-o "$output_path" \
|
||
"$url" 2>&1) || true
|
||
|
||
# Check if download succeeded
|
||
if [[ -f "$output_path" ]] && [[ -s "$output_path" ]]; then
|
||
local size
|
||
size=$(du -h "$output_path" | cut -f1)
|
||
print_success "Downloaded ${BOLD_WHITE}${filename}${RESET} (${size})"
|
||
return 0
|
||
else
|
||
print_error "Failed to download ${filename} (HTTP: ${http_code})"
|
||
rm -f "$output_path" 2>/dev/null || true
|
||
return 1
|
||
fi
|
||
}
|
||
|
||
# ============================================================================
|
||
# LINK FUNCTIONS
|
||
# ============================================================================
|
||
|
||
link_model() {
|
||
local repo_id="$1"
|
||
local model_type="$2"
|
||
local format="$3"
|
||
|
||
local filename="${repo_id}.${format}"
|
||
local source_path="${CACHE_DIR}/${filename}"
|
||
local target_dir="${OUTPUT_DIR}/${model_type}"
|
||
local link_path="${target_dir}/${filename}"
|
||
|
||
print_detail "Source: ${CYAN}${source_path}${RESET}"
|
||
print_detail "Target: ${CYAN}${link_path}${RESET}"
|
||
|
||
# Check if source exists
|
||
if [[ ! -f "$source_path" ]]; then
|
||
print_warning "Source file not found: ${filename}"
|
||
return 1
|
||
fi
|
||
|
||
# Dry-run mode
|
||
if [[ "$DRY_RUN" == true ]]; then
|
||
print_info "DRY-RUN: Would link ${BOLD_WHITE}${filename}${RESET} → ${model_type}/"
|
||
return 0
|
||
fi
|
||
|
||
# Create target directory
|
||
mkdir -p "$target_dir"
|
||
|
||
# Remove existing symlink
|
||
if [[ -L "$link_path" ]]; then
|
||
rm -f "$link_path"
|
||
elif [[ -e "$link_path" ]]; then
|
||
print_warning "File exists (not a symlink): ${filename}"
|
||
return 1
|
||
fi
|
||
|
||
# Create symlink
|
||
ln -s "$source_path" "$link_path"
|
||
print_success "Linked: ${BOLD_WHITE}${filename}${RESET} ${LINK_ICON} ${model_type}/"
|
||
return 0
|
||
}
|
||
|
||
# ============================================================================
|
||
# VERIFY FUNCTIONS
|
||
# ============================================================================
|
||
|
||
verify_model() {
|
||
local repo_id="$1"
|
||
local model_type="$2"
|
||
local format="$3"
|
||
|
||
local filename="${repo_id}.${format}"
|
||
local cache_path="${CACHE_DIR}/${filename}"
|
||
local link_path="${OUTPUT_DIR}/${model_type}/${filename}"
|
||
|
||
local cache_status="${BOLD_RED}${CROSS_MARK}${RESET}"
|
||
local link_status="${BOLD_RED}${CROSS_MARK}${RESET}"
|
||
|
||
# Check cache
|
||
if [[ -f "$cache_path" ]]; then
|
||
local size
|
||
size=$(du -h "$cache_path" | cut -f1)
|
||
cache_status="${BOLD_GREEN}${CHECK_MARK}${RESET} (${size})"
|
||
fi
|
||
|
||
# Check symlink
|
||
if [[ -L "$link_path" ]]; then
|
||
if [[ -e "$link_path" ]]; then
|
||
link_status="${BOLD_GREEN}${CHECK_MARK}${RESET}"
|
||
else
|
||
link_status="${BOLD_YELLOW}${WARNING}${RESET} (broken)"
|
||
fi
|
||
fi
|
||
|
||
echo -e " ${BOLD_WHITE}${repo_id}${RESET}"
|
||
echo -e " Cache: ${cache_status}"
|
||
echo -e " Link: ${link_status} → ${model_type}/"
|
||
}
|
||
|
||
# ============================================================================
|
||
# MAIN WORKFLOW
|
||
# ============================================================================
|
||
|
||
process_models() {
|
||
local action="$1"
|
||
|
||
print_section "Processing Models (${action})"
|
||
|
||
local total
|
||
total=$(get_model_count "$CONFIG_FILE")
|
||
|
||
if [[ "$total" == "0" ]]; then
|
||
print_warning "No models found in configuration"
|
||
return 0
|
||
fi
|
||
|
||
local processed=0
|
||
local succeeded=0
|
||
local failed=0
|
||
local skipped=0
|
||
|
||
for ((i=0; i<total; i++)); do
|
||
local repo_id version_id model_id description category model_type format
|
||
repo_id=$(get_model_field "$CONFIG_FILE" "$i" "repo_id")
|
||
version_id=$(get_model_field "$CONFIG_FILE" "$i" "version_id")
|
||
category=$(get_model_field "$CONFIG_FILE" "$i" "category")
|
||
model_type=$(get_model_field "$CONFIG_FILE" "$i" "type")
|
||
format=$(get_model_field "$CONFIG_FILE" "$i" "format")
|
||
description=$(get_model_field "$CONFIG_FILE" "$i" "description")
|
||
|
||
# Apply filters
|
||
if ! matches_filters "$repo_id" "$category"; then
|
||
skipped=$((skipped+1))
|
||
continue
|
||
fi
|
||
|
||
processed=$((processed+1))
|
||
echo ""
|
||
print_step "$processed" "$((total-skipped))" "${BOLD_MAGENTA}${repo_id}${RESET}"
|
||
|
||
local result=0
|
||
case "$action" in
|
||
download)
|
||
download_model "$repo_id" "$version_id" "$format" "$description" || result=1
|
||
;;
|
||
link)
|
||
link_model "$repo_id" "$model_type" "$format" || result=1
|
||
;;
|
||
verify)
|
||
verify_model "$repo_id" "$model_type" "$format"
|
||
;;
|
||
esac
|
||
|
||
if [[ $result -eq 0 ]]; then
|
||
succeeded=$((succeeded+1))
|
||
else
|
||
failed=$((failed+1))
|
||
fi
|
||
done
|
||
|
||
echo ""
|
||
if [[ $skipped -gt 0 ]]; then
|
||
print_info "Skipped ${skipped} model(s) (filtered)"
|
||
fi
|
||
print_info "Summary: ${BOLD_GREEN}${succeeded} succeeded${RESET}, ${BOLD_RED}${failed} failed${RESET}"
|
||
}
|
||
|
||
show_help() {
|
||
cat << 'EOF'
|
||
CivitAI Model Downloader - A Beautiful CLI Tool
|
||
|
||
Usage: artifact_civitai_download.sh [COMMAND] [OPTIONS]
|
||
|
||
Commands:
|
||
download Download models to cache directory (default)
|
||
link Create symlinks from cache to output directory
|
||
verify Verify model status (cache and symlinks)
|
||
|
||
Options:
|
||
-c, --config FILE Configuration YAML file (required)
|
||
--cache-dir DIR Cache directory
|
||
--output-dir DIR Output/installation directory
|
||
--category CAT1,CAT2 Filter by category (comma-separated)
|
||
--repo-id ID1,ID2 Filter by repo_id (comma-separated)
|
||
--auth-token TOKEN CivitAI API token (or set CIVITAI_API_KEY env var)
|
||
-n, --dry-run Show what would be done without making changes
|
||
-h, --help Show this help message
|
||
|
||
Environment Variables:
|
||
CIVITAI_API_KEY CivitAI API token
|
||
CACHE_DIR Override default cache directory
|
||
OUTPUT_DIR Override default output directory
|
||
|
||
Examples:
|
||
# Download all models from config
|
||
./artifact_civitai_download.sh download -c models_civitai.yaml
|
||
|
||
# Download only checkpoints
|
||
./artifact_civitai_download.sh download -c models_civitai.yaml --category checkpoints
|
||
|
||
# Create symlinks for downloaded models
|
||
./artifact_civitai_download.sh link -c models_civitai.yaml
|
||
|
||
# Verify all models
|
||
./artifact_civitai_download.sh verify -c models_civitai.yaml
|
||
|
||
# Dry-run to preview operations
|
||
./artifact_civitai_download.sh download -c models_civitai.yaml --dry-run
|
||
|
||
EOF
|
||
}
|
||
|
||
main() {
|
||
# Check for yq
|
||
check_yq
|
||
|
||
# Load token from .env files
|
||
load_env_token
|
||
|
||
# Display banner
|
||
print_banner "${ROCKET} CivitAI Model Downloader ${SPARKLES}"
|
||
|
||
# Show dry-run warning
|
||
if [[ "$DRY_RUN" == true ]]; then
|
||
echo -e "${BOLD_YELLOW}${WARNING} DRY-RUN MODE - No changes will be made ${WARNING}${RESET}\n"
|
||
fi
|
||
|
||
# Validate configuration
|
||
print_section "Configuration"
|
||
|
||
if [[ -z "$CONFIG_FILE" ]]; then
|
||
print_error "Configuration file required. Use -c/--config"
|
||
exit 1
|
||
fi
|
||
|
||
if [[ ! -f "$CONFIG_FILE" ]]; then
|
||
print_error "Configuration file not found: $CONFIG_FILE"
|
||
exit 1
|
||
fi
|
||
|
||
print_success "Config: ${CYAN}${CONFIG_FILE}${RESET}"
|
||
print_success "Cache: ${CYAN}${CACHE_DIR}${RESET}"
|
||
print_success "Output: ${CYAN}${OUTPUT_DIR}${RESET}"
|
||
print_success "Command: ${BOLD_MAGENTA}${COMMAND}${RESET}"
|
||
|
||
# Check API key for download
|
||
if [[ "$COMMAND" == "download" ]]; then
|
||
if [[ -z "$CIVITAI_API_KEY" ]]; then
|
||
print_error "CIVITAI_API_KEY not set"
|
||
print_info "Set via environment variable or --auth-token"
|
||
print_info "Get your API key: https://civitai.com/user/account"
|
||
exit 1
|
||
fi
|
||
print_success "API Key: ${DIM}${CIVITAI_API_KEY:0:8}...${RESET}"
|
||
fi
|
||
|
||
# Show filters if set
|
||
[[ -n "$CATEGORY_FILTER" ]] && print_info "Category filter: ${BOLD_WHITE}${CATEGORY_FILTER}${RESET}"
|
||
[[ -n "$REPO_ID_FILTER" ]] && print_info "Repo ID filter: ${BOLD_WHITE}${REPO_ID_FILTER}${RESET}"
|
||
|
||
# Process based on command
|
||
process_models "$COMMAND"
|
||
|
||
# Final message
|
||
echo ""
|
||
print_banner "${SPARKLES} Complete ${SPARKLES}"
|
||
}
|
||
|
||
# ============================================================================
|
||
# ARGUMENT PARSING
|
||
# ============================================================================
|
||
|
||
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
|
||
;;
|
||
--category)
|
||
CATEGORY_FILTER="$2"
|
||
shift 2
|
||
;;
|
||
--repo-id)
|
||
REPO_ID_FILTER="$2"
|
||
shift 2
|
||
;;
|
||
--auth-token)
|
||
CIVITAI_API_KEY="$2"
|
||
shift 2
|
||
;;
|
||
-n|--dry-run)
|
||
DRY_RUN=true
|
||
shift
|
||
;;
|
||
download|link|verify)
|
||
COMMAND="$1"
|
||
shift
|
||
;;
|
||
-h|--help)
|
||
show_help
|
||
exit 0
|
||
;;
|
||
-*)
|
||
print_error "Unknown option: $1"
|
||
echo "Use --help for usage information"
|
||
exit 1
|
||
;;
|
||
*)
|
||
# Positional argument - treat as config file
|
||
if [[ -z "$CONFIG_FILE" ]]; then
|
||
CONFIG_FILE="$1"
|
||
fi
|
||
shift
|
||
;;
|
||
esac
|
||
done
|
||
|
||
# Run main
|
||
main
|