diff --git a/artifact_huggingface_download.sh b/artifact_huggingface_download.sh index 24912c7..1a7f689 100755 --- a/artifact_huggingface_download.sh +++ b/artifact_huggingface_download.sh @@ -3,9 +3,38 @@ # ComfyUI Model Downloader - A Beautiful CLI Tool # Downloads AI models from HuggingFace and creates symlinks to ComfyUI directories # -# Usage: ./artifact_comfyui_download.sh [COMMAND] [options] +# Usage: ./artifact_huggingface_download.sh [COMMAND] [OPTIONS] # -# Commands: download, link, both (default) +# Commands: +# download Download models from HuggingFace to cache directory +# link Create symlinks from cache to ComfyUI models directory +# both Download and link (default) +# verify Verify symlinks in ComfyUI models directory +# +# Options: +# -c, --config FILE Path to YAML configuration file +# --cache-dir DIR HuggingFace cache directory (default: auto-detect) +# --comfyui-dir DIR ComfyUI installation directory +# --filter-repo REPO... Only process specific repositories +# --category CAT1[,CAT2] Filter by category (comma-separated for multiple) +# --cleanup, --clean Remove unused cache files (link/both only) +# --dry-run, -n Show what would be done without making changes +# +# Examples: +# # Download and link all models from config +# ./artifact_huggingface_download.sh both -c models.yaml +# +# # Download only specific categories +# ./artifact_huggingface_download.sh download -c models.yaml --category image_models,video_models +# +# # Link with cleanup (remove unused cache files) +# ./artifact_huggingface_download.sh link -c models.yaml --cleanup +# +# # Dry-run to preview operations +# ./artifact_huggingface_download.sh both -c models.yaml --dry-run +# +# # Process only specific repositories +# ./artifact_huggingface_download.sh both -c models.yaml --filter-repo black-forest-labs/FLUX.1-schnell # set -euo pipefail @@ -111,6 +140,11 @@ fi # Default command COMMAND="both" +# Feature flags +CATEGORY_FILTER="" # Empty = all categories, or comma-separated list +CLEANUP_MODE=false # Remove unused files from HuggingFace cache +DRY_RUN=false # Simulate operations without making changes + # HuggingFace token from environment or .env file # Initialize HF_TOKEN if not set HF_TOKEN="${HF_TOKEN:-}" @@ -232,6 +266,49 @@ except Exception as e: EOPYAML } +# Validate and get categories to process +validate_and_get_categories() { + local config_file="$1" + + # Get all available categories + local all_categories + all_categories=$(parse_yaml "$config_file" "categories") + + # If no filter specified, return all categories + if [[ -z "$CATEGORY_FILTER" ]]; then + echo "$all_categories" + return 0 + fi + + # Split comma-separated categories and validate each + local requested_categories + IFS=',' read -ra requested_categories <<< "$CATEGORY_FILTER" + + local validated_categories=() + for requested in "${requested_categories[@]}"; do + # Trim whitespace + requested=$(echo "$requested" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//') + + # Check if category exists + if echo "$all_categories" | grep -q "^${requested}$"; then + validated_categories+=("$requested") + else + print_error "Invalid category: '$requested'" + echo "" + echo "Available categories:" + while IFS= read -r cat; do + echo " - $cat" + done <<< "$all_categories" + exit 1 + fi + done + + # Print validated categories (one per line) + for cat in "${validated_categories[@]}"; do + echo "$cat" + done +} + # Parse file mappings for a specific model parse_file_mappings() { local yaml_file="$1" @@ -334,6 +411,12 @@ validate_config() { print_info "Verify mode: HuggingFace token not required" fi + # Validate flag combinations + if [[ "$CLEANUP_MODE" == true ]] && [[ ! "$COMMAND" =~ ^(link|both)$ ]]; then + print_error "--cleanup can only be used with 'link' or 'both' commands" + exit 1 + fi + # Cache directory if [[ "$COMMAND" == "download" ]] || [[ "$COMMAND" == "both" ]]; then if [[ ! -d "$CACHE_DIR" ]]; then @@ -452,6 +535,12 @@ download_model() { print_detail "Description: ${description}" print_detail "Size: ${BOLD_YELLOW}${size_gb}GB${RESET}" + # Dry-run mode: skip actual download + if [[ "$DRY_RUN" == true ]]; then + print_info "DRY-RUN: Would download ${BOLD_WHITE}${repo_id}${RESET} (~${size_gb} GB)" + return 0 + fi + # Download using Python python3 - <