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