diff --git a/__init__.py b/__init__.py index 6ee9240..f428a7c 100644 --- a/__init__.py +++ b/__init__.py @@ -45,6 +45,7 @@ _SUBPACKAGES: list[str] = [ "qwen", # 第 4 节：Qwen 一体化节点 "latent", # 第 4 节：Latent 空间操作 "image", # 第 5 节：图像处理与后处理 + "repair", # 第 5 节（修复子域）：图像修复 "upscale", # 第 5 节（超分子域）：图像超分 "mask", # 第 6 节：遮罩操作 "controlnet", # 第 7 节：ControlNet / 控制引导 diff --git a/repair/__init__.py b/repair/__init__.py new file mode 100644 index 0000000..23a59ef --- /dev/null +++ b/repair/__init__.py @@ -0,0 +1,26 @@ +""" +node.repair +=========== +功能域：图像修复（Repair）节点 +""" +from .repair_finish import ( + NODE_CLASS_MAPPINGS as FINISH_NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as FINISH_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, +) + + +NODE_CLASS_MAPPINGS = { + **PREPARE_NODE_CLASS_MAPPINGS, + **FINISH_NODE_CLASS_MAPPINGS, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + **PREPARE_NODE_DISPLAY_NAME_MAPPINGS, + **FINISH_NODE_DISPLAY_NAME_MAPPINGS, +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/repair/repair_finish.py b/repair/repair_finish.py new file mode 100644 index 0000000..9caa69b --- /dev/null +++ b/repair/repair_finish.py @@ -0,0 +1,46 @@ +from __future__ import annotations + + +_EDGE_FIX_CHOICES = ["off", "feather", "color_blend"] +_PREVIEW_MODE_CHOICES = ["final", "preview", "split"] + + +class LLSSimpleRepairFinish: + CATEGORY = "LLS/Image Repair" + FUNCTION = "finish" + RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_NAMES = ("final_image", "preview_image") + DESCRIPTION = "Composite repaired content back into the original image." + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "original_image": ("IMAGE",), + "generated_image": ("IMAGE",), + "repair_info": ("LLS_REPAIR_INFO",), + "feather": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 256.0, "step": 0.5}), + "color_match": ("BOOLEAN", {"default": True}), + "brightness_match": ("BOOLEAN", {"default": True}), + "blend_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "restore_unmasked_area": ("BOOLEAN", {"default": True}), + "edge_fix": (_EDGE_FIX_CHOICES, {"default": "feather"}), + "preview_mode": (_PREVIEW_MODE_CHOICES, {"default": "final"}), + }, + "optional": { + "work_mask": ("MASK",), + "sample_info": ("STRING",), + }, + } + + def finish(self, *args, **kwargs): + raise RuntimeError("[LLS] LLS Simple Repair Finish is not implemented yet.") + + +NODE_CLASS_MAPPINGS = { + "LLSSimpleRepairFinish": LLSSimpleRepairFinish, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LLSSimpleRepairFinish": "LLS Simple Repair Finish", +} diff --git a/repair/repair_prepare.py b/repair/repair_prepare.py new file mode 100644 index 0000000..9556515 --- /dev/null +++ b/repair/repair_prepare.py @@ -0,0 +1,61 @@ +from __future__ import annotations + + +_REPAIR_SCOPE_CHOICES = ["masked_area", "bounding_box", "full_image"] +_REPAIR_KERNEL_CHOICES = ["gaussian", "box", "dilate"] +_TASK_HINT_CHOICES = ["general", "cleanup", "object_removal", "detail_fix"] +_RESIZE_MODE_CHOICES = ["fit", "pad", "stretch"] +_CANVAS_FILL_CHOICES = ["original", "mask_mean", "solid"] + + +class LLSSimpleRepairPrepare: + CATEGORY = "LLS/Image Repair" + FUNCTION = "prepare" + RETURN_TYPES = ("LATENT", "IMAGE", "MASK", "LLS_REPAIR_INFO", "FLOAT") + RETURN_NAMES = ("latent", "work_image", "work_mask", "repair_info", "recommended_denoise") + DESCRIPTION = "Prepare repair inputs, work area, and repair metadata." + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "mask": ("MASK",), + "vae": ("VAE",), + "repair_scope": (_REPAIR_SCOPE_CHOICES, {"default": "masked_area"}), + "repair_kernel": (_REPAIR_KERNEL_CHOICES, {"default": "gaussian"}), + "task_hint": (_TASK_HINT_CHOICES, {"default": "general"}), + "mask_grow": ("INT", {"default": 24, "min": 0, "max": 2048, "step": 1}), + "mask_blur": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 256.0, "step": 0.5}), + "mask_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "invert_mask": ("BOOLEAN", {"default": False}), + "crop_context": ("BOOLEAN", {"default": True}), + "crop_context_factor": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 8.0, "step": 0.1}), + "min_size": ("INT", {"default": 256, "min": 64, "max": 8192, "step": 8}), + "max_size": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8}), + "resize_mode": (_RESIZE_MODE_CHOICES, {"default": "fit"}), + "expand_left": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1}), + "expand_right": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1}), + "expand_top": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1}), + "expand_bottom": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1}), + "canvas_fill": (_CANVAS_FILL_CHOICES, {"default": "original"}), + "auto_recommend": ("BOOLEAN", {"default": True}), + }, + "optional": { + "model_info": ("STRING",), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + }, + } + + def prepare(self, *args, **kwargs): + raise RuntimeError("[LLS] LLS Simple Repair Prepare is not implemented yet.") + + +NODE_CLASS_MAPPINGS = { + "LLSSimpleRepairPrepare": LLSSimpleRepairPrepare, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LLSSimpleRepairPrepare": "LLS Simple Repair Prepare", +} diff --git a/tests/test_repair_helpers.py b/tests/test_repair_helpers.py new file mode 100644 index 0000000..47fcf82 --- /dev/null +++ b/tests/test_repair_helpers.py @@ -0,0 +1,23 @@ +import importlib.util +import pathlib +import sys + + +ROOT = pathlib.Path(__file__).resolve().parents[1] +MODULE_NAME = "lls_node_test_repair" + + +def load_plugin_package(): + for name in list(sys.modules): + if name == MODULE_NAME or name.startswith(f"{MODULE_NAME}."): + sys.modules.pop(name) + + spec = importlib.util.spec_from_file_location( + MODULE_NAME, + ROOT / "__init__.py", + submodule_search_locations=[str(ROOT)], + ) + module = importlib.util.module_from_spec(spec) + sys.modules[MODULE_NAME] = module + spec.loader.exec_module(module) + return module diff --git a/tests/test_repair_registration.py b/tests/test_repair_registration.py new file mode 100644 index 0000000..58789da --- /dev/null +++ b/tests/test_repair_registration.py @@ -0,0 +1,95 @@ +import unittest + +from test_repair_helpers import load_plugin_package + + +class TestRepairRegistration(unittest.TestCase): + def test_plugin_registers_repair_nodes_with_display_names(self): + plugin = load_plugin_package() + + self.assertIn("LLSSimpleRepairPrepare", plugin.NODE_CLASS_MAPPINGS) + self.assertIn("LLSSimpleRepairFinish", plugin.NODE_CLASS_MAPPINGS) + self.assertEqual( + plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSSimpleRepairPrepare"], + "LLS Simple Repair Prepare", + ) + self.assertEqual( + plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSSimpleRepairFinish"], + "LLS Simple Repair Finish", + ) + + def test_prepare_node_schema(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"] + schema = node_cls.INPUT_TYPES() + + self.assertEqual(node_cls.CATEGORY, "LLS/Image Repair") + self.assertEqual(node_cls.FUNCTION, "prepare") + self.assertEqual( + node_cls.RETURN_TYPES, + ("LATENT", "IMAGE", "MASK", "LLS_REPAIR_INFO", "FLOAT"), + ) + + required = schema["required"] + optional = schema["optional"] + + for field in ( + "image", + "mask", + "vae", + "repair_scope", + "repair_kernel", + "task_hint", + "mask_grow", + "mask_blur", + "mask_threshold", + "invert_mask", + "crop_context", + "crop_context_factor", + "min_size", + "max_size", + "resize_mode", + "expand_left", + "expand_right", + "expand_top", + "expand_bottom", + "canvas_fill", + "auto_recommend", + ): + self.assertIn(field, required) + + for field in ("model_info", "positive", "negative"): + self.assertIn(field, optional) + + def test_finish_node_schema(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairFinish"] + schema = node_cls.INPUT_TYPES() + + self.assertEqual(node_cls.CATEGORY, "LLS/Image Repair") + self.assertEqual(node_cls.FUNCTION, "finish") + self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE", "IMAGE")) + + required = schema["required"] + optional = schema["optional"] + + for field in ( + "original_image", + "generated_image", + "repair_info", + "feather", + "color_match", + "brightness_match", + "blend_strength", + "restore_unmasked_area", + "edge_fix", + "preview_mode", + ): + self.assertIn(field, required) + + for field in ("work_mask", "sample_info"): + self.assertIn(field, optional) + + +if __name__ == "__main__": + unittest.main()