feat(download): add configurable base model download exclusions

2026-03-26 23:48:52 -03:00 · 2026-03-26 23:06:14 +08:00
parent 5b065b47d4
commit a5191414cc
22 changed files with 988 additions and 4 deletions
--- a/tests/frontend/managers/settingsManager.downloadSkipBaseModels.test.js
+++ b/tests/frontend/managers/settingsManager.downloadSkipBaseModels.test.js
@@ -0,0 +1,152 @@
+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_skip_base_models: [],
+        }),
+    };
+});
+
+vi.mock('../../../static/js/api/modelApiFactory.js', () => ({
+    resetAndReload: vi.fn(),
+}));
+
+vi.mock('../../../static/js/utils/constants.js', () => ({
+    DOWNLOAD_PATH_TEMPLATES: {},
+    DEFAULT_PATH_TEMPLATES: {},
+    MAPPABLE_BASE_MODELS: ['Flux.1 D', 'Pony', 'SDXL 1.0', 'Other'],
+    PATH_TEMPLATE_PLACEHOLDERS: {},
+    DEFAULT_PRIORITY_TAG_CONFIG: {
+        lora: 'character, style',
+        checkpoint: 'base, guide',
+        embedding: 'hint',
+    },
+}));
+
+vi.mock('../../../static/js/utils/i18nHelpers.js', () => ({
+    translate: (key, params, fallback) => {
+        if (key === 'settings.downloadSkipBaseModels.summary.none') {
+            return 'None selected';
+        }
+        if (key === 'settings.downloadSkipBaseModels.summary.count') {
+            return `${params?.count ?? 0} selected`;
+        }
+        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(),
+    },
+}));
+
+vi.mock('../../../static/js/components/SidebarManager.js', () => ({
+    sidebarManager: {
+        setSidebarEnabled: vi.fn().mockResolvedValue(),
+    },
+}));
+
+import { SettingsManager } from '../../../static/js/managers/SettingsManager.js';
+import { state } from '../../../static/js/state/index.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 appendDownloadSkipUi = () => {
+    document.body.innerHTML = `
+        <button id="downloadSkipBaseModelsToggle" aria-expanded="false">
+            <span id="downloadSkipBaseModelsSummary"></span>
+            <span class="base-model-skip-toggle-label"></span>
+        </button>
+        <div id="downloadSkipBaseModelsPanel" hidden>
+            <input id="downloadSkipBaseModelsSearch" />
+            <button id="downloadSkipBaseModelsClear" type="button">Clear</button>
+            <div id="downloadSkipBaseModelsContainer"></div>
+            <div id="downloadSkipBaseModelsEmpty" hidden></div>
+        </div>
+        <div id="downloadSkipBaseModelsError"></div>
+    `;
+};
+
+describe('SettingsManager download skip base models UI', () => {
+    beforeEach(() => {
+        document.body.innerHTML = '';
+        vi.clearAllMocks();
+        state.global.settings = {
+            download_skip_base_models: [],
+        };
+    });
+
+    it('renders a compact summary for selected base models', () => {
+        appendDownloadSkipUi();
+        state.global.settings.download_skip_base_models = ['Flux.1 D', 'Pony'];
+        const manager = createManager();
+
+        manager.renderDownloadSkipBaseModels();
+
+        expect(document.getElementById('downloadSkipBaseModelsSummary').textContent).toBe('Flux.1 D, Pony');
+        expect(document.querySelectorAll('#downloadSkipBaseModelsContainer input')).toHaveLength(3);
+    });
+
+    it('filters the list using the search input and shows an empty state', () => {
+        appendDownloadSkipUi();
+        state.global.settings.download_skip_base_models = ['Flux.1 D'];
+        const manager = createManager();
+        const searchInput = document.getElementById('downloadSkipBaseModelsSearch');
+
+        searchInput.value = 'pony';
+        manager.renderDownloadSkipBaseModels();
+
+        expect(document.querySelectorAll('#downloadSkipBaseModelsContainer input')).toHaveLength(1);
+        expect(document.querySelector('#downloadSkipBaseModelsContainer input').value).toBe('Pony');
+
+        searchInput.value = 'zzz';
+        manager.renderDownloadSkipBaseModels();
+
+        expect(document.querySelectorAll('#downloadSkipBaseModelsContainer input')).toHaveLength(0);
+        expect(document.getElementById('downloadSkipBaseModelsEmpty').hidden).toBe(false);
+    });
+});
--- a/tests/routes/test_base_model_routes_smoke.py
+++ b/tests/routes/test_base_model_routes_smoke.py
@@ -719,3 +719,42 @@ def test_auto_organize_conflict_when_running(mock_service):
            await client.close()

    asyncio.run(scenario())
