diff --git a/README.md b/README.md index 97497e5..fd8d2aa 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,13 @@ pip install -r requirements.txt 2. 从 ComfyUI `input` 目录选择图片,或通过图片控件上传图片。 3. 使用节点输出连接后续节点,或直接查看输出文本。 +## 节点参数说明 + +- `parameter_index`:用于 A1111 / Forge 多组参数 fallback 时选择第几组参数。普通使用保持 `0`。 +- `prefer_cached_text`:优先使用 `ShowText` 等节点保存的缓存结果,避免把 Gemini / OpenAI 的 `system_instruction`、`prompt template` 当成最终提示词。普通使用保持 `true`。 +- `include_raw_json`:是否在内部解析结果中保留原始 `prompt` / `workflow` JSON。普通使用保持 `true`。 +- `max_depth`:graph 反向追踪最大深度,用于防止循环。普通使用保持 `40`;极复杂工作流可以适当加大。 + ## 输出端口 主节点输出: @@ -124,6 +131,8 @@ python -m pytest -q python -m compileall . ``` +`tests/` 目录是开发和回归测试用的。普通 ComfyUI 用户不需要手动运行这些测试;仓库保留它们是为了在后续兼容更多 custom nodes 时,防止破坏已有功能。 + ## 手动 ComfyUI 验证 1. 重启 ComfyUI。 diff --git a/README_EN.md b/README_EN.md index 4f8b25d..79a6e0f 100644 --- a/README_EN.md +++ b/README_EN.md @@ -45,6 +45,13 @@ Restart ComfyUI after installing. 2. Select an image from the ComfyUI `input` directory, or upload one through the image widget. 3. Use the node outputs in your workflow, or inspect the text outputs directly. +## Node Options + +- `parameter_index`: Selects which parameter block to use when falling back to multi-parameter A1111 / Forge metadata. Keep this at `0` for normal use. +- `prefer_cached_text`: Prefer cached outputs from nodes such as `ShowText`, which helps avoid treating Gemini / OpenAI `system_instruction` or prompt templates as final prompts. Keep this `true` for normal use. +- `include_raw_json`: Keeps the original `prompt` / `workflow` JSON in the internal parse result. Keep this `true` for normal use. +- `max_depth`: Maximum reverse graph traversal depth, used to prevent cycles. Keep this at `40` for normal use; increase it only for very complex workflows. + ## Outputs The main node outputs: @@ -124,6 +131,8 @@ python -m pytest -q python -m compileall . ``` +The `tests/` directory is for development and regression testing. Regular ComfyUI users do not need to run it manually; it is kept in the repository to prevent future custom node compatibility work from breaking existing behavior. + ## Manual ComfyUI Verification 1. Restart ComfyUI. diff --git a/smart_metadata_reader/image_io.py b/smart_metadata_reader/image_io.py index e336f47..a1958b6 100644 --- a/smart_metadata_reader/image_io.py +++ b/smart_metadata_reader/image_io.py @@ -20,6 +20,14 @@ def _torch_module() -> Any | None: return torch +def _numpy_module() -> Any | None: + try: + import numpy + except ImportError: + return None + return numpy + + def _image_tensor(image: Image.Image) -> Any: torch = _torch_module() width, height = image.size @@ -27,7 +35,12 @@ def _image_tensor(image: Image.Image) -> Any: return ShapeOnlyTensor((1, height, width, 3)) rgb_image = image.convert("RGB") - data = list(rgb_image.getdata()) + numpy = _numpy_module() + if numpy is not None and hasattr(torch, "from_numpy"): + array = numpy.asarray(rgb_image, dtype=numpy.float32) + return torch.from_numpy(array.reshape(1, height, width, 3)) / 255.0 + + data = list(rgb_image.tobytes()) return torch.tensor(data, dtype=torch.float32).reshape(1, height, width, 3) / 255.0 @@ -37,10 +50,17 @@ def _mask_tensor(image: Image.Image) -> Any: if torch is None: return ShapeOnlyTensor((1, height, width)) + numpy = _numpy_module() if "A" in image.getbands(): alpha = image.getchannel("A") - mask_data = [1.0 - (value / 255.0) for value in alpha.getdata()] + if numpy is not None and hasattr(torch, "from_numpy"): + alpha_array = numpy.asarray(alpha, dtype=numpy.float32) + mask_array = (1.0 - (alpha_array / 255.0)).reshape(1, height, width) + return torch.from_numpy(mask_array) + mask_data = [1.0 - (value / 255.0) for value in alpha.tobytes()] else: + if numpy is not None and hasattr(torch, "from_numpy"): + return torch.from_numpy(numpy.zeros((1, height, width), dtype=numpy.float32)) mask_data = [0.0] * (width * height) return torch.tensor(mask_data, dtype=torch.float32).reshape(1, height, width) diff --git a/tests/test_image_io.py b/tests/test_image_io.py new file mode 100644 index 0000000..ccea076 --- /dev/null +++ b/tests/test_image_io.py @@ -0,0 +1,116 @@ +import numpy as np +from PIL import Image + +from smart_metadata_reader import image_io + + +class FakeTensor: + def __init__(self, array): + self.array = np.array(array) + + @property + def shape(self): + return self.array.shape + + def reshape(self, *shape): + return FakeTensor(self.array.reshape(*shape)) + + def float(self): + return FakeTensor(self.array.astype(np.float32)) + + def __truediv__(self, value): + return FakeTensor(self.array / value) + + def tolist(self): + return self.array.tolist() + + +class FakeTorch: + float32 = np.float32 + + @staticmethod + def tensor(data, dtype=None): + return FakeTensor(np.array(data, dtype=dtype)) + + @staticmethod + def from_numpy(array): + return FakeTensor(array) + + +def use_fake_torch(monkeypatch): + monkeypatch.setattr(image_io, "_torch_module", lambda: FakeTorch) + + +def test_rgb_png_outputs_image_and_empty_mask(tmp_path, monkeypatch): + use_fake_torch(monkeypatch) + path = tmp_path / "rgb.png" + Image.new("RGB", (2, 1), (128, 64, 32)).save(path) + + image, mask, width, height = image_io.load_image_and_mask(path) + + assert (width, height) == (2, 1) + assert image.shape == (1, 1, 2, 3) + assert mask.shape == (1, 1, 2) + assert np.allclose(mask.array, 0.0) + assert np.allclose(image.array[0, 0, 0], [128 / 255.0, 64 / 255.0, 32 / 255.0]) + + +def test_rgba_png_outputs_alpha_mask(tmp_path, monkeypatch): + use_fake_torch(monkeypatch) + path = tmp_path / "rgba.png" + image = Image.new("RGBA", (3, 1)) + image.putdata([(255, 0, 0, 255), (0, 255, 0, 0), (0, 0, 255, 128)]) + image.save(path) + + loaded, mask, width, height = image_io.load_image_and_mask(path) + + assert (width, height) == (3, 1) + assert loaded.shape == (1, 1, 3, 3) + assert mask.shape == (1, 1, 3) + assert np.allclose(mask.array[0, 0], [0.0, 1.0, 1.0 - (128 / 255.0)]) + + +def test_large_image_conversion_keeps_expected_shape(tmp_path, monkeypatch): + use_fake_torch(monkeypatch) + path = tmp_path / "large.png" + Image.new("RGB", (640, 512), (10, 20, 30)).save(path) + + image, mask, width, height = image_io.load_image_and_mask(path) + + assert (width, height) == (640, 512) + assert image.shape == (1, 512, 640, 3) + assert mask.shape == (1, 512, 640) + + +def test_numpy_path_does_not_use_getdata(tmp_path, monkeypatch): + use_fake_torch(monkeypatch) + path = tmp_path / "numpy.png" + Image.new("RGBA", (2, 2), (10, 20, 30, 128)).save(path) + + def broken_getdata(self): + raise AssertionError("numpy path should not call getdata") + + monkeypatch.setattr(Image.Image, "getdata", broken_getdata) + + image, mask, width, height = image_io.load_image_and_mask(path) + + assert (width, height) == (2, 2) + assert image.shape == (1, 2, 2, 3) + assert mask.shape == (1, 2, 2) + + +def test_numpy_unavailable_fallback_shape_matches_numpy_path(tmp_path, monkeypatch): + path = tmp_path / "fallback.png" + Image.new("RGBA", (4, 3), (10, 20, 30, 128)).save(path) + + use_fake_torch(monkeypatch) + numpy_image, numpy_mask, width, height = image_io.load_image_and_mask(path) + + monkeypatch.setattr(image_io, "_numpy_module", lambda: None, raising=False) + fallback_image, fallback_mask, fallback_width, fallback_height = image_io.load_image_and_mask(path) + + assert (width, height) == (fallback_width, fallback_height) + assert fallback_image.shape == numpy_image.shape + assert fallback_mask.shape == numpy_mask.shape + assert np.allclose(fallback_image.array, numpy_image.array) + assert np.allclose(fallback_mask.array, numpy_mask.array)