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
|
#!/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
|
||||||
|
|
||||||
if [[ -z "$HF_TOKEN" ]]; then
|
# HF_TOKEN only required for download and both commands
|
||||||
print_error "HF_TOKEN not set. Please set it in .env file or environment."
|
if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
|
||||||
exit 1
|
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
|
fi
|
||||||
print_success "HuggingFace token configured: ${DIM}${HF_TOKEN:0:10}...${RESET}"
|
|
||||||
|
|
||||||
if [[ ! -d "$CACHE_DIR" ]]; then
|
# Cache directory
|
||||||
print_info "Creating cache directory: ${CYAN}${CACHE_DIR}${RESET}"
|
if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then
|
||||||
mkdir -p "$CACHE_DIR"
|
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
|
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
|
# 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
|
||||||
;;
|
;;
|
||||||
-*)
|
-*)
|
||||||
|
|||||||
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