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:
2025-11-22 17:20:26 +01:00
parent 1f94435d36
commit 520e6ce260

View File

@@ -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" <<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() {
print_section "Checking Dependencies"
@@ -406,6 +445,7 @@ link_model() {
local repo_id="$1"
local model_type="$2"
local filename_filter="$3"
local file_mappings="$4" # Optional: explicit source|dest mappings
print_detail "Linking to: ${CYAN}${COMFYUI_DIR}/${model_type}/${RESET}"
@@ -429,34 +469,73 @@ link_model() {
return 1
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
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