From b8aa7184bd37c5f6674bbe695bbf31fb2511b8a7 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Wed, 13 Aug 2025 19:23:37 +0800 Subject: [PATCH] feat: update download path template handling for model types and migrate old settings --- py/routes/base_model_routes.py | 6 +++--- py/services/download_manager.py | 13 +++++++---- py/services/settings_manager.py | 31 +++++++++++++++++++++++++++ py/utils/utils.py | 11 +++++++--- static/js/managers/SettingsManager.js | 4 ++-- 5 files changed, 53 insertions(+), 12 deletions(-) diff --git a/py/routes/base_model_routes.py b/py/routes/base_model_routes.py index d4522871..37680bc9 100644 --- a/py/routes/base_model_routes.py +++ b/py/routes/base_model_routes.py @@ -756,8 +756,8 @@ class BaseModelRoutes(ABC): 'error': 'No model roots configured' }, status=400) - # Check if flat structure is configured - path_template = settings.get('download_path_template', '{base_model}/{first_tag}') + # Check if flat structure is configured for this model type + path_template = settings.get_download_path_template(self.service.model_type) is_flat_structure = not path_template # Prepare results tracking @@ -832,7 +832,7 @@ class BaseModelRoutes(ABC): target_dir = current_root else: # Calculate new relative path based on settings - new_relative_path = calculate_relative_path_for_model(model) + new_relative_path = calculate_relative_path_for_model(model, self.service.model_type) # If no relative path calculated (insufficient metadata), skip if not new_relative_path: diff --git a/py/services/download_manager.py b/py/services/download_manager.py index 55e8715d..1ee46b5b 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -258,7 +258,7 @@ class DownloadManager: save_dir = default_path # Calculate relative path using template - relative_path = self._calculate_relative_path(version_info) + relative_path = self._calculate_relative_path(version_info, model_type) # Update save directory with relative path if provided if relative_path: @@ -331,17 +331,18 @@ class DownloadManager: return {'success': False, 'error': f"Early access restriction: {str(e)}. Please ensure you have purchased early access and are logged in to Civitai."} return {'success': False, 'error': str(e)} - def _calculate_relative_path(self, version_info: Dict) -> str: + def _calculate_relative_path(self, version_info: Dict, model_type: str = 'lora') -> str: """Calculate relative path using template from settings Args: version_info: Version info from Civitai API + model_type: Type of model ('lora', 'checkpoint', 'embedding') Returns: Relative path string """ - # Get path template from settings, default to '{base_model}/{first_tag}' - path_template = settings.get('download_path_template', '{base_model}/{first_tag}') + # Get path template from settings for specific model type + path_template = settings.get_download_path_template(model_type) # If template is empty, return empty path (flat structure) if not path_template: @@ -350,6 +351,9 @@ class DownloadManager: # Get base model name base_model = version_info.get('baseModel', '') + # Get author from creator data + author = version_info.get('creator', {}).get('username', 'Anonymous') + # Apply mapping if available base_model_mappings = settings.get('base_model_path_mappings', {}) mapped_base_model = base_model_mappings.get(base_model, base_model) @@ -372,6 +376,7 @@ class DownloadManager: formatted_path = path_template formatted_path = formatted_path.replace('{base_model}', mapped_base_model) formatted_path = formatted_path.replace('{first_tag}', first_tag) + formatted_path = formatted_path.replace('{author}', author) return formatted_path diff --git a/py/services/settings_manager.py b/py/services/settings_manager.py index d9e2d866..a6a6e24a 100644 --- a/py/services/settings_manager.py +++ b/py/services/settings_manager.py @@ -9,6 +9,7 @@ class SettingsManager: def __init__(self): self.settings_file = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'settings.json') self.settings = self._load_settings() + self._migrate_download_path_template() self._auto_set_default_roots() self._check_environment_variables() @@ -22,6 +23,24 @@ class SettingsManager: logger.error(f"Error loading settings: {e}") return self._get_default_settings() + def _migrate_download_path_template(self): + """Migrate old download_path_template to new download_path_templates""" + old_template = self.settings.get('download_path_template') + templates = self.settings.get('download_path_templates') + + # If old template exists and new templates don't exist, migrate + if old_template is not None and not templates: + logger.info("Migrating download_path_template to download_path_templates") + self.settings['download_path_templates'] = { + 'lora': old_template, + 'checkpoint': old_template, + 'embedding': old_template + } + # Remove old setting + del self.settings['download_path_template'] + self._save_settings() + logger.info("Migration completed") + def _auto_set_default_roots(self): """Auto set default root paths if only one folder is present and default is empty.""" folder_paths = self.settings.get('folder_paths', {}) @@ -81,4 +100,16 @@ class SettingsManager: except Exception as e: logger.error(f"Error saving settings: {e}") + def get_download_path_template(self, model_type: str) -> str: + """Get download path template for specific model type + + Args: + model_type: The type of model ('lora', 'checkpoint', 'embedding') + + Returns: + Template string for the model type, defaults to '{base_model}/{first_tag}' + """ + templates = self.settings.get('download_path_templates', {}) + return templates.get(model_type, '{base_model}/{first_tag}') + settings = SettingsManager() diff --git a/py/utils/utils.py b/py/utils/utils.py index fe97c5da..2c67c645 100644 --- a/py/utils/utils.py +++ b/py/utils/utils.py @@ -132,17 +132,18 @@ def calculate_recipe_fingerprint(loras): return fingerprint -def calculate_relative_path_for_model(model_data: Dict) -> str: +def calculate_relative_path_for_model(model_data: Dict, model_type: str = 'lora') -> str: """Calculate relative path for existing model using template from settings Args: model_data: Model data from scanner cache + model_type: Type of model ('lora', 'checkpoint', 'embedding') Returns: Relative path string (empty string for flat structure) """ - # Get path template from settings, default to '{base_model}/{first_tag}' - path_template = settings.get('download_path_template', '{base_model}/{first_tag}') + # Get path template from settings for specific model type + path_template = settings.get_download_path_template(model_type) # If template is empty, return empty path (flat structure) if not path_template: @@ -154,9 +155,12 @@ def calculate_relative_path_for_model(model_data: Dict) -> str: # For CivitAI models, prefer civitai data only if 'id' exists; for non-CivitAI models, use model_data directly if civitai_data and civitai_data.get('id') is not None: base_model = civitai_data.get('baseModel', '') + # Get author from civitai creator data + author = civitai_data.get('creator', {}).get('username', 'Anonymous') else: # Fallback to model_data fields for non-CivitAI models base_model = model_data.get('base_model', '') + author = 'Anonymous' # Default for non-CivitAI models model_tags = model_data.get('tags', []) @@ -182,6 +186,7 @@ def calculate_relative_path_for_model(model_data: Dict) -> str: formatted_path = path_template formatted_path = formatted_path.replace('{base_model}', mapped_base_model) formatted_path = formatted_path.replace('{first_tag}', first_tag) + formatted_path = formatted_path.replace('{author}', author) return formatted_path diff --git a/static/js/managers/SettingsManager.js b/static/js/managers/SettingsManager.js index 6203b17d..dfb0f3c3 100644 --- a/static/js/managers/SettingsManager.js +++ b/static/js/managers/SettingsManager.js @@ -92,7 +92,7 @@ export class SettingsManager { // Ensure all model types have templates Object.keys(DEFAULT_PATH_TEMPLATES).forEach(modelType => { - if (!state.global.settings.download_path_templates[modelType]) { + if (typeof state.global.settings.download_path_templates[modelType] === 'undefined') { state.global.settings.download_path_templates[modelType] = DEFAULT_PATH_TEMPLATES[modelType]; } }); @@ -586,7 +586,7 @@ export class SettingsManager { // Find matching preset const matchingPreset = this.findMatchingPreset(template); - if (matchingPreset) { + if (matchingPreset !== null) { presetSelect.value = matchingPreset; if (customRow) customRow.style.display = 'none'; } else {