From f82ac946faaad75874fa7381201d7d755ba755ae Mon Sep 17 00:00:00 2001 From: CyberShadowX Date: Thu, 28 May 2026 11:19:41 +0800 Subject: [PATCH] Fix input image listing for ComfyUI loader --- nodes.py | 60 +++++++++++++++------ tests/test_node_contract.py | 101 ++++++++++++++++++++++++++++++++++-- 2 files changed, 142 insertions(+), 19 deletions(-) diff --git a/nodes.py b/nodes.py index b257485..13da34b 100644 --- a/nodes.py +++ b/nodes.py @@ -18,7 +18,7 @@ except ImportError: folder_paths = None # type: ignore -IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tif", ".tiff"} +IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"} class SmartMetadataReader: @@ -68,18 +68,21 @@ class SmartMetadataReader: @classmethod def IS_CHANGED(cls, image: str, **kwargs: Any) -> str: del kwargs - image_path = _resolve_image_path(image) + image_path = _resolve_image_path_safely(image) return _file_digest(image_path) @classmethod def VALIDATE_INPUTS(cls, image: str, **kwargs: Any) -> bool | str: del kwargs if folder_paths is not None and hasattr(folder_paths, "exists_annotated_filepath"): - if folder_paths.exists_annotated_filepath(image): - return True - return f"Invalid image file: {image}" + try: + if folder_paths.exists_annotated_filepath(image): + return True + return f"Invalid image file: {image}" + except Exception: + pass - image_path = _resolve_image_path(image) + image_path = _resolve_image_path_safely(image) if Path(image_path).exists(): return True return f"Invalid image file: {image_path}" @@ -139,18 +142,29 @@ class SmartMetadataReader: def _input_images() -> list[str]: - if folder_paths is not None and hasattr(folder_paths, "get_filename_list"): - return list(folder_paths.get_filename_list("input")) if folder_paths is None or not hasattr(folder_paths, "get_input_directory"): return [] - input_dir = Path(folder_paths.get_input_directory()) - if not input_dir.exists(): + + try: + input_dir_value = folder_paths.get_input_directory() + except Exception: return [] - return sorted( - entry.name - for entry in input_dir.iterdir() - if entry.is_file() and entry.suffix.lower() in IMAGE_EXTENSIONS - ) + if not input_dir_value: + return [] + + input_dir = Path(input_dir_value) + if not input_dir.is_dir(): + return [] + + images: list[str] = [] + for root, _dirs, files in os.walk(input_dir): + for filename in files: + if Path(filename).suffix.lower() not in IMAGE_EXTENSIONS: + continue + full_path = Path(root) / filename + relative_path = os.path.relpath(full_path, input_dir) + images.append(relative_path.replace(os.sep, "/")) + return sorted(images) def _resolve_image_path(image: str) -> str: @@ -161,6 +175,22 @@ def _resolve_image_path(image: str) -> str: return os.fspath(image) +def _resolve_image_path_safely(image: str) -> str: + try: + return _resolve_image_path(image) + except Exception: + pass + + if folder_paths is not None and hasattr(folder_paths, "get_input_directory"): + try: + input_dir = folder_paths.get_input_directory() + except Exception: + input_dir = None + if input_dir: + return str(Path(input_dir) / image) + return os.fspath(image) + + def _load_image_or_raise(image_path: str) -> tuple[Any, Any, int, int]: try: return load_image_and_mask(image_path) diff --git a/tests/test_node_contract.py b/tests/test_node_contract.py index aadc3ff..02f3e64 100644 --- a/tests/test_node_contract.py +++ b/tests/test_node_contract.py @@ -229,25 +229,48 @@ def test_plugin_package_import_works_without_repo_root_on_sys_path(tmp_path): assert "package import ok" in completed.stdout -def test_input_types_prefers_comfyui_filename_list(monkeypatch): +def test_input_types_ignores_unusable_comfyui_filename_list_and_scans_input_dir( + tmp_path, + monkeypatch, +): import nodes + (tmp_path / "root.png").write_bytes(b"fake") + subdir = tmp_path / "sub" + subdir.mkdir() + (subdir / "a.jpg").write_bytes(b"fake") + (subdir / "b.webp").write_bytes(b"fake") + (subdir / "note.txt").write_text("ignore me") + + def broken_filename_list(category): + raise KeyError(category) + monkeypatch.setattr( nodes, "folder_paths", - SimpleNamespace(get_filename_list=lambda category: ["z.png", "nested/a.webp"]), + SimpleNamespace( + get_filename_list=broken_filename_list, + get_input_directory=lambda: str(tmp_path), + ), ) input_types = nodes.SmartMetadataReader.INPUT_TYPES() - assert input_types["required"]["image"][0] == ["z.png", "nested/a.webp"] + assert input_types["required"]["image"][0] == [ + "root.png", + "sub/a.jpg", + "sub/b.webp", + ] def test_smart_metadata_reader_input_types_include_uploadable_image(tmp_path, monkeypatch): import nodes (tmp_path / "a.png").write_bytes(b"fake") + (tmp_path / "c.jpg").write_bytes(b"fake") + (tmp_path / "d.jpeg").write_bytes(b"fake") (tmp_path / "b.webp").write_bytes(b"fake") + (tmp_path / "ignored.bmp").write_bytes(b"fake") (tmp_path / "note.txt").write_text("ignore me") monkeypatch.setattr( @@ -259,7 +282,7 @@ def test_smart_metadata_reader_input_types_include_uploadable_image(tmp_path, mo input_types = nodes.SmartMetadataReader.INPUT_TYPES() required = input_types["required"] - assert required["image"][0] == ["a.png", "b.webp"] + assert required["image"][0] == ["a.png", "b.webp", "c.jpg", "d.jpeg"] assert required["image"][1]["image_upload"] is True assert required["parameter_index"][0] == "INT" assert required["parameter_index"][1]["default"] == 0 @@ -272,6 +295,39 @@ def test_smart_metadata_reader_input_types_include_uploadable_image(tmp_path, mo assert required["max_depth"][1]["max"] == 200 +def test_smart_metadata_reader_input_types_empty_or_missing_input_dir_is_safe( + tmp_path, + monkeypatch, +): + import nodes + + monkeypatch.setattr( + nodes, + "folder_paths", + SimpleNamespace(get_input_directory=lambda: str(tmp_path)), + ) + assert nodes.SmartMetadataReader.INPUT_TYPES()["required"]["image"][0] == [] + + monkeypatch.setattr( + nodes, + "folder_paths", + SimpleNamespace(get_input_directory=lambda: str(tmp_path / "missing")), + ) + assert nodes.SmartMetadataReader.INPUT_TYPES()["required"]["image"][0] == [] + + +def test_smart_metadata_reader_input_types_handles_input_directory_errors(monkeypatch): + import nodes + + monkeypatch.setattr( + nodes, + "folder_paths", + SimpleNamespace(get_input_directory=lambda: (_ for _ in ()).throw(KeyError("input"))), + ) + + assert nodes.SmartMetadataReader.INPUT_TYPES()["required"]["image"][0] == [] + + def test_smart_metadata_reader_return_contract(): import nodes @@ -477,6 +533,21 @@ def test_smart_metadata_reader_is_changed_tracks_file_content(tmp_path, monkeypa assert first != second +def test_smart_metadata_reader_is_changed_handles_missing_or_bad_paths(monkeypatch): + import nodes + + monkeypatch.setattr( + nodes, + "folder_paths", + SimpleNamespace(get_annotated_filepath=lambda image: (_ for _ in ()).throw(KeyError("input"))), + ) + + value = nodes.SmartMetadataReader.IS_CHANGED("missing.png") + + assert isinstance(value, str) + assert value.startswith("missing:") + + def test_smart_metadata_reader_validate_inputs_checks_image_exists(tmp_path, monkeypatch): import nodes @@ -496,6 +567,28 @@ def test_smart_metadata_reader_validate_inputs_checks_image_exists(tmp_path, mon assert str(missing) in invalid +def test_smart_metadata_reader_validate_inputs_falls_back_when_exists_helper_errors( + tmp_path, + monkeypatch, +): + import nodes + + existing = tmp_path / "exists.png" + existing.write_bytes(b"fake") + + monkeypatch.setattr( + nodes, + "folder_paths", + SimpleNamespace( + exists_annotated_filepath=lambda image: (_ for _ in ()).throw(KeyError("input")), + get_input_directory=lambda: str(tmp_path), + ), + ) + + assert nodes.SmartMetadataReader.VALIDATE_INPUTS("exists.png") is True + assert "Invalid image file" in nodes.SmartMetadataReader.VALIDATE_INPUTS("missing.png") + + def test_smart_metadata_reader_validate_inputs_uses_comfyui_exists_helper(monkeypatch): import nodes