diff --git a/repair/repair_prepare.py b/repair/repair_prepare.py index 15b7c2f..4b29961 100644 --- a/repair/repair_prepare.py +++ b/repair/repair_prepare.py @@ -1,5 +1,22 @@ from __future__ import annotations +from .repair_utils import ( + build_canvas_info, + compute_crop_box, + expand_canvas_image, + expand_canvas_mask, + get_image_size, + get_mask_metrics, + make_noise_mask, + normalize_model_info, + recommend_denoise, + resize_image_to, + resize_mask_to, + resolve_repair_kernel, + resolve_repair_scope, + resolve_work_size, +) + _REPAIR_SCOPE_CHOICES = ["auto", "region", "crop", "canvas"] _REPAIR_KERNEL_CHOICES = ["auto", "latent_mask", "vae_inpaint", "native_fill"] @@ -61,8 +78,149 @@ class LLSSimpleRepairPrepare: }, } - def prepare(self, *args, **kwargs): - raise RuntimeError("[LLS] LLS Simple Repair Prepare is not implemented yet.") + def prepare( + self, + 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, + model_info=None, + positive=None, + negative=None, + ): + del canvas_fill, positive, negative + + if vae is None: + raise RuntimeError("[LLS] Missing VAE for LLS Simple Repair Prepare.") + + 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) + canvas_expand = [ + max(0, int(expand_left)), + max(0, int(expand_right)), + max(0, int(expand_top)), + max(0, int(expand_bottom)), + ] + + effective_scope = resolve_repair_scope( + repair_scope, + mask_area_ratio=mask_area_ratio, + mask_bbox=mask_bbox, + image_size=original_size, + canvas_expand=tuple(canvas_expand), + ) + effective_kernel, warnings = resolve_repair_kernel( + repair_kernel, + scope=effective_scope, + task_hint=task_hint, + mask_area_ratio=mask_area_ratio, + model_info=normalized_model, + ) + warnings = list(warnings) + if effective_kernel == "native_fill": + effective_kernel = "vae_inpaint" + warnings.append( + "native_fill requested but runtime sampler support is not implemented yet; using vae_inpaint" + ) + + recommended = recommend_denoise(task_hint, effective_scope, effective_kernel, auto_recommend) + repair_info = { + "repair_scope": effective_scope, + "repair_kernel": effective_kernel, + "task_hint": str(task_hint or "auto"), + "original_size": original_size, + "work_size": original_size, + "crop_box": None, + "crop_scale": 1.0, + "canvas_expand": canvas_expand, + "mask_grow": int(mask_grow), + "mask_blur": float(mask_blur), + "mask_threshold": float(mask_threshold), + "invert_mask": bool(invert_mask), + "recommended_denoise": recommended, + "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, + "warnings": warnings, + } + + if effective_scope == "region": + work_image = image + work_mask = resize_mask_to(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) + return latent, work_image, work_mask, repair_info, recommended + + if effective_scope == "crop": + if mask_bbox is None or mask_area_ratio <= 0.0: + raise RuntimeError("[LLS] crop scope requires a non-empty mask.") + + crop_box = compute_crop_box(mask_bbox, original_size, int(crop_context), float(crop_context_factor)) + crop_width = max(1, crop_box[2] - crop_box[0]) + crop_height = max(1, crop_box[3] - crop_box[1]) + work_width, work_height, crop_scale = resolve_work_size( + (crop_width, crop_height), + int(min_size), + int(max_size), + resize_mode, + ) + work_image = resize_image_to(image, work_width, work_height) + work_mask = resize_mask_to(mask, work_width, work_height) + latent_samples = vae.encode(work_image) + latent = {"samples": latent_samples, "source": "repair_prepare_crop"} + if effective_kernel == "latent_mask": + latent["noise_mask"] = make_noise_mask(work_mask, latent_samples) + + repair_info["crop_box"] = list(crop_box) + repair_info["crop_scale"] = crop_scale + repair_info["work_size"] = (work_width, work_height) + return latent, work_image, work_mask, repair_info, recommended + + if effective_scope == "canvas": + canvas_info = build_canvas_info( + original_size, + int(expand_left), + int(expand_right), + int(expand_top), + 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) + 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"]) + return latent, work_image, work_mask, repair_info, recommended + + raise RuntimeError(f"[LLS] Unsupported repair scope '{effective_scope}'.") NODE_CLASS_MAPPINGS = { diff --git a/repair/repair_utils.py b/repair/repair_utils.py index e1dbf20..dbf0986 100644 --- a/repair/repair_utils.py +++ b/repair/repair_utils.py @@ -248,6 +248,89 @@ def build_canvas_info( } +def get_image_size(image: Any) -> tuple[int, int]: + shape = getattr(image, "shape", None) + if not isinstance(shape, (list, tuple)) or len(shape) != 4: + raise RuntimeError("[LLS] image must have shape [batch, height, width, channels].") + + _, height, width, _ = shape + width = _coerce_int(width, 0) + height = _coerce_int(height, 0) + if width <= 0 or height <= 0: + raise RuntimeError("[LLS] image width and height must be positive.") + return width, height + + +def get_mask_metrics(mask: Any, image_size: tuple[int, int]) -> tuple[tuple[int, int, int, int] | None, float]: + 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].") + + raw_bbox = getattr(mask, "mask_bbox", None) + mask_bbox = None + 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)) + 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) + 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 resizing.") + + +def resize_mask_to(mask: Any, width: int, height: int) -> 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, "resized"): + return mask.resized(width, height) + _, 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 resizing.") + + +def expand_canvas_image(image: Any, width: int, height: int) -> Any: + if hasattr(image, "canvas_expanded"): + return image.canvas_expanded(width, height) + 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: + 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) + _, 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 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: + raise RuntimeError("[LLS] latent samples must have shape [batch, channels, height, width].") + + _, _, latent_height, latent_width = latent_shape + latent_width = _coerce_int(latent_width, 0) + latent_height = _coerce_int(latent_height, 0) + if latent_width <= 0 or latent_height <= 0: + raise RuntimeError("[LLS] latent samples must have positive height and width.") + return resize_mask_to(mask, latent_width, latent_height) + + 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 47fcf82..4d634be 100644 --- a/tests/test_repair_helpers.py +++ b/tests/test_repair_helpers.py @@ -7,6 +7,57 @@ ROOT = pathlib.Path(__file__).resolve().parents[1] MODULE_NAME = "lls_node_test_repair" +class FakeTensor: + def __init__(self, shape, label="image"): + self.shape = tuple(shape) + self.label = label + + def resized(self, width, height): + batch, _, _, channels = self.shape + return FakeTensor((batch, height, width, channels), label=f"{self.label}:resized") + + def canvas_expanded(self, width, height): + batch, _, _, channels = self.shape + return FakeTensor((batch, height, width, channels), label=f"{self.label}:canvas") + + +class FakeMask: + def __init__(self, shape, mask_bbox=None, mask_area_ratio=0.0, label="mask"): + self.shape = tuple(shape) + self.mask_bbox = mask_bbox + self.mask_area_ratio = float(mask_area_ratio) + self.label = label + + def resized(self, width, height): + batch, _, _ = self.shape + return FakeMask( + (batch, height, width), + mask_bbox=self.mask_bbox, + mask_area_ratio=self.mask_area_ratio, + label=f"{self.label}:resized", + ) + + def canvas_expanded(self, width, height): + batch, _, _ = self.shape + return FakeMask( + (batch, height, width), + mask_bbox=self.mask_bbox, + mask_area_ratio=self.mask_area_ratio, + label=f"{self.label}:canvas", + ) + + +class FakeLatentTensor: + def __init__(self, shape): + self.shape = tuple(shape) + + +class FakeVAE: + def encode(self, image): + batch, height, width, _ = image.shape + return FakeLatentTensor((batch, 4, height // 8, width // 8)) + + def load_plugin_package(): for name in list(sys.modules): if name == MODULE_NAME or name.startswith(f"{MODULE_NAME}."): diff --git a/tests/test_repair_prepare.py b/tests/test_repair_prepare.py new file mode 100644 index 0000000..62f0d8c --- /dev/null +++ b/tests/test_repair_prepare.py @@ -0,0 +1,173 @@ +import unittest + +try: + from .test_repair_helpers import FakeMask, FakeVAE, FakeTensor, load_plugin_package +except ImportError: # pragma: no cover - discovery mode imports from top level + from test_repair_helpers import FakeMask, FakeVAE, FakeTensor, load_plugin_package + + +class TestRepairPrepare(unittest.TestCase): + def test_prepare_region_with_latent_mask_attaches_noise_mask(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]() + + image = FakeTensor((1, 1024, 1024, 3), label="region-image") + mask = FakeMask( + (1, 1024, 1024), + mask_bbox=(128, 128, 384, 384), + mask_area_ratio=0.05, + label="region-mask", + ) + + latent, work_image, work_mask, repair_info, recommended_denoise = node.prepare( + image=image, + mask=mask, + vae=FakeVAE(), + repair_scope="region", + repair_kernel="latent_mask", + task_hint="repair", + mask_grow=24, + mask_blur=8.0, + mask_threshold=0.5, + 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.assertEqual(work_image.shape, (1, 1024, 1024, 3)) + self.assertEqual(work_mask.shape, (1, 1024, 1024)) + self.assertIn("noise_mask", latent) + self.assertEqual(repair_info["repair_scope"], "region") + self.assertEqual(repair_info["repair_kernel"], "latent_mask") + self.assertEqual(recommended_denoise, 0.45) + + def test_prepare_crop_with_vae_inpaint_populates_crop_metadata(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]() + + image = FakeTensor((1, 768, 1024, 3), label="crop-image") + mask = FakeMask( + (1, 768, 1024), + mask_bbox=(200, 120, 420, 360), + mask_area_ratio=0.06, + label="crop-mask", + ) + + latent, work_image, work_mask, repair_info, recommended_denoise = node.prepare( + image=image, + mask=mask, + vae=FakeVAE(), + repair_scope="crop", + repair_kernel="vae_inpaint", + task_hint="content", + mask_grow=24, + mask_blur=8.0, + mask_threshold=0.5, + 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.assertEqual(repair_info["repair_scope"], "crop") + self.assertIsNotNone(repair_info["crop_box"]) + self.assertGreater(repair_info["crop_scale"], 0.0) + self.assertEqual(repair_info["work_size"], (work_image.shape[2], work_image.shape[1])) + self.assertEqual(recommended_denoise, 0.65) + self.assertEqual(latent["source"], "repair_prepare_crop") + self.assertEqual(work_mask.shape, work_image.shape[:3]) + + def test_prepare_canvas_with_vae_inpaint_supports_empty_mask(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]() + + image = FakeTensor((1, 640, 640, 3), label="canvas-image") + mask = FakeMask((1, 640, 640), mask_bbox=None, mask_area_ratio=0.0, label="canvas-mask") + + latent, work_image, work_mask, repair_info, recommended_denoise = node.prepare( + image=image, + mask=mask, + vae=FakeVAE(), + repair_scope="auto", + repair_kernel="vae_inpaint", + task_hint="fill", + mask_grow=24, + mask_blur=8.0, + mask_threshold=0.5, + invert_mask=False, + crop_context=64, + crop_context_factor=1.5, + min_size=256, + max_size=1024, + resize_mode="fit", + expand_left=128, + expand_right=0, + expand_top=0, + expand_bottom=0, + canvas_fill="edge", + auto_recommend="enabled", + model_info={"model_family": "SDXL", "model_role": "base"}, + ) + + self.assertEqual(work_image.shape, (1, 640, 768, 3)) + self.assertEqual(work_mask.shape, (1, 640, 768)) + self.assertEqual(repair_info["repair_scope"], "canvas") + self.assertEqual(repair_info["canvas_expand"], [128, 0, 0, 0]) + self.assertEqual(recommended_denoise, 0.90) + self.assertEqual(latent["source"], "repair_prepare_canvas") + + def test_prepare_requires_vae(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]() + + image = FakeTensor((1, 1024, 1024, 3), label="missing-vae-image") + mask = FakeMask((1, 1024, 1024), mask_bbox=(0, 0, 64, 64), mask_area_ratio=0.01) + + with self.assertRaisesRegex(RuntimeError, "Missing VAE"): + node.prepare( + image=image, + mask=mask, + vae=None, + repair_scope="region", + repair_kernel="latent_mask", + task_hint="repair", + mask_grow=24, + mask_blur=8.0, + mask_threshold=0.5, + 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"}, + ) + + +if __name__ == "__main__": + unittest.main()