Improve workflow analyzer widget category hints and test coverage

This commit is contained in:
Dariusz L
2026-07-24 14:43:18 +02:00
parent 449c2c8c68
commit 03c545e297
2 changed files with 229 additions and 13 deletions
+43 -13
View File
@@ -325,10 +325,10 @@ def get_widget_name_candidates(node: Dict[str, Any], widget_index: int) -> List[
widget_inputs = _get_widget_inputs(node)
if widget_index < len(widget_inputs):
candidates.extend(_widget_item_name_candidates(widget_inputs[widget_index]))
inputs = node.get("inputs", [])
if isinstance(inputs, list) and widget_index < len(inputs):
candidates.extend(_widget_item_name_candidates(inputs[widget_index]))
elif not widget_inputs:
inputs = node.get("inputs", [])
if isinstance(inputs, list) and widget_index < len(inputs):
candidates.extend(_widget_item_name_candidates(inputs[widget_index]))
return unique_ordered_strings(candidates)
@@ -400,6 +400,7 @@ def _summarize_choice_info_for_log(info_by_key: Dict[Any, Dict[str, Any]]) -> Di
def _summarize_dynamic_hints_for_log(hints: Dict[str, Any]) -> Dict[str, Any]:
return {
"widget_names": hints.get("widget_names", []),
"by_name": hints.get("by_name", {}),
"by_index": hints.get("by_index", {}),
"choice_info_by_name": _summarize_choice_info_for_log(
@@ -604,6 +605,7 @@ def _iter_widget_input_type_entries(
def _empty_dynamic_hints() -> Dict[str, Any]:
return {
"widget_names": [],
"by_name": {},
"by_index": {},
"choice_info_by_name": {},
@@ -680,6 +682,9 @@ def _merge_dynamic_hint_entry(
choice_info: Dict[str, Any],
) -> None:
if normalized_name:
hints["widget_names"] = _ordered_unique_categories(
[*list(hints.get("widget_names") or []), normalized_name]
)
_merge_category_hints(hints["by_name"], normalized_name, categories)
if choice_info:
_merge_choice_info(
@@ -717,7 +722,10 @@ def _build_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]:
def traced_get_filename_list(category: Any, *args: Any, **kwargs: Any) -> List[str]:
category_name = str(category or "").strip()
sentinel = f"{_DYNAMIC_CATEGORY_SENTINEL_PREFIX}{len(sentinel_to_category)}"
sentinel = (
f"{_DYNAMIC_CATEGORY_SENTINEL_PREFIX}"
f"{len(sentinel_to_category)}.safetensors"
)
sentinel_to_category[sentinel] = category_name
return [sentinel]
@@ -813,7 +821,7 @@ def _build_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]:
def get_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]:
if not node_type:
return {"by_name": {}, "by_index": {}}
return _empty_dynamic_hints()
with _DYNAMIC_NODE_WIDGET_CATEGORY_LOCK:
cached = _DYNAMIC_NODE_WIDGET_CATEGORY_CACHE.get(node_type)
@@ -837,9 +845,17 @@ def get_dynamic_widget_category_hints(
categories: List[str] = []
by_name = hints.get("by_name", {})
normalized_candidates = [
normalize_widget_name(candidate)
for candidate in get_widget_name_candidates(node, widget_index)
]
if isinstance(by_name, dict):
for candidate in get_widget_name_candidates(node, widget_index):
categories.extend(by_name.get(normalize_widget_name(candidate), []))
for candidate in normalized_candidates:
categories.extend(by_name.get(candidate, []))
known_widget_names = set(hints.get("widget_names") or [])
if any(candidate in known_widget_names for candidate in normalized_candidates):
return _ordered_unique_categories(categories)
by_index = hints.get("by_index", {})
if isinstance(by_index, dict):
@@ -891,11 +907,19 @@ def get_dynamic_widget_choice_info(
info: Dict[str, Any] = {"source": "unknown", "choices": []}
by_name = hints.get("choice_info_by_name", {})
normalized_candidates = [
normalize_widget_name(candidate)
for candidate in get_widget_name_candidates(node, widget_index)
]
if isinstance(by_name, dict):
for candidate in get_widget_name_candidates(node, widget_index):
candidate_info = by_name.get(normalize_widget_name(candidate), {})
for candidate in normalized_candidates:
candidate_info = by_name.get(candidate, {})
info = _merge_widget_choice_info_values(info, candidate_info)
known_widget_names = set(hints.get("widget_names") or [])
if any(candidate in known_widget_names for candidate in normalized_candidates):
return info
by_index = hints.get("choice_info_by_index", {})
if isinstance(by_index, dict):
info = _merge_widget_choice_info_values(info, by_index.get(widget_index, {}))
@@ -942,6 +966,9 @@ def static_or_hybrid_choice_looks_like_model(
MODEL_WIDGET_PLACEHOLDERS = {
"",
".none",
".use_ckpt_clip",
".use_ckpt_vae",
"none",
"[none]",
"null",
@@ -986,11 +1013,12 @@ def get_node_model_widget_category_hint(
def get_node_output_category_hint(node: Dict[str, Any]) -> Optional[str]:
"""Return a model category hint from strongly typed loader outputs."""
"""Return a category only when strongly typed model outputs are unambiguous."""
outputs = node.get("outputs", [])
if not isinstance(outputs, list):
return None
categories: List[str] = []
for output in outputs:
if not isinstance(output, dict):
continue
@@ -998,8 +1026,10 @@ def get_node_output_category_hint(node: Dict[str, Any]) -> Optional[str]:
token = str(output.get(key, "") or "").strip().upper()
category = MODEL_OUTPUT_TYPE_TO_CATEGORY.get(token)
if category:
return category
return None
categories.append(category)
unique_categories = _ordered_unique_categories(categories)
return unique_categories[0] if len(unique_categories) == 1 else None
def get_model_widget_category_hint(
+186
View File
@@ -1,8 +1,10 @@
import os
import tempfile
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from core import workflow_analyzer
from core.scanner import scan_directory
from core.workflow_analyzer import analyze_workflow_models, identify_missing_models
@@ -81,6 +83,190 @@ class WorkflowAnalyzerCaseSensitivityTests(unittest.TestCase):
class WorkflowAnalyzerCategoryHintTests(unittest.TestCase):
def test_named_widget_does_not_inherit_conflicting_index_category(self):
cases = [
("easy fullLoader", 7, "resolution", "512 x 512"),
("ttN pipeLoader_v2", 8, "negative", "Negative"),
("WanVideoVAELoader", 1, "precision", "bf16"),
]
for node_type, widget_index, widget_name, widget_value in cases:
with self.subTest(node_type=node_type, widget_name=widget_name):
dynamic_hints = {
"widget_names": [widget_name, "vae_name"],
"by_name": {"vae_name": ["vae"]},
"by_index": {widget_index: ["vae"]},
"choice_info_by_name": {},
"choice_info_by_index": {
widget_index: {"source": "folder_paths", "choices": []}
},
}
widget_inputs = [
{
"name": f"widget_{index}",
"type": "STRING",
"widget": {"name": f"widget_{index}"},
}
for index in range(widget_index)
]
widget_inputs.append(
{
"name": widget_name,
"type": "STRING",
"widget": {"name": widget_name},
}
)
workflow = {
"nodes": [
{
"id": 500,
"type": node_type,
"inputs": [
{"name": "model_override", "type": "MODEL"},
*widget_inputs,
],
"widgets_values": [
*([""] * widget_index),
widget_value,
],
"outputs": [],
}
]
}
with patch(
"core.workflow_analyzer.get_dynamic_node_widget_category_hints",
return_value=dynamic_hints,
):
refs = analyze_workflow_models(workflow, available_models=[])
self.assertEqual([], refs)
def test_dynamic_sentinel_survives_model_extension_filter(self):
fake_folder_paths = SimpleNamespace()
def get_filename_list(category):
return [f"{category}.safetensors"]
fake_folder_paths.get_filename_list = get_filename_list
class FilteredMultiCategoryLoader:
@classmethod
def INPUT_TYPES(cls):
choices = (
fake_folder_paths.get_filename_list("checkpoints")
+ fake_folder_paths.get_filename_list("diffusion_models")
)
return {
"required": {
"model_name": (
[
choice
for choice in choices
if choice.endswith((".ckpt", ".safetensors"))
],
)
}
}
with (
patch.object(
workflow_analyzer,
"_get_comfy_node_class",
return_value=FilteredMultiCategoryLoader,
),
patch.object(
workflow_analyzer,
"_get_folder_paths_module",
return_value=fake_folder_paths,
),
):
hints = workflow_analyzer._build_dynamic_node_widget_category_hints(
"FilteredMultiCategoryLoader"
)
self.assertEqual(
["checkpoints", "diffusion_models"],
hints["by_name"]["model_name"],
)
def test_res4lyf_placeholders_are_not_model_references(self):
category_by_index = {
0: ["checkpoints", "diffusion_models"],
1: [],
2: ["text_encoders"],
3: ["text_encoders"],
4: ["vae"],
5: ["clip_vision"],
6: ["style_models"],
}
workflow = {
"nodes": [
{
"id": 501,
"type": "FluxLoader",
"widgets": [
{"name": "model_name"},
{"name": "weight_dtype"},
{"name": "clip_name1"},
{"name": "clip_name2_opt"},
{"name": "vae_name"},
{"name": "clip_vision_name"},
{"name": "style_model_name"},
],
"widgets_values": [
".none",
"default",
".use_ckpt_clip",
".none",
".use_ckpt_vae",
".none",
".none",
],
"outputs": [
{"type": "MODEL", "links": []},
{"type": "CLIP", "links": []},
{"type": "VAE", "links": []},
{"type": "CLIP_VISION", "links": []},
{"type": "STYLE_MODEL", "links": []},
],
}
]
}
with patch(
"core.workflow_analyzer.get_dynamic_widget_category_hints",
side_effect=lambda _node, index: category_by_index[index],
):
refs = analyze_workflow_models(workflow, available_models=[])
self.assertEqual([], refs)
def test_ambiguous_model_outputs_do_not_override_widget_name_category(self):
workflow = {
"nodes": [
{
"id": 502,
"type": "FluxLoader",
"widgets": [{"name": "model_name"}],
"widgets_values": ["flux.safetensors"],
"outputs": [
{"type": "CLIP_VISION", "links": []},
{"type": "STYLE_MODEL", "links": []},
],
}
]
}
with patch(
"core.workflow_analyzer.get_dynamic_widget_category_hints",
return_value=[],
):
refs = analyze_workflow_models(workflow, available_models=[])
self.assertEqual(1, len(refs))
self.assertEqual("diffusion_models", refs[0]["category"])
def test_show_anything_cached_model_text_is_not_a_model_reference(self):
with tempfile.TemporaryDirectory() as tmpdir:
model_dir = os.path.join(tmpdir, "KREA2")