mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
feat(settings): filename templates for download and bulk rename (#1071)
Add per-model-type filename templates ({model_name}, {version_name},
{base_model}, {author}, {first_tag}, {hash_short}, {original_name}) so
downloaded files get informative names instead of e.g. V1.safetensors.
Empty template keeps the current filename (opt-in, off by default).
- apply template automatically after downloads; rename conflicts keep
the original name and never fail the download
- record original_file_name in metadata on rename for traceability
- bulk apply via GET|POST /api/lm/{prefix}/apply-filename-template with
WebSocket progress, sharing the auto-organize lock
- settings UI lives in the new Organization tab with validation, live
preview, and per-type 'apply to library' actions
This commit is contained in:
@@ -35,6 +35,8 @@ vi.mock('../../../static/js/utils/constants.js', () => ({
|
||||
DEFAULT_PATH_TEMPLATES: {},
|
||||
MAPPABLE_BASE_MODELS: ['Flux.1 D', 'Pony', 'SDXL 1.0', 'Other'],
|
||||
PATH_TEMPLATE_PLACEHOLDERS: {},
|
||||
FILENAME_TEMPLATE_PLACEHOLDERS: [],
|
||||
DEFAULT_FILENAME_TEMPLATES: { lora: '', checkpoint: '', embedding: '' },
|
||||
DEFAULT_PRIORITY_TAG_CONFIG: {
|
||||
lora: 'character, style',
|
||||
checkpoint: 'base, guide',
|
||||
|
||||
@@ -0,0 +1,276 @@
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
|
||||
vi.mock('../../../static/js/managers/ModalManager.js', () => ({
|
||||
modalManager: {
|
||||
closeModal: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/uiHelpers.js', () => ({
|
||||
showToast: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/state/index.js', () => {
|
||||
const settings = {};
|
||||
return {
|
||||
state: {
|
||||
global: {
|
||||
settings,
|
||||
},
|
||||
},
|
||||
createDefaultSettings: () => ({
|
||||
language: 'en',
|
||||
download_filename_templates: { lora: '', checkpoint: '', embedding: '' },
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock('../../../static/js/api/modelApiFactory.js', () => ({
|
||||
resetAndReload: vi.fn(),
|
||||
getModelApiClient: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/constants.js', () => ({
|
||||
DOWNLOAD_PATH_TEMPLATES: {},
|
||||
DEFAULT_PATH_TEMPLATES: {},
|
||||
MAPPABLE_BASE_MODELS: [],
|
||||
PATH_TEMPLATE_PLACEHOLDERS: [],
|
||||
FILENAME_TEMPLATE_PLACEHOLDERS: [
|
||||
'{model_name}',
|
||||
'{version_name}',
|
||||
'{base_model}',
|
||||
'{author}',
|
||||
'{first_tag}',
|
||||
'{hash_short}',
|
||||
'{original_name}',
|
||||
],
|
||||
DEFAULT_FILENAME_TEMPLATES: { lora: '', checkpoint: '', embedding: '' },
|
||||
DEFAULT_PRIORITY_TAG_CONFIG: {
|
||||
lora: 'character, style',
|
||||
checkpoint: 'base, guide',
|
||||
embedding: 'hint',
|
||||
},
|
||||
getMappableBaseModelsDynamic: () => [],
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/i18nHelpers.js', () => ({
|
||||
translate: (key, params, fallback) => {
|
||||
if (params && fallback) {
|
||||
return fallback.replace(/\{(\w+)\}/g, (match, name) => params[name] ?? match);
|
||||
}
|
||||
return fallback ?? '';
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/i18n/index.js', () => ({
|
||||
i18n: {
|
||||
getCurrentLocale: () => 'en',
|
||||
setLanguage: vi.fn().mockResolvedValue(),
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/components/shared/ModelCard.js', () => ({
|
||||
configureModelCardVideo: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/managers/BannerService.js', () => ({
|
||||
bannerService: {
|
||||
registerBanner: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
import { SettingsManager } from '../../../static/js/managers/SettingsManager.js';
|
||||
import { state } from '../../../static/js/state/index.js';
|
||||
import { showToast } from '../../../static/js/utils/uiHelpers.js';
|
||||
import { resetAndReload, getModelApiClient } from '../../../static/js/api/modelApiFactory.js';
|
||||
|
||||
const createManager = () => {
|
||||
const initSettingsSpy = vi
|
||||
.spyOn(SettingsManager.prototype, 'initializeSettings')
|
||||
.mockResolvedValue();
|
||||
const initializeSpy = vi
|
||||
.spyOn(SettingsManager.prototype, 'initialize')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
const manager = new SettingsManager();
|
||||
|
||||
initSettingsSpy.mockRestore();
|
||||
initializeSpy.mockRestore();
|
||||
|
||||
return manager;
|
||||
};
|
||||
|
||||
const appendFilenameTemplateUi = (modelType = 'lora') => {
|
||||
document.body.innerHTML = `
|
||||
<input id="${modelType}FilenameTemplate" />
|
||||
<div id="${modelType}FilenameValidation"></div>
|
||||
<div id="${modelType}FilenamePreview"></div>
|
||||
<button id="${modelType}ApplyFilenameTemplate" type="button"></button>
|
||||
`;
|
||||
};
|
||||
|
||||
describe('SettingsManager filename templates', () => {
|
||||
beforeEach(() => {
|
||||
document.body.innerHTML = '';
|
||||
vi.clearAllMocks();
|
||||
state.global.settings = {
|
||||
download_filename_templates: { lora: '', checkpoint: '', embedding: '' },
|
||||
};
|
||||
});
|
||||
|
||||
it('treats an empty template as valid (keep original filename)', () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
|
||||
expect(manager.validateFilenameTemplate('lora', '')).toBe(true);
|
||||
|
||||
const validation = document.getElementById('loraFilenameValidation');
|
||||
expect(validation.classList.contains('valid')).toBe(true);
|
||||
expect(validation.textContent).toContain('keep original filename');
|
||||
});
|
||||
|
||||
it('rejects templates with path separators or OS-illegal characters', () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
|
||||
expect(manager.validateFilenameTemplate('lora', '{base_model}/{model_name}')).toBe(false);
|
||||
expect(manager.validateFilenameTemplate('lora', 'a:b')).toBe(false);
|
||||
|
||||
const validation = document.getElementById('loraFilenameValidation');
|
||||
expect(validation.classList.contains('invalid')).toBe(true);
|
||||
});
|
||||
|
||||
it('rejects unknown placeholders', () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
|
||||
expect(manager.validateFilenameTemplate('lora', '{bogus}-{model_name}')).toBe(false);
|
||||
|
||||
const validation = document.getElementById('loraFilenameValidation');
|
||||
expect(validation.textContent).toContain('{bogus}');
|
||||
});
|
||||
|
||||
it('accepts a template using only known placeholders', () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
|
||||
const template = '{base_model}-{model_name}-{version_name}-{hash_short}';
|
||||
expect(manager.validateFilenameTemplate('lora', template)).toBe(true);
|
||||
expect(document.getElementById('loraFilenameValidation').classList.contains('valid')).toBe(true);
|
||||
});
|
||||
|
||||
it('saves a valid template via saveSetting with the merged dict', () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
manager.saveSetting = vi.fn().mockResolvedValue();
|
||||
|
||||
manager.updateFilenameTemplate('lora', '{model_name}');
|
||||
|
||||
expect(state.global.settings.download_filename_templates.lora).toBe('{model_name}');
|
||||
expect(manager.saveSetting).toHaveBeenCalledWith(
|
||||
'download_filename_templates',
|
||||
{ lora: '{model_name}', checkpoint: '', embedding: '' },
|
||||
);
|
||||
});
|
||||
|
||||
it('does not save an invalid template', () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
manager.saveSetting = vi.fn().mockResolvedValue();
|
||||
|
||||
manager.updateFilenameTemplate('lora', '{unknown_placeholder}');
|
||||
|
||||
expect(state.global.settings.download_filename_templates.lora).toBe('');
|
||||
expect(manager.saveSetting).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('previews the original filename when the template is empty', () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
|
||||
manager.updateFilenamePreview('lora', '');
|
||||
|
||||
expect(document.getElementById('loraFilenamePreview').textContent).toBe('V1.safetensors');
|
||||
});
|
||||
|
||||
it('renders a preview with example placeholder values', () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
|
||||
manager.updateFilenamePreview('lora', '{base_model}-{model_name}-{version_name}-{hash_short}');
|
||||
|
||||
expect(document.getElementById('loraFilenamePreview').textContent)
|
||||
.toBe('Flux.1 D-model-name-v3-a1b2c3d4e5.safetensors');
|
||||
});
|
||||
|
||||
it('disables the apply button when the template is empty', () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
|
||||
manager.updateFilenameTemplateApplyButton('lora', '');
|
||||
expect(document.getElementById('loraApplyFilenameTemplate').disabled).toBe(true);
|
||||
|
||||
manager.updateFilenameTemplateApplyButton('lora', '{model_name}');
|
||||
expect(document.getElementById('loraApplyFilenameTemplate').disabled).toBe(false);
|
||||
});
|
||||
|
||||
it('merges backend download_filename_templates over defaults', () => {
|
||||
const manager = createManager();
|
||||
|
||||
const merged = manager.mergeSettingsWithDefaults({
|
||||
download_filename_templates: { lora: '{model_name}' },
|
||||
});
|
||||
expect(merged.download_filename_templates).toEqual({
|
||||
lora: '{model_name}',
|
||||
checkpoint: '',
|
||||
embedding: '',
|
||||
});
|
||||
|
||||
const fromString = manager.mergeSettingsWithDefaults({
|
||||
download_filename_templates: '{"checkpoint":"{hash_short}"}',
|
||||
});
|
||||
expect(fromString.download_filename_templates).toEqual({
|
||||
lora: '',
|
||||
checkpoint: '{hash_short}',
|
||||
embedding: '',
|
||||
});
|
||||
});
|
||||
|
||||
it('shows an info toast and does nothing when applying an empty template', async () => {
|
||||
appendFilenameTemplateUi();
|
||||
const manager = createManager();
|
||||
vi.stubGlobal('confirm', vi.fn(() => true));
|
||||
|
||||
await manager.applyFilenameTemplate('lora');
|
||||
|
||||
expect(showToast).toHaveBeenCalledWith('settings.filenameTemplates.emptyTemplateInfo', {}, 'info');
|
||||
expect(getModelApiClient).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('applies the template through the model API client and reloads', async () => {
|
||||
appendFilenameTemplateUi();
|
||||
state.global.settings.download_filename_templates.lora = '{model_name}';
|
||||
const manager = createManager();
|
||||
vi.stubGlobal('confirm', vi.fn(() => true));
|
||||
const apiClient = { applyFilenameTemplate: vi.fn().mockResolvedValue() };
|
||||
getModelApiClient.mockReturnValue(apiClient);
|
||||
|
||||
await manager.applyFilenameTemplate('lora');
|
||||
|
||||
expect(getModelApiClient).toHaveBeenCalledWith('loras');
|
||||
expect(apiClient.applyFilenameTemplate).toHaveBeenCalledWith();
|
||||
expect(resetAndReload).toHaveBeenCalledWith(true);
|
||||
});
|
||||
|
||||
it('does not apply when the confirm dialog is declined', async () => {
|
||||
appendFilenameTemplateUi();
|
||||
state.global.settings.download_filename_templates.lora = '{model_name}';
|
||||
const manager = createManager();
|
||||
vi.stubGlobal('confirm', vi.fn(() => false));
|
||||
|
||||
await manager.applyFilenameTemplate('lora');
|
||||
|
||||
expect(getModelApiClient).not.toHaveBeenCalled();
|
||||
expect(resetAndReload).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -38,6 +38,8 @@ vi.mock('../../../static/js/utils/constants.js', () => ({
|
||||
DEFAULT_PATH_TEMPLATES: {},
|
||||
MAPPABLE_BASE_MODELS: [],
|
||||
PATH_TEMPLATE_PLACEHOLDERS: {},
|
||||
FILENAME_TEMPLATE_PLACEHOLDERS: [],
|
||||
DEFAULT_FILENAME_TEMPLATES: { lora: '', checkpoint: '', embedding: '' },
|
||||
DEFAULT_PRIORITY_TAG_CONFIG: {
|
||||
lora: 'character, style',
|
||||
checkpoint: 'base, guide',
|
||||
|
||||
@@ -36,6 +36,8 @@ vi.mock('../../../static/js/utils/constants.js', () => ({
|
||||
DEFAULT_PATH_TEMPLATES: {},
|
||||
MAPPABLE_BASE_MODELS: [],
|
||||
PATH_TEMPLATE_PLACEHOLDERS: {},
|
||||
FILENAME_TEMPLATE_PLACEHOLDERS: [],
|
||||
DEFAULT_FILENAME_TEMPLATES: { lora: '', checkpoint: '', embedding: '' },
|
||||
DEFAULT_PRIORITY_TAG_CONFIG: {},
|
||||
getMappableBaseModelsDynamic: () => [],
|
||||
}));
|
||||
|
||||
@@ -34,6 +34,8 @@ vi.mock('../../../static/js/utils/constants.js', () => ({
|
||||
DEFAULT_PATH_TEMPLATES: {},
|
||||
MAPPABLE_BASE_MODELS: [],
|
||||
PATH_TEMPLATE_PLACEHOLDERS: {},
|
||||
FILENAME_TEMPLATE_PLACEHOLDERS: [],
|
||||
DEFAULT_FILENAME_TEMPLATES: { lora: '', checkpoint: '', embedding: '' },
|
||||
DEFAULT_PRIORITY_TAG_CONFIG: {},
|
||||
getMappableBaseModelsDynamic: () => [],
|
||||
}));
|
||||
|
||||
@@ -23,6 +23,7 @@ from py.services.metadata_sync_service import MetadataSyncService
|
||||
from py.services.model_file_service import AutoOrganizeResult
|
||||
from py.services.model_update_service import ModelVersionRecord
|
||||
from py.services.service_registry import ServiceRegistry
|
||||
from py.services.use_cases import FilenameTemplateUseCase
|
||||
from py.services.websocket_manager import ws_manager
|
||||
from py.utils.exif_utils import ExifUtils
|
||||
from py.utils.metadata_manager import MetadataManager
|
||||
@@ -126,9 +127,11 @@ async def create_test_client(service) -> TestClient[Any, Any]:
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_ws_manager_state():
|
||||
ws_manager.cleanup_auto_organize_progress()
|
||||
ws_manager.cleanup_filename_template_progress()
|
||||
ws_manager._download_progress.clear()
|
||||
yield
|
||||
ws_manager.cleanup_auto_organize_progress()
|
||||
ws_manager.cleanup_filename_template_progress()
|
||||
ws_manager._download_progress.clear()
|
||||
|
||||
|
||||
@@ -762,6 +765,105 @@ def test_auto_organize_conflict_when_running(mock_service):
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_apply_filename_template_route_emits_progress(
|
||||
mock_service, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
async def fake_execute(self, file_paths=None, progress_callback=None):
|
||||
result = AutoOrganizeResult()
|
||||
result.total = 1
|
||||
result.processed = 1
|
||||
result.success_count = 1
|
||||
result.operation_type = "filename_template"
|
||||
if progress_callback is not None:
|
||||
await progress_callback.on_progress(
|
||||
{"type": "filename_template_progress", "status": "started"}
|
||||
)
|
||||
await progress_callback.on_progress(
|
||||
{"type": "filename_template_progress", "status": "completed"}
|
||||
)
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(FilenameTemplateUseCase, "execute", fake_execute)
|
||||
|
||||
async def scenario():
|
||||
client = await create_test_client(mock_service)
|
||||
try:
|
||||
response = await client.post(
|
||||
"/api/lm/test-models/apply-filename-template",
|
||||
json={"file_paths": ["/tmp/a.safetensors"]},
|
||||
)
|
||||
payload = await response.json()
|
||||
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
assert payload["summary"]["operation_type"] == "filename_template"
|
||||
|
||||
progress = ws_manager.get_filename_template_progress()
|
||||
assert progress is not None
|
||||
assert progress["status"] == "completed"
|
||||
# Auto-organize progress state must stay untouched.
|
||||
assert ws_manager.get_auto_organize_progress() is None
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_apply_filename_template_get_parses_query_file_paths(
|
||||
mock_service, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
captured = {}
|
||||
|
||||
async def fake_execute(self, file_paths=None, progress_callback=None):
|
||||
captured["file_paths"] = file_paths
|
||||
result = AutoOrganizeResult()
|
||||
result.operation_type = "filename_template"
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(FilenameTemplateUseCase, "execute", fake_execute)
|
||||
|
||||
async def scenario():
|
||||
client = await create_test_client(mock_service)
|
||||
try:
|
||||
response = await client.get(
|
||||
"/api/lm/test-models/apply-filename-template",
|
||||
params={"file_paths": "/tmp/a.safetensors, /tmp/b.safetensors"},
|
||||
)
|
||||
payload = await response.json()
|
||||
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
assert captured["file_paths"] == [
|
||||
"/tmp/a.safetensors",
|
||||
"/tmp/b.safetensors",
|
||||
]
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_apply_filename_template_conflict_when_running(mock_service):
|
||||
async def scenario():
|
||||
client = await create_test_client(mock_service)
|
||||
try:
|
||||
await ws_manager.broadcast_filename_template_progress(
|
||||
{"type": "filename_template_progress", "status": "started"}
|
||||
)
|
||||
|
||||
response = await client.post("/api/lm/test-models/apply-filename-template")
|
||||
payload = await response.json()
|
||||
|
||||
assert response.status == 409
|
||||
assert payload == {
|
||||
"success": False,
|
||||
"error": "Another library operation is already running. Please wait for it to complete.",
|
||||
}
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_download_model_returns_skipped_success(mock_service, download_manager_stub):
|
||||
async def scenario():
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Tests for the post-download filename template rename phase."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from py.services.download_manager import DownloadManager
|
||||
from py.services.service_registry import ServiceRegistry
|
||||
from py.services.settings_manager import get_settings_manager
|
||||
|
||||
|
||||
class DummyScanner:
|
||||
def __init__(self, root: Path):
|
||||
self._root = root
|
||||
self.model_type = "lora"
|
||||
self.updates = []
|
||||
|
||||
def get_model_roots(self):
|
||||
return [str(self._root)]
|
||||
|
||||
async def update_single_model_cache(self, original_path, new_path, metadata):
|
||||
self.updates.append((original_path, new_path, metadata))
|
||||
return True
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def download_manager() -> DownloadManager:
|
||||
return DownloadManager()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_recipe_scanner(monkeypatch: pytest.MonkeyPatch):
|
||||
async def _no_scanner():
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(ServiceRegistry, "get_recipe_scanner", _no_scanner)
|
||||
|
||||
|
||||
def _set_template(template: str, model_type: str = "lora") -> None:
|
||||
manager = get_settings_manager()
|
||||
templates = dict(manager.settings.get("download_filename_templates") or {})
|
||||
templates[model_type] = template
|
||||
manager.settings["download_filename_templates"] = templates
|
||||
|
||||
|
||||
def _write_model(root: Path, stem: str, model_name: str, sha256: str) -> Path:
|
||||
model_path = root / f"{stem}.safetensors"
|
||||
model_path.write_bytes(b"model")
|
||||
metadata_path = root / f"{stem}.metadata.json"
|
||||
metadata_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"file_name": stem,
|
||||
"file_path": model_path.as_posix(),
|
||||
"model_name": model_name,
|
||||
"sha256": sha256,
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
)
|
||||
)
|
||||
return model_path
|
||||
|
||||
|
||||
async def test_download_rename_applies_filename_template(
|
||||
tmp_path: Path, download_manager: DownloadManager
|
||||
):
|
||||
_set_template("{model_name}-{hash_short}")
|
||||
model_path = _write_model(tmp_path, "V1", "My Model", "abcdef0123456789")
|
||||
download_manager._active_downloads["dl1"] = {"file_path": model_path.as_posix()}
|
||||
|
||||
downloaded_metadata = [
|
||||
{
|
||||
"file_path": model_path.as_posix(),
|
||||
"file_name": "V1",
|
||||
"model_name": "My Model",
|
||||
"sha256": "abcdef0123456789",
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
]
|
||||
|
||||
await download_manager._apply_download_filename_template(
|
||||
scanner=DummyScanner(tmp_path),
|
||||
model_type="lora",
|
||||
downloaded_metadata=downloaded_metadata,
|
||||
download_id="dl1",
|
||||
)
|
||||
|
||||
new_path = tmp_path / "My Model-abcdef0123.safetensors"
|
||||
assert new_path.exists()
|
||||
assert not model_path.exists()
|
||||
|
||||
new_metadata = json.loads(
|
||||
(tmp_path / "My Model-abcdef0123.metadata.json").read_text()
|
||||
)
|
||||
assert new_metadata["original_file_name"] == "V1"
|
||||
|
||||
assert (
|
||||
download_manager._active_downloads["dl1"]["file_path"]
|
||||
== new_path.as_posix()
|
||||
)
|
||||
|
||||
|
||||
async def test_download_rename_keeps_original_on_conflict(
|
||||
tmp_path: Path, download_manager: DownloadManager
|
||||
):
|
||||
_set_template("{model_name}-{hash_short}")
|
||||
model_path = _write_model(tmp_path, "V1", "My Model", "abcdef0123456789")
|
||||
# Conflicting target already exists.
|
||||
(tmp_path / "My Model-abcdef0123.safetensors").write_bytes(b"other")
|
||||
|
||||
downloaded_metadata = [
|
||||
{
|
||||
"file_path": model_path.as_posix(),
|
||||
"file_name": "V1",
|
||||
"model_name": "My Model",
|
||||
"sha256": "abcdef0123456789",
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
]
|
||||
|
||||
# Must not raise: a rename conflict never fails the download.
|
||||
await download_manager._apply_download_filename_template(
|
||||
scanner=DummyScanner(tmp_path),
|
||||
model_type="lora",
|
||||
downloaded_metadata=downloaded_metadata,
|
||||
download_id=None,
|
||||
)
|
||||
|
||||
assert model_path.exists()
|
||||
|
||||
|
||||
async def test_download_rename_noop_without_template(
|
||||
tmp_path: Path, download_manager: DownloadManager
|
||||
):
|
||||
_set_template("")
|
||||
model_path = _write_model(tmp_path, "V1", "My Model", "abcdef0123456789")
|
||||
|
||||
await download_manager._apply_download_filename_template(
|
||||
scanner=DummyScanner(tmp_path),
|
||||
model_type="lora",
|
||||
downloaded_metadata=[{"file_path": model_path.as_posix()}],
|
||||
download_id=None,
|
||||
)
|
||||
|
||||
assert model_path.exists()
|
||||
assert (tmp_path / "V1.metadata.json").exists()
|
||||
@@ -424,6 +424,52 @@ async def test_rename_model_preserves_extension(tmp_path: Path):
|
||||
assert payload["file_name"] == new_name
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rename_model_records_original_file_name(tmp_path: Path):
|
||||
old_name = "V1"
|
||||
new_name = "flux-my-model-v3"
|
||||
|
||||
model_path = tmp_path / f"{old_name}.safetensors"
|
||||
model_path.write_bytes(b"model")
|
||||
|
||||
metadata_path = tmp_path / f"{old_name}.metadata.json"
|
||||
metadata_payload = {
|
||||
"file_name": old_name,
|
||||
"file_path": model_path.as_posix(),
|
||||
}
|
||||
metadata_path.write_text(json.dumps(metadata_payload))
|
||||
|
||||
async def metadata_loader(path: str):
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
return json.load(handle)
|
||||
|
||||
service = ModelLifecycleService(
|
||||
scanner=DummyScanner(),
|
||||
metadata_manager=PassthroughMetadataManager(),
|
||||
metadata_loader=metadata_loader,
|
||||
)
|
||||
|
||||
await service.rename_model(
|
||||
file_path=model_path.as_posix(),
|
||||
new_file_name=new_name,
|
||||
)
|
||||
|
||||
saved_metadata = json.loads((tmp_path / f"{new_name}.metadata.json").read_text())
|
||||
assert saved_metadata["original_file_name"] == old_name
|
||||
|
||||
# A second rename keeps the very first recorded name.
|
||||
second_name = "flux-my-model-v4"
|
||||
await service.rename_model(
|
||||
file_path=(tmp_path / f"{new_name}.safetensors").as_posix(),
|
||||
new_file_name=second_name,
|
||||
)
|
||||
|
||||
saved_metadata = json.loads(
|
||||
(tmp_path / f"{second_name}.metadata.json").read_text()
|
||||
)
|
||||
assert saved_metadata["original_file_name"] == old_name
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rename_model_with_dotted_basename(tmp_path: Path):
|
||||
old_name = "model.v1"
|
||||
|
||||
@@ -19,6 +19,7 @@ from py.services.use_cases import (
|
||||
DownloadModelEarlyAccessError,
|
||||
DownloadModelUseCase,
|
||||
DownloadModelValidationError,
|
||||
FilenameTemplateUseCase,
|
||||
ImportExampleImagesUseCase,
|
||||
ImportExampleImagesValidationError,
|
||||
)
|
||||
@@ -33,7 +34,7 @@ from py.utils.example_images_processor import (
|
||||
ExampleImagesValidationError,
|
||||
)
|
||||
from py.utils.metadata_manager import MetadataManager
|
||||
from tests.conftest import MockModelService, MockScanner
|
||||
from tests.conftest import MockCache, MockModelService, MockScanner
|
||||
|
||||
|
||||
class StubLockProvider:
|
||||
@@ -502,4 +503,179 @@ async def test_import_example_images_use_case_propagates_generic_error() -> None
|
||||
request = DummyJsonRequest({"model_hash": "abc", "file_paths": ["/tmp/file"]})
|
||||
|
||||
with pytest.raises(ExampleImagesImportError):
|
||||
await use_case.execute(request) # pyright: ignore[reportArgumentType]
|
||||
await use_case.execute(request) # pyright: ignore[reportArgumentType]
|
||||
|
||||
|
||||
class StubLifecycleService:
|
||||
def __init__(self, scanner: Optional[MockScanner] = None) -> None:
|
||||
self.renames: List[Dict[str, str]] = []
|
||||
self.error: Optional[Exception] = None
|
||||
self.cancel_on_rename = False
|
||||
self._scanner = scanner
|
||||
|
||||
async def rename_model(self, *, file_path: str, new_file_name: str) -> Dict[str, Any]:
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
self.renames.append({"file_path": file_path, "new_file_name": new_file_name})
|
||||
if self.cancel_on_rename and self._scanner is not None:
|
||||
self._scanner.cancel_task()
|
||||
return {"success": True, "new_file_path": file_path}
|
||||
|
||||
|
||||
def _filename_template_model(
|
||||
file_path: str,
|
||||
model_name: str,
|
||||
sha256: str = "abcdef0123456789",
|
||||
) -> Dict[str, Any]:
|
||||
return {
|
||||
"file_path": file_path,
|
||||
"file_name": file_path.rsplit("/", 1)[-1].rsplit(".", 1)[0],
|
||||
"model_name": model_name,
|
||||
"sha256": sha256,
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
|
||||
|
||||
def _set_filename_template(template: str, model_type: str = "lora") -> None:
|
||||
from py.services.settings_manager import get_settings_manager
|
||||
|
||||
manager = get_settings_manager()
|
||||
templates = dict(manager.settings.get("download_filename_templates") or {})
|
||||
templates[model_type] = template
|
||||
manager.settings["download_filename_templates"] = templates
|
||||
|
||||
|
||||
def _make_filename_template_use_case(
|
||||
scanner: MockScanner,
|
||||
lifecycle: StubLifecycleService,
|
||||
lock_provider: Optional[StubLockProvider] = None,
|
||||
) -> FilenameTemplateUseCase:
|
||||
return FilenameTemplateUseCase(
|
||||
scanner=scanner,
|
||||
lifecycle_service=lifecycle, # pyright: ignore[reportArgumentType]
|
||||
lock_provider=lock_provider or StubLockProvider(),
|
||||
model_type="lora",
|
||||
)
|
||||
|
||||
|
||||
async def test_filename_template_use_case_renames_models() -> None:
|
||||
_set_filename_template("{model_name}-{hash_short}")
|
||||
scanner = MockScanner(cache=MockCache([
|
||||
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
||||
_filename_template_model("/library/beta.safetensors", "Beta"),
|
||||
]))
|
||||
lifecycle = StubLifecycleService()
|
||||
progress = ProgressCollector()
|
||||
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
||||
|
||||
result = await use_case.execute(progress_callback=progress)
|
||||
|
||||
assert result.status == "success"
|
||||
assert result.operation_type == "filename_template"
|
||||
assert result.total == 2
|
||||
assert result.success_count == 2
|
||||
assert result.failure_count == 0
|
||||
assert lifecycle.renames == [
|
||||
{"file_path": "/library/alpha.safetensors", "new_file_name": "Alpha-abcdef0123"},
|
||||
{"file_path": "/library/beta.safetensors", "new_file_name": "Beta-abcdef0123"},
|
||||
]
|
||||
statuses = [event["status"] for event in progress.events]
|
||||
assert statuses[0] == "started"
|
||||
assert statuses[-1] == "completed"
|
||||
assert all(event["type"] == "filename_template_progress" for event in progress.events)
|
||||
|
||||
|
||||
async def test_filename_template_use_case_skips_unchanged_names() -> None:
|
||||
_set_filename_template("{model_name}-{hash_short}")
|
||||
scanner = MockScanner(cache=MockCache([
|
||||
_filename_template_model("/library/Alpha-abcdef0123.safetensors", "Alpha"),
|
||||
]))
|
||||
lifecycle = StubLifecycleService()
|
||||
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
||||
|
||||
result = await use_case.execute(progress_callback=None)
|
||||
|
||||
assert result.success_count == 0
|
||||
assert result.skipped_count == 1
|
||||
assert lifecycle.renames == []
|
||||
|
||||
|
||||
async def test_filename_template_use_case_skips_all_when_template_empty() -> None:
|
||||
_set_filename_template("")
|
||||
scanner = MockScanner(cache=MockCache([
|
||||
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
||||
]))
|
||||
lifecycle = StubLifecycleService()
|
||||
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
||||
|
||||
result = await use_case.execute(progress_callback=None)
|
||||
|
||||
assert result.skipped_count == 1
|
||||
assert lifecycle.renames == []
|
||||
|
||||
|
||||
async def test_filename_template_use_case_counts_conflicts_as_failures() -> None:
|
||||
_set_filename_template("{model_name}")
|
||||
scanner = MockScanner(cache=MockCache([
|
||||
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
||||
_filename_template_model("/library/beta.safetensors", "Beta"),
|
||||
]))
|
||||
lifecycle = StubLifecycleService()
|
||||
lifecycle.error = ValueError("A file with this name already exists")
|
||||
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
||||
|
||||
result = await use_case.execute(progress_callback=None)
|
||||
|
||||
assert result.status == "success"
|
||||
assert result.failure_count == 2
|
||||
assert result.success_count == 0
|
||||
assert len(result.results) == 2
|
||||
|
||||
|
||||
async def test_filename_template_use_case_honours_cancellation() -> None:
|
||||
_set_filename_template("{model_name}-{hash_short}")
|
||||
scanner = MockScanner(cache=MockCache([
|
||||
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
||||
_filename_template_model("/library/beta.safetensors", "Beta"),
|
||||
]))
|
||||
lifecycle = StubLifecycleService(scanner=scanner)
|
||||
lifecycle.cancel_on_rename = True
|
||||
progress = ProgressCollector()
|
||||
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
||||
|
||||
result = await use_case.execute(progress_callback=progress)
|
||||
|
||||
assert result.status == "cancelled"
|
||||
assert len(lifecycle.renames) == 1
|
||||
assert progress.events[-1]["status"] == "cancelled"
|
||||
|
||||
|
||||
async def test_filename_template_use_case_filters_file_paths() -> None:
|
||||
_set_filename_template("{model_name}-{hash_short}")
|
||||
scanner = MockScanner(cache=MockCache([
|
||||
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
||||
_filename_template_model("/library/beta.safetensors", "Beta"),
|
||||
]))
|
||||
lifecycle = StubLifecycleService()
|
||||
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
||||
|
||||
result = await use_case.execute(
|
||||
file_paths=["/library/beta.safetensors"], progress_callback=None
|
||||
)
|
||||
|
||||
assert result.total == 1
|
||||
assert lifecycle.renames == [
|
||||
{"file_path": "/library/beta.safetensors", "new_file_name": "Beta-abcdef0123"}
|
||||
]
|
||||
|
||||
|
||||
async def test_filename_template_use_case_rejects_when_lock_held() -> None:
|
||||
_set_filename_template("{model_name}")
|
||||
scanner = MockScanner(cache=MockCache())
|
||||
lifecycle = StubLifecycleService()
|
||||
lock_provider = StubLockProvider()
|
||||
lock_provider.running = True
|
||||
use_case = _make_filename_template_use_case(scanner, lifecycle, lock_provider)
|
||||
|
||||
with pytest.raises(AutoOrganizeInProgressError):
|
||||
await use_case.execute(progress_callback=None)
|
||||
@@ -3,6 +3,7 @@ import pytest
|
||||
from py.services.settings_manager import SettingsManager, get_settings_manager
|
||||
from py.services.service_registry import ServiceRegistry
|
||||
from py.utils.utils import (
|
||||
calculate_filename_for_model,
|
||||
calculate_recipe_fingerprint,
|
||||
calculate_relative_path_for_model,
|
||||
get_lora_info,
|
||||
@@ -164,6 +165,145 @@ def test_calculate_recipe_fingerprint_empty_input():
|
||||
assert calculate_recipe_fingerprint([]) == ""
|
||||
|
||||
|
||||
def _set_filename_templates(isolated_settings, template, model_types=("lora", "checkpoint", "embedding")):
|
||||
isolated_settings["download_filename_templates"] = {
|
||||
model_type: template for model_type in model_types
|
||||
}
|
||||
|
||||
|
||||
def test_calculate_filename_returns_empty_without_template(isolated_settings):
|
||||
model_data = {"model_name": "Model", "file_path": "/models/V1.safetensors"}
|
||||
|
||||
assert calculate_filename_for_model(model_data, "lora") == ""
|
||||
|
||||
|
||||
def test_calculate_filename_substitutes_all_placeholders(isolated_settings):
|
||||
_set_filename_templates(
|
||||
isolated_settings,
|
||||
"{base_model}-{model_name}-{version_name}-{author}-{first_tag}-{hash_short}-{original_name}",
|
||||
)
|
||||
|
||||
model_data = {
|
||||
"model_name": "My Model",
|
||||
"base_model": "SDXL",
|
||||
"tags": ["Style"],
|
||||
"sha256": "ABCDEF0123456789",
|
||||
"file_path": "/models/V1.safetensors",
|
||||
"civitai": {"id": 1, "name": "v3", "creator": {"username": "Author"}},
|
||||
}
|
||||
|
||||
result = calculate_filename_for_model(model_data, "lora")
|
||||
|
||||
assert result == "SDXL-My Model-v3-Author-style-abcdef0123-V1"
|
||||
|
||||
|
||||
def test_calculate_filename_hash_short_empty_when_unknown(isolated_settings):
|
||||
_set_filename_templates(isolated_settings, "{model_name}-{hash_short}")
|
||||
|
||||
model_data = {
|
||||
"model_name": "My Model",
|
||||
"file_path": "/models/V1.safetensors",
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
|
||||
# Missing hash leaves an empty segment; the dangling separator collapses.
|
||||
assert calculate_filename_for_model(model_data, "lora") == "My Model"
|
||||
|
||||
|
||||
def test_calculate_filename_missing_metadata_produces_empty_segments(isolated_settings):
|
||||
_set_filename_templates(isolated_settings, "{base_model}-{model_name}")
|
||||
|
||||
model_data = {
|
||||
"model_name": "My Model",
|
||||
"base_model": "",
|
||||
"tags": [],
|
||||
"file_path": "/models/V1.safetensors",
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
|
||||
assert calculate_filename_for_model(model_data, "lora") == "My Model"
|
||||
|
||||
|
||||
def test_calculate_filename_rejects_path_separators(isolated_settings):
|
||||
_set_filename_templates(isolated_settings, "{base_model}/{model_name}")
|
||||
|
||||
model_data = {
|
||||
"model_name": "My Model",
|
||||
"base_model": "SDXL",
|
||||
"file_path": "/models/V1.safetensors",
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
|
||||
assert calculate_filename_for_model(model_data, "lora") == ""
|
||||
|
||||
_set_filename_templates(isolated_settings, "{base_model}\\{model_name}")
|
||||
assert calculate_filename_for_model(model_data, "lora") == ""
|
||||
|
||||
|
||||
def test_calculate_filename_strips_illegal_characters(isolated_settings):
|
||||
_set_filename_templates(isolated_settings, '{model_name}:"custom"')
|
||||
|
||||
model_data = {
|
||||
"model_name": "My:Model*",
|
||||
"file_path": "/models/V1.safetensors",
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
|
||||
result = calculate_filename_for_model(model_data, "lora")
|
||||
|
||||
assert result == "My_Modelcustom"
|
||||
|
||||
|
||||
def test_calculate_filename_empty_result_returns_empty(isolated_settings):
|
||||
_set_filename_templates(isolated_settings, "{base_model}-{first_tag}")
|
||||
|
||||
model_data = {
|
||||
"base_model": "",
|
||||
"tags": [],
|
||||
"file_path": "/models/V1.safetensors",
|
||||
}
|
||||
|
||||
assert calculate_filename_for_model(model_data, "lora") == ""
|
||||
|
||||
|
||||
def test_calculate_filename_uses_base_model_mapping(isolated_settings):
|
||||
_set_filename_templates(isolated_settings, "{base_model}-{model_name}")
|
||||
isolated_settings["base_model_path_mappings"] = {"SDXL": "sdxl-mapped"}
|
||||
|
||||
model_data = {
|
||||
"model_name": "Model",
|
||||
"base_model": "SDXL",
|
||||
"file_path": "/models/V1.safetensors",
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
|
||||
assert calculate_filename_for_model(model_data, "lora") == "sdxl-mapped-Model"
|
||||
|
||||
|
||||
def test_calculate_filename_embedding_replaces_spaces(isolated_settings):
|
||||
_set_filename_templates(isolated_settings, "{base_model} {model_name}")
|
||||
|
||||
model_data = {
|
||||
"model_name": "My Model",
|
||||
"base_model": "Base Model",
|
||||
"file_path": "/models/V1.safetensors",
|
||||
"civitai": {"id": 1},
|
||||
}
|
||||
|
||||
assert calculate_filename_for_model(model_data, "embedding") == "Base_Model_My_Model"
|
||||
|
||||
|
||||
def test_calculate_filename_original_name_falls_back_to_file_name(isolated_settings):
|
||||
_set_filename_templates(isolated_settings, "{original_name}-{hash_short}")
|
||||
|
||||
model_data = {
|
||||
"file_name": "legacy-name",
|
||||
"sha256": "0123456789abcdef",
|
||||
}
|
||||
|
||||
assert calculate_filename_for_model(model_data, "lora") == "legacy-name-0123456789"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"original, expected",
|
||||
[
|
||||
|
||||
Reference in New Issue
Block a user