From da071e84526bb1df4ee70ebce8af3e3e0351d099 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Mon, 24 Aug 2026 23:17:10 +0800 Subject: [PATCH] feat(versions): add file-variant badge and hide download button for in-library versions (#1058) --- locales/de.json | 2 +- locales/en.json | 2 +- locales/es.json | 2 +- locales/fr.json | 2 +- locales/he.json | 2 +- locales/ja.json | 2 +- locales/ko.json | 2 +- locales/ru.json | 2 +- locales/zh-CN.json | 2 +- locales/zh-TW.json | 2 +- py/routes/handlers/model_handlers.py | 3 + py/services/model_update_service.py | 48 ++- .../css/components/modal/download-modal.css | 1 + .../js/components/shared/ModelVersionsTab.js | 124 ++++---- .../modelVersionsTab.downloadButton.test.js | 285 ++++++++++++++++++ tests/routes/test_model_update_handler.py | 18 ++ tests/services/test_model_update_service.py | 123 ++++++++ 17 files changed, 557 insertions(+), 65 deletions(-) create mode 100644 tests/frontend/components/modelVersionsTab.downloadButton.test.js diff --git a/locales/de.json b/locales/de.json index a4963a3b..0b38aa3f 100644 --- a/locales/de.json +++ b/locales/de.json @@ -1616,7 +1616,7 @@ "actions": { "download": "Herunterladen", "downloadTooltip": "Diese Version herunterladen", - "downloadRemainingTooltip": "Verbleibende Dateien dieser Version herunterladen", + "downloadChooseFilesTooltip": "Auswählen, welche Dateien heruntergeladen werden sollen", "downloadEarlyAccessTooltip": "Diese Early-Access-Version von Civitai herunterladen", "downloadPaidTooltip": "Diese bezahlte Version von Civitai herunterladen", "downloadNotAllowedTooltip": "Diese Version ist nur für die On-Site-Generierung auf Civitai verfügbar", diff --git a/locales/en.json b/locales/en.json index 2e2eff37..23a46aa8 100644 --- a/locales/en.json +++ b/locales/en.json @@ -1616,7 +1616,7 @@ "actions": { "download": "Download", "downloadTooltip": "Download this version", - "downloadRemainingTooltip": "Download remaining files of this version", + "downloadChooseFilesTooltip": "Choose which files to download", "downloadEarlyAccessTooltip": "Download this early access version from Civitai", "downloadPaidTooltip": "Download this paid version from Civitai", "downloadNotAllowedTooltip": "This version is only available for on-site generation on Civitai", diff --git a/locales/es.json b/locales/es.json index 8c7ca0f7..300a41dd 100644 --- a/locales/es.json +++ b/locales/es.json @@ -1616,7 +1616,7 @@ "actions": { "download": "Descargar", "downloadTooltip": "Descargar esta versión", - "downloadRemainingTooltip": "Descargar los archivos restantes de esta versión", + "downloadChooseFilesTooltip": "Elegir qué archivos descargar", "downloadEarlyAccessTooltip": "Descargar esta versión de acceso temprano desde Civitai", "downloadPaidTooltip": "Descargar esta versión de pago desde Civitai", "downloadNotAllowedTooltip": "Esta versión solo está disponible para generación en el sitio de Civitai", diff --git a/locales/fr.json b/locales/fr.json index a88a7570..3244bc68 100644 --- a/locales/fr.json +++ b/locales/fr.json @@ -1616,7 +1616,7 @@ "actions": { "download": "Télécharger", "downloadTooltip": "Télécharger cette version", - "downloadRemainingTooltip": "Télécharger les fichiers restants de cette version", + "downloadChooseFilesTooltip": "Choisir les fichiers à télécharger", "downloadEarlyAccessTooltip": "Télécharger cette version en accès anticipé depuis Civitai", "downloadPaidTooltip": "Télécharger cette version payante depuis Civitai", "downloadNotAllowedTooltip": "Cette version n'est disponible que pour la génération sur le site Civitai", diff --git a/locales/he.json b/locales/he.json index df596458..fe43ea7e 100644 --- a/locales/he.json +++ b/locales/he.json @@ -1616,7 +1616,7 @@ "actions": { "download": "הורדה", "downloadTooltip": "הורד את הגרסה הזו", - "downloadRemainingTooltip": "הורד את הקבצים הנותרים של גרסה זו", + "downloadChooseFilesTooltip": "בחר אילו קבצים להוריד", "downloadEarlyAccessTooltip": "הורד את גרסת ה-Early Access הזו מ-Civitai", "downloadPaidTooltip": "הורד את הגרסה בתשלום הזו מ-Civitai", "downloadNotAllowedTooltip": "גרסה זו זמינה רק ליצירה באתר Civitai", diff --git a/locales/ja.json b/locales/ja.json index 7f5bf71c..a774bfa8 100644 --- a/locales/ja.json +++ b/locales/ja.json @@ -1616,7 +1616,7 @@ "actions": { "download": "ダウンロード", "downloadTooltip": "このバージョンをダウンロード", - "downloadRemainingTooltip": "このバージョンの残りのファイルをダウンロード", + "downloadChooseFilesTooltip": "ダウンロードするファイルを選択", "downloadEarlyAccessTooltip": "Civitai からこの早期アクセス版をダウンロード", "downloadPaidTooltip": "Civitai からこの有料バージョンをダウンロード", "downloadNotAllowedTooltip": "このバージョンはCivitaiサイト内でのみ利用可能で、ダウンロードはできません", diff --git a/locales/ko.json b/locales/ko.json index 343c52ee..0bda407a 100644 --- a/locales/ko.json +++ b/locales/ko.json @@ -1616,7 +1616,7 @@ "actions": { "download": "다운로드", "downloadTooltip": "이 버전 다운로드", - "downloadRemainingTooltip": "이 버전의 나머지 파일 다운로드", + "downloadChooseFilesTooltip": "다운로드할 파일 선택", "downloadEarlyAccessTooltip": "Civitai에서 이 얼리 액세스 버전 다운로드", "downloadPaidTooltip": "Civitai에서 이 유료 버전 다운로드", "downloadNotAllowedTooltip": "이 버전은 Civitai 사이트 내에서만 사용 가능하며 다운로드할 수 없습니다", diff --git a/locales/ru.json b/locales/ru.json index 0964f34d..03df82f5 100644 --- a/locales/ru.json +++ b/locales/ru.json @@ -1616,7 +1616,7 @@ "actions": { "download": "Скачать", "downloadTooltip": "Скачать эту версию", - "downloadRemainingTooltip": "Скачать оставшиеся файлы этой версии", + "downloadChooseFilesTooltip": "Выбрать файлы для скачивания", "downloadEarlyAccessTooltip": "Скачать эту версию раннего доступа с Civitai", "downloadPaidTooltip": "Скачать эту платную версию с Civitai", "downloadNotAllowedTooltip": "Эта версия доступна только для генерации на сайте Civitai", diff --git a/locales/zh-CN.json b/locales/zh-CN.json index 7be918bf..66ff68b2 100644 --- a/locales/zh-CN.json +++ b/locales/zh-CN.json @@ -1616,7 +1616,7 @@ "actions": { "download": "下载", "downloadTooltip": "下载此版本", - "downloadRemainingTooltip": "下载此版本的剩余文件", + "downloadChooseFilesTooltip": "选择要下载的文件", "downloadEarlyAccessTooltip": "从 Civitai 下载此抢先体验版本", "downloadPaidTooltip": "从 Civitai 下载此付费版本", "downloadNotAllowedTooltip": "此版本仅在 Civitai 站内可用，无法下载", diff --git a/locales/zh-TW.json b/locales/zh-TW.json index e9e4a779..d5d96348 100644 --- a/locales/zh-TW.json +++ b/locales/zh-TW.json @@ -1616,7 +1616,7 @@ "actions": { "download": "下載", "downloadTooltip": "下載此版本", - "downloadRemainingTooltip": "下載此版本的剩餘檔案", + "downloadChooseFilesTooltip": "選擇要下載的檔案", "downloadEarlyAccessTooltip": "從 Civitai 下載此搶先體驗版本", "downloadPaidTooltip": "從 Civitai 下載此付費版本", "downloadNotAllowedTooltip": "此版本僅在 Civitai 站內可用，無法下載", diff --git a/py/routes/handlers/model_handlers.py b/py/routes/handlers/model_handlers.py index 22c00ee6..4f1b30a7 100644 --- a/py/routes/handlers/model_handlers.py +++ b/py/routes/handlers/model_handlers.py @@ -3120,6 +3120,9 @@ class ModelUpdateHandler: "paidAccess": paid_access_payload, "filePath": context.get("file_path"), "fileName": context.get("file_name"), + # Weight-file variant count (None when unknown); lets the UI hide + # the download affordance for single-file in-library versions. + "fileCount": getattr(version, "file_count", None), } async def _build_version_context( diff --git a/py/services/model_update_service.py b/py/services/model_update_service.py index 08686784..3c3b12bc 100644 --- a/py/services/model_update_service.py +++ b/py/services/model_update_service.py @@ -18,6 +18,7 @@ from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence from .errors import RateLimitError, ResourceNotFoundError from .settings_manager import get_settings_manager from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration +from ..utils.constants import MODEL_WEIGHT_FILE_TYPES from ..utils.civitai_utils import rewrite_preview_url from ..utils.preview_selection import resolve_mature_threshold, select_preview_media @@ -77,6 +78,10 @@ class ModelVersionRecord: usage_control: Optional[str] = None # "Download", "Generation", "InternalGeneration" paid_access: Optional[str] = None # JSON string of the CivitAI paidAccess DTO is_paid: bool = False # True when paidAccess.permanent is True (permanent paid gate) + # Number of downloadable weight files for the version (None when unknown, + # e.g. records persisted before this field existed or locally-synthesized + # entries). Mirrors the frontend isModelWeightFile() filter. + file_count: Optional[int] = None @dataclass @@ -273,6 +278,7 @@ class ModelUpdateService: usage_control TEXT, paid_access TEXT, is_paid INTEGER NOT NULL DEFAULT 0, + file_count INTEGER, PRIMARY KEY (model_id, version_id), FOREIGN KEY(model_id) REFERENCES model_update_status(model_id) ON DELETE CASCADE ); @@ -520,6 +526,10 @@ class ModelUpdateService: "ALTER TABLE model_update_versions " "ADD COLUMN is_paid INTEGER NOT NULL DEFAULT 0" ), + "file_count": ( + "ALTER TABLE model_update_versions " + "ADD COLUMN file_count INTEGER" + ), } for column, statement in migrations.items(): @@ -623,6 +633,7 @@ class ModelUpdateService: is_early_access INTEGER NOT NULL DEFAULT 0, paid_access TEXT, is_paid INTEGER NOT NULL DEFAULT 0, + file_count INTEGER, PRIMARY KEY (model_id, version_id), FOREIGN KEY(model_id) REFERENCES model_update_status(model_id) ON DELETE CASCADE ) @@ -644,6 +655,7 @@ class ModelUpdateService: "is_early_access", "paid_access", "is_paid", + "file_count", ] defaults = { "sort_index": "0", @@ -658,6 +670,7 @@ class ModelUpdateService: "is_early_access": "0", "paid_access": "NULL", "is_paid": "0", + "file_count": "NULL", } select_parts = [] @@ -1504,6 +1517,7 @@ class ModelUpdateService: ) ignore_map = {version.version_id: version.should_ignore for version in existing.versions} if existing else {} preview_map = {version.version_id: version.preview_url for version in existing.versions} if existing else {} + file_count_map = {version.version_id: version.file_count for version in existing.versions} if existing else {} sort_map = {version.version_id: version.sort_index for version in existing.versions} if existing else {} existing_map = {version.version_id: version for version in existing.versions} if existing else {} @@ -1528,6 +1542,11 @@ class ModelUpdateService: usage_control=remote_version.usage_control, paid_access=remote_version.paid_access, is_paid=remote_version.is_paid, + file_count=( + remote_version.file_count + if remote_version.file_count is not None + else file_count_map.get(version_id) + ), ) ) @@ -1620,6 +1639,7 @@ class ModelUpdateService: base_model = _normalize_string(entry.get("baseModel")) released_at = _normalize_string(entry.get("publishedAt") or entry.get("createdAt")) size_bytes = self._extract_size_bytes(entry.get("files")) + file_count = self._extract_file_count(entry.get("files")) preview_url = self._extract_preview_url(entry.get("images")) early_access_ends_at = _normalize_string(entry.get("earlyAccessEndsAt")) @@ -1655,6 +1675,7 @@ class ModelUpdateService: usage_control=usage_control, paid_access=paid_access_json, is_paid=is_paid, + file_count=file_count, ) @staticmethod @@ -1683,6 +1704,25 @@ class ModelUpdateService: return None return {"permanent": permanent, "endsAt": ends_at} + @staticmethod + def _extract_file_count(files) -> Optional[int]: + """Count downloadable weight files in a version entry's ``files`` list. + + Returns None when the payload carries no files array (unknown), so + callers can distinguish "no weight files" from "no data". + """ + + if not isinstance(files, list): + return None + count = 0 + for entry in files: + if not isinstance(entry, Mapping): + continue + entry_type = entry.get("type") + if isinstance(entry_type, str) and entry_type in MODEL_WEIGHT_FILE_TYPES: + count += 1 + return count + def _extract_size_bytes(self, files) -> Optional[int]: if not isinstance(files, Iterable): return None @@ -1795,7 +1835,7 @@ class ModelUpdateService: f""" SELECT model_id, version_id, sort_index, name, base_model, released_at, size_bytes, preview_url, is_in_library, should_ignore, early_access_ends_at, - is_early_access, usage_control, paid_access, is_paid + is_early_access, usage_control, paid_access, is_paid, file_count FROM model_update_versions WHERE model_id IN ({placeholders}) ORDER BY model_id ASC, sort_index ASC, version_id ASC @@ -1826,6 +1866,7 @@ class ModelUpdateService: usage_control=row["usage_control"], paid_access=row["paid_access"], is_paid=bool(row["is_paid"]), + file_count=_normalize_int(row["file_count"]), ) ) @@ -1888,8 +1929,8 @@ class ModelUpdateService: INSERT INTO model_update_versions ( version_id, model_id, sort_index, name, base_model, released_at, size_bytes, preview_url, is_in_library, should_ignore, early_access_ends_at, - is_early_access, usage_control, paid_access, is_paid - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + is_early_access, usage_control, paid_access, is_paid, file_count + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( version.version_id, @@ -1907,6 +1948,7 @@ class ModelUpdateService: version.usage_control, paid_access_value, 1 if version.is_paid else 0, + version.file_count, ), ) conn.commit() diff --git a/static/css/components/modal/download-modal.css b/static/css/components/modal/download-modal.css index a3b69700..49f606c0 100644 --- a/static/css/components/modal/download-modal.css +++ b/static/css/components/modal/download-modal.css @@ -514,6 +514,7 @@ background: oklch(var(--lora-accent) / 0.18); color: var(--lora-accent); font-size: inherit; + font-family: inherit; font-weight: 600; cursor: pointer; transition: var(--transition-base); diff --git a/static/js/components/shared/ModelVersionsTab.js b/static/js/components/shared/ModelVersionsTab.js index 7c188a4e..f0d6c3bd 100644 --- a/static/js/components/shared/ModelVersionsTab.js +++ b/static/js/components/shared/ModelVersionsTab.js @@ -575,50 +575,60 @@ function renderRow(version, options) { const actions = []; const canDownload = isDownloadAllowed(version); const downloadIcon = isEarlyAccess ? ' ' : ''; - let downloadTitle; - if (!canDownload) { - downloadTitle = translate( - 'modals.model.versions.actions.downloadNotAllowedTooltip', - {}, - 'This version is only available for on-site generation on Civitai' - ); - } else if (version.isInLibrary) { - // In-library versions may still have undownloaded weight files; the - // download modal's file dialog decides what remains (#1058). - downloadTitle = translate( - 'modals.model.versions.actions.downloadRemainingTooltip', - {}, - 'Download remaining files of this version' - ); - } else if (isPaidPermanent(version)) { - downloadTitle = translate( - 'modals.model.versions.actions.downloadPaidTooltip', - {}, - 'Download this paid version from Civitai' - ); - } else if (isEarlyAccess) { - downloadTitle = translate( - 'modals.model.versions.actions.downloadEarlyAccessTooltip', - {}, - 'Download this early access version from Civitai' - ); - } else { - downloadTitle = translate( - 'modals.model.versions.actions.downloadTooltip', - {}, - 'Download this version' - ); - } - actions.push(buildActionButton( - downloadLabel, - canDownload ? 'version-action-primary' : 'version-action-disabled', - canDownload ? 'download' : '', - { - title: downloadTitle, - iconMarkup: downloadIcon, - disabled: !canDownload, + // The Download button always fetches the default (primary) file, keeping + // the single-file experience for users who don't care about variants. + // In-library versions hide it: their default file already exists locally, + // and multi-file versions use the "N files" badge below for the remaining + // variants instead (#1058). fileCount is null for records persisted before + // the field existed; default to the single-file behavior in that case. + const fileCount = typeof version.fileCount === 'number' ? version.fileCount : null; + const showDownload = !version.isInLibrary; + if (showDownload) { + let downloadTitle; + if (!canDownload) { + downloadTitle = translate( + 'modals.model.versions.actions.downloadNotAllowedTooltip', + {}, + 'This version is only available for on-site generation on Civitai' + ); + } else if (isPaidPermanent(version)) { + downloadTitle = translate( + 'modals.model.versions.actions.downloadPaidTooltip', + {}, + 'Download this paid version from Civitai' + ); + } else if (isEarlyAccess) { + downloadTitle = translate( + 'modals.model.versions.actions.downloadEarlyAccessTooltip', + {}, + 'Download this early access version from Civitai' + ); + } else { + downloadTitle = translate( + 'modals.model.versions.actions.downloadTooltip', + {}, + 'Download this version' + ); } - )); + actions.push(buildActionButton( + downloadLabel, + canDownload ? 'version-action-primary' : 'version-action-disabled', + canDownload ? 'download' : '', + { + title: downloadTitle, + iconMarkup: downloadIcon, + disabled: !canDownload, + } + )); + } + + // Multi-file versions get an explicit entry into the download modal's + // file-selection step, mirroring the version step's file badge (#1058). + const fileSelectionBadge = fileCount !== null && fileCount > 1 + ? `` + : ''; if (version.isInLibrary && version.filePath) { actions.push(buildActionButton( deleteLabel, @@ -696,6 +706,7 @@ function renderRow(version, options) {
${badges.join('')}
${buildMetaMarkup(version, { showEarlyAccess: true })} + ${fileSelectionBadge}
@@ -1429,15 +1440,9 @@ export function initVersionsTab({ button.disabled = true; try { - // In-library versions may still have undownloaded weight files - // (#1058). The tab payload has no per-file state, so open the - // download modal's file dialog, which refetches the full version - // payload and shows what remains. - if (version.isInLibrary) { - await downloadManager.openFileSelectionForVersion(modelType, modelId, versionId); - return; - } - + // The Download button only renders for versions not in the library + // and always fetches the default (primary) file. Multi-file + // variants are reached through the "N files" badge instead. const pathInfo = await resolveDownloadPathFromCurrentVersion(); const resolveTemplatePath = shouldResolveTemplatePath(version, pathInfo); const success = await downloadManager.downloadVersionWithDefaults(modelType, modelId, versionId, { @@ -1516,6 +1521,21 @@ export function initVersionsTab({ return; } + // File-selection badge: enter the download modal's file step directly. + // Must run before the row-click navigation below (rows are clickable). + const filesBadge = event.target.closest('[data-version-files]'); + if (filesBadge) { + event.preventDefault(); + event.stopPropagation(); + const row = filesBadge.closest('.model-version-row'); + if (!row) { + return; + } + const versionId = Number(row.dataset.versionId); + await downloadManager.openFileSelectionForVersion(modelType, modelId, versionId); + return; + } + const row = event.target.closest('.model-version-row.is-clickable'); const civitaiLink = event.target.closest('.version-civitai-link'); if (civitaiLink) { diff --git a/tests/frontend/components/modelVersionsTab.downloadButton.test.js b/tests/frontend/components/modelVersionsTab.downloadButton.test.js new file mode 100644 index 00000000..5510ebac --- /dev/null +++ b/tests/frontend/components/modelVersionsTab.downloadButton.test.js @@ -0,0 +1,285 @@ +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(); +const openFileSelectionForVersion = vi.fn(); + +vi.mock(DOWNLOAD_MANAGER_MODULE, () => ({ + downloadManager: { + downloadVersionWithDefaults, + openFileSelectionForVersion, + }, +})); + +vi.mock(UI_HELPERS_MODULE, () => ({ + showToast: vi.fn(), + openCivitaiUrl: vi.fn(), +})); + +const stateMock = { + global: { + settings: { + autoplay_on_hover: false, + version_grouping: 'any', + }, + }, +}; +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(), +})); + +function buildRecord(versions) { + return { + success: true, + record: { + shouldIgnore: false, + inLibraryVersionIds: versions.filter(v => v.isInLibrary).map(v => v.versionId), + versions, + }, + }; +} + +async function renderVersions(record) { + const { initVersionsTab } = await import(MODEL_VERSIONS_MODULE); + const controller = initVersionsTab({ + modalId: 'model-versions-modal', + modelType: 'loras', + modelId: 123, + currentVersionId: null, + }); + await controller.load(); +} + +function downloadButtonFor(versionId) { + return document.querySelector( + `.model-version-row[data-version-id="${versionId}"] [data-version-action="download"]` + ); +} + +function filesBadgeFor(versionId) { + return document.querySelector( + `.model-version-row[data-version-id="${versionId}"] [data-version-files]` + ); +} + +describe('ModelVersionsTab download button visibility', () => { + let getModelApiClient; + let fetchModelUpdateVersions; + + beforeEach(async () => { + vi.resetModules(); + downloadVersionWithDefaults.mockReset(); + downloadVersionWithDefaults.mockResolvedValue(true); + openFileSelectionForVersion.mockReset(); + openFileSelectionForVersion.mockResolvedValue(undefined); + document.body.innerHTML = ` +
+
+
+
+
+ `; + ({ getModelApiClient } = await import(API_FACTORY_MODULE)); + fetchModelUpdateVersions = vi.fn(); + getModelApiClient.mockReturnValue({ + fetchModelUpdateVersions, + fetchModelRoots: vi.fn(), + setModelUpdateIgnore: vi.fn(), + setVersionUpdateIgnore: vi.fn(), + deleteModel: vi.fn(), + }); + }); + + afterEach(() => { + document.body.innerHTML = ''; + }); + + it('hides the download button for a single-file in-library version', async () => { + fetchModelUpdateVersions.mockResolvedValue(buildRecord([ + { + versionId: 10, + name: 'v1.0', + baseModel: 'Illustrious', + isInLibrary: true, + shouldIgnore: false, + filePath: '/models/loras/file.safetensors', + fileCount: 1, + }, + ])); + + await renderVersions(); + + expect(downloadButtonFor(10)).toBeFalsy(); + expect(filesBadgeFor(10)).toBeFalsy(); + // The delete affordance must remain for in-library versions. + expect(document.querySelector( + '.model-version-row[data-version-id="10"] [data-version-action="delete"]' + )).toBeTruthy(); + }); + + it('shows the files badge instead of a download button for a multi-file in-library version', async () => { + fetchModelUpdateVersions.mockResolvedValue(buildRecord([ + { + versionId: 10, + name: 'v1.0', + baseModel: 'Illustrious', + isInLibrary: true, + shouldIgnore: false, + filePath: '/models/loras/file.safetensors', + fileCount: 3, + }, + ])); + + await renderVersions(); + + expect(downloadButtonFor(10)).toBeFalsy(); + const badge = filesBadgeFor(10); + expect(badge).toBeTruthy(); + expect(badge.textContent).toContain('3 files'); + expect(badge.getAttribute('title')).toBe('Choose which files to download'); + + badge.click(); + await new Promise(resolve => setTimeout(resolve, 0)); + + expect(openFileSelectionForVersion).toHaveBeenCalledWith('loras', 123, 10); + }); + + it('hides the download button when fileCount is unknown for an in-library version', async () => { + fetchModelUpdateVersions.mockResolvedValue(buildRecord([ + { + versionId: 10, + name: 'v1.0', + baseModel: 'Illustrious', + isInLibrary: true, + shouldIgnore: false, + filePath: '/models/loras/file.safetensors', + }, + ])); + + await renderVersions(); + + expect(downloadButtonFor(10)).toBeFalsy(); + expect(filesBadgeFor(10)).toBeFalsy(); + }); + + it('shows the download button for versions not in the library', async () => { + fetchModelUpdateVersions.mockResolvedValue(buildRecord([ + { + versionId: 11, + name: 'v1.1', + baseModel: 'Illustrious', + isInLibrary: false, + shouldIgnore: false, + fileCount: 1, + }, + { + versionId: 12, + name: 'v1.2', + baseModel: 'Illustrious', + isInLibrary: false, + shouldIgnore: false, + }, + ])); + + await renderVersions(); + + expect(downloadButtonFor(11)).toBeTruthy(); + expect(downloadButtonFor(12)).toBeTruthy(); + // Single-file and unknown-count versions get no files badge. + expect(filesBadgeFor(11)).toBeFalsy(); + expect(filesBadgeFor(12)).toBeFalsy(); + }); + + it('keeps the default-file download button and offers the files badge for a multi-file version not in the library', async () => { + fetchModelUpdateVersions.mockResolvedValue(buildRecord([ + { + versionId: 11, + name: 'v1.1', + baseModel: 'Illustrious', + isInLibrary: false, + shouldIgnore: false, + fileCount: 2, + }, + ])); + + await renderVersions(); + + // The Download button stays bound to the default (primary) file. + const button = downloadButtonFor(11); + expect(button).toBeTruthy(); + expect(button.getAttribute('title')).toBe('Download this version'); + + button.click(); + await new Promise(resolve => setTimeout(resolve, 0)); + + expect(downloadVersionWithDefaults).toHaveBeenCalledWith( + 'loras', 123, 11, + expect.objectContaining({ versionName: 'v1.1' }) + ); + expect(openFileSelectionForVersion).not.toHaveBeenCalled(); + + // The badge is the advanced entry into the file-selection step. + const badge = filesBadgeFor(11); + expect(badge).toBeTruthy(); + expect(badge.textContent).toContain('2 files'); + + badge.click(); + await new Promise(resolve => setTimeout(resolve, 0)); + + expect(openFileSelectionForVersion).toHaveBeenCalledWith('loras', 123, 11); + }); + + it('keeps the direct default download for a single-file version not in the library', async () => { + fetchModelUpdateVersions.mockResolvedValue(buildRecord([ + { + versionId: 11, + name: 'v1.1', + baseModel: 'Illustrious', + isInLibrary: false, + shouldIgnore: false, + fileCount: 1, + }, + ])); + + await renderVersions(); + + expect(filesBadgeFor(11)).toBeFalsy(); + downloadButtonFor(11).click(); + await new Promise(resolve => setTimeout(resolve, 0)); + + expect(downloadVersionWithDefaults).toHaveBeenCalledWith( + 'loras', 123, 11, + expect.objectContaining({ versionName: 'v1.1' }) + ); + expect(openFileSelectionForVersion).not.toHaveBeenCalled(); + }); +}); diff --git a/tests/routes/test_model_update_handler.py b/tests/routes/test_model_update_handler.py index cd5e40ed..3be9ff81 100644 --- a/tests/routes/test_model_update_handler.py +++ b/tests/routes/test_model_update_handler.py @@ -693,3 +693,21 @@ async def test_enrich_early_access_details_skips_permanent_paid(monkeypatch): enriched_map = {v.version_id: v for v in enriched.versions} assert enriched_map[200].early_access_ends_at == "2099-01-01T00:00:00.000Z" assert enriched_map[100].early_access_ends_at is None + + +def test_serialize_version_includes_file_count(): + version = ModelVersionRecord( + version_id=11, name="v11", base_model=None, released_at=None, size_bytes=None, + preview_url=None, is_in_library=True, should_ignore=False, file_count=2, + ) + serialized = ModelUpdateHandler._serialize_version(version, None) + assert serialized["fileCount"] == 2 + + +def test_serialize_version_file_count_defaults_to_none(): + version = ModelVersionRecord( + version_id=12, name="v12", base_model=None, released_at=None, size_bytes=None, + preview_url=None, is_in_library=False, should_ignore=False, + ) + serialized = ModelUpdateHandler._serialize_version(version, None) + assert serialized["fileCount"] is None diff --git a/tests/services/test_model_update_service.py b/tests/services/test_model_update_service.py index f2008801..04e72ad7 100644 --- a/tests/services/test_model_update_service.py +++ b/tests/services/test_model_update_service.py @@ -798,3 +798,126 @@ def test_build_record_from_remote_preserves_paid_fields(tmp_path): rebuilt = record.versions[0] assert rebuilt.paid_access == '{"permanent": true, "endsAt": null}' assert rebuilt.is_paid is True + + +def test_extract_file_count_counts_weight_files(tmp_path): + """file_count counts only weight-type files; a missing files array stays + None (unknown) so the UI can distinguish it from "no weight files".""" + db_path = tmp_path / "updates.sqlite" + service = ModelUpdateService(str(db_path)) + + response = { + "modelVersions": [ + { + "id": 42, + "files": [ + {"sizeKB": 100, "type": "Model", "primary": True}, + {"sizeKB": 10, "type": "Training Data"}, + {"sizeKB": 50, "type": "Pruned Model"}, + ], + "images": [], + }, + {"id": 43, "images": []}, + {"id": 44, "files": [], "images": []}, + ] + } + + versions = service._extract_versions(response) + assert versions is not None + assert versions[0].file_count == 2 + assert versions[1].file_count is None + assert versions[2].file_count == 0 + + +@pytest.mark.asyncio +async def test_refresh_persists_file_count(tmp_path): + db_path = tmp_path / "updates.sqlite" + service = ModelUpdateService(str(db_path), ttl_seconds=3600) + raw_data = [{"civitai": {"modelId": 1, "id": 11}}] + scanner = DummyScanner(raw_data) + provider = DummyProvider( + { + "modelVersions": [ + { + "id": 11, + "name": "v1", + "baseModel": "SD15", + "files": [ + {"sizeKB": 1024, "type": "Model", "primary": True}, + {"sizeKB": 2048, "type": "Model"}, + {"sizeKB": 128, "type": "Training Data"}, + ], + "images": [], + } + ] + } + ) + + await service.refresh_for_model_type("lora", scanner, provider) + record = await service.get_record("lora", 1) + + assert record is not None + assert record.versions[0].file_count == 2 + + +def test_build_record_from_remote_preserves_file_count(tmp_path): + """A remote payload without files data must not clobber the previously + persisted file_count; a populated payload wins.""" + db_path = tmp_path / "updates.sqlite" + service = ModelUpdateService(str(db_path)) + + existing = make_record( + ModelVersionRecord( + version_id=7, + name="v7", + base_model=None, + released_at=None, + size_bytes=None, + preview_url=None, + is_in_library=True, + should_ignore=False, + file_count=3, + ) + ) + remote_without_count = ModelVersionRecord( + version_id=7, + name="v7", + base_model=None, + released_at=None, + size_bytes=None, + preview_url=None, + is_in_library=False, + should_ignore=False, + file_count=None, + ) + + record = service._build_record_from_remote( + model_type="lora", + model_id=999, + local_versions=[7], + remote_versions=[remote_without_count], + existing=existing, + timestamp=1.0, + ) + assert record.versions[0].file_count == 3 + + remote_with_count = ModelVersionRecord( + version_id=7, + name="v7", + base_model=None, + released_at=None, + size_bytes=None, + preview_url=None, + is_in_library=False, + should_ignore=False, + file_count=1, + ) + record = service._build_record_from_remote( + model_type="lora", + model_id=999, + local_versions=[7], + remote_versions=[remote_with_count], + existing=existing, + timestamp=2.0, + ) + assert record.versions[0].file_count == 1