diff --git a/artifact_comfyui_download.sh b/artifact_comfyui_download.sh index 37961b9..0966f56 100755 --- a/artifact_comfyui_download.sh +++ b/artifact_comfyui_download.sh @@ -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 - <