From f749dd0d52ed9a21244cc3cbfc925cce67eeaff5 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Wed, 16 Jul 2025 17:07:13 +0800 Subject: [PATCH] feat: Add YAML configuration for path mappings to customize model download directories --- path_mappings.yaml | 71 +++++++++++++++++++++++++++++++++ py/services/download_manager.py | 40 ++++++++++++++++++- 2 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 path_mappings.yaml diff --git a/path_mappings.yaml b/path_mappings.yaml new file mode 100644 index 00000000..60f434a8 --- /dev/null +++ b/path_mappings.yaml @@ -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: /flux//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: //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" diff --git a/py/services/download_manager.py b/py/services/download_manager.py index 3e3d1e25..d53256a8 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -1,6 +1,7 @@ import logging import os import asyncio +import yaml from typing import Dict from ..utils.models import LoraMetadata, CheckpointMetadata from ..utils.constants import CARD_PREVIEW_WIDTH, VALID_LORA_TYPES, CIVITAI_MODEL_TAGS @@ -33,6 +34,36 @@ class DownloadManager: self._initialized = True 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): """Lazily initialize CivitaiClient from registry""" @@ -131,6 +162,9 @@ class DownloadManager: model_tags = version_info.get('model', {}).get('tags', []) 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 prioritized_tag = None for civitai_tag in CIVITAI_MODEL_TAGS: @@ -143,9 +177,11 @@ class DownloadManager: prioritized_tag = model_tags[0] 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: - relative_path = base_model + relative_path = mapped_base_model # Update save directory with relative path if provided if relative_path: