feat: add ComfyUI symlink functionality and subcommands

Adds subcommand structure (download, link, both) to the ComfyUI downloader script, allowing users to download models from HuggingFace and/or create symlinks to ComfyUI model directories. Features include:

- New subcommands: download (download only), link (symlink only), both (default)
- ComfyUI directory configuration (--comfyui-dir, default: ~/ComfyUI/models)
- Smart symlink creation to appropriate ComfyUI subdirectories (checkpoints, vae, loras, controlnet, upscale_models, etc.)
- YAML configuration extended with 'type' and 'filename' fields for precise model organization
- Automatic ComfyUI subdirectory creation
- Graceful handling of existing symlinks and files
- HF_TOKEN validation only when needed (download/both commands)
- Example configuration file (comfyui_models.example.yaml) demonstrating proper setup

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
2025-11-22 02:20:46 +01:00
parent 394432ac7d
commit a310a08e15
2 changed files with 261 additions and 20 deletions

View File

@@ -1,9 +1,11 @@
#!/bin/bash #!/bin/bash
# #
# ComfyUI Model Downloader - A Beautiful CLI Tool # ComfyUI Model Downloader - A Beautiful CLI Tool
# Downloads AI models from HuggingFace with style and grace # Downloads AI models from HuggingFace and creates symlinks to ComfyUI directories
# #
# Usage: ./artifact_comfyui_download.sh [options] # Usage: ./artifact_comfyui_download.sh [COMMAND] [options]
#
# Commands: download, link, both (default)
# #
set -euo pipefail set -euo pipefail
@@ -98,6 +100,12 @@ fi
# Default cache directory (use HuggingFace default) # Default cache directory (use HuggingFace default)
CACHE_DIR="${CACHE_DIR:-${HOME}/.cache/huggingface}" CACHE_DIR="${CACHE_DIR:-${HOME}/.cache/huggingface}"
# Default ComfyUI models directory
COMFYUI_DIR="${COMFYUI_DIR:-${HOME}/ComfyUI/models}"
# Default command
COMMAND="both"
# HuggingFace token from environment or .env file # HuggingFace token from environment or .env file
if [[ -f "${PROJECT_ROOT}/.env" ]]; then if [[ -f "${PROJECT_ROOT}/.env" ]]; then
HF_TOKEN=$(grep ^HF_TOKEN "${PROJECT_ROOT}/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true) HF_TOKEN=$(grep ^HF_TOKEN "${PROJECT_ROOT}/.env" | cut -d'=' -f2- | tr -d '"' | tr -d "'" || true)
@@ -196,7 +204,9 @@ try:
description = model.get('description', '') description = model.get('description', '')
size_gb = model.get('size_gb', 0) size_gb = model.get('size_gb', 0)
essential = model.get('essential', False) essential = model.get('essential', False)
print(f"{repo_id}|{description}|{size_gb}|{essential}") model_type = model.get('type', 'checkpoints')
filename = model.get('filename', '')
print(f"{repo_id}|{description}|{size_gb}|{essential}|{model_type}|{filename}")
else: else:
sys.exit(1) sys.exit(1)
except Exception as e: except Exception as e:
@@ -244,6 +254,9 @@ check_dependencies() {
validate_config() { validate_config() {
print_section "Validating Configuration" print_section "Validating Configuration"
# Show current command
print_info "Command: ${BOLD_CYAN}${COMMAND}${RESET}"
if [[ -n "$CONFIG_FILE" ]]; then if [[ -n "$CONFIG_FILE" ]]; then
if [[ ! -f "$CONFIG_FILE" ]]; then if [[ ! -f "$CONFIG_FILE" ]]; then
print_error "Configuration file not found: $CONFIG_FILE" print_error "Configuration file not found: $CONFIG_FILE"
@@ -254,17 +267,77 @@ validate_config() {
print_warning "No configuration file specified" print_warning "No configuration file specified"
fi fi
# HF_TOKEN only required for download and both commands
if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
if [[ -z "$HF_TOKEN" ]]; then if [[ -z "$HF_TOKEN" ]]; then
print_error "HF_TOKEN not set. Please set it in .env file or environment." print_error "HF_TOKEN not set. Please set it in .env file or environment."
exit 1 exit 1
fi fi
print_success "HuggingFace token configured: ${DIM}${HF_TOKEN:0:10}...${RESET}" 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 if [[ ! -d "$CACHE_DIR" ]]; then
print_info "Creating cache directory: ${CYAN}${CACHE_DIR}${RESET}" print_info "Creating cache directory: ${CYAN}${CACHE_DIR}${RESET}"
mkdir -p "$CACHE_DIR" mkdir -p "$CACHE_DIR"
fi fi
print_success "Cache directory ready: ${CYAN}${CACHE_DIR}${RESET}" 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 - <<EOPYFINDFIND
import os
import sys
from pathlib import Path
cache_dir = '${CACHE_DIR}'
repo_id = '${repo_id}'
filename_filter = '${filename_filter}'
# HuggingFace cache structure: cache_dir/hub/models--org--name/snapshots/hash/
cache_path = Path(cache_dir) / 'hub'
repo_path = repo_id.replace('/', '--')
model_dir = cache_path / f'models--{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 a single model
@@ -311,8 +384,62 @@ EOPYDOWNLOAD
fi fi
} }
# Download models by category # Create symlink for a model
download_category() { link_model() {
local repo_id="$1"
local model_type="$2"
local filename_filter="$3"
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
# 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
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 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): ${filename}"
continue
fi
# Create symlink
ln -s "$source_file" "$link_path"
print_detail "${LINK} Linked: ${DIM}${filename}${RESET}"
((linked_count++))
fi
done <<< "$model_files"
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="$1"
local category_display="$2" local category_display="$2"
@@ -333,13 +460,31 @@ download_category() {
local succeeded=0 local succeeded=0
local failed=0 local failed=0
while IFS='|' read -r repo_id description size_gb essential; do while IFS='|' read -r repo_id description size_gb essential model_type filename; do
((current++)) ((current++))
echo "" echo ""
print_step "$current" "$total_models" "${BOLD_MAGENTA}${description}${RESET}" print_step "$current" "$total_models" "${BOLD_MAGENTA}${description}${RESET}"
if download_model "$repo_id" "$description" "$size_gb"; then 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
if ! link_model "$repo_id" "$model_type" "$filename"; then
success=false
fi
fi
fi
if $success; then
((succeeded++)) ((succeeded++))
else else
((failed++)) ((failed++))
@@ -417,13 +562,13 @@ main() {
local total_succeeded=0 local total_succeeded=0
local total_failed=0 local total_failed=0
# Download each category # Process each category
while IFS= read -r category; do while IFS= read -r category; do
# Get category display name (capitalize and add spaces) # Get category display name (capitalize and add spaces)
local category_display local category_display
category_display=$(echo "$category" | sed 's/_/ /g' | sed 's/\b\(.\)/\u\1/g') category_display=$(echo "$category" | sed 's/_/ /g' | sed 's/\b\(.\)/\u\1/g')
download_category "$category" "$category_display" process_category "$category" "$category_display"
# Update counters (this is simplified, you'd need to track actual numbers) # Update counters (this is simplified, you'd need to track actual numbers)
((total_succeeded++)) ((total_succeeded++))
@@ -451,16 +596,32 @@ while [[ $# -gt 0 ]]; do
CACHE_DIR="$2" CACHE_DIR="$2"
shift 2 shift 2
;; ;;
--comfyui-dir)
COMFYUI_DIR="$2"
shift 2
;;
download|link|both)
COMMAND="$1"
shift
;;
-h|--help) -h|--help)
echo "Usage: $0 [CONFIG_FILE] [options]" echo "Usage: $0 [COMMAND] [options]"
echo "" echo ""
echo "Arguments:" echo "Commands:"
echo " CONFIG_FILE Configuration file path (optional, can also use -c)" 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 ""
echo "Options:" echo "Options:"
echo " -c, --config FILE Configuration file (default: NONE)" echo " -c, --config FILE Configuration file (default: NONE)"
echo " --cache-dir DIR Cache directory (default: ~/.cache/huggingface)" echo " --cache-dir DIR Cache directory (default: ~/.cache/huggingface)"
echo " --comfyui-dir DIR ComfyUI models directory (default: ~/ComfyUI/models)"
echo " -h, --help Show this help message" 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 exit 0
;; ;;
-*) -*)

View File

@@ -0,0 +1,80 @@
# ComfyUI Models Configuration Example
#
# This file defines which models to download from HuggingFace
# and where to symlink them in your ComfyUI installation.
#
# Model types correspond to ComfyUI subdirectories:
# - checkpoints: Stable Diffusion checkpoints
# - vae: VAE models
# - loras: LoRA models
# - controlnet: ControlNet models
# - clip: CLIP models
# - clip_vision: CLIP Vision models
# - upscale_models: Upscaler models
# - embeddings: Textual inversion embeddings
# - hypernetworks: Hypernetwork models
# - style_models: Style transfer models
settings:
cache_dir: ~/.cache/huggingface
parallel_downloads: 1
model_categories:
# Stable Diffusion Checkpoints
checkpoints:
- repo_id: runwayml/stable-diffusion-v1-5
description: Stable Diffusion v1.5
size_gb: 4
essential: true
type: checkpoints
filename: "" # Empty means all model files
- repo_id: stabilityai/stable-diffusion-xl-base-1.0
description: SDXL Base 1.0
size_gb: 7
essential: true
type: checkpoints
filename: ""
# VAE Models
vae:
- repo_id: stabilityai/sd-vae-ft-mse-original
description: SD VAE ft MSE
size_gb: 0.3
essential: true
type: vae
filename: ""
# LoRA Models
loras:
- repo_id: latent-consistency/lcm-lora-sdv1-5
description: LCM LoRA for SD v1.5
size_gb: 0.1
essential: false
type: loras
filename: ""
# ControlNet Models
controlnet:
- repo_id: lllyasviel/control_v11p_sd15_canny
description: ControlNet Canny
size_gb: 1.4
essential: false
type: controlnet
filename: ""
- repo_id: lllyasviel/control_v11p_sd15_openpose
description: ControlNet OpenPose
size_gb: 1.4
essential: false
type: controlnet
filename: ""
# Upscale Models
upscale_models:
- repo_id: ai-forever/Real-ESRGAN
description: Real-ESRGAN x4
size_gb: 0.1
essential: false
type: upscale_models
filename: "RealESRGAN_x4plus.pth"