+
+
+
+def test_download_model_returns_skipped_success(mock_service, download_manager_stub):
+    async def scenario():
+        download_manager_stub.last_progress_snapshot = None
+
+        async def fake_download(**kwargs):
+            download_manager_stub.calls.append(kwargs)
+            return {
+                "success": True,
+                "skipped": True,
+                "status": "skipped",
+                "reason": "base_model_excluded",
+                "message": "Skipped by settings",
+                "base_model": "SDXL 1.0",
+                "file_name": "demo.safetensors",
+            }
+
+        download_manager_stub.download_from_civitai = fake_download
+
+        client = await create_test_client(mock_service)
+        try:
+            response = await client.post(
+                "/api/lm/download-model",
+                json={"model_version_id": 123},
+            )
+            payload = await response.json()
+
+            assert response.status == 200
+            assert payload["success"] is True
+            assert payload["skipped"] is True
+            assert payload["reason"] == "base_model_excluded"
+            assert payload["base_model"] == "SDXL 1.0"
+            assert payload["file_name"] == "demo.safetensors"
+        finally:
+            await client.close()
+
+    asyncio.run(scenario())
--- a/tests/services/test_download_manager_basic.py
+++ b/tests/services/test_download_manager_basic.py
@@ -38,6 +38,7 @@ def isolate_settings(monkeypatch, tmp_path):
                "embedding": "{base_model}/{first_tag}",
            },
            "base_model_path_mappings": {"BaseModel": "MappedModel"},
+            "download_skip_base_models": [],
        }
    )
    monkeypatch.setattr(manager, "settings", default_settings)
@@ -443,3 +444,49 @@ def test_distribute_preview_to_entries_keeps_existing_file(tmp_path):

    assert targets[0] == str(existing_preview)
    assert Path(targets[1]).read_bytes() == b"preview"
+
+
+
+@pytest.mark.asyncio
+async def test_download_skips_excluded_base_model(monkeypatch, scanners, metadata_provider):
+    manager = DownloadManager()
+    get_settings_manager().settings["download_skip_base_models"] = ["SDXL 1.0"]
+
+    metadata_provider.get_model_version = AsyncMock(
+        return_value={
+            "id": 42,
+            "model": {"type": "LoRA", "tags": ["fantasy"]},
+            "baseModel": "SDXL 1.0",
+            "creator": {"username": "Author"},
+            "files": [
+                {
+                    "type": "Model",
+                    "primary": True,
+                    "downloadUrl": "https://example.invalid/file.safetensors",
+                    "name": "file.safetensors",
+                }
+            ],
+        }
+    )
+
+    execute_download = AsyncMock()
+    monkeypatch.setattr(
+        DownloadManager, "_execute_download", execute_download, raising=False
+    )
+
+    result = await manager.download_from_civitai(
+        model_version_id=99,
+        use_default_paths=True,
+        progress_callback=None,
+        source=None,
+    )
+
+    assert result["success"] is True
+    assert result["skipped"] is True
+    assert result["status"] == "skipped"
+    assert result["reason"] == "base_model_excluded"
+    assert result["base_model"] == "SDXL 1.0"
+    assert result["file_name"] == "file.safetensors"
+    assert "file.safetensors" in result["message"]
+    execute_download.assert_not_called()
+    assert manager._active_downloads[result["download_id"]]["status"] == "skipped"
--- a/tests/services/test_settings_manager.py
+++ b/tests/services/test_settings_manager.py
@@ -605,3 +605,28 @@ def test_delete_library_switches_active(manager, tmp_path):
    manager.delete_library("other")

    assert manager.get_active_library_name() == "default"
+
+
+
+def test_download_skip_base_models_are_normalized(manager):
+    manager.settings["download_skip_base_models"] = [
+        "SDXL 1.0",
+        "Invalid",
+        "SDXL 1.0",
+        "Pony",
+        "Other",
+    ]
+
+    result = manager.get_download_skip_base_models()
+
+    assert result == ["SDXL 1.0", "Pony"]
+    assert manager.settings["download_skip_base_models"] == ["SDXL 1.0", "Pony"]
+
+
+def test_setting_download_skip_base_models_normalizes_string_input(manager):
+    manager.set(
+        "download_skip_base_models",
+        "SDXL 1.0, Pony; Invalid\nSDXL 1.0"
+    )
+
+    assert manager.get("download_skip_base_models") == ["SDXL 1.0", "Pony"]