feat: add repair node scaffolding

2026-05-20 18:40:03 +08:00
parent fc06d58432
commit a4ea3f2089
6 changed files with 252 additions and 0 deletions
@@ -45,6 +45,7 @@ _SUBPACKAGES: list[str] = [
    "qwen",           # 第 4 节：Qwen 一体化节点
    "latent",         # 第 4 节：Latent 空间操作
    "image",          # 第 5 节：图像处理与后处理
+    "repair",         # 第 5 节（修复子域）：图像修复
    "upscale",        # 第 5 节（超分子域）：图像超分
    "mask",           # 第 6 节：遮罩操作
    "controlnet",     # 第 7 节：ControlNet / 控制引导
@@ -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"]
@@ -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",
+}
@@ -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",
+}
@@ -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
@@ -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()