diff --git a/repair/__init__.py b/repair/__init__.py index 23a59ef..cd066c0 100644 --- a/repair/__init__.py +++ b/repair/__init__.py @@ -13,14 +13,25 @@ from .repair_prepare import ( ) -NODE_CLASS_MAPPINGS = { - **PREPARE_NODE_CLASS_MAPPINGS, - **FINISH_NODE_CLASS_MAPPINGS, -} +def _merge_registry_maps(*maps): + merged = {} + for mapping in maps: + overlap = set(merged).intersection(mapping) + if overlap: + duplicate_keys = ", ".join(sorted(overlap)) + raise RuntimeError(f"[LLS] Duplicate repair node registration keys: {duplicate_keys}") + merged.update(mapping) + return merged -NODE_DISPLAY_NAME_MAPPINGS = { - **PREPARE_NODE_DISPLAY_NAME_MAPPINGS, - **FINISH_NODE_DISPLAY_NAME_MAPPINGS, -} + +NODE_CLASS_MAPPINGS = _merge_registry_maps( + PREPARE_NODE_CLASS_MAPPINGS, + FINISH_NODE_CLASS_MAPPINGS, +) + +NODE_DISPLAY_NAME_MAPPINGS = _merge_registry_maps( + 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 index 9caa69b..017f5d7 100644 --- a/repair/repair_finish.py +++ b/repair/repair_finish.py @@ -1,8 +1,10 @@ from __future__ import annotations -_EDGE_FIX_CHOICES = ["off", "feather", "color_blend"] -_PREVIEW_MODE_CHOICES = ["final", "preview", "split"] +_COLOR_MATCH_CHOICES = ["disabled", "mean_std", "histogram_simple"] +_BRIGHTNESS_MATCH_CHOICES = ["disabled", "enabled"] +_EDGE_FIX_CHOICES = ["none", "soft", "strong"] +_PREVIEW_MODE_CHOICES = ["final", "compare", "mask", "before_after"] class LLSSimpleRepairFinish: @@ -20,11 +22,11 @@ class LLSSimpleRepairFinish: "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}), + "color_match": (_COLOR_MATCH_CHOICES, {"default": "disabled"}), + "brightness_match": (_BRIGHTNESS_MATCH_CHOICES, {"default": "enabled"}), "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"}), + "edge_fix": (_EDGE_FIX_CHOICES, {"default": "soft"}), "preview_mode": (_PREVIEW_MODE_CHOICES, {"default": "final"}), }, "optional": { diff --git a/repair/repair_prepare.py b/repair/repair_prepare.py index 9556515..15b7c2f 100644 --- a/repair/repair_prepare.py +++ b/repair/repair_prepare.py @@ -1,11 +1,24 @@ 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"] +_REPAIR_SCOPE_CHOICES = ["auto", "region", "crop", "canvas"] +_REPAIR_KERNEL_CHOICES = ["auto", "latent_mask", "vae_inpaint", "native_fill"] +_TASK_HINT_CHOICES = [ + "auto", + "repair", + "remove", + "replace", + "fill", + "appearance", + "content", + "structure", + "dehaze", + "deshadow", + "recolor", +] _RESIZE_MODE_CHOICES = ["fit", "pad", "stretch"] -_CANVAS_FILL_CHOICES = ["original", "mask_mean", "solid"] +_AUTO_RECOMMEND_CHOICES = ["enabled", "disabled"] +_CANVAS_FILL_CHOICES = ["edge", "blur", "black", "white", "neutral"] class LLSSimpleRepairPrepare: @@ -22,14 +35,14 @@ class LLSSimpleRepairPrepare: "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"}), + "repair_scope": (_REPAIR_SCOPE_CHOICES, {"default": "auto"}), + "repair_kernel": (_REPAIR_KERNEL_CHOICES, {"default": "auto"}), + "task_hint": (_TASK_HINT_CHOICES, {"default": "auto"}), "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": ("INT", {"default": 64, "min": 0, "max": 512}), "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}), @@ -38,8 +51,8 @@ class LLSSimpleRepairPrepare: "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}), + "canvas_fill": (_CANVAS_FILL_CHOICES, {"default": "edge"}), + "auto_recommend": (_AUTO_RECOMMEND_CHOICES, {"default": "enabled"}), }, "optional": { "model_info": ("STRING",), diff --git a/tests/test_repair_registration.py b/tests/test_repair_registration.py index 58789da..d547aa3 100644 --- a/tests/test_repair_registration.py +++ b/tests/test_repair_registration.py @@ -1,6 +1,9 @@ import unittest -from test_repair_helpers import load_plugin_package +try: + from .test_repair_helpers import load_plugin_package +except ImportError: # pragma: no cover - discovery mode imports from top level + from test_repair_helpers import load_plugin_package class TestRepairRegistration(unittest.TestCase): @@ -29,6 +32,10 @@ class TestRepairRegistration(unittest.TestCase): node_cls.RETURN_TYPES, ("LATENT", "IMAGE", "MASK", "LLS_REPAIR_INFO", "FLOAT"), ) + self.assertEqual( + node_cls.RETURN_NAMES, + ("latent", "work_image", "work_mask", "repair_info", "recommended_denoise"), + ) required = schema["required"] optional = schema["optional"] @@ -60,6 +67,51 @@ class TestRepairRegistration(unittest.TestCase): for field in ("model_info", "positive", "negative"): self.assertIn(field, optional) + self.assertEqual(required["image"], ("IMAGE",)) + self.assertEqual(required["mask"], ("MASK",)) + self.assertEqual(required["vae"], ("VAE",)) + self.assertEqual(optional["model_info"], ("STRING",)) + self.assertEqual(optional["positive"], ("CONDITIONING",)) + self.assertEqual(optional["negative"], ("CONDITIONING",)) + self.assertEqual( + required["repair_scope"], + (["auto", "region", "crop", "canvas"], {"default": "auto"}), + ) + self.assertEqual( + required["repair_kernel"], + (["auto", "latent_mask", "vae_inpaint", "native_fill"], {"default": "auto"}), + ) + self.assertEqual( + required["task_hint"], + ( + [ + "auto", + "repair", + "remove", + "replace", + "fill", + "appearance", + "content", + "structure", + "dehaze", + "deshadow", + "recolor", + ], + {"default": "auto"}, + ), + ) + self.assertEqual( + required["crop_context"], + ("INT", {"default": 64, "min": 0, "max": 512}), + ) + self.assertEqual( + required["auto_recommend"], + (["enabled", "disabled"], {"default": "enabled"}), + ) + self.assertEqual( + required["canvas_fill"], + (["edge", "blur", "black", "white", "neutral"], {"default": "edge"}), + ) def test_finish_node_schema(self): plugin = load_plugin_package() @@ -69,6 +121,7 @@ class TestRepairRegistration(unittest.TestCase): self.assertEqual(node_cls.CATEGORY, "LLS/Image Repair") self.assertEqual(node_cls.FUNCTION, "finish") self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE", "IMAGE")) + self.assertEqual(node_cls.RETURN_NAMES, ("final_image", "preview_image")) required = schema["required"] optional = schema["optional"] @@ -89,6 +142,27 @@ class TestRepairRegistration(unittest.TestCase): for field in ("work_mask", "sample_info"): self.assertIn(field, optional) + self.assertEqual(required["original_image"], ("IMAGE",)) + self.assertEqual(required["generated_image"], ("IMAGE",)) + self.assertEqual(required["repair_info"], ("LLS_REPAIR_INFO",)) + self.assertEqual(optional["work_mask"], ("MASK",)) + self.assertEqual(optional["sample_info"], ("STRING",)) + self.assertEqual( + required["color_match"], + (["disabled", "mean_std", "histogram_simple"], {"default": "disabled"}), + ) + self.assertEqual( + required["brightness_match"], + (["disabled", "enabled"], {"default": "enabled"}), + ) + self.assertEqual( + required["edge_fix"], + (["none", "soft", "strong"], {"default": "soft"}), + ) + self.assertEqual( + required["preview_mode"], + (["final", "compare", "mask", "before_after"], {"default": "final"}), + ) if __name__ == "__main__":