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:
@@ -1,9 +1,11 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# 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
|
||||
@@ -98,6 +100,12 @@ fi
|
||||
# Default cache directory (use HuggingFace default)
|
||||
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
|
||||
if [[ -f "${PROJECT_ROOT}/.env" ]]; then
|
||||
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', '')
|
||||
size_gb = model.get('size_gb', 0)
|
||||
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:
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
@@ -244,6 +254,9 @@ check_dependencies() {
|
||||
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"
|
||||
@@ -254,17 +267,77 @@ validate_config() {
|
||||
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
|
||||
# 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
|
||||
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"
|
||||
# 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
|
||||
print_success "Cache directory ready: ${CYAN}${CACHE_DIR}${RESET}"
|
||||
|
||||
# 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
|
||||
@@ -311,8 +384,62 @@ EOPYDOWNLOAD
|
||||
fi
|
||||
}
|
||||
|
||||
# Download models by category
|
||||
download_category() {
|
||||
# Create symlink for a model
|
||||
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_display="$2"
|
||||
|
||||
@@ -333,13 +460,31 @@ download_category() {
|
||||
local succeeded=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++))
|
||||
|
||||
echo ""
|
||||
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++))
|
||||
else
|
||||
((failed++))
|
||||
@@ -417,13 +562,13 @@ main() {
|
||||
local total_succeeded=0
|
||||
local total_failed=0
|
||||
|
||||
# Download each category
|
||||
# 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')
|
||||
|
||||
download_category "$category" "$category_display"
|
||||
process_category "$category" "$category_display"
|
||||
|
||||
# Update counters (this is simplified, you'd need to track actual numbers)
|
||||
((total_succeeded++))
|
||||
@@ -451,16 +596,32 @@ while [[ $# -gt 0 ]]; do
|
||||
CACHE_DIR="$2"
|
||||
shift 2
|
||||
;;
|
||||
--comfyui-dir)
|
||||
COMFYUI_DIR="$2"
|
||||
shift 2
|
||||
;;
|
||||
download|link|both)
|
||||
COMMAND="$1"
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
echo "Usage: $0 [CONFIG_FILE] [options]"
|
||||
echo "Usage: $0 [COMMAND] [options]"
|
||||
echo ""
|
||||
echo "Arguments:"
|
||||
echo " CONFIG_FILE Configuration file path (optional, can also use -c)"
|
||||
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: ~/.cache/huggingface)"
|
||||
echo " --comfyui-dir DIR ComfyUI models directory (default: ~/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
|
||||
;;
|
||||
-*)
|
||||
|
||||
80
comfyui_models.example.yaml
Normal file
80
comfyui_models.example.yaml
Normal 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"
|
||||
Reference in New Issue
Block a user