From 23dd271e39c4a99e3e9c01eabfb8ca5f181ca92e Mon Sep 17 00:00:00 2001 From: Gin3601 <2405559452@qq.com> Date: Mon, 25 May 2026 09:39:23 +0800 Subject: [PATCH] Allow LLS Save Image to output mask-only --- README.md | 7 ++- image/nodes.py | 85 ++++++++++++++++------------ tests/test_loader_prompt_refactor.py | 77 ++++++++++++++++++++++++- 3 files changed, 128 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index fb43e02..55c2b97 100644 --- a/README.md +++ b/README.md @@ -191,7 +191,7 @@ node/ ### `LLS Save Image` -`LLS Save Image` 现在除了保存或预览 `image`，也支持接收一个可选的 `mask` 输入。连上 `mask` 之后，节点会额外输出一份独立的黑白遮罩图，不会覆盖原图结果，也不会做彩色叠加。 +`LLS Save Image` 现在支持独立保存或预览 `image` 和 `mask`。`image`、`mask` 两个端口都不是必填，但至少要连接其中一个。连上哪个端口，就输出哪个结果；两个都连时，就同时输出两份结果。 **推荐工作流：** @@ -202,8 +202,9 @@ node/ **行为说明：** -- `output_mode = save` 时：保存原图，同时额外保存一份 `_mask` -- `output_mode = preview_only` 时：预览原图，同时额外预览一份黑白 mask +- 只连接 `image`：只保存或预览原图 +- 只连接 `mask`：只保存或预览黑白 mask +- 同时连接 `image + mask`：两者都输出，mask 在保存模式下使用 `_mask` - `image` 与 `mask` 各自独立输出，互不影响 ### `LLS Simple Mask Draw` diff --git a/image/nodes.py b/image/nodes.py index 647f63f..2392af1 100644 --- a/image/nodes.py +++ b/image/nodes.py @@ -356,12 +356,12 @@ class LLSSaveImage: def INPUT_TYPES(cls): return { "required": { - "image": ("IMAGE",), "filename_prefix": ("STRING", {"default": "LLS"}), "output_mode": (_SAVE_OUTPUT_MODES, {"default": "save"}), "save_metadata": ("BOOLEAN", {"default": True}), }, "optional": { + "image": ("IMAGE",), "mask": ("MASK",), "prompt_info": ("STRING", {"forceInput": True}), "latent_info": ("STRING", {"forceInput": True}), @@ -376,6 +376,8 @@ class LLSSaveImage: } def _merge_ui_results(self, primary: dict, secondary: dict | None = None) -> dict: + if not primary: + return secondary or {} if not secondary: return primary @@ -454,7 +456,7 @@ class LLSSaveImage: def save( self, - image, + image=None, filename_prefix: str = "LLS", output_mode: str = "save", save_metadata: bool = True, @@ -467,53 +469,62 @@ class LLSSaveImage: prompt=None, extra_pnginfo=None, ): + if image is None and mask is None: + raise RuntimeError("[LLS] At least one of IMAGE or MASK must be connected for LLS Save Image.") + if output_mode == "preview_only": if comfy_core_nodes is None or not hasattr(comfy_core_nodes, "PreviewImage"): raise RuntimeError("[LLS] ComfyUI core PreviewImage node is not available.") previewer = comfy_core_nodes.PreviewImage() - image_result = previewer.save_images( - image, - prompt=prompt, - extra_pnginfo=None, - ) - if mask is None: - return image_result - mask_result = previewer.save_images( - mask_to_image(mask), - prompt=prompt, - extra_pnginfo=None, - ) + image_result = None + if image is not None: + image_result = previewer.save_images( + image, + prompt=prompt, + extra_pnginfo=None, + ) + + mask_result = None + if mask is not None: + mask_result = previewer.save_images( + mask_to_image(mask), + prompt=prompt, + extra_pnginfo=None, + ) return self._merge_ui_results(image_result, mask_result) if comfy_core_nodes is None or not hasattr(comfy_core_nodes, "SaveImage"): raise RuntimeError("[LLS] ComfyUI core SaveImage node is not available.") saver = comfy_core_nodes.SaveImage() - merged_extra_pnginfo = dict(extra_pnginfo or {}) - if save_metadata: - merged_extra_pnginfo["lls_metadata"] = self._build_metadata( - image=image, - prompt_info=prompt_info, - latent_info=latent_info, - sample_info=sample_info, - decode_info=decode_info, - upscale_info=upscale_info, + image_result = None + if image is not None: + merged_extra_pnginfo = dict(extra_pnginfo or {}) + if save_metadata: + merged_extra_pnginfo["lls_metadata"] = self._build_metadata( + image=image, + prompt_info=prompt_info, + latent_info=latent_info, + sample_info=sample_info, + decode_info=decode_info, + upscale_info=upscale_info, + ) + image_result = saver.save_images( + image, + filename_prefix=filename_prefix, + prompt=prompt, + extra_pnginfo=merged_extra_pnginfo, ) - image_result = saver.save_images( - image, - filename_prefix=filename_prefix, - prompt=prompt, - extra_pnginfo=merged_extra_pnginfo, - ) - if mask is None: - return image_result - mask_result = saver.save_images( - mask_to_image(mask), - filename_prefix=f"{filename_prefix}_mask", - prompt=prompt, - extra_pnginfo={}, - ) + mask_result = None + if mask is not None: + mask_filename_prefix = filename_prefix if image is None else f"{filename_prefix}_mask" + mask_result = saver.save_images( + mask_to_image(mask), + filename_prefix=mask_filename_prefix, + prompt=prompt, + extra_pnginfo={}, + ) return self._merge_ui_results(image_result, mask_result) diff --git a/tests/test_loader_prompt_refactor.py b/tests/test_loader_prompt_refactor.py index 4931142..b3a889d 100644 --- a/tests/test_loader_prompt_refactor.py +++ b/tests/test_loader_prompt_refactor.py @@ -977,13 +977,14 @@ class TestLoaderPromptRefactor(unittest.TestCase): required = node_cls.INPUT_TYPES()["required"] optional = node_cls.INPUT_TYPES()["optional"] - self.assertEqual(required["image"][0], "IMAGE") self.assertEqual(required["filename_prefix"][0], "STRING") self.assertEqual(required["output_mode"][0], ["save", "preview_only"]) self.assertEqual(required["save_metadata"][0], "BOOLEAN") + self.assertNotIn("image", required) self.assertNotIn("model", optional) self.assertNotIn("clip", optional) self.assertNotIn("vae", optional) + self.assertEqual(optional["image"], ("IMAGE",)) self.assertEqual(optional["mask"], ("MASK",)) self.assertEqual(optional["prompt_info"], ("STRING", {"forceInput": True})) self.assertEqual(optional["latent_info"], ("STRING", {"forceInput": True})) @@ -1092,6 +1093,35 @@ class TestLoaderPromptRefactor(unittest.TestCase): self.assertEqual(len(result["ui"]["images"]), 2) self.assertIsNone(result["ui"]["images"][1]["extra_pnginfo"]) + def test_save_image_preview_only_supports_mask_without_image(self): + load_plugin_package() + from lls_node_test_refactor.image import nodes as image_nodes + + CoreNodesStub.save_calls = [] + CoreNodesStub.preview_calls = [] + + with mock.patch.object(image_nodes, "comfy_core_nodes", CoreNodesStub()): + with mock.patch.object( + image_nodes, + "mask_to_image", + side_effect=lambda mask: FakeTensor((mask.shape[0], mask.shape[1], mask.shape[2], 3)), + create=True, + ): + node = image_nodes.LLSSaveImage() + result = node.save( + image=None, + mask=FakeTensor((1, 4, 4)), + filename_prefix="LLS", + save_metadata=True, + output_mode="preview_only", + ) + + self.assertEqual(CoreNodesStub.save_calls, []) + self.assertEqual(len(CoreNodesStub.preview_calls), 1) + self.assertEqual(tuple(CoreNodesStub.preview_calls[0]["images"].shape), (1, 4, 4, 3)) + self.assertEqual(len(result["ui"]["images"]), 1) + self.assertEqual(result["ui"]["images"][0]["type"], "temp") + def test_save_image_save_mode_emits_separate_mask_file_without_lls_metadata(self): load_plugin_package() from lls_node_test_refactor.image import nodes as image_nodes @@ -1124,6 +1154,51 @@ class TestLoaderPromptRefactor(unittest.TestCase): self.assertEqual(tuple(CoreNodesStub.save_calls[1]["images"].shape), (1, 4, 4, 3)) self.assertEqual(len(result["ui"]["images"]), 2) + def test_save_image_save_mode_supports_mask_without_image(self): + load_plugin_package() + from lls_node_test_refactor.image import nodes as image_nodes + + CoreNodesStub.save_calls = [] + CoreNodesStub.preview_calls = [] + + with mock.patch.object(image_nodes, "comfy_core_nodes", CoreNodesStub()): + with mock.patch.object( + image_nodes, + "mask_to_image", + side_effect=lambda mask: FakeTensor((mask.shape[0], mask.shape[1], mask.shape[2], 3)), + create=True, + ): + node = image_nodes.LLSSaveImage() + result = node.save( + image=None, + mask=FakeTensor((1, 4, 4)), + filename_prefix="LLS", + save_metadata=True, + ) + + self.assertEqual(len(CoreNodesStub.preview_calls), 0) + self.assertEqual(len(CoreNodesStub.save_calls), 1) + self.assertEqual(CoreNodesStub.save_calls[0]["filename_prefix"], "LLS") + self.assertEqual(CoreNodesStub.save_calls[0]["extra_pnginfo"], {}) + self.assertEqual(tuple(CoreNodesStub.save_calls[0]["images"].shape), (1, 4, 4, 3)) + self.assertEqual(len(result["ui"]["images"]), 1) + + def test_save_image_rejects_missing_image_and_mask(self): + load_plugin_package() + from lls_node_test_refactor.image import nodes as image_nodes + + with mock.patch.object(image_nodes, "comfy_core_nodes", CoreNodesStub()): + node = image_nodes.LLSSaveImage() + + with self.assertRaisesRegex(RuntimeError, "At least one of IMAGE or MASK"): + node.save( + image=None, + mask=None, + filename_prefix="LLS", + save_metadata=True, + output_mode="preview_only", + ) + def test_generation_config_uses_model_inference_without_context(self): plugin = load_plugin_package() node_cls = plugin.NODE_CLASS_MAPPINGS["LLSGenerationConfig"]