feat: Add YAML configuration for path mappings to customize model download directories

This commit is contained in:
Will Miao
2025-07-16 17:07:13 +08:00
parent 48a8a42108
commit f749dd0d52
2 changed files with 109 additions and 2 deletions

71
path_mappings.yaml Normal file
View File

@@ -0,0 +1,71 @@
# Path mappings configuration for ComfyUI-Lora-Manager
# This file allows you to customize how base models and Civitai tags map to directories when downloading models
# Base model mappings
# Format: "Original Base Model": "Custom Directory Name"
#
# Example: If you change "Flux.1 D": "flux"
# Then models with base model "Flux.1 D" will be stored in a directory named "flux"
# So the final path would be: <model_directory>/flux/<tag>/model_file.safetensors
base_models:
"SD 1.4": "SD 1.4"
"SD 1.5": "SD 1.5"
"SD 1.5 LCM": "SD 1.5 LCM"
"SD 1.5 Hyper": "SD 1.5 Hyper"
"SD 2.0": "SD 2.0"
"SD 2.1": "SD 2.1"
"SDXL 1.0": "SDXL 1.0"
"SD 3": "SD 3"
"SD 3.5": "SD 3.5"
"SD 3.5 Medium": "SD 3.5 Medium"
"SD 3.5 Large": "SD 3.5 Large"
"SD 3.5 Large Turbo": "SD 3.5 Large Turbo"
"Pony": "Pony"
"Flux.1 S": "Flux.1 S"
"Flux.1 D": "Flux.1 D"
"Flux.1 Kontext": "Flux.1 Kontext"
"AuraFlow": "AuraFlow"
"SDXL Lightning": "SDXL Lightning"
"SDXL Hyper": "SDXL Hyper"
"Stable Cascade": "Stable Cascade"
"SVD": "SVD"
"PixArt a": "PixArt a"
"PixArt E": "PixArt E"
"Hunyuan 1": "Hunyuan 1"
"Hunyuan Video": "Hunyuan Video"
"Lumina": "Lumina"
"Kolors": "Kolors"
"Illustrious": "Illustrious"
"Mochi": "Mochi"
"LTXV": "LTXV"
"CogVideoX": "CogVideoX"
"NoobAI": "NoobAI"
"Wan Video": "Wan Video"
"Wan Video 1.3B t2v": "Wan Video 1.3B t2v"
"Wan Video 14B t2v": "Wan Video 14B t2v"
"Wan Video 14B i2v 480p": "Wan Video 14B i2v 480p"
"Wan Video 14B i2v 720p": "Wan Video 14B i2v 720p"
"HiDream": "HiDream"
"Other": "Other"
# Civitai model tag mappings
# Format: "Original Tag": "Custom Directory Name"
#
# Example: If you change "character": "characters"
# Then models with tag "character" will be stored in a directory named "characters"
# So the final path would be: <model_directory>/<base_model>/characters/model_file.safetensors
model_tags:
"character": "character"
"style": "style"
"concept": "concept"
"clothing": "clothing"
"base model": "base model"
"poses": "poses"
"background": "background"
"tool": "tool"
"vehicle": "vehicle"
"buildings": "buildings"
"objects": "objects"
"assets": "assets"
"animal": "animal"
"action": "action"

View File

@@ -1,6 +1,7 @@
import logging import logging
import os import os
import asyncio import asyncio
import yaml
from typing import Dict from typing import Dict
from ..utils.models import LoraMetadata, CheckpointMetadata from ..utils.models import LoraMetadata, CheckpointMetadata
from ..utils.constants import CARD_PREVIEW_WIDTH, VALID_LORA_TYPES, CIVITAI_MODEL_TAGS from ..utils.constants import CARD_PREVIEW_WIDTH, VALID_LORA_TYPES, CIVITAI_MODEL_TAGS
@@ -33,6 +34,36 @@ class DownloadManager:
self._initialized = True self._initialized = True
self._civitai_client = None # Will be lazily initialized self._civitai_client = None # Will be lazily initialized
self._path_mappings = self._load_path_mappings()
def _load_path_mappings(self):
"""Load path mappings from YAML configuration"""
path_mappings = {
'base_models': {},
'model_tags': {}
}
# Path to the configuration file
config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'path_mappings.yaml')
try:
if os.path.exists(config_path):
with open(config_path, 'r', encoding='utf-8') as f:
mappings = yaml.safe_load(f)
if mappings and isinstance(mappings, dict):
if 'base_models' in mappings and isinstance(mappings['base_models'], dict):
path_mappings['base_models'] = mappings['base_models']
if 'model_tags' in mappings and isinstance(mappings['model_tags'], dict):
path_mappings['model_tags'] = mappings['model_tags']
logger.info(f"Loaded path mappings from {config_path}")
else:
logger.info(f"Path mappings configuration file not found at {config_path}, using default mappings")
except Exception as e:
logger.error(f"Error loading path mappings: {e}", exc_info=True)
return path_mappings
async def _get_civitai_client(self): async def _get_civitai_client(self):
"""Lazily initialize CivitaiClient from registry""" """Lazily initialize CivitaiClient from registry"""
@@ -131,6 +162,9 @@ class DownloadManager:
model_tags = version_info.get('model', {}).get('tags', []) model_tags = version_info.get('model', {}).get('tags', [])
if base_model: if base_model:
# Apply base model mapping if available
mapped_base_model = self._path_mappings['base_models'].get(base_model, base_model)
# Find the first Civitai model tag that exists in model_tags # Find the first Civitai model tag that exists in model_tags
prioritized_tag = None prioritized_tag = None
for civitai_tag in CIVITAI_MODEL_TAGS: for civitai_tag in CIVITAI_MODEL_TAGS:
@@ -143,9 +177,11 @@ class DownloadManager:
prioritized_tag = model_tags[0] prioritized_tag = model_tags[0]
if prioritized_tag: if prioritized_tag:
relative_path = os.path.join(base_model, prioritized_tag) # Apply tag mapping if available
mapped_tag = self._path_mappings['model_tags'].get(prioritized_tag, prioritized_tag)
relative_path = os.path.join(mapped_base_model, mapped_tag)
else: else:
relative_path = base_model relative_path = mapped_base_model
# Update save directory with relative path if provided # Update save directory with relative path if provided
if relative_path: if relative_path: