diff --git a/artifact_comfyui_download.sh b/artifact_comfyui_download.sh index b22405a..f328f32 100755 --- a/artifact_comfyui_download.sh +++ b/artifact_comfyui_download.sh @@ -227,6 +227,45 @@ except Exception as e: EOPYAML } +# Parse file mappings for a specific model +parse_file_mappings() { + local yaml_file="$1" + local category="$2" + local repo_id="$3" + + python3 - "$yaml_file" "$category" "$repo_id" < "musicgen-medium" - local model_name=$(echo "$repo_id" | sed 's/.*\///') - local linked_count=0 - while IFS= read -r source_file; do - if [[ -f "$source_file" ]]; then - local filename=$(basename "$source_file") - # Add model name prefix to filename for better organization - # e.g., "pytorch_model.bin" -> "musicgen-medium-pytorch_model.bin" - local prefixed_filename="${model_name}-${filename}" - local link_path="${target_dir}/${prefixed_filename}" + # If explicit file mappings are provided, use them + if [[ -n "$file_mappings" ]]; then + print_detail "Using explicit file mappings from YAML" + + while IFS='|' read -r source_pattern dest_filename; do + if [[ -z "$source_pattern" ]]; then + continue + fi + + # Find the file matching the source pattern in model_files + local source_file + source_file=$(echo "$model_files" | grep -F "/$source_pattern" | head -n1) + + if [[ -z "$source_file" ]] || [[ ! -f "$source_file" ]]; then + print_warning "Source file not found: ${source_pattern}" + continue + fi + + local link_path="${target_dir}/${dest_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): ${prefixed_filename}" + print_warning "File already exists (not a symlink): ${dest_filename}" continue fi # Create symlink ln -s "$source_file" "$link_path" - print_detail "${LINK} Linked: ${DIM}${prefixed_filename}${RESET}" + print_detail "${LINK} Linked: ${DIM}${dest_filename}${RESET}" linked_count=$((linked_count+1)) - fi - done <<< "$model_files" + done <<< "$file_mappings" + else + # Fallback: use automatic prefixing for files without explicit mappings + print_detail "No file mappings found, using automatic prefixing" + + # Extract model name from repo_id for prefixing filenames + # e.g., "facebook/musicgen-medium" -> "musicgen-medium" + local model_name=$(echo "$repo_id" | sed 's/.*\///') + + while IFS= read -r source_file; do + if [[ -f "$source_file" ]]; then + local filename=$(basename "$source_file") + + # Add model name prefix to filename for better organization + # e.g., "pytorch_model.bin" -> "musicgen-medium-pytorch_model.bin" + local prefixed_filename="${model_name}-${filename}" + local link_path="${target_dir}/${prefixed_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): ${prefixed_filename}" + continue + fi + + # Create symlink + ln -s "$source_file" "$link_path" + print_detail "${LINK} Linked: ${DIM}${prefixed_filename}${RESET}" + linked_count=$((linked_count+1)) + fi + done <<< "$model_files" + fi # Also link checkpoint files to checkpoints/ directory if they exist # This handles models like SDXL that have both diffusers and checkpoint formats @@ -551,7 +630,12 @@ process_category() { # 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 + # Parse file mappings from YAML for this model + local file_mappings + file_mappings=$(parse_file_mappings "$CONFIG_FILE" "$category" "$repo_id") + + # Pass file mappings to link_model (empty string if no mappings found) + if ! link_model "$repo_id" "$model_type" "$filename" "$file_mappings"; then success=false fi fi