543 lines
18 KiB
Python
543 lines
18 KiB
Python
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from core import resolver as resolver_core
|
|
from core.contracts import ModelReference, Resolution
|
|
from core.workflow import analysis, references, subgraphs, traversal
|
|
from core.workflow_updater import update_workflow_nodes
|
|
|
|
|
|
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_workflow_node_traversal_preserves_node_order_and_scope():
|
|
top_level_node = {"id": 1, "type": "TopLevel"}
|
|
subgraph_node = {"id": 2, "type": "Inner"}
|
|
workflow = {
|
|
"nodes": [top_level_node, "invalid-top-level-node"],
|
|
"definitions": {
|
|
"subgraphs": [
|
|
None,
|
|
{
|
|
"id": "group-1",
|
|
"name": "Group one",
|
|
"nodes": [subgraph_node, "invalid-subgraph-node"],
|
|
},
|
|
]
|
|
},
|
|
}
|
|
|
|
contexts = list(traversal.iter_workflow_nodes_with_scope(workflow))
|
|
|
|
assert [context.node for context in contexts] == [
|
|
top_level_node,
|
|
"invalid-top-level-node",
|
|
subgraph_node,
|
|
"invalid-subgraph-node",
|
|
]
|
|
assert contexts[0].is_top_level is True
|
|
assert contexts[0].subgraph_id == ""
|
|
assert contexts[1].is_top_level is True
|
|
assert contexts[2].is_top_level is False
|
|
assert contexts[2].subgraph_id == "group-1"
|
|
assert contexts[2].subgraph_name == "Group one"
|
|
assert contexts[3].subgraph_id == "group-1"
|
|
|
|
|
|
def test_workflow_cache_helpers_share_scope_keys_and_normalized_hashing():
|
|
node = {"id": 2, "type": "Inner", "pos": [10, 20]}
|
|
|
|
assert analysis._get_workflow_node_cache_key(
|
|
node,
|
|
subgraph_id="group-1",
|
|
is_top_level=False,
|
|
) == analysis._get_workflow_scope_key(
|
|
2,
|
|
subgraph_id="group-1",
|
|
is_top_level=False,
|
|
)
|
|
assert analysis._get_workflow_node_cache_key(
|
|
node,
|
|
is_top_level=True,
|
|
) == ("top", "", "2")
|
|
assert analysis._get_workflow_node_fingerprint(node) == (
|
|
analysis._hash_normalized_workflow_value(node)
|
|
)
|
|
assert analysis._get_workflow_node_fingerprint(
|
|
{"type": "Inner", "id": 2, "pos": [999, 999]}
|
|
) == analysis._get_workflow_node_fingerprint(node)
|
|
|
|
|
|
def test_active_workflow_node_traversal_follows_nested_referenced_subgraphs_only():
|
|
workflow = {
|
|
"nodes": [
|
|
{"id": "top", "type": "subgraph-1", "widgets_values": []},
|
|
"invalid-top-level-node",
|
|
],
|
|
"definitions": {
|
|
"subgraphs": [
|
|
{
|
|
"id": "subgraph-1",
|
|
"name": "Active group",
|
|
"nodes": [
|
|
{"id": "inner", "type": "subgraph-2"},
|
|
"invalid-inner-node",
|
|
],
|
|
},
|
|
{
|
|
"id": "subgraph-2",
|
|
"name": "Nested active group",
|
|
"nodes": [{"id": "nested", "type": "NestedLoader"}],
|
|
},
|
|
{
|
|
"id": "unused",
|
|
"name": "Unused group",
|
|
"nodes": [{"id": "unused-node", "type": "UnusedLoader"}],
|
|
},
|
|
]
|
|
},
|
|
}
|
|
|
|
active_contexts = list(
|
|
traversal.iter_active_workflow_nodes_with_scope(workflow)
|
|
)
|
|
|
|
assert [context.node["id"] for context in active_contexts] == [
|
|
"top",
|
|
"inner",
|
|
"nested",
|
|
]
|
|
assert active_contexts[1].subgraph_id == "subgraph-1"
|
|
assert active_contexts[2].subgraph_id == "subgraph-2"
|
|
|
|
|
|
def test_extract_workflow_urls_includes_all_serialized_subgraph_nodes():
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": "top",
|
|
"type": "TopLoader",
|
|
"properties": {
|
|
"models": [
|
|
{
|
|
"name": "top.safetensors",
|
|
"url": "https://civitai.com/api/download/models/1",
|
|
"directory": "checkpoints",
|
|
}
|
|
]
|
|
},
|
|
}
|
|
],
|
|
"definitions": {
|
|
"subgraphs": [
|
|
{
|
|
"id": "unused",
|
|
"nodes": [
|
|
{
|
|
"id": "subgraph-node",
|
|
"type": "SubgraphLoader",
|
|
"properties": {
|
|
"models": [
|
|
{
|
|
"name": "subgraph.safetensors",
|
|
"url": "https://civitai.com/api/download/models/2",
|
|
"directory": "loras",
|
|
}
|
|
]
|
|
},
|
|
}
|
|
],
|
|
}
|
|
]
|
|
},
|
|
}
|
|
|
|
url_map = resolver_core.extract_workflow_urls(workflow)
|
|
|
|
assert url_map["top.safetensors"] == {
|
|
"url": "https://civitai.com/api/download/models/1",
|
|
"model_url": "https://civitai.com/api/download/models/1",
|
|
"directory": "checkpoints",
|
|
"node_type": "TopLoader",
|
|
"source": "node_properties",
|
|
}
|
|
assert url_map["subgraph.safetensors"] == {
|
|
"url": "https://civitai.com/api/download/models/2",
|
|
"model_url": "https://civitai.com/api/download/models/2",
|
|
"directory": "loras",
|
|
"node_type": "SubgraphLoader",
|
|
"source": "node_properties",
|
|
}
|
|
|
|
|
|
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 = ModelReference.from_mapping(
|
|
{"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",
|
|
}
|
|
|
|
reference = subgraphs._apply_instance_promoted_widget_context(
|
|
reference,
|
|
context,
|
|
available_models,
|
|
)
|
|
|
|
assert reference.extra_value("promoted_inner_node_id") == 10
|
|
assert reference.category == "checkpoints"
|
|
assert reference.extra_value("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 = ModelReference.from_mapping(
|
|
{
|
|
"subgraph_id": "subgraph-1",
|
|
"node_id": 10,
|
|
"widget_index": 99,
|
|
"widget_name": "ckpt_name",
|
|
"original_path": "second",
|
|
}
|
|
)
|
|
|
|
reference = subgraphs._apply_promoted_widget_locator(reference, contexts)
|
|
|
|
assert reference.extra_value("locate_node_id") == 2
|
|
assert reference.extra_value("locate_subgraph_name") == "Two"
|
|
assert reference.extra_value("locate_is_top_level") is True
|
|
assert reference.extra_value("locate_via_promoted_widget") is True
|
|
|
|
|
|
def test_subgraph_input_reference_targets_instance_for_update_and_location():
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 1,
|
|
"type": "subgraph-1",
|
|
"title": "Model subgraph",
|
|
"inputs": [
|
|
{
|
|
"name": "ckpt_name",
|
|
"type": "COMBO",
|
|
"widget": {"name": "ckpt_name"},
|
|
"link": None,
|
|
}
|
|
],
|
|
"outputs": [{"links": [20]}],
|
|
"properties": {"proxyWidgets": [["10", "ckpt_name"]]},
|
|
# Some serialized subgraph instances have no values here;
|
|
# the connected inner widget still contains the old value.
|
|
"widgets_values": [],
|
|
}
|
|
],
|
|
"definitions": {
|
|
"subgraphs": [
|
|
{
|
|
"id": "subgraph-1",
|
|
"name": "Model subgraph",
|
|
"inputs": [
|
|
{"name": "ckpt_name", "type": "COMBO", "linkIds": [7]}
|
|
],
|
|
"nodes": [
|
|
{
|
|
"id": 10,
|
|
"type": "CheckpointLoaderSimple",
|
|
"inputs": [
|
|
{
|
|
"name": "ckpt_name",
|
|
"type": "COMBO",
|
|
"widget": {"name": "ckpt_name"},
|
|
"link": 7,
|
|
}
|
|
],
|
|
"outputs": [{"type": "MODEL", "links": [20]}],
|
|
"widgets_values": ["inner-stale.safetensors"],
|
|
}
|
|
],
|
|
}
|
|
]
|
|
},
|
|
}
|
|
|
|
refs = analysis.analyze_workflow_models(workflow, available_models=[])
|
|
|
|
assert len(refs) == 1
|
|
assert refs[0].node_id == 1
|
|
assert refs[0].widget_index == 0
|
|
assert refs[0].subgraph_id == "subgraph-1"
|
|
assert refs[0].is_top_level is True
|
|
assert refs[0].extra_value("subgraph_path") is None
|
|
assert refs[0].extra_value("locate_node_id") == 1
|
|
assert refs[0].extra_value("locate_subgraph_id") == ""
|
|
assert refs[0].extra_value("promoted_inner_node_id") == 10
|
|
|
|
no_proxy_workflow = {
|
|
**workflow,
|
|
"nodes": [
|
|
{
|
|
**workflow["nodes"][0],
|
|
"properties": {},
|
|
}
|
|
],
|
|
}
|
|
no_proxy_refs = analysis.analyze_workflow_models(
|
|
no_proxy_workflow, available_models=[]
|
|
)
|
|
assert no_proxy_refs[0].node_id == 1
|
|
assert no_proxy_refs[0].extra_value("promoted_inner_node_id") == 10
|
|
|
|
resolution = Resolution.from_mapping(
|
|
{
|
|
"node_id": refs[0].node_id,
|
|
"widget_index": refs[0].widget_index,
|
|
"resolved_path": "replacement.safetensors",
|
|
"category": "checkpoints",
|
|
"subgraph_id": refs[0].subgraph_id,
|
|
"is_top_level": refs[0].is_top_level,
|
|
"promoted_widget_name": refs[0].extra_value("promoted_widget_name"),
|
|
}
|
|
)
|
|
update_workflow_nodes(workflow, [resolution])
|
|
|
|
assert workflow["nodes"][0]["widgets_values"] == ["replacement.safetensors"]
|
|
assert workflow["definitions"]["subgraphs"][0]["nodes"][0]["widgets_values"] == [
|
|
"inner-stale.safetensors"
|
|
]
|
|
|
|
|
|
def test_reused_subgraph_instance_refreshes_promoted_model_context():
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 1,
|
|
"type": "subgraph-1",
|
|
"inputs": [
|
|
{
|
|
"name": "ckpt_name",
|
|
"type": "COMBO",
|
|
"widget": {"name": "ckpt_name"},
|
|
"link": None,
|
|
}
|
|
],
|
|
"properties": {"proxyWidgets": [["10", "ckpt_name"]]},
|
|
"widgets_values": ["old.safetensors"],
|
|
}
|
|
],
|
|
"definitions": {
|
|
"subgraphs": [
|
|
{
|
|
"id": "subgraph-1",
|
|
"name": "Promoted loader",
|
|
"nodes": [_loader_node(value="old.safetensors")],
|
|
}
|
|
]
|
|
},
|
|
}
|
|
previous_node_cache = {}
|
|
analysis.analyze_workflow_models(
|
|
workflow,
|
|
available_models=[],
|
|
node_cache_out=previous_node_cache,
|
|
)
|
|
|
|
updated_workflow = {
|
|
**workflow,
|
|
"nodes": [
|
|
{
|
|
**workflow["nodes"][0],
|
|
"widgets_values": ["new.safetensors"],
|
|
}
|
|
],
|
|
}
|
|
# Simulate a serialized proxy widget change whose node cache entry was
|
|
# retained while the promoted context changed.
|
|
previous_node_cache[("top", "", "1")]["fingerprint"] = (
|
|
analysis._get_workflow_node_fingerprint(updated_workflow["nodes"][0])
|
|
)
|
|
|
|
refs = analysis.analyze_workflow_models(
|
|
updated_workflow,
|
|
available_models=[],
|
|
previous_node_cache=previous_node_cache,
|
|
)
|
|
|
|
assert len(refs) == 1
|
|
assert refs[0].node_id == 1
|
|
assert refs[0].original_path == "new.safetensors"
|
|
assert refs[0].extra_value("promoted_inner_node_id") == 10
|
|
|
|
|
|
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),
|
|
)
|