From 7b2a108596bfd3225701c243cbfb0476ba4ae944 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Fri, 2 Oct 2026 10:34:59 +0800 Subject: [PATCH] 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. --- .../handlers/download_routing_handlers.py | 5 ++ py/services/download_manager.py | 6 ++- py/utils/constants.py | 30 ++++++++++-- .../routes/test_download_routing_handlers.py | 15 ++++++ tests/services/test_download_manager_basic.py | 48 +++++++++++++++++++ tests/services/test_download_routing.py | 23 ++++++++- 6 files changed, 122 insertions(+), 5 deletions(-) diff --git a/py/routes/handlers/download_routing_handlers.py b/py/routes/handlers/download_routing_handlers.py index fffeee5a..8107afc5 100644 --- a/py/routes/handlers/download_routing_handlers.py +++ b/py/routes/handlers/download_routing_handlers.py @@ -55,6 +55,11 @@ class DownloadRoutingHandler: 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: from ...services.settings_manager import get_settings_manager diff --git a/py/services/download_manager.py b/py/services/download_manager.py index fe767116..62fff77a 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -1523,7 +1523,11 @@ class DownloadManager: return {"success": False, "error": "Failed to fetch model metadata"} 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" elif model_type_from_info in VALID_LORA_TYPES: model_type = "lora" diff --git a/py/utils/constants.py b/py/utils/constants.py index bf671271..9a9341ca 100644 --- a/py/utils/constants.py +++ b/py/utils/constants.py @@ -352,12 +352,15 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset( "PixArt E", # Video diffusion models "CogVideoX", + "Flux 3 Video", "Hunyuan Video", "LTXV", "LTXV2", "LTXV 2.3", + "LTXV 2.5", "Mochi", "SVD", + "SVD XT", "Wan Video", "Wan Video 1.3B t2v", "Wan Video 14B t2v", @@ -368,12 +371,19 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset( "Wan Video 2.2 T2V-A14B", "Wan Video 2.5 T2V", "Wan Video 2.5 I2V", + "Wan Video 2.7", + "Wan Video 3.0", # Other diffusion models + "Boogu", "Ernie", "Ernie Turbo", + "HiDream-O1", + "Ming Image Design 0.1", + "Ming Image Design Layer 0.1", "MiniMax H3", "Nucleus", "Qwen", + "Wan Image 2.7", "ZImageBase", "ZImageTurbo", # 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 # 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). +# model by default (see py/services/download_routing.py). Cross-checked +# against CivitAI's official baseModelRecords (packages/civitai-shared +# 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( [ # Stable Diffusion 1.x @@ -397,7 +410,10 @@ CHECKPOINT_BASE_MODELS = frozenset( "SD 1.5 Hyper", # Stable Diffusion 2.x "SD 2.0", + "SD 2.0 768", "SD 2.1", + "SD 2.1 768", + "SD 2.1 Unclip", # Stable Diffusion 3.x "SD 3", "SD 3.5", @@ -405,12 +421,20 @@ CHECKPOINT_BASE_MODELS = frozenset( "SD 3.5 Large", "SD 3.5 Large Turbo", # SDXL and its full-checkpoint derivatives + "SDXL 0.9", "SDXL 1.0", + "SDXL 1.0 LCM", "SDXL Lightning", "SDXL Hyper", + "SDXL Turbo", + "SDXL Distilled", "Pony", "Illustrious", "NoobAI", + # Other full-checkpoint families + "Playground v2", + # Stable Cascade is unCLIP-style but loads via CheckpointLoader + "Stable Cascade", ] ) diff --git a/tests/routes/test_download_routing_handlers.py b/tests/routes/test_download_routing_handlers.py index 4ea57b34..05695dfa 100644 --- a/tests/routes/test_download_routing_handlers.py +++ b/tests/routes/test_download_routing_handlers.py @@ -109,6 +109,21 @@ async def test_unknown_base_model_honors_checkpoint_routing_setting(): 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 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 76894b4a..39066cb6 100644 --- a/tests/services/test_download_manager_basic.py +++ b/tests/services/test_download_manager_basic.py @@ -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 async def test_download_falls_back_to_rest_name_when_mini_fails( monkeypatch, scanners, metadata_provider, tmp_path diff --git a/tests/services/test_download_routing.py b/tests/services/test_download_routing.py index 843ba16b..b118b0a6 100644 --- a/tests/services/test_download_routing.py +++ b/tests/services/test_download_routing.py @@ -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(): """CivitAI model 2877206: type "Checkpoint", baseModel "MiniMax H3", 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): """CHECKPOINT_BASE_MODELS members never route to unet, even when the unknown-base-model default is diffusion."""