From 3aa32120df3191bc91d35d7802ad5cf10d25473a Mon Sep 17 00:00:00 2001 From: Will Miao Date: Fri, 2 Oct 2026 09:58:01 +0800 Subject: [PATCH] fix(downloads): route unknown checkpoint baseModels to diffusion models by default CivitAI labels new DiT architectures (MiniMax H3, future Flux/Wan/Qwen variants) as model.type "Checkpoint" with plain "Model" file entries, so the DIFFUSION_MODEL_BASE_MODELS allowlist could never keep up and such downloads were mis-routed to the checkpoint roots (e.g. model 2877206 / version 3374439). The set of true full-checkpoint families is closed, so the baseModel fallback is inverted: 1. file type UNet/Diffusion Model -> unet (unchanged) 2. baseModel in DIFFUSION_MODEL_BASE_MODELS (now incl. MiniMax H3) -> unet 3. baseModel in new CHECKPOINT_BASE_MODELS (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) -> checkpoint 4. unknown/empty baseModel -> new unknown_base_model_routing setting, defaulting to diffusion models The setting is exposed under Settings > Downloads, validated in SettingsManager, and threaded into both the download manager and the download routing endpoint so they keep agreeing. --- locales/de.json | 8 +++ locales/en.json | 8 +++ locales/es.json | 8 +++ locales/fr.json | 8 +++ locales/he.json | 8 +++ locales/ja.json | 8 +++ locales/ko.json | 8 +++ locales/ru.json | 8 +++ locales/zh-CN.json | 8 +++ locales/zh-TW.json | 8 +++ .../handlers/download_routing_handlers.py | 5 ++ py/services/download_manager.py | 3 ++ py/services/download_routing.py | 24 ++++++++- py/services/settings_manager.py | 15 ++++++ py/utils/constants.py | 34 ++++++++++++ static/js/managers/SettingsManager.js | 5 ++ static/js/state/index.js | 1 + .../components/modals/settings/general.html | 4 ++ .../routes/test_download_routing_handlers.py | 36 +++++++++++++ tests/services/test_download_manager_basic.py | 2 +- tests/services/test_download_routing.py | 54 +++++++++++++++++-- tests/services/test_settings_manager.py | 15 ++++++ 22 files changed, 270 insertions(+), 8 deletions(-) diff --git a/locales/de.json b/locales/de.json index 59834869..ab2ec09c 100644 --- a/locales/de.json +++ b/locales/de.json @@ -353,6 +353,14 @@ "placeholder": "Leer lassen, um aria2c aus dem PATH zu verwenden" }, "aria2HelpLink": "Erfahren Sie, wie Sie das aria2-Download-Backend einrichten", + "unknownBaseModelRouting": { + "label": "[TODO: Translate] Unknown base model routing", + "help": "[TODO: Translate] Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "[TODO: Translate] Diffusion models (recommended)", + "checkpoint": "[TODO: Translate] Checkpoints" + } + }, "civitaiHostBanner": { "title": "CivitAI-Host-Einstellung verfügbar", "content": "CivitAI verwendet jetzt civitai.com für SFW-Inhalte und civitai.red für uneingeschränkte Inhalte. In den Einstellungen können Sie ändern, welche Seite standardmäßig geöffnet wird.", diff --git a/locales/en.json b/locales/en.json index 96da919b..4b33f3b1 100644 --- a/locales/en.json +++ b/locales/en.json @@ -353,6 +353,14 @@ "placeholder": "Leave empty to use aria2c from PATH" }, "aria2HelpLink": "Learn how to set up the aria2 download backend", + "unknownBaseModelRouting": { + "label": "Unknown base model routing", + "help": "Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "Diffusion models (recommended)", + "checkpoint": "Checkpoints" + } + }, "civitaiHostBanner": { "title": "CivitAI host preference available", "content": "CivitAI now uses civitai.com for SFW content and civitai.red for unrestricted content. You can change which site opens by default in Settings.", diff --git a/locales/es.json b/locales/es.json index d301e746..0606532c 100644 --- a/locales/es.json +++ b/locales/es.json @@ -353,6 +353,14 @@ "placeholder": "Déjalo vacío para usar aria2c desde el PATH" }, "aria2HelpLink": "Aprende a configurar el backend de descarga aria2", + "unknownBaseModelRouting": { + "label": "[TODO: Translate] Unknown base model routing", + "help": "[TODO: Translate] Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "[TODO: Translate] Diffusion models (recommended)", + "checkpoint": "[TODO: Translate] Checkpoints" + } + }, "civitaiHostBanner": { "title": "Preferencia de host de CivitAI disponible", "content": "CivitAI ahora usa civitai.com para contenido SFW y civitai.red para contenido sin restricciones. Puedes cambiar en Ajustes qué sitio se abre por defecto.", diff --git a/locales/fr.json b/locales/fr.json index 8ad8cc02..054560b5 100644 --- a/locales/fr.json +++ b/locales/fr.json @@ -353,6 +353,14 @@ "placeholder": "Laisser vide pour utiliser aria2c depuis le PATH" }, "aria2HelpLink": "Apprenez à configurer le backend de téléchargement aria2", + "unknownBaseModelRouting": { + "label": "[TODO: Translate] Unknown base model routing", + "help": "[TODO: Translate] Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "[TODO: Translate] Diffusion models (recommended)", + "checkpoint": "[TODO: Translate] Checkpoints" + } + }, "civitaiHostBanner": { "title": "Préférence d’hôte CivitAI disponible", "content": "CivitAI utilise désormais civitai.com pour le contenu SFW et civitai.red pour le contenu sans restriction. Vous pouvez modifier dans les paramètres le site ouvert par défaut.", diff --git a/locales/he.json b/locales/he.json index cfca1ccd..d821934c 100644 --- a/locales/he.json +++ b/locales/he.json @@ -353,6 +353,14 @@ "placeholder": "השאר ריק כדי להשתמש ב-aria2c מתוך ה-PATH" }, "aria2HelpLink": "למד כיצד להגדיר את מנוע ההורדה aria2", + "unknownBaseModelRouting": { + "label": "[TODO: Translate] Unknown base model routing", + "help": "[TODO: Translate] Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "[TODO: Translate] Diffusion models (recommended)", + "checkpoint": "[TODO: Translate] Checkpoints" + } + }, "civitaiHostBanner": { "title": "העדפת מארח CivitAI זמינה", "content": "CivitAI משתמש כעת ב-civitai.com עבור תוכן SFW וב-civitai.red עבור תוכן ללא הגבלות. ניתן לשנות בהגדרות איזה אתר ייפתח כברירת מחדל.", diff --git a/locales/ja.json b/locales/ja.json index 14ebd64c..f145ebc1 100644 --- a/locales/ja.json +++ b/locales/ja.json @@ -353,6 +353,14 @@ "placeholder": "空欄のままにすると PATH 上の aria2c を使用します" }, "aria2HelpLink": "aria2 ダウンロードバックエンドの設定方法", + "unknownBaseModelRouting": { + "label": "[TODO: Translate] Unknown base model routing", + "help": "[TODO: Translate] Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "[TODO: Translate] Diffusion models (recommended)", + "checkpoint": "[TODO: Translate] Checkpoints" + } + }, "civitaiHostBanner": { "title": "CivitAI ホスト設定を利用できます", "content": "CivitAI は現在、SFW コンテンツには civitai.com、制限なしコンテンツには civitai.red を使用しています。設定で既定で開くサイトを変更できます。", diff --git a/locales/ko.json b/locales/ko.json index c221ff42..ed2b7a2c 100644 --- a/locales/ko.json +++ b/locales/ko.json @@ -353,6 +353,14 @@ "placeholder": "비워 두면 PATH의 aria2c를 사용합니다" }, "aria2HelpLink": "aria2 다운로드 백엔드 설정 방법 알아보기", + "unknownBaseModelRouting": { + "label": "[TODO: Translate] Unknown base model routing", + "help": "[TODO: Translate] Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "[TODO: Translate] Diffusion models (recommended)", + "checkpoint": "[TODO: Translate] Checkpoints" + } + }, "civitaiHostBanner": { "title": "CivitAI 호스트 기본 설정 사용 가능", "content": "이제 CivitAI는 SFW 콘텐츠에 civitai.com을, 무제한 콘텐츠에 civitai.red를 사용합니다. 설정에서 기본으로 열 사이트를 변경할 수 있습니다.", diff --git a/locales/ru.json b/locales/ru.json index 94b6cfce..90530258 100644 --- a/locales/ru.json +++ b/locales/ru.json @@ -353,6 +353,14 @@ "placeholder": "Оставьте пустым, чтобы использовать aria2c из PATH" }, "aria2HelpLink": "Узнайте, как настроить сервер загрузки aria2", + "unknownBaseModelRouting": { + "label": "[TODO: Translate] Unknown base model routing", + "help": "[TODO: Translate] Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "[TODO: Translate] Diffusion models (recommended)", + "checkpoint": "[TODO: Translate] Checkpoints" + } + }, "civitaiHostBanner": { "title": "Доступна настройка хоста CivitAI", "content": "Теперь CivitAI использует civitai.com для контента SFW и civitai.red для контента без ограничений. В настройках можно изменить, какой сайт открывать по умолчанию.", diff --git a/locales/zh-CN.json b/locales/zh-CN.json index df3fd027..cf38b02f 100644 --- a/locales/zh-CN.json +++ b/locales/zh-CN.json @@ -353,6 +353,14 @@ "placeholder": "留空则使用 PATH 中的 aria2c" }, "aria2HelpLink": "了解如何配置 aria2 下载后端", + "unknownBaseModelRouting": { + "label": "[TODO: Translate] Unknown base model routing", + "help": "[TODO: Translate] Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "[TODO: Translate] Diffusion models (recommended)", + "checkpoint": "[TODO: Translate] Checkpoints" + } + }, "civitaiHostBanner": { "title": "已提供 CivitAI 站点偏好设置", "content": "CivitAI 现在使用 civitai.com 提供 SFW 内容，使用 civitai.red 提供无限制内容。你可以在设置中更改默认打开的站点。", diff --git a/locales/zh-TW.json b/locales/zh-TW.json index 589e3c9a..8adbc8fa 100644 --- a/locales/zh-TW.json +++ b/locales/zh-TW.json @@ -353,6 +353,14 @@ "placeholder": "留空則使用 PATH 中的 aria2c" }, "aria2HelpLink": "了解如何設定 aria2 下載後端", + "unknownBaseModelRouting": { + "label": "[TODO: Translate] Unknown base model routing", + "help": "[TODO: Translate] Decides where checkpoint downloads go when CivitAI reports a base model that is neither a known checkpoint (SD 1.x/2.x/3.x, SDXL, Pony, Illustrious, NoobAI) nor a known diffusion model. New diffusion architectures appear frequently, so routing them to diffusion models is usually correct.", + "options": { + "diffusionModel": "[TODO: Translate] Diffusion models (recommended)", + "checkpoint": "[TODO: Translate] Checkpoints" + } + }, "civitaiHostBanner": { "title": "已提供 CivitAI 站點偏好設定", "content": "CivitAI 現在使用 civitai.com 提供 SFW 內容，使用 civitai.red 提供無限制內容。您可以在設定中變更預設開啟的站點。", diff --git a/py/routes/handlers/download_routing_handlers.py b/py/routes/handlers/download_routing_handlers.py index 20ecb7ef..fffeee5a 100644 --- a/py/routes/handlers/download_routing_handlers.py +++ b/py/routes/handlers/download_routing_handlers.py @@ -96,10 +96,15 @@ class DownloadRoutingHandler: } ) + from ...services.settings_manager import get_settings_manager + is_diffusion = is_diffusion_model_download( model_type, file_types=(str(t) for t in file_types), base_model=base_model, + unknown_base_model_default=get_settings_manager().get( + "unknown_base_model_routing", "diffusion_model" + ), ) return web.json_response( { diff --git a/py/services/download_manager.py b/py/services/download_manager.py index b2cf40ec..fe767116 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -1670,6 +1670,9 @@ class DownloadManager: model_type, file_types=(f.get("type", "") for f in version_info.get("files", [])), base_model=base_model_value, + unknown_base_model_default=get_settings_manager().get( + "unknown_base_model_routing", "diffusion_model" + ), ) # Existence check after the metadata fetch (#1058): diff --git a/py/services/download_routing.py b/py/services/download_routing.py index a18c2080..175ca9a4 100644 --- a/py/services/download_routing.py +++ b/py/services/download_routing.py @@ -13,6 +13,7 @@ import logging from typing import Iterable, Optional from ..utils.constants import ( + CHECKPOINT_BASE_MODELS, CIVITAI_FILE_TYPE_TO_OTHER_SUB_TYPE, CIVITAI_TYPE_TO_OTHER_SUB_TYPE, DIFFUSION_MODEL_BASE_MODELS, @@ -24,17 +25,27 @@ logger = logging.getLogger(__name__) # model (loaded via UNETLoader in ComfyUI) rather than a full checkpoint. DIFFUSION_FILE_TYPES = frozenset({"UNet", "Diffusion Model"}) +# Allowed values for the "unknown_base_model_routing" setting / the +# unknown_base_model_default parameter below. +ROUTING_DIFFUSION_MODEL = "diffusion_model" +ROUTING_CHECKPOINT = "checkpoint" + def is_diffusion_model_download( model_type: str, file_types: Iterable[str] = (), base_model: str = "", + unknown_base_model_default: str = ROUTING_DIFFUSION_MODEL, ) -> bool: """Return True when a download should be routed to the unet roots. Only applies to downloads initiated from the checkpoint library. Priority: (1) any file has type "UNet" or "Diffusion Model" (the more - direct signal from CivitAI), (2) baseModel is a known diffusion model. + direct signal from CivitAI), (2) baseModel is a known diffusion model, + (3) baseModel is a known full checkpoint -> not diffusion, (4) unknown + or empty baseModel -> the ``unknown_base_model_default`` setting, which + defaults to diffusion because the set of true checkpoint families is + closed while new DiT base models appear all the time. """ if model_type != "checkpoint": return False @@ -54,7 +65,16 @@ def is_diffusion_model_download( ) return True - return False + if base_model in CHECKPOINT_BASE_MODELS: + return False + + is_diffusion = unknown_base_model_default != ROUTING_CHECKPOINT + logger.info( + "baseModel '%s' is unknown, routing to %s folder (unknown_base_model_routing)", + base_model, + "unet" if is_diffusion else "checkpoint", + ) + return is_diffusion def resolve_other_download_sub_type( diff --git a/py/services/settings_manager.py b/py/services/settings_manager.py index 0e874776..df564d30 100644 --- a/py/services/settings_manager.py +++ b/py/services/settings_manager.py @@ -130,6 +130,9 @@ DEFAULT_SETTINGS: Dict[str, Any] = { "metadata_refresh_skip_paths": [], "skip_previously_downloaded_model_versions": False, "download_skip_base_models": [], + # Routing target for checkpoint downloads whose baseModel is neither a + # known diffusion model nor a known full checkpoint (CHECKPOINT_BASE_MODELS). + "unknown_base_model_routing": "diffusion_model", "backup_auto_enabled": True, "backup_retention_count": 5, "use_new_license_icons": True, @@ -1628,6 +1631,16 @@ class SettingsManager: return normalized return "alongside" + @staticmethod + def _normalize_unknown_base_model_routing(value: Any) -> str: + """Return a valid unknown-base-model routing target, falling back to ``diffusion_model``.""" + + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in ("diffusion_model", "checkpoint"): + return normalized + return "diffusion_model" + def _refresh_sidecar_storage_config(self) -> None: """Rebuild dependent config state after sidecar storage settings change.""" @@ -1840,6 +1853,8 @@ class SettingsManager: self._migrate_recipes_directory(current_recipes_dir, target_recipes_dir) elif key == "sidecar_storage_mode": value = self._normalize_sidecar_storage_mode(value) + elif key == "unknown_base_model_routing": + value = self._normalize_unknown_base_model_routing(value) elif key == "sidecar_storage_path": value = self._normalize_recipes_path_value(value) self.settings[key] = value diff --git a/py/utils/constants.py b/py/utils/constants.py index 6a283ced..bf671271 100644 --- a/py/utils/constants.py +++ b/py/utils/constants.py @@ -371,6 +371,7 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset( # Other diffusion models "Ernie", "Ernie Turbo", + "MiniMax H3", "Nucleus", "Qwen", "ZImageBase", @@ -380,6 +381,39 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset( ] ) +# baseModel values from CivitAI that are true full checkpoints (loaded via +# CheckpointLoaderSimple in ComfyUI). New DiT families appear on CivitAI all +# the time, so download routing inverts the fallback: anything NOT in this +# closed set (and not a known diffusion model) is treated as a diffusion +# model by default (see py/services/download_routing.py). +# "Pony V7" is deliberately excluded: it is not an SDXL-derivative full +# checkpoint, so it follows the unknown-base-model default (diffusion). +CHECKPOINT_BASE_MODELS = frozenset( + [ + # Stable Diffusion 1.x + "SD 1.4", + "SD 1.5", + "SD 1.5 LCM", + "SD 1.5 Hyper", + # Stable Diffusion 2.x + "SD 2.0", + "SD 2.1", + # Stable Diffusion 3.x + "SD 3", + "SD 3.5", + "SD 3.5 Medium", + "SD 3.5 Large", + "SD 3.5 Large Turbo", + # SDXL and its full-checkpoint derivatives + "SDXL 1.0", + "SDXL Lightning", + "SDXL Hyper", + "Pony", + "Illustrious", + "NoobAI", + ] +) + # Supported baseModel values for download exclusion settings. # Keep this aligned with static/js/utils/constants.js, excluding the generic "Other" value. SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS = frozenset( diff --git a/static/js/managers/SettingsManager.js b/static/js/managers/SettingsManager.js index 808e4692..e6763d97 100644 --- a/static/js/managers/SettingsManager.js +++ b/static/js/managers/SettingsManager.js @@ -1053,6 +1053,11 @@ export class SettingsManager { downloadBackendSelect.value = state.global.settings.download_backend || 'python'; } + const unknownBaseModelRoutingSelect = document.getElementById('unknownBaseModelRouting'); + if (unknownBaseModelRoutingSelect) { + unknownBaseModelRoutingSelect.value = state.global.settings.unknown_base_model_routing || 'diffusion_model'; + } + const aria2cPathInput = document.getElementById('aria2cPath'); if (aria2cPathInput) { aria2cPathInput.value = state.global.settings.aria2c_path || ''; diff --git a/static/js/state/index.js b/static/js/state/index.js index f50e5ac8..acf76dd3 100644 --- a/static/js/state/index.js +++ b/static/js/state/index.js @@ -60,6 +60,7 @@ const DEFAULT_SETTINGS_BASE = Object.freeze({ metadata_refresh_skip_paths: [], skip_previously_downloaded_model_versions: false, download_skip_base_models: [], + unknown_base_model_routing: 'diffusion_model', backup_auto_enabled: true, backup_retention_count: 5, sidecar_storage_mode: 'alongside', diff --git a/templates/components/modals/settings/general.html b/templates/components/modals/settings/general.html index f29282f9..66514b3a 100644 --- a/templates/components/modals/settings/general.html +++ b/templates/components/modals/settings/general.html @@ -146,6 +146,10 @@ + {{ sm.setting_select('unknownBaseModelRouting', 'unknown_base_model_routing', 'settings.unknownBaseModelRouting.label', [ + ('diffusion_model', 'settings.unknownBaseModelRouting.options.diffusionModel'), + ('checkpoint', 'settings.unknownBaseModelRouting.options.checkpoint'), + ], 'settings.unknownBaseModelRouting.help') }} diff --git a/tests/routes/test_download_routing_handlers.py b/tests/routes/test_download_routing_handlers.py index 4e1474b6..4ea57b34 100644 --- a/tests/routes/test_download_routing_handlers.py +++ b/tests/routes/test_download_routing_handlers.py @@ -73,6 +73,42 @@ async def test_regular_checkpoint_stays_on_checkpoint_root(): assert payload["root_kind"] == "checkpoint" +@pytest.mark.asyncio +async def test_unknown_base_model_routes_to_unet_by_default(): + handler = DownloadRoutingHandler() + response = await handler.get_download_routing( + FakeRequest( + { + "model_type": "checkpoint", + "base_model": "Brand New Arch", + "file_types": ["Model"], + } + ) + ) + payload = json.loads(response.text) + assert payload["is_diffusion_model"] is True + assert payload["root_kind"] == "unet" + + +@pytest.mark.asyncio +async def test_unknown_base_model_honors_checkpoint_routing_setting(): + get_settings_manager().settings["unknown_base_model_routing"] = "checkpoint" + + handler = DownloadRoutingHandler() + response = await handler.get_download_routing( + FakeRequest( + { + "model_type": "checkpoint", + "base_model": "Brand New Arch", + "file_types": ["Model"], + } + ) + ) + payload = json.loads(response.text) + assert payload["is_diffusion_model"] is False + assert payload["root_kind"] == "checkpoint" + + @pytest.mark.asyncio async def test_lora_is_never_diffusion(): handler = DownloadRoutingHandler() diff --git a/tests/services/test_download_manager_basic.py b/tests/services/test_download_manager_basic.py index 97fc974b..76894b4a 100644 --- a/tests/services/test_download_manager_basic.py +++ b/tests/services/test_download_manager_basic.py @@ -2209,7 +2209,7 @@ async def test_download_falls_back_to_rest_name_when_mini_fails( metadata_provider.payload = { "id": 42, "model": {"type": "Checkpoint", "tags": ["fantasy"]}, - "baseModel": "BaseModel", + "baseModel": "SDXL 1.0", "creator": {"username": "Author"}, "files": [ { diff --git a/tests/services/test_download_routing.py b/tests/services/test_download_routing.py index 07d78bd1..843ba16b 100644 --- a/tests/services/test_download_routing.py +++ b/tests/services/test_download_routing.py @@ -20,9 +20,47 @@ def test_base_model_fallback_routes_to_unet(): ) -def test_regular_checkpoint_stays_on_checkpoint_roots(): +def test_minimax_h3_routes_to_unet(): + """CivitAI model 2877206: type "Checkpoint", baseModel "MiniMax H3", + all file entries typed "Model".""" + assert is_diffusion_model_download( + "checkpoint", file_types=["Model"], base_model="MiniMax H3" + ) + + +@pytest.mark.parametrize("base_model", ["SDXL 1.0", "Illustrious", "SD 1.5", "SD 3.5 Large"]) +def test_known_checkpoint_base_models_stay_on_checkpoint_roots(base_model): + """CHECKPOINT_BASE_MODELS members never route to unet, even when the + unknown-base-model default is diffusion.""" assert not is_diffusion_model_download( - "checkpoint", file_types=["Model"], base_model="SDXL 1.0" + "checkpoint", file_types=["Model"], base_model=base_model + ) + + +def test_unknown_base_model_defaults_to_unet(): + """New DiT families appear faster than the allowlist can track them, so + unknown baseModels route to unet by default.""" + assert is_diffusion_model_download( + "checkpoint", file_types=["Model"], base_model="Brand New Arch" + ) + + +def test_unknown_base_model_honors_checkpoint_setting(): + assert not is_diffusion_model_download( + "checkpoint", + file_types=["Model"], + base_model="Brand New Arch", + unknown_base_model_default="checkpoint", + ) + + +def test_unknown_base_model_checkpoint_setting_still_loses_to_file_type(): + """The file-type signal stays first regardless of the setting.""" + assert is_diffusion_model_download( + "checkpoint", + file_types=["Diffusion Model"], + base_model="SDXL 1.0", + unknown_base_model_default="checkpoint", ) @@ -35,9 +73,15 @@ def test_non_checkpoint_types_never_route_to_unet(): ) -def test_empty_inputs_stay_on_checkpoint_roots(): - assert not is_diffusion_model_download("checkpoint") - assert not is_diffusion_model_download("checkpoint", file_types=[], base_model="") +def test_empty_base_model_follows_unknown_default(): + assert is_diffusion_model_download("checkpoint") + assert is_diffusion_model_download("checkpoint", file_types=[], base_model="") + assert not is_diffusion_model_download( + "checkpoint", + file_types=[], + base_model="", + unknown_base_model_default="checkpoint", + ) from py.services.download_routing import resolve_other_download_sub_type diff --git a/tests/services/test_settings_manager.py b/tests/services/test_settings_manager.py index f06391ae..4bdbd44d 100644 --- a/tests/services/test_settings_manager.py +++ b/tests/services/test_settings_manager.py @@ -1300,6 +1300,21 @@ def test_setting_download_skip_base_models_normalizes_string_input(manager): assert manager.get("download_skip_base_models") == ["SDXL 1.0", "Pony"] +def test_unknown_base_model_routing_defaults_to_diffusion_model(manager): + assert manager.get("unknown_base_model_routing") == "diffusion_model" + + +def test_unknown_base_model_routing_is_normalized(manager): + manager.set("unknown_base_model_routing", "CHECKPOINT") + assert manager.get("unknown_base_model_routing") == "checkpoint" + + manager.set("unknown_base_model_routing", "bogus") + assert manager.get("unknown_base_model_routing") == "diffusion_model" + + manager.set("unknown_base_model_routing", None) + assert manager.get("unknown_base_model_routing") == "diffusion_model" + + def test_skip_previously_downloaded_model_versions_defaults_false(manager): assert manager.get_skip_previously_downloaded_model_versions() is False