Improve workflow analyzer widget category hints and test coverage
This commit is contained in:
+43
-13
@@ -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(
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user