diff --git a/py/services/download_manager.py b/py/services/download_manager.py index b8042414..cce8cc84 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -415,8 +415,10 @@ class DownloadManager: base_model_mappings = settings_manager.get('base_model_path_mappings', {}) mapped_base_model = base_model_mappings.get(base_model, base_model) + model_info = version_info.get('model') or {} + # Get model tags - model_tags = version_info.get('model', {}).get('tags', []) + model_tags = model_info.get('tags', []) first_tag = settings_manager.resolve_priority_tag_for_model(model_tags, model_type) @@ -425,6 +427,8 @@ class DownloadManager: 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) + formatted_path = formatted_path.replace('{model_name}', model_info.get('name', '')) + formatted_path = formatted_path.replace('{version_name}', version_info.get('name', '')) if model_type == 'embedding': formatted_path = formatted_path.replace(' ', '_') diff --git a/py/utils/utils.py b/py/utils/utils.py index c77136c6..fa0d0fc0 100644 --- a/py/utils/utils.py +++ b/py/utils/utils.py @@ -175,10 +175,18 @@ def calculate_relative_path_for_model(model_data: Dict, model_type: str = 'lora' first_tag = 'no tags' # Default if no tags available # Format the template with available data + model_name = model_data.get('model_name', '') + version_name = '' + + if isinstance(civitai_data, dict): + version_name = civitai_data.get('name') or '' + 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) + formatted_path = formatted_path.replace('{model_name}', model_name) + formatted_path = formatted_path.replace('{version_name}', version_name) if model_type == 'embedding': formatted_path = formatted_path.replace(' ', '_') diff --git a/static/js/utils/constants.js b/static/js/utils/constants.js index e26a0347..3f0e5aeb 100644 --- a/static/js/utils/constants.js +++ b/static/js/utils/constants.js @@ -117,7 +117,9 @@ export const DOWNLOAD_PATH_TEMPLATES = { export const PATH_TEMPLATE_PLACEHOLDERS = [ '{base_model}', '{author}', - '{first_tag}' + '{first_tag}', + '{model_name}', + '{version_name}' ]; // Default templates for each model type diff --git a/templates/components/modals/settings_modal.html b/templates/components/modals/settings_modal.html index dd91f240..1c44d55a 100644 --- a/templates/components/modals/settings_modal.html +++ b/templates/components/modals/settings_modal.html @@ -301,6 +301,8 @@ {base_model} {author} {first_tag} + {model_name} + {version_name} diff --git a/tests/services/test_download_manager.py b/tests/services/test_download_manager.py index eedca8d7..d5acff88 100644 --- a/tests/services/test_download_manager.py +++ b/tests/services/test_download_manager.py @@ -324,6 +324,22 @@ def test_embedding_relative_path_replaces_spaces(): assert relative_path == "Base_Model/tag_with_space" +def test_relative_path_supports_model_and_version_placeholders(): + manager = DownloadManager() + settings_manager = get_settings_manager() + settings_manager.settings["download_path_templates"]["lora"] = "{model_name}/{version_name}" + + version_info = { + "baseModel": "BaseModel", + "name": "Version One", + "model": {"name": "Fancy Model", "tags": []}, + } + + relative_path = manager._calculate_relative_path(version_info, "lora") + + assert relative_path == "Fancy Model/Version One" + + async def test_execute_download_retries_urls(monkeypatch, tmp_path): manager = DownloadManager() diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index ca371494..952c5944 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -53,6 +53,21 @@ def test_calculate_relative_path_for_model_uses_mappings_and_defaults(isolated_s assert relative_path == "SDXL-mapped/no tags/Creator" +def test_calculate_relative_path_supports_model_and_version(isolated_settings): + isolated_settings["download_path_templates"]["lora"] = "{model_name}/{version_name}" + + model_data = { + "model_name": "Fancy Model", + "base_model": "SDXL", + "tags": ["tag"], + "civitai": {"id": 1, "name": "Version One", "creator": {"username": "Creator"}}, + } + + relative_path = calculate_relative_path_for_model(model_data, "lora") + + assert relative_path == "Fancy Model/Version One" + + def test_calculate_recipe_fingerprint_filters_and_sorts(): loras = [ {"hash": "ABC", "strength": 0.1234},