mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-03 00:35:32 -03:00
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:
@@ -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
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
+27
-3
@@ -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",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user