From dd8b0533f46943953d8cb6df3aa2f8d27bef4120 Mon Sep 17 00:00:00 2001 From: Gin3601 Date: Fri, 22 May 2026 11:51:24 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E=E5=86=85=E5=A2=9E=E5=A4=96?= =?UTF-8?q?=E9=83=A8=E8=8A=82=E7=82=B9=E5=92=8C=E4=BF=AE=E6=94=B9=E5=9B=BE?= =?UTF-8?q?=E5=83=8F=E4=BF=9D=E5=AD=98=E8=8A=82=E7=82=B9=EF=BC=88=E6=94=AF?= =?UTF-8?q?=E6=8C=81=E4=BF=9D=E5=AD=98=E5=92=8C=E9=A2=84=E8=A7=88=E5=8F=8C?= =?UTF-8?q?=E6=A8=A1=E5=BC=8F=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- image/nodes.py | 13 +++ repair/__init__.py | 6 ++ repair/native_inpaint.py | 107 ++++++++++++++++++++++ tests/test_loader_prompt_refactor.py | 55 ++++++++++++ tests/test_repair_native_inpaint.py | 127 +++++++++++++++++++++++++++ 5 files changed, 308 insertions(+) create mode 100644 repair/native_inpaint.py create mode 100644 tests/test_repair_native_inpaint.py diff --git a/image/nodes.py b/image/nodes.py index 79b7bf8..db53042 100644 --- a/image/nodes.py +++ b/image/nodes.py @@ -40,6 +40,7 @@ else: _RESIZE_MODES = ["keep_aspect", "crop_center", "stretch", "none"] _SIZE_SOURCES = ["input_image", "custom", "model_recommended"] +_SAVE_OUTPUT_MODES = ["save", "preview_only"] def _round_to_multiple(value: int, multiple: int) -> int: @@ -351,6 +352,7 @@ class LLSSaveImage: "required": { "image": ("IMAGE",), "filename_prefix": ("STRING", {"default": "LLS"}), + "output_mode": (_SAVE_OUTPUT_MODES, {"default": "save"}), "save_metadata": ("BOOLEAN", {"default": True}), }, "optional": { @@ -429,6 +431,7 @@ class LLSSaveImage: self, image, filename_prefix: str = "LLS", + output_mode: str = "save", save_metadata: bool = True, model=None, clip=None, @@ -441,6 +444,16 @@ class LLSSaveImage: prompt=None, extra_pnginfo=None, ): + 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() + return previewer.save_images( + image, + prompt=prompt, + extra_pnginfo=None, + ) + if comfy_core_nodes is None or not hasattr(comfy_core_nodes, "SaveImage"): raise RuntimeError("[LLS] ComfyUI core SaveImage node is not available.") diff --git a/repair/__init__.py b/repair/__init__.py index cd066c0..afd72ea 100644 --- a/repair/__init__.py +++ b/repair/__init__.py @@ -7,6 +7,10 @@ from .repair_finish import ( NODE_CLASS_MAPPINGS as FINISH_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as FINISH_NODE_DISPLAY_NAME_MAPPINGS, ) +from .native_inpaint import ( + NODE_CLASS_MAPPINGS as NATIVE_INPAINT_NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as NATIVE_INPAINT_NODE_DISPLAY_NAME_MAPPINGS, +) from .repair_prepare import ( NODE_CLASS_MAPPINGS as PREPARE_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as PREPARE_NODE_DISPLAY_NAME_MAPPINGS, @@ -25,11 +29,13 @@ def _merge_registry_maps(*maps): NODE_CLASS_MAPPINGS = _merge_registry_maps( + NATIVE_INPAINT_NODE_CLASS_MAPPINGS, PREPARE_NODE_CLASS_MAPPINGS, FINISH_NODE_CLASS_MAPPINGS, ) NODE_DISPLAY_NAME_MAPPINGS = _merge_registry_maps( + NATIVE_INPAINT_NODE_DISPLAY_NAME_MAPPINGS, PREPARE_NODE_DISPLAY_NAME_MAPPINGS, FINISH_NODE_DISPLAY_NAME_MAPPINGS, ) diff --git a/repair/native_inpaint.py b/repair/native_inpaint.py new file mode 100644 index 0000000..700944a --- /dev/null +++ b/repair/native_inpaint.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from ..utils.model_info import info_to_json, resolve_edit_capabilities +from . import runtime as repair_runtime + + +_PATCH_MODE_CHOICES = ["auto", "disabled", "differential_diffusion"] + + +class LLSNativeInpaintConditioning: + CATEGORY = "LLS/Image Repair" + FUNCTION = "encode" + RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING", "LATENT", "STRING") + RETURN_NAMES = ("model", "positive", "negative", "latent", "inpaint_info") + DESCRIPTION = ( + "Thin wrapper around ComfyUI InpaintModelConditioning. " + "Keeps the standard KSampler path and can auto-apply DifferentialDiffusion for FLUX." + ) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "vae": ("VAE",), + "image": ("IMAGE",), + "mask": ("MASK",), + "patch_mode": (_PATCH_MODE_CHOICES, {"default": "auto"}), + "patch_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "noise_mask": ("BOOLEAN", {"default": True}), + }, + "optional": { + "model_info": ("STRING",), + }, + } + + def encode( + self, + model, + positive, + negative, + vae, + image, + mask, + patch_mode, + patch_strength, + noise_mask, + model_info=None, + ): + capabilities = resolve_edit_capabilities(model=model, model_info=model_info) + + apply_patch = False + if patch_mode == "differential_diffusion": + apply_patch = True + elif patch_mode == "auto" and str(capabilities["model_family"]).startswith("FLUX"): + apply_patch = True + + patched_model = model + applied_patch = "" + if apply_patch: + patched_model = repair_runtime.apply_differential_diffusion( + model, + strength=float(patch_strength), + ) + applied_patch = "differential_diffusion" + + positive_out, negative_out, latent = repair_runtime.encode_inpaint_conditioning( + positive, + negative, + vae, + image, + mask, + noise_mask=bool(noise_mask), + ) + + inpaint_info = info_to_json( + { + "model_family": capabilities["model_family"], + "model_role": capabilities["model_role"], + "profile_id": capabilities["profile_id"], + "backend_type": capabilities["backend_type"], + "sampler_strategy": capabilities["sampler_strategy"], + "loader_strategy": capabilities["loader_strategy"], + "patch_mode": patch_mode, + "applied_patch": applied_patch, + "patch_strength": float(patch_strength), + "noise_mask": bool(noise_mask), + } + ) + return ( + patched_model, + positive_out, + negative_out, + latent, + inpaint_info, + ) + + +NODE_CLASS_MAPPINGS = { + "LLSNativeInpaintConditioning": LLSNativeInpaintConditioning, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LLSNativeInpaintConditioning": "LLS Native Inpaint Conditioning", +} diff --git a/tests/test_loader_prompt_refactor.py b/tests/test_loader_prompt_refactor.py index 119fd58..1b4d679 100644 --- a/tests/test_loader_prompt_refactor.py +++ b/tests/test_loader_prompt_refactor.py @@ -90,6 +90,9 @@ class ComfySDStub: class CoreNodesStub: + save_calls = [] + preview_calls = [] + class VAELoader: @staticmethod def vae_list(_self): @@ -108,6 +111,14 @@ class CoreNodesStub: class SaveImage: def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + CoreNodesStub.save_calls.append( + { + "images": images, + "filename_prefix": filename_prefix, + "prompt": prompt, + "extra_pnginfo": extra_pnginfo, + } + ) return { "ui": { "images": [ @@ -122,6 +133,29 @@ class CoreNodesStub: } } + class PreviewImage: + def save_images(self, images, prompt=None, extra_pnginfo=None): + CoreNodesStub.preview_calls.append( + { + "images": images, + "prompt": prompt, + "extra_pnginfo": extra_pnginfo, + } + ) + return { + "ui": { + "images": [ + { + "filename": "preview.png", + "subfolder": "", + "type": "temp", + "extra_pnginfo": extra_pnginfo, + "prompt": prompt, + } + ] + } + } + class RecordingClip: def __init__(self, token_keys, calls=None): @@ -829,6 +863,27 @@ class TestLoaderPromptRefactor(unittest.TestCase): self.assertEqual(metadata["steps"], 20) self.assertEqual(metadata["upscale_mode"], "none") + def test_save_image_preview_only_uses_native_preview_node_without_saving(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()): + node = image_nodes.LLSSaveImage() + result = node.save( + image=FakeTensor((1, 512, 512, 3)), + filename_prefix="LLS", + save_metadata=True, + output_mode="preview_only", + ) + + self.assertEqual(CoreNodesStub.save_calls, []) + self.assertEqual(len(CoreNodesStub.preview_calls), 1) + self.assertEqual(result["ui"]["images"][0]["type"], "temp") + self.assertIsNone(result["ui"]["images"][0]["extra_pnginfo"]) + def test_generation_config_uses_model_inference_without_context(self): plugin = load_plugin_package() node_cls = plugin.NODE_CLASS_MAPPINGS["LLSGenerationConfig"] diff --git a/tests/test_repair_native_inpaint.py b/tests/test_repair_native_inpaint.py new file mode 100644 index 0000000..0e13b11 --- /dev/null +++ b/tests/test_repair_native_inpaint.py @@ -0,0 +1,127 @@ +import json +import unittest +from unittest import mock + +try: + from .test_repair_helpers import FakeMask, FakeModel, FakeTensor, import_plugin_submodule, load_plugin_package +except ImportError: # pragma: no cover - discovery mode imports from top level + from test_repair_helpers import FakeMask, FakeModel, FakeTensor, import_plugin_submodule, load_plugin_package + + +class TestRepairNativeInpaint(unittest.TestCase): + def test_node_registers_with_expected_schema(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSNativeInpaintConditioning"] + schema = node_cls.INPUT_TYPES() + + self.assertEqual(plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSNativeInpaintConditioning"], "LLS Native Inpaint Conditioning") + self.assertEqual(node_cls.CATEGORY, "LLS/Image Repair") + self.assertEqual(node_cls.FUNCTION, "encode") + self.assertEqual( + node_cls.RETURN_TYPES, + ("MODEL", "CONDITIONING", "CONDITIONING", "LATENT", "STRING"), + ) + self.assertEqual( + node_cls.RETURN_NAMES, + ("model", "positive", "negative", "latent", "inpaint_info"), + ) + self.assertEqual(schema["required"]["patch_mode"][0], ["auto", "disabled", "differential_diffusion"]) + self.assertEqual(schema["optional"]["model_info"], ("STRING",)) + + def test_node_wraps_inpaint_model_conditioning_without_flux_patch(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSNativeInpaintConditioning"]() + repair_runtime = import_plugin_submodule(plugin, "repair.runtime") + model = FakeModel( + family="SDXL", + model_role="inpaint", + supports_inpaint_native=True, + supports_image_edit_native=False, + preferred_edit_backend="sdxl", + profile_id="sdxl_inpaint", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", + ) + positive = [["pos", {}]] + negative = [["neg", {}]] + latent = {"samples": "latent"} + + with mock.patch.object( + repair_runtime, + "encode_inpaint_conditioning", + return_value=(positive, negative, latent), + ) as encode_inpaint, mock.patch.object( + repair_runtime, + "apply_differential_diffusion", + ) as patch_model: + model_out, positive_out, negative_out, latent_out, info = node.encode( + model=model, + positive=positive, + negative=negative, + vae=object(), + image=FakeTensor((1, 1024, 1024, 3)), + mask=FakeMask((1, 1024, 1024), mask_bbox=(64, 64, 256, 256), mask_area_ratio=0.04), + patch_mode="auto", + patch_strength=1.0, + noise_mask=True, + ) + + encode_inpaint.assert_called_once() + patch_model.assert_not_called() + self.assertIs(model_out, model) + self.assertIs(positive_out, positive) + self.assertIs(negative_out, negative) + self.assertIs(latent_out, latent) + payload = json.loads(info) + self.assertEqual(payload["model_family"], "SDXL") + self.assertEqual(payload["applied_patch"], "") + + def test_node_auto_applies_flux_differential_diffusion(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSNativeInpaintConditioning"]() + repair_runtime = import_plugin_submodule(plugin, "repair.runtime") + model = FakeModel( + family="FLUX_DEV", + model_role="fill", + supports_inpaint_native=False, + supports_image_edit_native=True, + preferred_edit_backend="flux", + profile_id="flux_edit", + backend_type="flux_edit", + sampler_strategy="flux_guided", + loader_strategy="flux_split_or_bundle", + ) + patched_model = object() + + with mock.patch.object( + repair_runtime, + "apply_differential_diffusion", + return_value=patched_model, + ) as patch_model, mock.patch.object( + repair_runtime, + "encode_inpaint_conditioning", + return_value=([["pos", {}]], [["neg", {}]], {"samples": "latent"}), + ): + model_out, _positive_out, _negative_out, _latent_out, info = node.encode( + model=model, + positive=[["pos", {}]], + negative=[["neg", {}]], + vae=object(), + image=FakeTensor((1, 1024, 1024, 3)), + mask=FakeMask((1, 1024, 1024), mask_bbox=(64, 64, 256, 256), mask_area_ratio=0.04), + patch_mode="auto", + patch_strength=0.75, + noise_mask=True, + ) + + patch_model.assert_called_once_with(model, strength=0.75) + self.assertIs(model_out, patched_model) + payload = json.loads(info) + self.assertEqual(payload["model_family"], "FLUX_DEV") + self.assertEqual(payload["applied_patch"], "differential_diffusion") + self.assertEqual(payload["patch_strength"], 0.75) + + +if __name__ == "__main__": + unittest.main()