From d77044b929a18b71a185da7f16ebd70bc67bc090 Mon Sep 17 00:00:00 2001 From: Gin3601 <2405559452@qq.com> Date: Wed, 20 May 2026 18:55:46 +0800 Subject: [PATCH] feat: add repair utility helpers --- repair/repair_utils.py | 313 +++++++++++++++++++++++++++++++++++++ tests/test_repair_utils.py | 163 +++++++++++++++++++ 2 files changed, 476 insertions(+) create mode 100644 repair/repair_utils.py create mode 100644 tests/test_repair_utils.py diff --git a/repair/repair_utils.py b/repair/repair_utils.py new file mode 100644 index 0000000..8ff74b2 --- /dev/null +++ b/repair/repair_utils.py @@ -0,0 +1,313 @@ +from __future__ import annotations + +from typing import Any + +from ..utils.model_info import parse_jsonish_info + + +REPAIR_INFO_TYPE = "LLS_REPAIR_INFO" +GUIDANCE_STACK_TYPE = "LLS_GUIDANCE_STACK" +SUPPORTED_SAMPLER_FAMILIES = {"SD1.5", "SDXL", "SDXL_TURBO", "FLUX_DEV", "FLUX_SCHNELL"} + +_REPAIR_SCOPE_VALUES = {"auto", "region", "crop", "canvas"} +_REPAIR_KERNEL_VALUES = {"auto", "latent_mask", "vae_inpaint", "native_fill"} +_TASK_HINT_LATENT_MASK = {"repair", "appearance", "dehaze", "deshadow", "recolor"} +_TASK_DENOISE_VALUES = { + "repair": 0.45, + "appearance": 0.50, + "dehaze": 0.55, + "deshadow": 0.55, + "recolor": 0.55, + "structure": 0.60, + "content": 0.65, + "replace": 0.72, + "remove": 0.88, + "fill": 0.90, + "auto": 0.55, +} +_RESIZE_MODE_ALIASES = { + "fit": "keep_aspect", + "keep_aspect": "keep_aspect", + "pad": "force_square", + "force_square": "force_square", + "stretch": "ranged_size", + "ranged_size": "ranged_size", +} + + +def normalize_model_info(model_info: dict[str, Any] | str | None) -> dict[str, Any]: + raw = parse_jsonish_info(model_info) + return { + "model_family": str(raw.get("model_family") or raw.get("family") or "UNKNOWN"), + "model_role": str(raw.get("model_role") or raw.get("role") or "unknown"), + "supports_inpaint_native": bool(raw.get("supports_inpaint_native", False)), + } + + +def normalize_repair_info(repair_info: dict[str, Any] | str | None) -> dict[str, Any]: + raw = parse_jsonish_info(repair_info) + info = dict(raw) + model_info = normalize_model_info(raw.get("model_info")) + + info["repair_scope"] = str(info.get("repair_scope") or "region") + info["repair_kernel"] = str(info.get("repair_kernel") or "vae_inpaint") + info["task_hint"] = str(info.get("task_hint") or "auto") + info["original_size"] = _normalize_size_pair(info.get("original_size"), default=(0, 0)) + info["work_size"] = _normalize_size_pair(info.get("work_size"), default=(0, 0)) + info["crop_box"] = _normalize_box(info.get("crop_box")) + info["crop_scale"] = _normalize_optional_float(info.get("crop_scale")) + info["canvas_expand"] = _normalize_expand(info.get("canvas_expand")) + info["mask_grow"] = _coerce_int(info.get("mask_grow"), 8) + info["mask_blur"] = _coerce_float(info.get("mask_blur"), 8.0) + info["mask_threshold"] = _coerce_float(info.get("mask_threshold"), 0.5) + info["invert_mask"] = bool(info.get("invert_mask", False)) + info["recommended_denoise"] = _coerce_float(info.get("recommended_denoise"), 0.55) + info["model_family"] = str(info.get("model_family") or model_info["model_family"]) + info["model_role"] = str(info.get("model_role") or model_info["model_role"]) + info["supports_inpaint_native"] = bool( + info.get("supports_inpaint_native", model_info["supports_inpaint_native"]) + ) + info["repair_payload_version"] = str(info.get("repair_payload_version") or "1.0") + info["warnings"] = _normalize_warning_list(info.get("warnings")) + return info + + +def resolve_repair_scope( + requested_scope: str, + *, + mask_area_ratio: float, + mask_bbox: tuple[int, int, int, int] | None, + image_size: tuple[int, int], + canvas_expand: tuple[int, int, int, int], +) -> str: + scope = str(requested_scope or "auto") + if scope != "auto": + if scope not in _REPAIR_SCOPE_VALUES: + raise RuntimeError(f"[LLS] Unsupported repair_scope '{requested_scope}'.") + return scope + + if any(_coerce_int(value, 0) > 0 for value in canvas_expand): + return "canvas" + + if mask_bbox is None: + raise RuntimeError("[LLS] mask is empty after preprocessing.") + + image_width, image_height = image_size + x1, y1, x2, y2 = mask_bbox + bbox_width_ratio = max(0, x2 - x1) / float(max(1, image_width)) + bbox_height_ratio = max(0, y2 - y1) / float(max(1, image_height)) + if mask_area_ratio <= 0.18: + return "crop" + if mask_area_ratio <= 0.35 and bbox_width_ratio <= 0.8 and bbox_height_ratio <= 0.8: + return "crop" + return "region" + + +def resolve_repair_kernel( + requested_kernel: str, + *, + scope: str, + task_hint: str, + mask_area_ratio: float, + model_info: dict[str, Any] | str | None, +) -> tuple[str, list[str]]: + normalized_model = normalize_model_info(model_info) + kernel = str(requested_kernel or "auto") + warnings: list[str] = [] + + if kernel != "auto": + if kernel not in _REPAIR_KERNEL_VALUES: + raise RuntimeError(f"[LLS] Unsupported repair_kernel '{requested_kernel}'.") + if kernel == "native_fill" and not normalized_model["supports_inpaint_native"]: + warnings.append("native_fill requested but unsupported; falling back to vae_inpaint") + return "vae_inpaint", warnings + return kernel, warnings + + if ( + normalized_model["model_role"] in {"inpaint", "fill", "edit"} + and normalized_model["supports_inpaint_native"] + ): + return "native_fill", warnings + + if scope != "canvas" and mask_area_ratio <= 0.20 and task_hint in _TASK_HINT_LATENT_MASK: + return "latent_mask", warnings + + return "vae_inpaint", warnings + + +def recommend_denoise(task_hint: str, scope: str, kernel: str, auto_recommend: str) -> float: + if auto_recommend == "disabled": + if scope == "canvas": + return 0.90 + if scope == "crop": + return 0.50 + if kernel == "latent_mask": + return 0.45 + return 0.55 + + base = _TASK_DENOISE_VALUES.get(str(task_hint or "auto"), 0.55) + if scope == "canvas": + return max(0.90, base) + if scope == "crop": + return max(0.30, min(0.65, base)) + return base + + +def clamp_box(box: tuple[int, int, int, int], image_size: tuple[int, int]) -> tuple[int, int, int, int]: + image_width, image_height = image_size + x1, y1, x2, y2 = box + x1 = max(0, min(_coerce_int(x1, 0), image_width)) + y1 = max(0, min(_coerce_int(y1, 0), image_height)) + x2 = max(x1, min(_coerce_int(x2, x1), image_width)) + y2 = max(y1, min(_coerce_int(y2, y1), image_height)) + return x1, y1, x2, y2 + + +def compute_crop_box( + mask_bbox: tuple[int, int, int, int], + image_size: tuple[int, int], + crop_context: int, + crop_context_factor: float, +) -> tuple[int, int, int, int]: + x1, y1, x2, y2 = mask_bbox + box_width = max(1, x2 - x1) + box_height = max(1, y2 - y1) + center_x = (x1 + x2) / 2.0 + center_y = (y1 + y2) / 2.0 + expanded_width = (box_width + (2 * max(0, crop_context))) * max(1.0, float(crop_context_factor)) + expanded_height = (box_height + (2 * max(0, crop_context))) * max(1.0, float(crop_context_factor)) + half_width = expanded_width / 2.0 + half_height = expanded_height / 2.0 + return clamp_box( + ( + round(center_x - half_width), + round(center_y - half_height), + round(center_x + half_width), + round(center_y + half_height), + ), + image_size, + ) + + +def resolve_work_size( + crop_size: tuple[int, int], + min_size: int, + max_size: int, + resize_mode: str, +) -> tuple[int, int, float]: + crop_width = max(1, _coerce_int(crop_size[0], 1)) + crop_height = max(1, _coerce_int(crop_size[1], 1)) + target_min = max(1, _coerce_int(min_size, crop_width)) + target_max = max(target_min, _coerce_int(max_size, target_min)) + mode = _RESIZE_MODE_ALIASES.get(str(resize_mode or "fit"), "keep_aspect") + + if mode == "force_square": + side = min(target_max, max(target_min, max(crop_width, crop_height))) + scale = side / float(max(crop_width, crop_height)) + return side, side, scale + + longest = max(crop_width, crop_height) + shortest = min(crop_width, crop_height) + + if mode == "ranged_size" and target_min <= shortest and longest <= target_max: + return crop_width, crop_height, 1.0 + + if mode == "keep_aspect" and target_min <= longest <= target_max: + return crop_width, crop_height, 1.0 + + if longest < target_min: + scale = target_min / float(longest) + elif longest > target_max: + scale = target_max / float(longest) + else: + scale = 1.0 + + width = max(1, int(round(crop_width * scale))) + height = max(1, int(round(crop_height * scale))) + return width, height, scale + + +def build_canvas_info( + image_size: tuple[int, int], + expand_left: int, + expand_right: int, + expand_top: int, + expand_bottom: int, +) -> dict[str, tuple[int, int] | tuple[int, int, int, int]]: + image_width, image_height = image_size + left = max(0, _coerce_int(expand_left, 0)) + right = max(0, _coerce_int(expand_right, 0)) + top = max(0, _coerce_int(expand_top, 0)) + bottom = max(0, _coerce_int(expand_bottom, 0)) + + return { + "work_size": (image_width + left + right, image_height + top + bottom), + "original_box": (left, top, left + image_width, top + image_height), + } + + +def resolve_adapter_mode(adapter_mode: str, model_family: str) -> str: + mode = str(adapter_mode or "auto") + if mode != "auto": + return mode + + family = str(model_family or "") + if family in {"SD1.5", "SDXL", "SDXL_TURBO"}: + return "sd_classic" + if family in {"FLUX", "FLUX_DEV", "FLUX_SCHNELL"}: + return "flux" + if family == "SD3": + return "sd3" + if family == "QWEN": + return "qwen" + if family == "ZIMAGE": + return "zimage" + return "auto" + + +def _coerce_int(value: Any, default: int) -> int: + try: + return int(value) + except (TypeError, ValueError): + return default + + +def _coerce_float(value: Any, default: float) -> float: + try: + return float(value) + except (TypeError, ValueError): + return default + + +def _normalize_box(value: Any) -> tuple[int, int, int, int] | None: + if value is None: + return None + if isinstance(value, (list, tuple)) and len(value) == 4: + return tuple(_coerce_int(item, 0) for item in value) + return None + + +def _normalize_size_pair(value: Any, *, default: tuple[int, int]) -> tuple[int, int]: + if isinstance(value, (list, tuple)) and len(value) == 2: + return tuple(_coerce_int(item, 0) for item in value) + return default + + +def _normalize_expand(value: Any) -> tuple[int, int, int, int]: + if isinstance(value, (list, tuple)) and len(value) == 4: + return tuple(_coerce_int(item, 0) for item in value) + return (0, 0, 0, 0) + + +def _normalize_optional_float(value: Any) -> float | None: + if value in (None, ""): + return None + return _coerce_float(value, 0.0) + + +def _normalize_warning_list(value: Any) -> list[str]: + if isinstance(value, list): + return [str(item) for item in value] + if value in (None, ""): + return [] + return [str(value)] diff --git a/tests/test_repair_utils.py b/tests/test_repair_utils.py new file mode 100644 index 0000000..e174093 --- /dev/null +++ b/tests/test_repair_utils.py @@ -0,0 +1,163 @@ +import importlib +import importlib.util +import pathlib +import sys +import unittest + + +ROOT = pathlib.Path(__file__).resolve().parents[1] +MODULE_NAME = "lls_node_test_repair_utils" + + +def load_repair_utils(): + 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 importlib.import_module(f"{MODULE_NAME}.repair.repair_utils") + + +class TestRepairUtils(unittest.TestCase): + def test_resolve_repair_scope_prefers_canvas_then_crop_then_region(self): + utils = load_repair_utils() + + self.assertEqual( + utils.resolve_repair_scope( + "auto", + mask_area_ratio=0.0, + mask_bbox=None, + image_size=(1024, 1024), + canvas_expand=(128, 0, 0, 0), + ), + "canvas", + ) + self.assertEqual( + utils.resolve_repair_scope( + "auto", + mask_area_ratio=0.08, + mask_bbox=(100, 100, 260, 260), + image_size=(1024, 1024), + canvas_expand=(0, 0, 0, 0), + ), + "crop", + ) + self.assertEqual( + utils.resolve_repair_scope( + "auto", + mask_area_ratio=0.62, + mask_bbox=(0, 0, 1024, 1024), + image_size=(1024, 1024), + canvas_expand=(0, 0, 0, 0), + ), + "region", + ) + + def test_resolve_repair_kernel_prefers_native_then_latent_then_vae(self): + utils = load_repair_utils() + + self.assertEqual( + utils.resolve_repair_kernel( + "auto", + scope="region", + task_hint="repair", + mask_area_ratio=0.05, + model_info={"model_role": "fill", "supports_inpaint_native": True}, + )[0], + "native_fill", + ) + self.assertEqual( + utils.resolve_repair_kernel( + "auto", + scope="region", + task_hint="appearance", + mask_area_ratio=0.05, + model_info={"model_role": "normal", "supports_inpaint_native": False}, + )[0], + "latent_mask", + ) + self.assertEqual( + utils.resolve_repair_kernel( + "auto", + scope="canvas", + task_hint="fill", + mask_area_ratio=0.50, + model_info={"model_role": "normal", "supports_inpaint_native": False}, + )[0], + "vae_inpaint", + ) + + def test_recommend_denoise_applies_required_rules(self): + utils = load_repair_utils() + + self.assertEqual(utils.recommend_denoise("repair", "region", "latent_mask", "enabled"), 0.45) + self.assertEqual(utils.recommend_denoise("fill", "canvas", "vae_inpaint", "enabled"), 0.90) + self.assertEqual(utils.recommend_denoise("replace", "crop", "vae_inpaint", "enabled"), 0.65) + self.assertEqual(utils.recommend_denoise("auto", "crop", "vae_inpaint", "disabled"), 0.50) + + def test_crop_and_canvas_geometry_helpers(self): + utils = load_repair_utils() + + self.assertEqual( + utils.clamp_box((-20, 10, 1200, 2048), (1024, 768)), + (0, 10, 1024, 768), + ) + + crop_box = utils.compute_crop_box( + mask_bbox=(900, 900, 1100, 1100), + image_size=(1024, 1024), + crop_context=64, + crop_context_factor=1.5, + ) + self.assertEqual(crop_box, (754, 754, 1024, 1024)) + + keep_aspect = utils.resolve_work_size( + crop_size=(288, 288), + min_size=512, + max_size=1024, + resize_mode="fit", + ) + self.assertEqual(keep_aspect, (512, 512, 512 / 288)) + + square = utils.resolve_work_size( + crop_size=(400, 200), + min_size=256, + max_size=320, + resize_mode="pad", + ) + self.assertEqual(square, (320, 320, 320 / 400)) + + canvas_info = utils.build_canvas_info((640, 480), 64, 32, 16, 8) + self.assertEqual(canvas_info["work_size"], (736, 504)) + self.assertEqual(canvas_info["original_box"], (64, 16, 704, 496)) + + def test_normalize_repair_info_backfills_required_defaults(self): + utils = load_repair_utils() + + info = utils.normalize_repair_info({"repair_scope": "region", "repair_kernel": "latent_mask"}) + self.assertEqual(info["repair_scope"], "region") + self.assertEqual(info["repair_kernel"], "latent_mask") + self.assertEqual(info["task_hint"], "auto") + self.assertEqual(info["original_size"], (0, 0)) + self.assertEqual(info["work_size"], (0, 0)) + self.assertEqual(info["canvas_expand"], (0, 0, 0, 0)) + self.assertEqual(info["mask_grow"], 8) + self.assertEqual(info["mask_blur"], 8.0) + self.assertEqual(info["mask_threshold"], 0.5) + self.assertFalse(info["invert_mask"]) + self.assertEqual(info["recommended_denoise"], 0.55) + self.assertEqual(info["model_family"], "UNKNOWN") + self.assertEqual(info["model_role"], "unknown") + self.assertEqual(info["repair_payload_version"], "1.0") + self.assertEqual(info["warnings"], []) + + +if __name__ == "__main__": + unittest.main()