mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-13 12:52:13 -03:00
Will Miao
5dc2109d7c
Add a download_root_rules setting: ordered per-type rules
({match_type: base_model, match: [...], root: path}) that replace the
default download root when 'use default paths' is on. The first matching
rule wins; unmatched downloads fall back to default_*_root. Rules are
consulted under the final download type, so a checkpoint routed to the
unet roots uses the unet rule set.
- settings_manager: normalize/get/set with root validation against the
configured roots (folder_paths + extra_folder_paths); invalid rules
are dropped with a warning
- download_manager: apply the resolved root in the use_default_paths
branches for lora/checkpoint(unet)/embedding downloads
- download routing handler: expose resolved_root so the download
confirmation modal preselects the rule-routed root
- settings UI: per-type rule editor in the Organization section with
base-model multi-select panels, root dropdowns, reorder and remove
274 lines
9.2 KiB
Python
274 lines
9.2 KiB
Python
"""Tests for the shared download routing decision."""
|
|
|
|
import pytest
|
|
|
|
from py.services.download_routing import is_diffusion_model_download
|
|
|
|
|
|
@pytest.mark.parametrize("file_type", ["UNet", "Diffusion Model"])
|
|
def test_file_type_signal_routes_to_unet(file_type):
|
|
assert is_diffusion_model_download(
|
|
"checkpoint", file_types=[file_type], base_model="SDXL 1.0"
|
|
)
|
|
|
|
|
|
def test_base_model_fallback_routes_to_unet():
|
|
"""The reported Anima case: file type is plain "Model", but the
|
|
baseModel is a known diffusion model."""
|
|
assert is_diffusion_model_download(
|
|
"checkpoint", file_types=["Model"], base_model="Anima"
|
|
)
|
|
|
|
|
|
@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"."""
|
|
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 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."""
|
|
assert not is_diffusion_model_download(
|
|
"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",
|
|
)
|
|
|
|
|
|
def test_non_checkpoint_types_never_route_to_unet():
|
|
assert not is_diffusion_model_download(
|
|
"lora", file_types=["UNet"], base_model="Anima"
|
|
)
|
|
assert not is_diffusion_model_download(
|
|
"embedding", file_types=["Diffusion Model"], base_model="Anima"
|
|
)
|
|
|
|
|
|
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
|
|
|
|
|
|
class TestResolveOtherDownloadSubType:
|
|
"""Fixed priority: explicit file pick > model.type > file.type fallback."""
|
|
|
|
def test_explicit_file_pick_wins_over_model_type(self):
|
|
"""User explicitly picked a VAE component file of a Checkpoint model —
|
|
the picked file type wins."""
|
|
assert (
|
|
resolve_other_download_sub_type(
|
|
"Checkpoint", file_types=["Model", "VAE"], selected_file_type="VAE"
|
|
)
|
|
== "vae"
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"selected,expected",
|
|
[
|
|
("VAE", "vae"),
|
|
("Upscaler", "upscaler"),
|
|
("Text Encoder", "text_encoder"),
|
|
("Vision Encoder", "clip_vision"),
|
|
("CLIPVision", "clip_vision"),
|
|
("ControlNet", "controlnet"),
|
|
],
|
|
)
|
|
def test_explicit_file_pick_maps_all_known_types(self, selected, expected):
|
|
assert (
|
|
resolve_other_download_sub_type("Other", selected_file_type=selected)
|
|
== expected
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_type,expected",
|
|
[
|
|
("VAE", "vae"),
|
|
("Upscaler", "upscaler"),
|
|
("TextEncoder", "text_encoder"),
|
|
("CLIP", "text_encoder"),
|
|
("CLIPVision", "clip_vision"),
|
|
("Controlnet", "controlnet"),
|
|
],
|
|
)
|
|
def test_model_type_mapping(self, model_type, expected):
|
|
assert resolve_other_download_sub_type(model_type) == expected
|
|
|
|
def test_model_type_beats_unmappable_file_pick(self):
|
|
"""An explicit pick whose file type does not map (e.g. plain 'Model')
|
|
falls through to model.type."""
|
|
assert (
|
|
resolve_other_download_sub_type(
|
|
"TextEncoder", selected_file_type="Model"
|
|
)
|
|
== "text_encoder"
|
|
)
|
|
|
|
def test_bundled_component_files_never_override_model_type(self):
|
|
"""Anti-misrouting: a TextEncoder model bundling a VAE component file
|
|
must stay text_encoder — file types are a fallback, not an override."""
|
|
assert (
|
|
resolve_other_download_sub_type(
|
|
"TextEncoder", file_types=["Model", "VAE"]
|
|
)
|
|
== "text_encoder"
|
|
)
|
|
assert (
|
|
resolve_other_download_sub_type(
|
|
"Controlnet", file_types=["Model", "Text Encoder"]
|
|
)
|
|
== "controlnet"
|
|
)
|
|
|
|
def test_file_type_fallback_when_model_type_unmapped(self):
|
|
"""model.type 'Other' (or retired values) maps to nothing, so the
|
|
first mappable file type decides."""
|
|
assert (
|
|
resolve_other_download_sub_type("Other", file_types=["Model", "Upscaler"])
|
|
== "upscaler"
|
|
)
|
|
|
|
def test_file_type_fallback_for_civarchive_payload(self):
|
|
"""CivArchive-shaped payload: same fields, same decision path."""
|
|
assert (
|
|
resolve_other_download_sub_type(
|
|
"Other",
|
|
file_types=["Config", "Text Encoder"],
|
|
)
|
|
== "text_encoder"
|
|
)
|
|
|
|
@pytest.mark.parametrize("model_type", ["Other", "", "SomethingNew"])
|
|
def test_undecidable_returns_none(self, model_type):
|
|
assert (
|
|
resolve_other_download_sub_type(model_type, file_types=["Model"]) is None
|
|
)
|
|
assert resolve_other_download_sub_type(model_type) is None
|
|
|
|
def test_model_type_matching_is_case_insensitive(self):
|
|
assert resolve_other_download_sub_type("vAe") == "vae"
|
|
|
|
|
|
from py.services.download_routing import resolve_download_root
|
|
|
|
|
|
def test_resolve_download_root_returns_default_without_rules():
|
|
assert resolve_download_root(None, "Flux.1 D", "/default") == "/default"
|
|
assert resolve_download_root([], "Flux.1 D", "/default") == "/default"
|
|
|
|
|
|
def test_resolve_download_root_returns_default_without_base_model():
|
|
rules = [{"match_type": "base_model", "match": ["Flux.1 D"], "root": "/nas"}]
|
|
assert resolve_download_root(rules, "", "/default") == "/default"
|
|
|
|
|
|
def test_resolve_download_root_first_match_wins():
|
|
rules = [
|
|
{"match_type": "base_model", "match": ["Flux.1 D", "Wan Video"], "root": "/nas"},
|
|
{"match_type": "base_model", "match": ["Flux.1 D"], "root": "/other"},
|
|
]
|
|
assert resolve_download_root(rules, "Flux.1 D", "/default") == "/nas"
|
|
assert resolve_download_root(rules, "Wan Video", "/default") == "/nas"
|
|
|
|
|
|
def test_resolve_download_root_falls_back_when_no_rule_matches():
|
|
rules = [{"match_type": "base_model", "match": ["Pony"], "root": "/nas"}]
|
|
assert resolve_download_root(rules, "Flux.1 D", "/default") == "/default"
|
|
|
|
|
|
def test_resolve_download_root_skips_malformed_rules():
|
|
rules = [
|
|
{"match_type": "tags", "match": ["Flux.1 D"], "root": "/nas"},
|
|
{"match_type": "base_model", "match": "Flux.1 D", "root": "/nas"},
|
|
"not-a-dict",
|
|
{"match_type": "base_model", "match": ["Flux.1 D"], "root": ""},
|
|
{"match_type": "base_model", "match": ["Flux.1 D"], "root": "/nas"},
|
|
]
|
|
assert resolve_download_root(rules, "Flux.1 D", "/default") == "/nas"
|
|
|
|
|
|
def test_resolve_download_root_rule_lists_are_per_type():
|
|
"""Checkpoint and unet rule sets are consulted independently."""
|
|
rules_by_type = {
|
|
"checkpoints": [
|
|
{"match_type": "base_model", "match": ["SDXL 1.0"], "root": "/ckpt-nas"}
|
|
],
|
|
"unet": [
|
|
{"match_type": "base_model", "match": ["Flux.1 D"], "root": "/unet-nas"}
|
|
],
|
|
}
|
|
assert (
|
|
resolve_download_root(rules_by_type["checkpoints"], "SDXL 1.0", "/default")
|
|
== "/ckpt-nas"
|
|
)
|
|
assert (
|
|
resolve_download_root(rules_by_type["unet"], "SDXL 1.0", "/default")
|
|
== "/default"
|
|
)
|
|
assert (
|
|
resolve_download_root(rules_by_type["unet"], "Flux.1 D", "/default")
|
|
== "/unet-nas"
|
|
)
|
|
assert (
|
|
resolve_download_root(rules_by_type["checkpoints"], "Flux.1 D", "/default")
|
|
== "/default"
|
|
)
|