feat: add repair prepare node
This commit is contained in:
+160
-2
@@ -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 = {
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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}."):
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user