fix(downloads): close routing gaps and map CivitAI ModelType.UNet to the checkpoint branch

Cross-checked both baseModel lists against CivitAI's official
baseModelRecords (packages/civitai-shared src/basemodel.constants.ts):

- CHECKPOINT_BASE_MODELS gains SD 2.0/2.1 768, SD 2.1 Unclip, SDXL 0.9 /
  1.0 LCM / Turbo / Distilled, Playground v2 and Stable Cascade
  (unCLIP-style but CheckpointLoader-loaded).
- DIFFUSION_MODEL_BASE_MODELS gains SVD XT, LTXV 2.5, Flux 3 Video,
  Wan Image 2.7, Wan Video 2.7 / 3.0, HiDream-O1, Boogu and the Ming
  Image Design families. API-only (Kling/Sora/Veo/Imagen...), 3D and
  audio baseModels are intentionally skipped.
- Pony V7 exclusion now backed by live-API evidence (model 1901521 is
  AuraFlow-architecture shipping .gguf variants).

CivitAI has no model-level diffusion ModelType: DiT models are uploaded
as "Checkpoint" or "UNet", with only uploader-chosen file types to tell
them apart. model.type "unet" previously fell through type derivation
and failed with 'not supported for download'; it now goes through the
checkpoint branch in both the download manager and the download routing
endpoint, so the standard chain (file type -> baseModel lists -> unknown
default) applies.
This commit is contained in:
Will Miao
2026-10-02 10:34:59 +08:00
parent 3aa32120df
commit 7b2a108596
6 changed files with 122 additions and 5 deletions
@@ -55,6 +55,11 @@ class DownloadRoutingHandler:
status=400, status=400,
) )
# CivitAI ModelType.UNet downloads go through the checkpoint branch,
# same as in the download manager.
if model_type.lower() == "unet":
model_type = "checkpoint"
if model_type.lower() in VALID_OTHER_CIVITAI_TYPES: if model_type.lower() in VALID_OTHER_CIVITAI_TYPES:
from ...services.settings_manager import get_settings_manager from ...services.settings_manager import get_settings_manager
+5 -1
View File
@@ -1523,7 +1523,11 @@ class DownloadManager:
return {"success": False, "error": "Failed to fetch model metadata"} return {"success": False, "error": "Failed to fetch model metadata"}
model_type_from_info = version_info.get("model", {}).get("type", "").lower() model_type_from_info = version_info.get("model", {}).get("type", "").lower()
if model_type_from_info == "checkpoint": # CivitAI ModelType has no model-level diffusion variant: DiT
# models are uploaded as "Checkpoint" or "UNet". Both go through
# the checkpoint branch so the standard diffusion routing chain
# (file type -> baseModel lists -> unknown default) applies.
if model_type_from_info in ("checkpoint", "unet"):
model_type = "checkpoint" model_type = "checkpoint"
elif model_type_from_info in VALID_LORA_TYPES: elif model_type_from_info in VALID_LORA_TYPES:
model_type = "lora" model_type = "lora"
+27 -3
View File
@@ -352,12 +352,15 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset(
"PixArt E", "PixArt E",
# Video diffusion models # Video diffusion models
"CogVideoX", "CogVideoX",
"Flux 3 Video",
"Hunyuan Video", "Hunyuan Video",
"LTXV", "LTXV",
"LTXV2", "LTXV2",
"LTXV 2.3", "LTXV 2.3",
"LTXV 2.5",
"Mochi", "Mochi",
"SVD", "SVD",
"SVD XT",
"Wan Video", "Wan Video",
"Wan Video 1.3B t2v", "Wan Video 1.3B t2v",
"Wan Video 14B t2v", "Wan Video 14B t2v",
@@ -368,12 +371,19 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset(
"Wan Video 2.2 T2V-A14B", "Wan Video 2.2 T2V-A14B",
"Wan Video 2.5 T2V", "Wan Video 2.5 T2V",
"Wan Video 2.5 I2V", "Wan Video 2.5 I2V",
"Wan Video 2.7",
"Wan Video 3.0",
# Other diffusion models # Other diffusion models
"Boogu",
"Ernie", "Ernie",
"Ernie Turbo", "Ernie Turbo",
"HiDream-O1",
"Ming Image Design 0.1",
"Ming Image Design Layer 0.1",
"MiniMax H3", "MiniMax H3",
"Nucleus", "Nucleus",
"Qwen", "Qwen",
"Wan Image 2.7",
"ZImageBase", "ZImageBase",
"ZImageTurbo", "ZImageTurbo",
# Krea 2 — loaded via UNETLoader in ComfyUI # Krea 2 — loaded via UNETLoader in ComfyUI
@@ -385,9 +395,12 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset(
# CheckpointLoaderSimple in ComfyUI). New DiT families appear on CivitAI all # CheckpointLoaderSimple in ComfyUI). New DiT families appear on CivitAI all
# the time, so download routing inverts the fallback: anything NOT in this # 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 # closed set (and not a known diffusion model) is treated as a diffusion
# model by default (see py/services/download_routing.py). # model by default (see py/services/download_routing.py). Cross-checked
# "Pony V7" is deliberately excluded: it is not an SDXL-derivative full # against CivitAI's official baseModelRecords (packages/civitai-shared
# checkpoint, so it follows the unknown-base-model default (diffusion). # src/basemodel.constants.ts), not only the download skip list.
# "Pony V7" is deliberately excluded: verified via the live API (model
# 1901521) to be AuraFlow-architecture shipping .gguf variants (UNETLoader),
# so it follows the unknown-base-model default (diffusion).
CHECKPOINT_BASE_MODELS = frozenset( CHECKPOINT_BASE_MODELS = frozenset(
[ [
# Stable Diffusion 1.x # Stable Diffusion 1.x
@@ -397,7 +410,10 @@ CHECKPOINT_BASE_MODELS = frozenset(
"SD 1.5 Hyper", "SD 1.5 Hyper",
# Stable Diffusion 2.x # Stable Diffusion 2.x
"SD 2.0", "SD 2.0",
"SD 2.0 768",
"SD 2.1", "SD 2.1",
"SD 2.1 768",
"SD 2.1 Unclip",
# Stable Diffusion 3.x # Stable Diffusion 3.x
"SD 3", "SD 3",
"SD 3.5", "SD 3.5",
@@ -405,12 +421,20 @@ CHECKPOINT_BASE_MODELS = frozenset(
"SD 3.5 Large", "SD 3.5 Large",
"SD 3.5 Large Turbo", "SD 3.5 Large Turbo",
# SDXL and its full-checkpoint derivatives # SDXL and its full-checkpoint derivatives
"SDXL 0.9",
"SDXL 1.0", "SDXL 1.0",
"SDXL 1.0 LCM",
"SDXL Lightning", "SDXL Lightning",
"SDXL Hyper", "SDXL Hyper",
"SDXL Turbo",
"SDXL Distilled",
"Pony", "Pony",
"Illustrious", "Illustrious",
"NoobAI", "NoobAI",
# Other full-checkpoint families
"Playground v2",
# Stable Cascade is unCLIP-style but loads via CheckpointLoader
"Stable Cascade",
] ]
) )
@@ -109,6 +109,21 @@ async def test_unknown_base_model_honors_checkpoint_routing_setting():
assert payload["root_kind"] == "checkpoint" assert payload["root_kind"] == "checkpoint"
@pytest.mark.asyncio
async def test_unet_model_type_uses_checkpoint_branch():
"""CivitAI ModelType.UNet has no branch of its own: it goes through the
checkpoint branch, where its UNet-typed files route to the unet roots."""
handler = DownloadRoutingHandler()
response = await handler.get_download_routing(
FakeRequest(
{"model_type": "unet", "base_model": "MiniMax H3", "file_types": ["UNet"]}
)
)
payload = json.loads(response.text)
assert payload["is_diffusion_model"] is True
assert payload["root_kind"] == "unet"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_lora_is_never_diffusion(): async def test_lora_is_never_diffusion():
handler = DownloadRoutingHandler() handler = DownloadRoutingHandler()
@@ -2200,6 +2200,54 @@ async def test_download_uses_raw_file_name_from_mini_endpoint(
] ]
@pytest.mark.asyncio
async def test_unet_model_type_download_lands_in_unet_roots(
monkeypatch, scanners, metadata_provider, tmp_path
):
"""CivitAI ModelType.UNet goes through the checkpoint branch: its
UNet-typed files route to the unet roots like any diffusion model."""
manager = DownloadManager()
unet_root = tmp_path / "unet"
get_settings_manager().settings["default_unet_root"] = str(unet_root)
metadata_provider.payload = {
"id": 42,
"model": {"type": "UNet", "tags": ["fantasy"]},
"baseModel": "MiniMax H3",
"creator": {"username": "Author"},
"files": [
{
"id": 1001,
"type": "UNet",
"primary": True,
"name": "minimax_h3.safetensors",
"downloadUrl": "https://example.invalid/file.safetensors",
}
],
}
captured = {}
async def fake_execute_download(self, **kwargs):
captured["file_path"] = kwargs["metadata"].file_path
return {"success": True}
monkeypatch.setattr(
DownloadManager, "_execute_download", fake_execute_download, raising=False
)
result = await manager.download_from_civitai(
model_version_id=42,
save_dir=str(tmp_path),
use_default_paths=True,
progress_callback=None,
source=None,
)
assert result["success"] is True, result
assert captured["file_path"].startswith(str(unet_root))
assert captured["file_path"].endswith("minimax_h3.safetensors")
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_download_falls_back_to_rest_name_when_mini_fails( async def test_download_falls_back_to_rest_name_when_mini_fails(
monkeypatch, scanners, metadata_provider, tmp_path monkeypatch, scanners, metadata_provider, tmp_path
+22 -1
View File
@@ -20,6 +20,15 @@ def test_base_model_fallback_routes_to_unet():
) )
@pytest.mark.parametrize("base_model", ["SVD XT", "Boogu", "Wan Video 3.0", "Flux 3 Video"])
def test_new_diffusion_families_route_to_unet(base_model):
"""DiT families from CivitAI's official baseModel list that plain
"Model"-typed files must not strand on the checkpoint roots."""
assert is_diffusion_model_download(
"checkpoint", file_types=["Model"], base_model=base_model
)
def test_minimax_h3_routes_to_unet(): def test_minimax_h3_routes_to_unet():
"""CivitAI model 2877206: type "Checkpoint", baseModel "MiniMax H3", """CivitAI model 2877206: type "Checkpoint", baseModel "MiniMax H3",
all file entries typed "Model".""" all file entries typed "Model"."""
@@ -28,7 +37,19 @@ def test_minimax_h3_routes_to_unet():
) )
@pytest.mark.parametrize("base_model", ["SDXL 1.0", "Illustrious", "SD 1.5", "SD 3.5 Large"]) @pytest.mark.parametrize(
"base_model",
[
"SDXL 1.0",
"Illustrious",
"SD 1.5",
"SD 2.1 768",
"SD 3.5 Large",
"SDXL Turbo",
"Playground v2",
"Stable Cascade",
],
)
def test_known_checkpoint_base_models_stay_on_checkpoint_roots(base_model): def test_known_checkpoint_base_models_stay_on_checkpoint_roots(base_model):
"""CHECKPOINT_BASE_MODELS members never route to unet, even when the """CHECKPOINT_BASE_MODELS members never route to unet, even when the
unknown-base-model default is diffusion.""" unknown-base-model default is diffusion."""