mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-14 09:43:22 -03:00
fix(update): resolve template path when updating to a different base model (#1059)
Version-tab updates reused the current version's folder, so updating a LoRA to a version with a different base model (e.g. Illustrious -> Anima) ignored the download path template and landed in the old version's directory. When the target version's base model differs from the current local version and a path template is configured, re-resolve the template under the same model root. The backend keeps an explicitly provided root when use_save_dir_as_root is set, so regular downloads still use the default root.
This commit is contained in:
Will Miao
@@ -0,0 +1,210 @@
|
||||
import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
|
||||
|
||||
const {
|
||||
MODEL_VERSIONS_MODULE,
|
||||
API_FACTORY_MODULE,
|
||||
DOWNLOAD_MANAGER_MODULE,
|
||||
UI_HELPERS_MODULE,
|
||||
STATE_MODULE,
|
||||
I18N_HELPERS_MODULE,
|
||||
UTILS_MODULE,
|
||||
} = vi.hoisted(() => ({
|
||||
MODEL_VERSIONS_MODULE: new URL('../../../static/js/components/shared/ModelVersionsTab.js', import.meta.url).pathname,
|
||||
API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
|
||||
DOWNLOAD_MANAGER_MODULE: new URL('../../../static/js/managers/DownloadManager.js', import.meta.url).pathname,
|
||||
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
|
||||
STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
|
||||
I18N_HELPERS_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
|
||||
UTILS_MODULE: new URL('../../../static/js/components/shared/utils.js', import.meta.url).pathname,
|
||||
}));
|
||||
|
||||
const downloadVersionWithDefaults = vi.fn();
|
||||
|
||||
vi.mock(DOWNLOAD_MANAGER_MODULE, () => ({
|
||||
downloadManager: {
|
||||
downloadVersionWithDefaults,
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock(UI_HELPERS_MODULE, () => ({
|
||||
showToast: vi.fn(),
|
||||
openCivitaiUrl: vi.fn(),
|
||||
}));
|
||||
|
||||
const stateMock = {
|
||||
global: {
|
||||
settings: {
|
||||
autoplay_on_hover: false,
|
||||
version_grouping: 'any',
|
||||
download_path_templates: {
|
||||
lora: '{base_model}/{first_tag}',
|
||||
checkpoint: '{base_model}/{first_tag}',
|
||||
embedding: '{base_model}/{first_tag}',
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
vi.mock(STATE_MODULE, () => ({
|
||||
state: stateMock,
|
||||
}));
|
||||
|
||||
vi.mock(I18N_HELPERS_MODULE, () => ({
|
||||
translate: vi.fn((_, __, fallback) => fallback ?? ''),
|
||||
}));
|
||||
|
||||
vi.mock(UTILS_MODULE, () => ({
|
||||
formatFileSize: vi.fn(() => '1 MB'),
|
||||
}));
|
||||
|
||||
vi.mock(API_FACTORY_MODULE, () => ({
|
||||
getModelApiClient: vi.fn(),
|
||||
}));
|
||||
|
||||
const LORA_ROOT = '/models/loras';
|
||||
|
||||
function buildRecord(targetBaseModel = 'Anima') {
|
||||
return {
|
||||
success: true,
|
||||
record: {
|
||||
shouldIgnore: false,
|
||||
inLibraryVersionIds: [10],
|
||||
versions: [
|
||||
{
|
||||
versionId: 10,
|
||||
name: 'v1.0',
|
||||
baseModel: 'Illustrious',
|
||||
sizeBytes: 1024,
|
||||
isInLibrary: true,
|
||||
shouldIgnore: false,
|
||||
filePath: `${LORA_ROOT}/Illustrious/works/file.safetensors`,
|
||||
},
|
||||
{
|
||||
versionId: 11,
|
||||
name: 'v1.1',
|
||||
baseModel: targetBaseModel,
|
||||
sizeBytes: 2048,
|
||||
isInLibrary: false,
|
||||
shouldIgnore: false,
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function renderAndClickDownload({ currentVersionId = 10, record = null } = {}) {
|
||||
const { initVersionsTab } = await import(MODEL_VERSIONS_MODULE);
|
||||
const controller = initVersionsTab({
|
||||
modalId: 'model-versions-modal',
|
||||
modelType: 'loras',
|
||||
modelId: 123,
|
||||
currentVersionId,
|
||||
});
|
||||
await controller.load();
|
||||
const downloadButton = document.querySelector(
|
||||
'.model-version-row[data-version-id="11"] [data-version-action="download"]'
|
||||
);
|
||||
downloadButton?.click();
|
||||
await new Promise(resolve => setTimeout(resolve, 0));
|
||||
return controller;
|
||||
}
|
||||
|
||||
describe('ModelVersionsTab update download path resolution', () => {
|
||||
let getModelApiClient;
|
||||
let fetchModelUpdateVersions;
|
||||
let fetchModelRoots;
|
||||
|
||||
beforeEach(async () => {
|
||||
vi.resetModules();
|
||||
downloadVersionWithDefaults.mockReset();
|
||||
downloadVersionWithDefaults.mockResolvedValue(true);
|
||||
document.body.innerHTML = `
|
||||
<div id="model-versions-modal">
|
||||
<div id="versions-tab">
|
||||
<div class="model-versions-tab"></div>
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
stateMock.global.settings.version_grouping = 'any';
|
||||
stateMock.global.settings.download_path_templates.lora = '{base_model}/{first_tag}';
|
||||
({ getModelApiClient } = await import(API_FACTORY_MODULE));
|
||||
fetchModelUpdateVersions = vi.fn();
|
||||
fetchModelRoots = vi.fn();
|
||||
fetchModelRoots.mockResolvedValue({ roots: [LORA_ROOT] });
|
||||
getModelApiClient.mockReturnValue({
|
||||
fetchModelUpdateVersions,
|
||||
fetchModelRoots,
|
||||
setModelUpdateIgnore: vi.fn(),
|
||||
setVersionUpdateIgnore: vi.fn(),
|
||||
deleteModel: vi.fn(),
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
document.body.innerHTML = '';
|
||||
});
|
||||
|
||||
it('keeps the current folder when the target version has the same base model', async () => {
|
||||
fetchModelUpdateVersions.mockResolvedValue(buildRecord('Illustrious'));
|
||||
|
||||
await renderAndClickDownload();
|
||||
|
||||
expect(downloadVersionWithDefaults).toHaveBeenCalledWith(
|
||||
'loras', 123, 11,
|
||||
expect.objectContaining({
|
||||
modelRoot: LORA_ROOT,
|
||||
targetFolder: 'Illustrious/works',
|
||||
useDefaultPaths: null,
|
||||
useSaveDirAsRoot: false,
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it('resolves the template path when the target base model differs and a template is configured', async () => {
|
||||
fetchModelUpdateVersions.mockResolvedValue(buildRecord());
|
||||
|
||||
await renderAndClickDownload();
|
||||
|
||||
expect(downloadVersionWithDefaults).toHaveBeenCalledWith(
|
||||
'loras', 123, 11,
|
||||
expect.objectContaining({
|
||||
modelRoot: LORA_ROOT,
|
||||
targetFolder: '',
|
||||
useDefaultPaths: true,
|
||||
useSaveDirAsRoot: true,
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it('keeps the current folder when the target base model differs but no template is configured', async () => {
|
||||
stateMock.global.settings.download_path_templates.lora = '';
|
||||
fetchModelUpdateVersions.mockResolvedValue(buildRecord());
|
||||
|
||||
await renderAndClickDownload();
|
||||
|
||||
expect(downloadVersionWithDefaults).toHaveBeenCalledWith(
|
||||
'loras', 123, 11,
|
||||
expect.objectContaining({
|
||||
modelRoot: LORA_ROOT,
|
||||
targetFolder: 'Illustrious/works',
|
||||
useDefaultPaths: null,
|
||||
useSaveDirAsRoot: false,
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it('falls back to default paths when no local version exists', async () => {
|
||||
fetchModelUpdateVersions.mockResolvedValue(buildRecord());
|
||||
|
||||
await renderAndClickDownload({ currentVersionId: null });
|
||||
|
||||
expect(downloadVersionWithDefaults).toHaveBeenCalledWith(
|
||||
'loras', 123, 11,
|
||||
expect.objectContaining({
|
||||
modelRoot: '',
|
||||
targetFolder: '',
|
||||
useDefaultPaths: null,
|
||||
useSaveDirAsRoot: false,
|
||||
})
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -233,6 +233,58 @@ async def test_successful_download_uses_defaults(
|
||||
assert captured["download_urls"] == ["https://example.invalid/file.safetensors"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_keeps_save_dir_when_use_save_dir_as_root(
|
||||
monkeypatch, scanners, metadata_provider, tmp_path
|
||||
):
|
||||
"""use_default_paths with use_save_dir_as_root resolves the template under
|
||||
the provided save_dir instead of switching to the default root."""
|
||||
manager = DownloadManager()
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_execute_download(
|
||||
self,
|
||||
*,
|
||||
download_urls,
|
||||
save_dir,
|
||||
metadata,
|
||||
version_info,
|
||||
relative_path,
|
||||
progress_callback,
|
||||
model_type,
|
||||
download_id,
|
||||
transfer_backend=None,
|
||||
):
|
||||
captured.update(
|
||||
{
|
||||
"save_dir": Path(save_dir),
|
||||
"relative_path": relative_path,
|
||||
"model_type": model_type,
|
||||
}
|
||||
)
|
||||
return {"success": True}
|
||||
|
||||
monkeypatch.setattr(
|
||||
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
||||
)
|
||||
|
||||
custom_root = tmp_path / "custom_root"
|
||||
result = await manager.download_from_civitai(
|
||||
model_version_id=99,
|
||||
save_dir=str(custom_root),
|
||||
use_default_paths=True,
|
||||
use_save_dir_as_root=True,
|
||||
progress_callback=None,
|
||||
source=None,
|
||||
)
|
||||
|
||||
assert result["success"] is True
|
||||
assert captured["relative_path"] == "MappedModel/fantasy"
|
||||
assert captured["save_dir"] == custom_root / "MappedModel" / "fantasy"
|
||||
assert captured["model_type"] == "lora"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_download_schedules_auto_example_images(
|
||||
monkeypatch, scanners, metadata_provider, tmp_path
|
||||
@@ -618,6 +670,7 @@ async def test_resume_download_restores_persisted_aria2_task(monkeypatch, tmp_pa
|
||||
use_default_paths=False,
|
||||
source=None,
|
||||
file_params=None,
|
||||
use_save_dir_as_root=False,
|
||||
):
|
||||
created.update(
|
||||
{
|
||||
@@ -1037,6 +1090,7 @@ async def test_download_uses_captured_backend_when_settings_change(
|
||||
transfer_backend="python",
|
||||
source=None,
|
||||
file_params=None,
|
||||
use_save_dir_as_root=False,
|
||||
):
|
||||
captured["transfer_backend"] = transfer_backend
|
||||
return {"success": True}
|
||||
|
||||
Reference in New Issue
Block a user