Skip to content

Commit da7d9b3

Browse files
committed
Skip non-model reference categories in workflow analyzer
Treat widget values under the configs and custom_nodes dynamic categories as non-model references so the workflow analyzer does not falsely mark them as missing models. Added a regression test to verify a checkpoint widget is still detected when a config widget is present.
1 parent 2dd5322 commit da7d9b3

2 files changed

Lines changed: 43 additions & 1 deletion

File tree

core/workflow_analyzer.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,7 @@
233233
_DYNAMIC_CATEGORY_SENTINEL_PREFIX = "__model_resolver_folder_category__"
234234
_DYNAMIC_NODE_WIDGET_CATEGORY_CACHE: Dict[str, Dict[str, Any]] = {}
235235
_DYNAMIC_NODE_WIDGET_CATEGORY_LOCK = threading.RLock()
236+
NON_MODEL_REFERENCE_CATEGORIES = {"custom_nodes", "configs"}
236237

237238
# These ComfyUI INPUT_TYPES entries become widgets in widgets_values. Typed graph
238239
# inputs like MODEL or CLIP are links, so they should not shift widget indexes.
@@ -1407,7 +1408,14 @@ def get_node_model_info(
14071408
# For each widget value, check if it looks like a model file or URN
14081409
for idx, value in enumerate(widgets_values):
14091410
widget_name = get_widget_name_hint(node, idx)
1410-
model_widget_folder_key_hints = get_dynamic_widget_category_hints(node, idx)
1411+
dynamic_category_hints = get_dynamic_widget_category_hints(node, idx)
1412+
if dynamic_category_hints and all(
1413+
normalize_download_category(category)
1414+
in NON_MODEL_REFERENCE_CATEGORIES
1415+
for category in dynamic_category_hints
1416+
):
1417+
continue
1418+
model_widget_folder_key_hints = dynamic_category_hints
14111419
model_widget_choice_info = get_dynamic_widget_choice_info(node, idx)
14121420
model_widget_category_hints = get_model_widget_category_hints(node, idx)
14131421
model_widget_category_hint = (

tests/test_workflow_analyzer.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -466,6 +466,40 @@ def test_diffusers_folder_models_are_scanned_as_folder_entries(self):
466466

467467

468468
class WorkflowCategoryHintsTests(unittest.TestCase):
469+
def test_non_model_config_widget_does_not_hide_checkpoint_widget(self):
470+
workflow = {
471+
"nodes": [
472+
{
473+
"id": 470,
474+
"type": "CheckpointLoader",
475+
"widgets": [
476+
{"name": "config_name"},
477+
{"name": "ckpt_name"},
478+
],
479+
"widgets_values": [
480+
"v1-inference.yaml",
481+
"model.safetensors",
482+
],
483+
"outputs": [{"type": "MODEL", "links": [1]}],
484+
}
485+
]
486+
}
487+
488+
def dynamic_category_hints(_node, widget_index):
489+
if widget_index == 0:
490+
return ["configs"]
491+
return ["checkpoints"]
492+
493+
with patch(
494+
"core.workflow_analyzer.get_dynamic_widget_category_hints",
495+
side_effect=dynamic_category_hints,
496+
):
497+
refs = analyze_workflow_models(workflow, available_models=[])
498+
499+
self.assertEqual(1, len(refs))
500+
self.assertEqual("model.safetensors", refs[0]["original_path"])
501+
self.assertEqual("checkpoints", refs[0]["category"])
502+
469503
def test_node_type_to_category_hints_is_populated(self):
470504
from core.workflow_analyzer import NODE_TYPE_TO_CATEGORY_HINTS
471505
# Verify standard loader mappings are correctly generated

0 commit comments

Comments
 (0)