diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 10cf9e2..21d652d 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -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( diff --git a/tests/test_workflow_analyzer.py b/tests/test_workflow_analyzer.py index 076ec76..7dfd3ae 100644 --- a/tests/test_workflow_analyzer.py +++ b/tests/test_workflow_analyzer.py @@ -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")