Optimize image loading and document node options

2026-05-28 12:19:58 +08:00
parent b51e6d22bb
commit c02daae0e3
4 changed files with 156 additions and 2 deletions
@@ -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。
@@ -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.
@@ -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)

@@ -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)