feat: add repair prepare node

This commit is contained in:
Gin3601
2026-05-20 19:11:39 +08:00
parent daf5abacf1
commit 0523ffa124
4 changed files with 467 additions and 2 deletions
+160 -2
View File
@@ -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 = {
+83
View File
@@ -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":
+51
View File
@@ -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}."):
+173
View File
@@ -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()