diff --git a/artifact_comfyui_download.sh b/artifact_comfyui_download.sh index a1fa8e6..3f8cd71 100755 --- a/artifact_comfyui_download.sh +++ b/artifact_comfyui_download.sh @@ -416,6 +416,10 @@ link_model() { mkdir -p "$target_dir" fi + # Clean existing symlinks for this repo in the target directory + # This ensures we start fresh and don't have stale links + find "$target_dir" -type l -lname "*${repo_id/\//-}*" -delete 2>/dev/null || true + # Find model files in cache local model_files model_files=$(find_model_files "$repo_id" "$filename_filter") @@ -446,6 +450,49 @@ link_model() { fi done <<< "$model_files" + # Also link checkpoint files to checkpoints/ directory if they exist + # This handles models like SDXL that have both diffusers and checkpoint formats + local checkpoint_dir="${COMFYUI_DIR}/checkpoints" + if [[ "$model_type" == "diffusers" ]]; then + # Clean existing checkpoint symlinks for this repo + find "$checkpoint_dir" -type l -lname "*${repo_id/\//-}*" -delete 2>/dev/null || true + + # Look for checkpoint .safetensors files in the snapshot root (not in subdirectories) + # These are the consolidated model files like sd_xl_base_1.0.safetensors + local checkpoint_files + checkpoint_files=$(echo "$model_files" | grep -v '/' | grep '\.safetensors$' || echo "$model_files" | grep -E 'snapshots/[^/]+/[^/]+\.safetensors$' | grep -vE '/(vae|text_encoder|transformer|unet)/' || true) + + if [[ -n "$checkpoint_files" ]]; then + mkdir -p "$checkpoint_dir" + while IFS= read -r ckpt_file; do + if [[ -f "$ckpt_file" ]] && [[ -n "$ckpt_file" ]]; then + local ckpt_name=$(basename "$ckpt_file") + + # Skip if it's clearly a component file (has parent directory in name) + if [[ "$ckpt_file" =~ /(vae|text_encoder|transformer|unet|tokenizer)/ ]]; then + continue + fi + + # Only link if the file is large enough to be a checkpoint (> 500MB) + local file_size=$(stat -c%s "$ckpt_file" 2>/dev/null || echo "0") + if [[ $file_size -gt 524288000 ]]; then + local ckpt_link="${checkpoint_dir}/${ckpt_name}" + + # Remove existing symlink if it exists + if [[ -L "$ckpt_link" ]]; then + rm -f "$ckpt_link" + fi + + # Create symlink to checkpoints directory + ln -s "$ckpt_file" "$ckpt_link" + print_detail "${LINK} Also linked to checkpoints: ${DIM}${ckpt_name}${RESET}" + linked_count=$((linked_count+1)) + fi + fi + done <<< "$checkpoint_files" + fi + fi + if [[ $linked_count -gt 0 ]]; then print_success "Linked ${linked_count} file(s) for ${BOLD_WHITE}${repo_id}${RESET}" return 0