feat: add support for explicit file mappings in YAML
- Added parse_file_mappings() function to extract file mappings from YAML - Modified link_model() to accept and use explicit source→dest mappings - Falls back to automatic prefixing if no mappings specified - Updated process_category() to parse and pass file mappings - This allows workflows to reference models by their expected names 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -227,6 +227,45 @@ except Exception as e:
|
|||||||
EOPYAML
|
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" <<EOPYTHON
|
||||||
|
import yaml
|
||||||
|
import sys
|
||||||
|
|
||||||
|
yaml_file = sys.argv[1]
|
||||||
|
category = sys.argv[2]
|
||||||
|
repo_id = sys.argv[3]
|
||||||
|
|
||||||
|
try:
|
||||||
|
with open(yaml_file, 'r') as f:
|
||||||
|
config = yaml.safe_load(f)
|
||||||
|
|
||||||
|
if category in config.get('model_categories', {}):
|
||||||
|
models = config['model_categories'][category]
|
||||||
|
for model in models:
|
||||||
|
if model.get('repo_id', '') == repo_id:
|
||||||
|
files = model.get('files', [])
|
||||||
|
if files:
|
||||||
|
for file_mapping in files:
|
||||||
|
source = file_mapping.get('source', '')
|
||||||
|
dest = file_mapping.get('dest', source)
|
||||||
|
if source:
|
||||||
|
print('{0}|{1}'.format(source, dest))
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
# No file mappings found
|
||||||
|
sys.exit(0)
|
||||||
|
except Exception as e:
|
||||||
|
print("ERROR: {0}".format(e), file=sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
EOPYTHON
|
||||||
|
}
|
||||||
|
|
||||||
# Check dependencies
|
# Check dependencies
|
||||||
check_dependencies() {
|
check_dependencies() {
|
||||||
print_section "Checking Dependencies"
|
print_section "Checking Dependencies"
|
||||||
@@ -406,6 +445,7 @@ link_model() {
|
|||||||
local repo_id="$1"
|
local repo_id="$1"
|
||||||
local model_type="$2"
|
local model_type="$2"
|
||||||
local filename_filter="$3"
|
local filename_filter="$3"
|
||||||
|
local file_mappings="$4" # Optional: explicit source|dest mappings
|
||||||
|
|
||||||
print_detail "Linking to: ${CYAN}${COMFYUI_DIR}/${model_type}/${RESET}"
|
print_detail "Linking to: ${CYAN}${COMFYUI_DIR}/${model_type}/${RESET}"
|
||||||
|
|
||||||
@@ -429,34 +469,73 @@ link_model() {
|
|||||||
return 1
|
return 1
|
||||||
fi
|
fi
|
||||||
|
|
||||||
# Extract model name from repo_id for prefixing filenames
|
|
||||||
# e.g., "facebook/musicgen-medium" -> "musicgen-medium"
|
|
||||||
local model_name=$(echo "$repo_id" | sed 's/.*\///')
|
|
||||||
|
|
||||||
local linked_count=0
|
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
|
# If explicit file mappings are provided, use them
|
||||||
# e.g., "pytorch_model.bin" -> "musicgen-medium-pytorch_model.bin"
|
if [[ -n "$file_mappings" ]]; then
|
||||||
local prefixed_filename="${model_name}-${filename}"
|
print_detail "Using explicit file mappings from YAML"
|
||||||
local link_path="${target_dir}/${prefixed_filename}"
|
|
||||||
|
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
|
# Remove existing symlink or file if it exists
|
||||||
if [[ -L "$link_path" ]]; then
|
if [[ -L "$link_path" ]]; then
|
||||||
rm -f "$link_path"
|
rm -f "$link_path"
|
||||||
elif [[ -e "$link_path" ]]; then
|
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
|
continue
|
||||||
fi
|
fi
|
||||||
|
|
||||||
# Create symlink
|
# Create symlink
|
||||||
ln -s "$source_file" "$link_path"
|
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))
|
linked_count=$((linked_count+1))
|
||||||
fi
|
done <<< "$file_mappings"
|
||||||
done <<< "$model_files"
|
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
|
# Also link checkpoint files to checkpoints/ directory if they exist
|
||||||
# This handles models like SDXL that have both diffusers and checkpoint formats
|
# 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'
|
# Link if command is 'link' or 'both'
|
||||||
if [[ "$COMMAND" == "link" ]] || [[ "$COMMAND" == "both" ]]; then
|
if [[ "$COMMAND" == "link" ]] || [[ "$COMMAND" == "both" ]]; then
|
||||||
if $success; 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
|
success=false
|
||||||
fi
|
fi
|
||||||
fi
|
fi
|
||||||
|
|||||||
Reference in New Issue
Block a user