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.
This commit is contained in:
Will Miao
2026-10-02 09:58:01 +08:00
parent 2193ec8f38
commit 3aa32120df
22 changed files with 270 additions and 8 deletions
@@ -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()
@@ -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": [
{
+49 -5
View File
@@ -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
+15
View File
@@ -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