217 lines
7.2 KiB
Python
217 lines
7.2 KiB
Python
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from core.workflow import references, subgraphs
|
|
|
|
|
|
def _loader_node(node_id=10, value="model.safetensors", link=None):
|
|
return {
|
|
"id": node_id,
|
|
"type": "CheckpointLoaderSimple",
|
|
"title": "Loader",
|
|
"widgets_values": [value],
|
|
"inputs": [
|
|
{
|
|
"name": "ckpt_name",
|
|
"widget": {"name": "ckpt_name"},
|
|
"link": link,
|
|
}
|
|
],
|
|
"outputs": [],
|
|
}
|
|
|
|
|
|
def test_subgraph_widget_indexes_support_named_inputs_and_links():
|
|
node = {
|
|
**_loader_node(),
|
|
"inputs": [
|
|
{"name": "model", "widget": None, "link": 1},
|
|
{"name": "ckpt_name", "widget": {"name": "ckpt_name"}},
|
|
],
|
|
"widgets_values": ["model.safetensors"],
|
|
}
|
|
|
|
assert subgraphs._get_widget_index_for_input_index(node, 0) is None
|
|
assert subgraphs._get_widget_index_for_input_index(node, 1) == 0
|
|
assert subgraphs._get_widget_index_by_name(node, "ckpt-name") == 0
|
|
assert subgraphs._get_node_by_id([node], "10") is node
|
|
|
|
subgraph = {
|
|
"inputs": [{"name": "ckpt_name", "linkIds": [7, "8"]}],
|
|
}
|
|
assert subgraphs._get_subgraph_input_link_ids(subgraph, "ckpt-name") == {
|
|
"7",
|
|
"8",
|
|
}
|
|
assert subgraphs._get_subgraph_input_link_ids(subgraph, "missing") == set()
|
|
|
|
|
|
def test_find_promoted_targets_prefers_explicit_node_then_linked_input():
|
|
explicit_node = _loader_node(node_id=10)
|
|
linked_node = _loader_node(node_id=11, link=7)
|
|
subgraph = {
|
|
"inputs": [{"name": "ckpt_name", "linkIds": [7]}],
|
|
"nodes": [explicit_node, linked_node],
|
|
}
|
|
|
|
with patch.object(
|
|
subgraphs.references,
|
|
"get_effective_model_category_hint",
|
|
return_value="checkpoints",
|
|
):
|
|
explicit = subgraphs._find_promoted_widget_targets(
|
|
subgraph, "10", "ckpt_name"
|
|
)
|
|
linked = subgraphs._find_promoted_widget_targets(
|
|
subgraph, "-1", "ckpt_name"
|
|
)
|
|
|
|
assert [(item["node_id"], item["widget_index"]) for item in explicit] == [(10, 0)]
|
|
assert [(item["node_id"], item["widget_index"]) for item in linked] == [(11, 0)]
|
|
assert linked[0]["category"] == "checkpoints"
|
|
|
|
|
|
def test_build_promoted_contexts_records_instance_and_inner_locators():
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 1,
|
|
"type": "subgraph-1",
|
|
"widgets_values": ["promoted.safetensors"],
|
|
"properties": {"proxyWidgets": [["10", "ckpt_name"]]},
|
|
}
|
|
]
|
|
}
|
|
subgraph = {
|
|
"id": "subgraph-1",
|
|
"name": "Promoted loader",
|
|
"nodes": [_loader_node()],
|
|
}
|
|
|
|
with patch.object(
|
|
subgraphs.references,
|
|
"get_effective_model_category_hint",
|
|
return_value="checkpoints",
|
|
):
|
|
contexts = subgraphs._build_promoted_widget_contexts(workflow, [subgraph])
|
|
|
|
instance = contexts["instance_widgets"][("1", 0)]
|
|
assert instance["subgraph_name"] == "Promoted loader"
|
|
assert instance["promoted_value"] == "promoted.safetensors"
|
|
assert contexts["inner_widgets"][("subgraph-1", "10", 0)][0]["node_id"] == 1
|
|
assert contexts["inner_widget_names"][("subgraph-1", "10", "ckpt_name")]
|
|
|
|
|
|
def test_apply_instance_promoted_context_resolves_existing_model(tmp_path):
|
|
model_path = tmp_path / "promoted.safetensors"
|
|
model_path.write_bytes(b"model")
|
|
available_models = [
|
|
{
|
|
"category": "checkpoints",
|
|
"relative_path": "promoted.safetensors",
|
|
"path": str(model_path),
|
|
}
|
|
]
|
|
reference = {"original_path": "promoted.safetensors", "exists": False}
|
|
context = {
|
|
"proxy_widget_name": "ckpt_name",
|
|
"node_id": 10,
|
|
"node_type": "CheckpointLoaderSimple",
|
|
"node_title": "Loader",
|
|
"widget_index": 0,
|
|
"widget_name": "ckpt_name",
|
|
"category": "checkpoints",
|
|
}
|
|
|
|
subgraphs._apply_instance_promoted_widget_context(
|
|
reference,
|
|
context,
|
|
available_models,
|
|
)
|
|
|
|
assert reference["promoted_inner_node_id"] == 10
|
|
assert reference["category"] == "checkpoints"
|
|
assert reference["full_path"] == str(model_path)
|
|
assert reference["exists"] is True
|
|
|
|
|
|
def test_apply_promoted_locator_supports_name_fallback_and_value_selection():
|
|
contexts = {
|
|
"inner_widgets": {},
|
|
"inner_widget_names": {
|
|
("subgraph-1", "10", "ckpt_name"): [
|
|
{"node_id": 1, "node_type": "Subgraph", "subgraph_name": "One", "promoted_value": "first"},
|
|
{"node_id": 2, "node_type": "Subgraph", "subgraph_name": "Two", "promoted_value": "second"},
|
|
]
|
|
},
|
|
}
|
|
reference = {
|
|
"subgraph_id": "subgraph-1",
|
|
"node_id": 10,
|
|
"widget_index": 99,
|
|
"widget_name": "ckpt_name",
|
|
"original_path": "second",
|
|
}
|
|
|
|
subgraphs._apply_promoted_widget_locator(reference, contexts)
|
|
|
|
assert reference["locate_node_id"] == 2
|
|
assert reference["locate_subgraph_name"] == "Two"
|
|
assert reference["locate_is_top_level"] is True
|
|
assert reference["locate_via_promoted_widget"] is True
|
|
|
|
|
|
def test_reference_matching_handles_placeholders_paths_and_scanner_records(tmp_path):
|
|
model_path = tmp_path / "folder" / "model.safetensors"
|
|
model_path.parent.mkdir()
|
|
model_path.write_bytes(b"model")
|
|
available_models = [
|
|
{
|
|
"category": "checkpoints",
|
|
"relative_path": "folder/model.safetensors",
|
|
"path": str(model_path),
|
|
}
|
|
]
|
|
|
|
assert references.is_model_filename("model.safetensors") is True
|
|
assert references.is_model_filename("urn:air:sd:checkpoint:civitai:1@2") is True
|
|
assert references.should_scan_as_model_reference("None", True) is False
|
|
assert references.should_scan_as_model_reference("plain-name", True) is True
|
|
assert references.should_scan_as_model_reference("plain-name", False) is False
|
|
assert references.is_static_or_hybrid_widget_choice(
|
|
"Folder\\Model.safetensors",
|
|
{"source": "static", "choices": ["folder/model.safetensors"]},
|
|
)
|
|
assert references.static_or_hybrid_choice_looks_like_model(
|
|
"folder/model.safetensors",
|
|
{"source": "hybrid", "choices": []},
|
|
)
|
|
|
|
resolved = references.try_resolve_model_path(
|
|
"folder/model.safetensors",
|
|
["checkpoints", "configs"],
|
|
available_models=available_models,
|
|
)
|
|
assert resolved == ("checkpoints", str(model_path))
|
|
assert references.try_resolve_model_path(
|
|
"folder/MODEL.safetensors",
|
|
["checkpoints"],
|
|
available_models=available_models,
|
|
) is None
|
|
|
|
|
|
def test_reference_resolution_uses_folder_paths_when_scanner_is_not_available(tmp_path):
|
|
model_path = tmp_path / "model.safetensors"
|
|
model_path.write_bytes(b"model")
|
|
fake_folder_paths = SimpleNamespace(
|
|
folder_names_and_paths={"checkpoints": ([], set()), "configs": ([], set())},
|
|
get_filename_list=lambda category: ["model.safetensors"],
|
|
get_full_path=lambda category, filename: str(model_path),
|
|
)
|
|
|
|
with patch.object(references, "folder_paths", fake_folder_paths):
|
|
assert references.try_resolve_model_path("model.safetensors") == (
|
|
"checkpoints",
|
|
str(model_path),
|
|
)
|