From 98cec3f757d52a9937d3bcc80d60da2a70401954 Mon Sep 17 00:00:00 2001 From: Gin3601 <2405559452@qq.com> Date: Wed, 20 May 2026 19:36:59 +0800 Subject: [PATCH] fix: finalize repair prepare masking --- repair/repair_prepare.py | 62 +++++-- repair/repair_utils.py | 312 ++++++++++++++++++++++++++++++++++- tests/test_repair_helpers.py | 199 ++++++++++++++++++++-- tests/test_repair_prepare.py | 50 ++++++ tests/test_repair_utils.py | 36 ++++ 5 files changed, 632 insertions(+), 27 deletions(-) diff --git a/repair/repair_prepare.py b/repair/repair_prepare.py index b79ef53..b220920 100644 --- a/repair/repair_prepare.py +++ b/repair/repair_prepare.py @@ -2,6 +2,7 @@ from __future__ import annotations from .repair_utils import ( build_canvas_info, + build_canvas_repair_mask, compute_crop_box, crop_image_to, crop_mask_to, @@ -10,7 +11,9 @@ from .repair_utils import ( get_image_size, get_mask_metrics, make_noise_mask, + merge_masks, normalize_model_info, + preprocess_mask, recommend_denoise, resize_image_to, resize_mask_to, @@ -107,7 +110,7 @@ class LLSSimpleRepairPrepare: positive=None, negative=None, ): - del canvas_fill, positive, negative + del positive, negative if vae is None: raise RuntimeError("[LLS] Missing VAE for LLS Simple Repair Prepare.") @@ -115,7 +118,15 @@ class LLSSimpleRepairPrepare: image_width, image_height = get_image_size(image) original_size = (image_width, image_height) normalized_model = normalize_model_info(model_info) - mask_bbox, mask_area_ratio = get_mask_metrics(mask, original_size) + processed_mask = preprocess_mask( + mask, + original_size, + invert_mask=bool(invert_mask), + mask_threshold=float(mask_threshold), + mask_grow=int(mask_grow), + mask_blur=float(mask_blur), + ) + mask_bbox, mask_area_ratio = get_mask_metrics(processed_mask, original_size) canvas_expand = [ max(0, int(expand_left)), max(0, int(expand_right)), @@ -162,19 +173,23 @@ class LLSSimpleRepairPrepare: "model_family": normalized_model["model_family"], "model_role": normalized_model["model_role"], "repair_payload_version": "1.0", - "has_mask": mask_bbox is not None and mask_area_ratio > 0.0, - "mask_area_ratio": mask_area_ratio, - "mask_bbox": list(mask_bbox) if mask_bbox is not None else None, + "has_mask": False, + "mask_area_ratio": 0.0, + "mask_bbox": None, "warnings": warnings, } if effective_scope == "region": work_image = image - work_mask = resize_mask_to(mask, image_width, image_height) + work_mask = resize_mask_to(processed_mask, image_width, image_height) latent_samples = vae.encode(work_image) latent = {"samples": latent_samples, "source": "repair_prepare_region"} if effective_kernel == "latent_mask": latent["noise_mask"] = make_noise_mask(work_mask, latent_samples) + final_bbox, final_area_ratio = get_mask_metrics(work_mask, original_size) + repair_info["has_mask"] = final_bbox is not None and final_area_ratio > 0.0 + repair_info["mask_bbox"] = list(final_bbox) if final_bbox is not None else None + repair_info["mask_area_ratio"] = final_area_ratio return latent, work_image, work_mask, repair_info, recommended if effective_scope == "crop": @@ -191,7 +206,7 @@ class LLSSimpleRepairPrepare: resize_mode, ) cropped_image = crop_image_to(image, crop_box) - cropped_mask = crop_mask_to(mask, crop_box) + cropped_mask = crop_mask_to(processed_mask, crop_box) work_image = resize_image_to(cropped_image, work_width, work_height) work_mask = resize_mask_to(cropped_mask, work_width, work_height) latent_samples = vae.encode(work_image) @@ -202,6 +217,10 @@ class LLSSimpleRepairPrepare: repair_info["crop_box"] = list(crop_box) repair_info["crop_scale"] = crop_scale repair_info["work_size"] = (work_width, work_height) + final_bbox, final_area_ratio = get_mask_metrics(work_mask, (work_width, work_height)) + repair_info["has_mask"] = final_bbox is not None and final_area_ratio > 0.0 + repair_info["mask_bbox"] = list(final_bbox) if final_bbox is not None else None + repair_info["mask_area_ratio"] = final_area_ratio return latent, work_image, work_mask, repair_info, recommended if effective_scope == "canvas": @@ -213,15 +232,38 @@ class LLSSimpleRepairPrepare: int(expand_bottom), ) work_width, work_height = canvas_info["work_size"] - work_image = expand_canvas_image(image, work_width, work_height) - work_mask = expand_canvas_mask(mask, work_width, work_height) + original_box = canvas_info["original_box"] + work_image = expand_canvas_image( + image, + work_width, + work_height, + fill_mode=canvas_fill, + original_box=original_box, + ) + user_mask = expand_canvas_mask( + processed_mask, + work_width, + work_height, + original_box=original_box, + ) + canvas_mask = build_canvas_repair_mask( + user_mask, + work_width, + work_height, + original_box=original_box, + ) + work_mask = merge_masks(user_mask, canvas_mask, (work_width, work_height)) latent_samples = vae.encode(work_image) latent = {"samples": latent_samples, "source": "repair_prepare_canvas"} if effective_kernel == "latent_mask": latent["noise_mask"] = make_noise_mask(work_mask, latent_samples) repair_info["work_size"] = (work_width, work_height) - repair_info["original_box_in_canvas"] = list(canvas_info["original_box"]) + repair_info["original_box_in_canvas"] = list(original_box) + final_bbox, final_area_ratio = get_mask_metrics(work_mask, (work_width, work_height)) + repair_info["has_mask"] = final_bbox is not None and final_area_ratio > 0.0 + repair_info["mask_bbox"] = list(final_bbox) if final_bbox is not None else None + repair_info["mask_area_ratio"] = final_area_ratio return latent, work_image, work_mask, repair_info, recommended raise RuntimeError(f"[LLS] Unsupported repair scope '{effective_scope}'.") diff --git a/repair/repair_utils.py b/repair/repair_utils.py index 64a47ab..6dd4675 100644 --- a/repair/repair_utils.py +++ b/repair/repair_utils.py @@ -4,6 +4,25 @@ from typing import Any from ..utils.model_info import canonicalize_family, parse_jsonish_info +try: + import torch +except Exception as exc: # pragma: no cover - optional runtime dependency + torch = None + _TORCH_ERR = exc + torch_nn_functional = None +else: # pragma: no cover - exercised only in ComfyUI-like runtime + import torch.nn.functional as torch_nn_functional + + _TORCH_ERR = None + +try: + import comfy.utils as comfy_utils +except Exception as exc: # pragma: no cover - optional runtime dependency + comfy_utils = None + _COMFY_UTILS_ERR = exc +else: # pragma: no cover - exercised only in ComfyUI-like runtime + _COMFY_UTILS_ERR = None + REPAIR_INFO_TYPE = "LLS_REPAIR_INFO" GUIDANCE_STACK_TYPE = "LLS_GUIDANCE_STACK" @@ -56,6 +75,7 @@ def normalize_repair_info(repair_info: dict[str, Any] | str | None) -> dict[str, 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["original_box_in_canvas"] = _normalize_box(info.get("original_box_in_canvas")) 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) @@ -271,14 +291,35 @@ def get_mask_metrics(mask: Any, image_size: tuple[int, int]) -> tuple[tuple[int, if isinstance(raw_bbox, (list, tuple)) and len(raw_bbox) == 4: mask_bbox = clamp_box(tuple(_coerce_int(value, 0) for value in raw_bbox), image_size) - mask_area_ratio = _coerce_float(getattr(mask, "mask_area_ratio", 0.0), 0.0) - mask_area_ratio = max(0.0, min(1.0, mask_area_ratio)) + raw_area_ratio = getattr(mask, "mask_area_ratio", None) + if raw_area_ratio is not None: + mask_area_ratio = _coerce_float(raw_area_ratio, 0.0) + mask_area_ratio = max(0.0, min(1.0, mask_area_ratio)) + return mask_bbox, mask_area_ratio + + if _is_torch_tensor(mask): + active_map = _normalize_mask_tensor(mask) > 0.0 + if active_map.ndim == 3: + active_map = active_map.any(dim=0) + if not bool(active_map.any().item()): + return None, 0.0 + points = active_map.nonzero(as_tuple=False) + y1 = int(points[:, 0].min().item()) + y2 = int(points[:, 0].max().item()) + 1 + x1 = int(points[:, 1].min().item()) + x2 = int(points[:, 1].max().item()) + 1 + mask_bbox = clamp_box((x1, y1, x2, y2), image_size) + return mask_bbox, float(active_map.float().mean().item()) + + mask_area_ratio = 0.0 return mask_bbox, mask_area_ratio def resize_image_to(image: Any, width: int, height: int) -> Any: if hasattr(image, "resized"): return image.resized(width, height) + if _is_torch_tensor(image): + return _resize_image_tensor(image, width, height) current_width, current_height = get_image_size(image) if (current_width, current_height) == (width, height): return image @@ -291,6 +332,8 @@ def resize_mask_to(mask: Any, width: int, height: int) -> Any: raise RuntimeError("[LLS] mask must have shape [batch, height, width].") if hasattr(mask, "resized"): return mask.resized(width, height) + if _is_torch_tensor(mask): + return _resize_mask_tensor(mask, width, height) _, current_height, current_width = shape if (_coerce_int(current_width, 0), _coerce_int(current_height, 0)) == (width, height): return mask @@ -301,6 +344,9 @@ def crop_image_to(image: Any, box: tuple[int, int, int, int]) -> Any: cropped_box = clamp_box(box, get_image_size(image)) if hasattr(image, "cropped"): return image.cropped(*cropped_box) + if _is_torch_tensor(image): + x1, y1, x2, y2 = cropped_box + return image[:, y1:y2, x1:x2, :] current_width, current_height = get_image_size(image) if cropped_box == (0, 0, current_width, current_height): @@ -317,33 +363,149 @@ def crop_mask_to(mask: Any, box: tuple[int, int, int, int]) -> Any: cropped_box = clamp_box(box, (_coerce_int(current_width, 0), _coerce_int(current_height, 0))) if hasattr(mask, "cropped"): return mask.cropped(*cropped_box) + if _is_torch_tensor(mask): + x1, y1, x2, y2 = cropped_box + return mask[:, y1:y2, x1:x2] if cropped_box == (0, 0, _coerce_int(current_width, 0), _coerce_int(current_height, 0)): return mask raise RuntimeError("[LLS] mask object does not support cropping.") -def expand_canvas_image(image: Any, width: int, height: int) -> Any: +def expand_canvas_image( + image: Any, + width: int, + height: int, + *, + fill_mode: str = "edge", + original_box: tuple[int, int, int, int] | None = None, +) -> Any: if hasattr(image, "canvas_expanded"): - return image.canvas_expanded(width, height) + return image.canvas_expanded(width, height, fill_mode=fill_mode, original_box=original_box) + if _is_torch_tensor(image): + return _expand_canvas_image_tensor(image, width, height, fill_mode=fill_mode, original_box=original_box) current_width, current_height = get_image_size(image) if (current_width, current_height) == (width, height): return image raise RuntimeError("[LLS] image object does not support canvas expansion.") -def expand_canvas_mask(mask: Any, width: int, height: int) -> Any: +def expand_canvas_mask( + mask: Any, + width: int, + height: int, + *, + original_box: tuple[int, int, int, int] | None = None, +) -> Any: shape = getattr(mask, "shape", None) if not isinstance(shape, (list, tuple)) or len(shape) != 3: raise RuntimeError("[LLS] mask must have shape [batch, height, width].") if hasattr(mask, "canvas_expanded"): - return mask.canvas_expanded(width, height) + return mask.canvas_expanded(width, height, original_box=original_box) + if _is_torch_tensor(mask): + return _expand_canvas_mask_tensor(mask, width, height, original_box=original_box) _, current_height, current_width = shape if (_coerce_int(current_width, 0), _coerce_int(current_height, 0)) == (width, height): return mask raise RuntimeError("[LLS] mask object does not support canvas expansion.") +def preprocess_mask( + mask: Any, + image_size: tuple[int, int], + *, + invert_mask: bool, + mask_threshold: float, + mask_grow: int, + mask_blur: float, +) -> Any: + width, height = image_size + processed = resize_mask_to(mask, width, height) + processed = normalize_mask_values(processed) + if invert_mask: + processed = invert_mask_values(processed, image_size) + processed = apply_mask_threshold(processed, mask_threshold) + if _coerce_int(mask_grow, 0) > 0: + processed = apply_mask_grow(processed, _coerce_int(mask_grow, 0), image_size) + if _coerce_float(mask_blur, 0.0) > 0.0: + processed = apply_mask_blur(processed, _coerce_float(mask_blur, 0.0), image_size) + return clamp_mask_values(processed) + + +def normalize_mask_values(mask: Any) -> Any: + if hasattr(mask, "normalized"): + return mask.normalized() + if _is_torch_tensor(mask): + return _normalize_mask_tensor(mask) + return mask + + +def invert_mask_values(mask: Any, image_size: tuple[int, int]) -> Any: + if hasattr(mask, "inverted"): + return mask.inverted(image_size) + if _is_torch_tensor(mask): + return 1.0 - _normalize_mask_tensor(mask) + raise RuntimeError("[LLS] mask object does not support invert preprocessing.") + + +def apply_mask_threshold(mask: Any, threshold: float) -> Any: + if hasattr(mask, "thresholded"): + return mask.thresholded(float(threshold)) + if _is_torch_tensor(mask): + normalized = _normalize_mask_tensor(mask) + return (normalized >= float(threshold)).to(normalized.dtype) + return mask + + +def apply_mask_grow(mask: Any, amount: int, image_size: tuple[int, int]) -> Any: + if hasattr(mask, "grown"): + return mask.grown(int(amount), image_size) + if _is_torch_tensor(mask): + return _grow_mask_tensor(mask, int(amount)) + return mask + + +def apply_mask_blur(mask: Any, radius: float, image_size: tuple[int, int]) -> Any: + if hasattr(mask, "blurred"): + return mask.blurred(float(radius), image_size) + if _is_torch_tensor(mask): + return _blur_mask_tensor(mask, float(radius)) + return mask + + +def clamp_mask_values(mask: Any) -> Any: + if _is_torch_tensor(mask): + return _normalize_mask_tensor(mask) + return mask + + +def build_canvas_repair_mask( + mask: Any, + width: int, + height: int, + *, + original_box: tuple[int, int, int, int], +) -> Any: + canvas_shape = (int(getattr(mask, "shape", [1])[0]), height, width) + canvas_region_factory = getattr(type(mask), "canvas_region", None) + if callable(canvas_region_factory): + return canvas_region_factory(canvas_shape, original_box, label="canvas-region") + if _is_torch_tensor(mask): + repair_mask = torch.ones(canvas_shape, dtype=_normalize_mask_tensor(mask).dtype, device=mask.device) + x1, y1, x2, y2 = original_box + repair_mask[:, y1:y2, x1:x2] = 0.0 + return repair_mask + raise RuntimeError("[LLS] mask object does not support canvas repair mask creation.") + + +def merge_masks(mask: Any, other: Any, image_size: tuple[int, int]) -> Any: + if hasattr(mask, "merged_with"): + return mask.merged_with(other, image_size) + if _is_torch_tensor(mask) and _is_torch_tensor(other): + return torch.maximum(_normalize_mask_tensor(mask), _normalize_mask_tensor(other)) + raise RuntimeError("[LLS] mask objects do not support merging.") + + def make_noise_mask(mask: Any, latent_samples: Any) -> Any: latent_shape = getattr(latent_samples, "shape", None) if not isinstance(latent_shape, (list, tuple)) or len(latent_shape) != 4: @@ -357,6 +519,144 @@ def make_noise_mask(mask: Any, latent_samples: Any) -> Any: return resize_mask_to(mask, latent_width, latent_height) +def _is_torch_tensor(value: Any) -> bool: + return bool(torch is not None and isinstance(value, torch.Tensor)) + + +def _resize_image_tensor(image: Any, width: int, height: int) -> Any: + current_width, current_height = get_image_size(image) + if (current_width, current_height) == (width, height): + return image + channel_first = image.movedim(-1, 1) + resized = _upscale_tensor(channel_first, width, height, mode="bilinear") + return resized.movedim(1, -1) + + +def _resize_mask_tensor(mask: Any, width: int, height: int) -> Any: + _, current_height, current_width = tuple(mask.shape) + if (int(current_width), int(current_height)) == (width, height): + return mask + channel_first = _normalize_mask_tensor(mask).unsqueeze(1) + resized = _upscale_tensor(channel_first, width, height, mode="bilinear") + return resized.squeeze(1) + + +def _upscale_tensor(channel_first: Any, width: int, height: int, *, mode: str) -> Any: + if comfy_utils is not None: + try: + return comfy_utils.common_upscale(channel_first, width, height, mode, "disabled") + except Exception: + pass + if torch_nn_functional is None: + raise RuntimeError( + "[LLS] torch or comfy.utils is required for tensor resizing outside fake-object tests." + ) from _TORCH_ERR + interpolate_kwargs = { + "size": (height, width), + "mode": mode, + } + if mode in {"bilinear", "bicubic"}: + interpolate_kwargs["align_corners"] = False + return torch_nn_functional.interpolate(channel_first, **interpolate_kwargs) + + +def _expand_canvas_mask_tensor( + mask: Any, + width: int, + height: int, + *, + original_box: tuple[int, int, int, int] | None, +) -> Any: + if original_box is None: + original_box = (0, 0, mask.shape[2], mask.shape[1]) + expanded = torch.zeros((mask.shape[0], height, width), dtype=_normalize_mask_tensor(mask).dtype, device=mask.device) + x1, y1, x2, y2 = original_box + expanded[:, y1:y2, x1:x2] = _normalize_mask_tensor(mask) + return expanded + + +def _expand_canvas_image_tensor( + image: Any, + width: int, + height: int, + *, + fill_mode: str, + original_box: tuple[int, int, int, int] | None, +) -> Any: + if original_box is None: + original_box = (0, 0, image.shape[2], image.shape[1]) + x1, y1, x2, y2 = original_box + left = x1 + top = y1 + right = max(0, width - x2) + bottom = max(0, height - y2) + + if fill_mode == "edge": + channel_first = image.movedim(-1, 1) + return torch_nn_functional.pad(channel_first, (left, right, top, bottom), mode="replicate").movedim(1, -1) + + if fill_mode == "blur": + channel_first = image.movedim(-1, 1) + padded = torch_nn_functional.pad(channel_first, (left, right, top, bottom), mode="replicate") + blur_radius = max(1, min(8, max(left, right, top, bottom))) + blurred = torch_nn_functional.avg_pool2d( + padded, + kernel_size=(blur_radius * 2) + 1, + stride=1, + padding=blur_radius, + ).movedim(1, -1) + blurred[:, y1:y2, x1:x2, :] = image + return blurred + + fill_value = 0.0 + if fill_mode == "white": + fill_value = 1.0 + elif fill_mode == "neutral": + fill_value = 0.5 + canvas = image.new_full((image.shape[0], height, width, image.shape[3]), fill_value) + canvas[:, y1:y2, x1:x2, :] = image + return canvas + + +def _normalize_mask_tensor(mask: Any) -> Any: + if not _is_torch_tensor(mask): + return mask + normalized = mask.float() + if normalized.numel() == 0: + return normalized + max_value = float(normalized.max().item()) + min_value = float(normalized.min().item()) + if max_value > 1.0 and max_value <= 255.0 and min_value >= 0.0: + normalized = normalized / 255.0 + return normalized.clamp(0.0, 1.0) + + +def _grow_mask_tensor(mask: Any, amount: int) -> Any: + if amount <= 0: + return _normalize_mask_tensor(mask) + if torch_nn_functional is None: + raise RuntimeError("[LLS] torch is required for tensor mask grow operations.") from _TORCH_ERR + normalized = _normalize_mask_tensor(mask).unsqueeze(1) + grown = torch_nn_functional.max_pool2d(normalized, kernel_size=(amount * 2) + 1, stride=1, padding=amount) + return grown.squeeze(1) + + +def _blur_mask_tensor(mask: Any, radius: float) -> Any: + blur_amount = max(0, int(round(radius))) + if blur_amount <= 0: + return _normalize_mask_tensor(mask) + if torch_nn_functional is None: + raise RuntimeError("[LLS] torch is required for tensor mask blur operations.") from _TORCH_ERR + normalized = _normalize_mask_tensor(mask).unsqueeze(1) + blurred = torch_nn_functional.avg_pool2d( + normalized, + kernel_size=(blur_amount * 2) + 1, + stride=1, + padding=blur_amount, + ) + return blurred.squeeze(1) + + def resolve_adapter_mode(adapter_mode: str, model_family: str) -> str: mode = str(adapter_mode or "auto") if mode != "auto": diff --git a/tests/test_repair_helpers.py b/tests/test_repair_helpers.py index da6aa83..1c75cd6 100644 --- a/tests/test_repair_helpers.py +++ b/tests/test_repair_helpers.py @@ -5,13 +5,16 @@ import sys ROOT = pathlib.Path(__file__).resolve().parents[1] MODULE_NAME = "lls_node_test_repair" +_UNSET = object() class FakeTensor: - def __init__(self, shape, label="image", crop_box=None): + def __init__(self, shape, label="image", crop_box=None, fill_mode=None, original_box=None): self.shape = tuple(shape) self.label = label self.crop_box = crop_box + self.fill_mode = fill_mode + self.original_box = original_box def cropped(self, x1, y1, x2, y2): batch, _, _, channels = self.shape @@ -20,6 +23,8 @@ class FakeTensor: (batch, max(0, y2 - y1), max(0, x2 - x1), channels), label=f"{self.label}:crop[{x1},{y1},{x2},{y2}]", crop_box=crop_box, + fill_mode=self.fill_mode, + original_box=self.original_box, ) def resized(self, width, height): @@ -28,14 +33,18 @@ class FakeTensor: (batch, height, width, channels), label=f"{self.label}:resized", crop_box=self.crop_box, + fill_mode=self.fill_mode, + original_box=self.original_box, ) - def canvas_expanded(self, width, height): + def canvas_expanded(self, width, height, fill_mode="edge", original_box=None): batch, _, _, channels = self.shape return FakeTensor( (batch, height, width, channels), - label=f"{self.label}:canvas", + label=f"{self.label}:canvas[{fill_mode}]", crop_box=self.crop_box, + fill_mode=fill_mode, + original_box=original_box, ) @@ -47,37 +56,205 @@ class FakeMask: self.label = label self.crop_box = crop_box + def _clone(self, shape=None, mask_bbox=_UNSET, mask_area_ratio=None, label=None, crop_box=None): + return FakeMask( + shape or self.shape, + mask_bbox=self.mask_bbox if mask_bbox is _UNSET else mask_bbox, + mask_area_ratio=self.mask_area_ratio if mask_area_ratio is None else mask_area_ratio, + label=self.label if label is None else label, + crop_box=self.crop_box if crop_box is None else crop_box, + ) + + def _current_size(self): + return self.shape[2], self.shape[1] + + @staticmethod + def _clamp_box(box, width, height): + if box is None: + return None + x1, y1, x2, y2 = box + x1 = max(0, min(int(x1), width)) + y1 = max(0, min(int(y1), height)) + x2 = max(x1, min(int(x2), width)) + y2 = max(y1, min(int(y2), height)) + if x2 <= x1 or y2 <= y1: + return None + return (x1, y1, x2, y2) + + @staticmethod + def _bbox_area_ratio(box, width, height): + if box is None or width <= 0 or height <= 0: + return 0.0 + return ((box[2] - box[0]) * (box[3] - box[1])) / float(width * height) + def cropped(self, x1, y1, x2, y2): batch, _, _ = self.shape crop_box = (x1, y1, x2, y2) + new_width = max(0, x2 - x1) + new_height = max(0, y2 - y1) + new_bbox = None + if self.mask_bbox is not None: + new_bbox = self._clamp_box( + ( + self.mask_bbox[0] - x1, + self.mask_bbox[1] - y1, + self.mask_bbox[2] - x1, + self.mask_bbox[3] - y1, + ), + new_width, + new_height, + ) return FakeMask( - (batch, max(0, y2 - y1), max(0, x2 - x1)), - mask_bbox=self.mask_bbox, - mask_area_ratio=self.mask_area_ratio, + (batch, new_height, new_width), + mask_bbox=new_bbox, + mask_area_ratio=self._bbox_area_ratio(new_bbox, new_width, new_height), label=f"{self.label}:crop[{x1},{y1},{x2},{y2}]", crop_box=crop_box, ) def resized(self, width, height): batch, _, _ = self.shape + source_width, source_height = self._current_size() + new_bbox = None + if self.mask_bbox is not None and source_width > 0 and source_height > 0: + scale_x = width / float(source_width) + scale_y = height / float(source_height) + new_bbox = self._clamp_box( + ( + round(self.mask_bbox[0] * scale_x), + round(self.mask_bbox[1] * scale_y), + round(self.mask_bbox[2] * scale_x), + round(self.mask_bbox[3] * scale_y), + ), + width, + height, + ) return FakeMask( (batch, height, width), - mask_bbox=self.mask_bbox, - mask_area_ratio=self.mask_area_ratio, + mask_bbox=new_bbox, + mask_area_ratio=self._bbox_area_ratio(new_bbox, width, height), label=f"{self.label}:resized", crop_box=self.crop_box, ) - def canvas_expanded(self, width, height): + def canvas_expanded(self, width, height, original_box=None): batch, _, _ = self.shape + if original_box is None: + original_box = (0, 0, self.shape[2], self.shape[1]) + offset_x, offset_y = original_box[0], original_box[1] + new_bbox = None + if self.mask_bbox is not None: + new_bbox = self._clamp_box( + ( + self.mask_bbox[0] + offset_x, + self.mask_bbox[1] + offset_y, + self.mask_bbox[2] + offset_x, + self.mask_bbox[3] + offset_y, + ), + width, + height, + ) return FakeMask( (batch, height, width), - mask_bbox=self.mask_bbox, - mask_area_ratio=self.mask_area_ratio, + mask_bbox=new_bbox, + mask_area_ratio=self._bbox_area_ratio(new_bbox, width, height), label=f"{self.label}:canvas", crop_box=self.crop_box, ) + def normalized(self): + return self._clone(label=f"{self.label}:normalize") + + def inverted(self, image_size): + width, height = image_size + inverted_ratio = max(0.0, min(1.0, 1.0 - self.mask_area_ratio)) + inverted_bbox = None if inverted_ratio <= 0.0 else (0, 0, width, height) + return self._clone( + shape=(self.shape[0], height, width), + mask_bbox=inverted_bbox, + mask_area_ratio=inverted_ratio, + label=f"{self.label}:invert", + ) + + def thresholded(self, threshold): + width, height = self._current_size() + if self.mask_bbox is None or self.mask_area_ratio <= 0.0: + return self._clone(label=f"{self.label}:threshold[{threshold}]") + shrink_ratio = max(0.0, min(1.0, 1.0 - max(0.0, float(threshold) - 0.5))) + if shrink_ratio <= 0.0: + return self._clone(mask_bbox=None, mask_area_ratio=0.0, label=f"{self.label}:threshold[{threshold}]") + x1, y1, x2, y2 = self.mask_bbox + cx = (x1 + x2) / 2.0 + cy = (y1 + y2) / 2.0 + half_width = max(1.0, (x2 - x1) * shrink_ratio / 2.0) + half_height = max(1.0, (y2 - y1) * shrink_ratio / 2.0) + new_bbox = self._clamp_box( + ( + round(cx - half_width), + round(cy - half_height), + round(cx + half_width), + round(cy + half_height), + ), + width, + height, + ) + return self._clone( + mask_bbox=new_bbox, + mask_area_ratio=self._bbox_area_ratio(new_bbox, width, height), + label=f"{self.label}:threshold[{threshold}]", + ) + + def grown(self, amount, image_size): + width, height = image_size + if self.mask_bbox is None: + return self._clone(shape=(self.shape[0], height, width), label=f"{self.label}:grow[{amount}]") + x1, y1, x2, y2 = self.mask_bbox + new_bbox = self._clamp_box((x1 - amount, y1 - amount, x2 + amount, y2 + amount), width, height) + return self._clone( + shape=(self.shape[0], height, width), + mask_bbox=new_bbox, + mask_area_ratio=self._bbox_area_ratio(new_bbox, width, height), + label=f"{self.label}:grow[{amount}]", + ) + + def blurred(self, radius, image_size): + blur_expand = max(1, int(round(float(radius) / 4.0))) + return self.grown(blur_expand, image_size)._clone(label=f"{self.label}:blur[{radius}]") + + def merged_with(self, other, image_size): + width, height = image_size + if self.mask_bbox is None: + return other._clone(shape=(other.shape[0], height, width), label=f"{self.label}:merge:{other.label}") + if other.mask_bbox is None: + return self._clone(shape=(self.shape[0], height, width), label=f"{self.label}:merge:{other.label}") + new_bbox = self._clamp_box( + ( + min(self.mask_bbox[0], other.mask_bbox[0]), + min(self.mask_bbox[1], other.mask_bbox[1]), + max(self.mask_bbox[2], other.mask_bbox[2]), + max(self.mask_bbox[3], other.mask_bbox[3]), + ), + width, + height, + ) + return self._clone( + shape=(self.shape[0], height, width), + mask_bbox=new_bbox, + mask_area_ratio=self._bbox_area_ratio(new_bbox, width, height), + label=f"{self.label}:merge:{other.label}", + ) + + @classmethod + def canvas_region(cls, shape, original_box, label="canvas-region"): + batch, height, width = shape + original_width = max(0, original_box[2] - original_box[0]) + original_height = max(0, original_box[3] - original_box[1]) + total_area = max(1, width * height) + original_area = max(0, original_width * original_height) + area_ratio = max(0.0, min(1.0, (total_area - original_area) / float(total_area))) + bbox = None if area_ratio <= 0.0 else (0, 0, width, height) + return cls(shape, mask_bbox=bbox, mask_area_ratio=area_ratio, label=label) + class FakeLatentTensor: def __init__(self, shape): diff --git a/tests/test_repair_prepare.py b/tests/test_repair_prepare.py index 4eefeab..24c5734 100644 --- a/tests/test_repair_prepare.py +++ b/tests/test_repair_prepare.py @@ -138,6 +138,56 @@ class TestRepairPrepare(unittest.TestCase): self.assertEqual(repair_info["canvas_expand"], [128, 0, 0, 0]) self.assertEqual(recommended_denoise, 0.90) self.assertEqual(latent["source"], "repair_prepare_canvas") + self.assertGreater(work_mask.mask_area_ratio, 0.0) + self.assertTrue(repair_info["has_mask"]) + self.assertIsNotNone(repair_info["mask_bbox"]) + self.assertEqual(work_image.fill_mode, "edge") + + def test_prepare_region_applies_mask_preprocessing_before_metrics(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]() + + image = FakeTensor((1, 512, 512, 3), label="preprocess-image") + mask = FakeMask( + (1, 512, 512), + mask_bbox=(200, 200, 280, 280), + mask_area_ratio=((280 - 200) * (280 - 200)) / float(512 * 512), + label="preprocess-mask", + ) + + _latent, _work_image, work_mask, repair_info, _recommended_denoise = node.prepare( + image=image, + mask=mask, + vae=FakeVAE(), + repair_scope="region", + repair_kernel="vae_inpaint", + task_hint="content", + mask_grow=24, + mask_blur=8.0, + mask_threshold=0.8, + invert_mask=False, + crop_context=64, + crop_context_factor=1.5, + min_size=256, + max_size=1024, + resize_mode="fit", + expand_left=0, + expand_right=0, + expand_top=0, + expand_bottom=0, + canvas_fill="edge", + auto_recommend="enabled", + model_info={"model_family": "SDXL", "model_role": "base"}, + ) + + self.assertIn("normalize", work_mask.label) + self.assertIn("threshold[0.8]", work_mask.label) + self.assertIn("grow[24]", work_mask.label) + self.assertIn("blur[8.0]", work_mask.label) + self.assertNotEqual(work_mask.mask_bbox, mask.mask_bbox) + self.assertGreater(work_mask.mask_area_ratio, mask.mask_area_ratio) + self.assertEqual(repair_info["mask_bbox"], list(work_mask.mask_bbox)) + self.assertEqual(repair_info["mask_area_ratio"], work_mask.mask_area_ratio) def test_prepare_requires_vae(self): plugin = load_plugin_package() diff --git a/tests/test_repair_utils.py b/tests/test_repair_utils.py index 6cb0076..0ad8055 100644 --- a/tests/test_repair_utils.py +++ b/tests/test_repair_utils.py @@ -4,6 +4,11 @@ import pathlib import sys import unittest +try: + from .test_repair_helpers import FakeMask +except ImportError: # pragma: no cover - discovery mode imports from top level + from test_repair_helpers import FakeMask + ROOT = pathlib.Path(__file__).resolve().parents[1] MODULE_NAME = "lls_node_test_repair_utils" @@ -216,6 +221,37 @@ class TestRepairUtils(unittest.TestCase): self.assertTrue(utils.normalize_repair_info({"invert_mask": "true"})["invert_mask"]) self.assertTrue(utils.normalize_repair_info({"invert_mask": 1})["invert_mask"]) + def test_preprocess_mask_applies_requested_operations(self): + utils = load_repair_utils() + + mask = FakeMask( + (1, 256, 256), + mask_bbox=(100, 100, 140, 140), + mask_area_ratio=((140 - 100) * (140 - 100)) / float(256 * 256), + label="preprocess-mask", + ) + processed = utils.preprocess_mask( + mask, + (256, 256), + invert_mask=False, + mask_threshold=0.8, + mask_grow=12, + mask_blur=6.0, + ) + + self.assertIn("normalize", processed.label) + self.assertIn("threshold[0.8]", processed.label) + self.assertIn("grow[12]", processed.label) + self.assertIn("blur[6.0]", processed.label) + self.assertNotEqual(processed.mask_bbox, mask.mask_bbox) + self.assertGreater(processed.mask_area_ratio, mask.mask_area_ratio) + + def test_normalize_repair_info_normalizes_original_box_in_canvas(self): + utils = load_repair_utils() + + info = utils.normalize_repair_info({"original_box_in_canvas": ["8", 16, 128, 144]}) + self.assertEqual(info["original_box_in_canvas"], (8, 16, 128, 144)) + if __name__ == "__main__": unittest.main()