nodes/Gin3601-LLS-node
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

修改MASK创建节点

Browse Source
This commit is contained in:
Gin3601
2026-05-22 18:30:48 +08:00
parent 70185976f7
commit 904775dc9b
9 changed files with 296 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+40 -16
View File
@@ -81,6 +81,7 @@ node/
`LLS-node` 现在提供一套面向本地工作流的修复链路，由以下节点组成：
- `LLS Simple Mask Create`
- `LLS Simple Mask Preview`
- `LLS Simple Mask Draw`
- `LLS Simple Repair Prepare`
- `LLS Simple KSampler`
@@ -90,28 +91,30 @@ node/
### `LLS Simple Mask Create`
`LLS Simple Mask Create` 是一个面向局部重绘的规则几何遮罩节点，用来根据输入图像尺寸快速创建 `rectangle`、`square`、`circle`、`ellipse` 四种基础 `mask`，并输出面积信息与叠加预览图。
`LLS Simple Mask Create` 是一个基础几何遮罩生成节点。它不依赖原图内容，只根据输出尺寸和几何参数直接创建 `rectangle`、`square`、`circle`、`ellipse` 四种基础 `mask`。
它适合放在手工绘制前或修复准备前：
它适合放在修复准备前或手工绘制前：
- `Load Image -> LLS Simple Mask Create -> Preview Image`
- `Load Image -> LLS Simple Mask Create -> LLS Simple Repair Prepare`
- `Load Image -> LLS Simple Mask Create -> LLS Simple Mask Draw`
- `Load Image -> LLS Simple Mask Create -> LLS Simple Mask Draw -> LLS Simple Repair Prepare`
- `LLS Simple Mask Create.mask_image -> Preview Image`
- `Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Mask Preview -> Preview Image`
- `Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Repair Prepare`
- `LLS Simple Mask Create.mask -> LLS Simple Mask Draw.input_mask`
**推荐连接方式：**
- `Load Image.image -> LLS Simple Mask Create.image`
- `LLS Simple Mask Create.image -> LLS Simple Repair Prepare.image`
- `LLS Simple Mask Create.mask_image -> Preview Image`
- `Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Mask Preview`
- `Load Image.image -> LLS Simple Repair Prepare.image`
- `LLS Simple Mask Create.mask -> LLS Simple Repair Prepare.mask`
- `LLS Simple Mask Create.mask -> LLS Simple Mask Draw.input_mask`
- `LLS Simple Mask Create.image -> LLS Simple Mask Draw.image`
- `Load Image.image -> LLS Simple Mask Draw.image`
**输入：**
| 名称 | 类型 | 说明 |
|------|------|------|
| `image` | `IMAGE` | 必填，原图输入，用于确定几何 mask 尺寸 |
| `image_width` | `INT` | 必填，输出遮罩宽度 |
| `image_height` | `INT` | 必填，输出遮罩高度 |
| `input_mask` | `MASK` | 可选，已有遮罩；可与当前几何 mask 做 `replace` / `union` / `subtract` / `intersect` |
**核心参数：**
@@ -127,16 +130,13 @@ node/
| `blur` | 整体边缘模糊 |
| `invert_mask` | 反转最终 mask |
| `combine_mode` | `replace` / `union` / `subtract` / `intersect` |
| `overlay_alpha` | 预览叠加透明度 |
| `overlay_color` | `red` / `green` / `blue` |
**输出：**
| 名称 | 类型 | 说明 |
|------|------|------|
| `image` | `IMAGE` | 原图透传 |
| `mask` | `MASK` | 最终遮罩，白色 / `1` 表示要重绘，黑色 / `0` 表示不重绘 |
| `preview_image` | `IMAGE` | 原图与半透明彩色 mask 叠加后的预览图 |
| `mask_image` | `IMAGE` | 将 `mask` 直接可视化为黑白图像，方便单独预览 |
| `area_info` | `LLS_MASK_INFO` | 包含面积、占比、bbox、坐标模式等信息的 dict |
`area_info` 主要包含：
@@ -165,12 +165,12 @@ node/
- `coordinate_mode = percent`
- `center_x = 0.5`
- `center_y = 0.5`
- `image_width = 1024`
- `image_height = 1024`
- `width = 0.3`
- `height = 0.3`
- `radius = 0.15`
- `combine_mode = replace`
- `overlay_alpha = 0.4`
- `overlay_color = red`
这会在图像中心创建一个大约占宽高 `30%` 的矩形 mask。
@@ -188,6 +188,30 @@ node/
- 先创建基础几何区域，再交给 `LLS Simple Mask Draw` 手动修边
- 为 `LLS Simple Repair Prepare` 提供稳定的初始局部重绘 mask
### `LLS Simple Mask Preview`
`LLS Simple Mask Preview` 只负责把 `image + mask` 做半透明叠加预览，不参与采样，也不修改原图。
**推荐工作流：**
- `Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Mask Preview -> Preview Image`
- `Load Image.image + LLS Simple Mask Draw.mask -> LLS Simple Mask Preview -> Preview Image`
**输入：**
| 名称 | 类型 | 说明 |
|------|------|------|
| `image` | `IMAGE` | 必填，原图输入 |
| `mask` | `MASK` | 必填，待预览的遮罩 |
| `overlay_alpha` | `FLOAT` | 叠加透明度 |
| `overlay_color` | `red / green / blue` | 叠加颜色 |
**输出：**
| 名称 | 类型 | 说明 |
|------|------|------|
| `preview_image` | `IMAGE` | 原图与半透明彩色 mask 叠加后的预览图 |
### `LLS Simple Mask Draw`
`LLS Simple Mask Draw` 是一个面向局部重绘的交互式遮罩输入节点，用来直接在输入图像上手动画 `mask`，并把结果输出给后续修复链路。
mask/mask_create.py
+30 -22
View File
@@ -4,29 +4,39 @@ from .mask_utils import (
COMBINE_MODE_CHOICES,
COORDINATE_MODE_CHOICES,
MASK_INFO_TYPE,
OVERLAY_COLOR_CHOICES,
SHAPE_TYPE_CHOICES,
apply_invert,
apply_softening,
build_area_info,
build_preview_image,
combine_masks,
create_reference_image,
create_shape_mask,
mask_to_image,
)
def _resolve_batch_for_mask_creation(input_mask) -> int:
if input_mask is None:
return 1
shape = getattr(input_mask, "shape", None)
if not isinstance(shape, (list, tuple)) or len(shape) != 3:
raise RuntimeError("[LLS] mask must have shape [batch, height, width].")
return max(1, int(shape[0]))
class LLSSimpleMaskCreate:
CATEGORY = "LLS/Mask"
FUNCTION = "create_mask"
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", MASK_INFO_TYPE)
RETURN_NAMES = ("image", "mask", "preview_image", "area_info")
DESCRIPTION = "Create a geometric repair mask, preview overlay, and area metadata from an input image."
RETURN_TYPES = ("MASK", "IMAGE", MASK_INFO_TYPE)
RETURN_NAMES = ("mask", "mask_image", "area_info")
DESCRIPTION = "Create a geometric mask image from width, height, and shape parameters."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"image_width": ("INT", {"default": 1024, "min": 1, "max": 8192, "step": 1}),
"image_height": ("INT", {"default": 1024, "min": 1, "max": 8192, "step": 1}),
"shape_type": (SHAPE_TYPE_CHOICES, {"default": "rectangle"}),
"coordinate_mode": (COORDINATE_MODE_CHOICES, {"default": "percent"}),
"center_x": ("FLOAT", {"default": 0.5, "step": 0.01}),
@@ -38,8 +48,6 @@ class LLSSimpleMaskCreate:
"blur": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 128.0, "step": 0.5}),
"invert_mask": ("BOOLEAN", {"default": False}),
"combine_mode": (COMBINE_MODE_CHOICES, {"default": "replace"}),
"overlay_alpha": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}),
"overlay_color": (OVERLAY_COLOR_CHOICES, {"default": "red"}),
},
"optional": {
"input_mask": ("MASK",),
@@ -48,7 +56,8 @@ class LLSSimpleMaskCreate:
def create_mask(
self,
image,
image_width,
image_height,
shape_type,
coordinate_mode,
center_x,
@@ -60,15 +69,19 @@ class LLSSimpleMaskCreate:
blur,
invert_mask,
combine_mode,
overlay_alpha,
overlay_color,
input_mask=None,
):
if image is None:
raise RuntimeError("[LLS] Missing image for LLS Simple Mask Create.")
batch = _resolve_batch_for_mask_creation(input_mask)
reference_image = create_reference_image(
image_width=int(image_width),
image_height=int(image_height),
batch=batch,
dtype=getattr(input_mask, "dtype", None),
device=getattr(input_mask, "device", None),
)
created_mask, geometry_info = create_shape_mask(
image,
reference_image,
shape_type=shape_type,
coordinate_mode=coordinate_mode,
center_x=center_x,
@@ -77,15 +90,10 @@ class LLSSimpleMaskCreate:
height=height,
radius=radius,
)
combined_mask = combine_masks(created_mask, input_mask, image, combine_mode)
combined_mask = combine_masks(created_mask, input_mask, reference_image, combine_mode)
final_mask = apply_invert(combined_mask, bool(invert_mask))
final_mask = apply_softening(final_mask, feather=float(feather), blur=float(blur))
preview_image = build_preview_image(
image,
final_mask,
overlay_alpha=float(overlay_alpha),
overlay_color=str(overlay_color or "red"),
)
mask_image = mask_to_image(final_mask)
area_info = build_area_info(
final_mask,
geometry_info,
@@ -94,7 +102,7 @@ class LLSSimpleMaskCreate:
invert_mask=bool(invert_mask),
combine_mode=str(combine_mode or "replace"),
)
return image, final_mask, preview_image, area_info
return final_mask, mask_image, area_info
NODE_CLASS_MAPPINGS = {"LLSSimpleMaskCreate": LLSSimpleMaskCreate}
mask/mask_preview.py
+39
View File
@@ -0,0 +1,39 @@
from __future__ import annotations
from .mask_utils import OVERLAY_COLOR_CHOICES, build_preview_image
class LLSSimpleMaskPreview:
CATEGORY = "LLS/Mask"
FUNCTION = "preview_mask"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("preview_image",)
DESCRIPTION = "Overlay a mask on top of an image for preview only."
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"mask": ("MASK",),
"overlay_alpha": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}),
"overlay_color": (OVERLAY_COLOR_CHOICES, {"default": "red"}),
}
}
def preview_mask(self, image, mask, overlay_alpha, overlay_color):
if image is None:
raise RuntimeError("[LLS] Missing image for LLS Simple Mask Preview.")
if mask is None:
raise RuntimeError("[LLS] Missing mask for LLS Simple Mask Preview.")
preview_image = build_preview_image(
image,
mask,
overlay_alpha=float(overlay_alpha),
overlay_color=str(overlay_color or "red"),
)
return (preview_image,)
NODE_CLASS_MAPPINGS = {"LLSSimpleMaskPreview": LLSSimpleMaskPreview}
NODE_DISPLAY_NAME_MAPPINGS = {"LLSSimpleMaskPreview": "LLS Simple Mask Preview"}
mask/mask_utils.py
+29
View File
@@ -26,6 +26,35 @@ _OVERLAY_COLORS = {
}
def create_reference_image(
image_width: int,
image_height: int,
*,
batch: int = 1,
dtype=None,
device=None,
):
_require_torch("[LLS] torch is required for mask generation.")
width = max(1, int(image_width))
height = max(1, int(image_height))
batch = max(1, int(batch))
kwargs = {"dtype": dtype or torch.float32}
if device is not None:
kwargs["device"] = device
return torch.zeros((batch, height, width, 3), **kwargs)
def mask_to_image(mask):
_require_torch("[LLS] torch is required for mask image generation.")
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].")
normalized = mask.clamp(0.0, 1.0)
return normalized.unsqueeze(-1).repeat(1, 1, 1, 3)
def get_image_size(image) -> tuple[int, int]:
shape = getattr(image, "shape", None)
if not isinstance(shape, (list, tuple)) or len(shape) != 4:
mask/nodes.py
+16 -1
View File
@@ -3,6 +3,21 @@ LLS / Mask
==========
功能域：遮罩（Mask）操作（对应功能分类总览第 6 节）
"""
from .mask_create import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from .mask_create import (
NODE_CLASS_MAPPINGS as MASK_CREATE_NODE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as MASK_CREATE_NODE_DISPLAY_NAME_MAPPINGS,
)
from .mask_preview import (
NODE_CLASS_MAPPINGS as MASK_PREVIEW_NODE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as MASK_PREVIEW_NODE_DISPLAY_NAME_MAPPINGS,
)
NODE_CLASS_MAPPINGS = {}
NODE_CLASS_MAPPINGS.update(MASK_CREATE_NODE_CLASS_MAPPINGS)
NODE_CLASS_MAPPINGS.update(MASK_PREVIEW_NODE_CLASS_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS.update(MASK_CREATE_NODE_DISPLAY_NAME_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(MASK_PREVIEW_NODE_DISPLAY_NAME_MAPPINGS)
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
tests/test_mask_create_docs.py
+6 -2
View File
@@ -11,11 +11,15 @@ class TestMaskCreateDocs(unittest.TestCase):
readme = (ROOT / "README.md").read_text(encoding="utf-8")
self.assertIn("LLS Simple Mask Create", readme)
self.assertIn("LLS Simple Mask Preview", readme)
self.assertIn("image_width", readme)
self.assertIn("image_height", readme)
self.assertIn("shape_type", readme)
self.assertIn("coordinate_mode", readme)
self.assertIn("area_info", readme)
self.assertIn("Load Image -> LLS Simple Mask Create -> LLS Simple Repair Prepare", readme)
self.assertIn("Load Image -> LLS Simple Mask Create -> LLS Simple Mask Draw", readme)
self.assertIn("LLS Simple Mask Create.mask -> LLS Simple Repair Prepare.mask", readme)
self.assertIn("LLS Simple Mask Create.mask_image -> Preview Image", readme)
self.assertIn("Load Image.image + LLS Simple Mask Create.mask -> LLS Simple Mask Preview", readme)
def test_readme_lists_supported_shapes_and_combine_modes(self):
readme = (ROOT / "README.md").read_text(encoding="utf-8")
tests/test_mask_create_node.py
+31 -67
View File
@@ -49,7 +49,8 @@ class TestMaskCreateNode(unittest.TestCase):
def _create(self, **overrides):
params = {
"image": make_image(width=100, height=80, color=0.2),
"image_width": 100,
"image_height": 80,
"shape_type": "rectangle",
"coordinate_mode": "percent",
"center_x": 0.5,
@@ -61,19 +62,16 @@ class TestMaskCreateNode(unittest.TestCase):
"blur": 0.0,
"invert_mask": False,
"combine_mode": "replace",
"overlay_alpha": 0.4,
"overlay_color": "red",
"input_mask": None,
}
params.update(overrides)
return self.node.create_mask(**params)
def test_rectangle_percent_mode_creates_expected_bbox_and_area_info(self):
image_out, mask, preview_image, area_info = self._create()
mask, mask_image, area_info = self._create()
self.assertEqual(tuple(mask.shape), (1, 80, 100))
self.assertEqual(tuple(preview_image.shape), (1, 80, 100, 3))
self.assertEqual(tuple(image_out.shape), (1, 80, 100, 3))
self.assertEqual(tuple(mask_image.shape), (1, 80, 100, 3))
self.assertEqual(area_info["image_size"], [100, 80])
self.assertEqual(area_info["shape_type"], "rectangle")
self.assertEqual(area_info["coordinate_mode"], "percent")
@@ -86,11 +84,13 @@ class TestMaskCreateNode(unittest.TestCase):
self.assertAlmostEqual(area_info["area_ratio"], 600.0 / 8000.0, places=4)
self.assertEqual(float(mask[0, 40, 50].item()), 1.0)
self.assertEqual(float(mask[0, 5, 5].item()), 0.0)
self.assertEqual(float(mask_image[0, 40, 50, 0].item()), 1.0)
self.assertEqual(float(mask_image[0, 5, 5, 0].item()), 0.0)
def test_square_pixel_mode_uses_width_as_side(self):
image = make_image(width=64, height=64, color=0.2)
_image_out, mask, _preview_image, area_info = self.node.create_mask(
image=image,
mask, _mask_image, area_info = self.node.create_mask(
image_width=64,
image_height=64,
shape_type="square",
coordinate_mode="pixel",
center_x=20.0,
@@ -102,8 +102,6 @@ class TestMaskCreateNode(unittest.TestCase):
blur=0.0,
invert_mask=False,
combine_mode="replace",
overlay_alpha=0.4,
overlay_color="red",
input_mask=None,
)
@@ -114,9 +112,9 @@ class TestMaskCreateNode(unittest.TestCase):
self.assertEqual(float(mask[0, 5, 5].item()), 0.0)
def test_circle_generates_centered_mask(self):
image = make_image(width=101, height=101, color=0.2)
_image_out, mask, _preview_image, area_info = self.node.create_mask(
image=image,
mask, _mask_image, area_info = self.node.create_mask(
image_width=101,
image_height=101,
shape_type="circle",
coordinate_mode="pixel",
center_x=50.0,
@@ -128,8 +126,6 @@ class TestMaskCreateNode(unittest.TestCase):
blur=0.0,
invert_mask=False,
combine_mode="replace",
overlay_alpha=0.4,
overlay_color="red",
input_mask=None,
)
@@ -140,9 +136,9 @@ class TestMaskCreateNode(unittest.TestCase):
self.assertEqual(float(mask[0, 10, 10].item()), 0.0)
def test_ellipse_generates_nonempty_mask(self):
image = make_image(width=120, height=90, color=0.2)
_image_out, mask, _preview_image, area_info = self.node.create_mask(
image=image,
mask, _mask_image, area_info = self.node.create_mask(
image_width=120,
image_height=90,
shape_type="ellipse",
coordinate_mode="pixel",
center_x=60.0,
@@ -154,8 +150,6 @@ class TestMaskCreateNode(unittest.TestCase):
blur=0.0,
invert_mask=False,
combine_mode="replace",
overlay_alpha=0.4,
overlay_color="red",
input_mask=None,
)
@@ -165,36 +159,12 @@ class TestMaskCreateNode(unittest.TestCase):
self.assertEqual(float(mask[0, 45, 60].item()), 1.0)
self.assertEqual(float(mask[0, 5, 5].item()), 0.0)
def test_overlay_image_uses_requested_color(self):
image = make_image(width=20, height=20, color=0.2)
_image_out, _mask, preview_image, _area_info = self.node.create_mask(
image=image,
shape_type="rectangle",
coordinate_mode="pixel",
center_x=10.0,
center_y=10.0,
width=10.0,
height=10.0,
radius=5.0,
feather=0.0,
blur=0.0,
invert_mask=False,
combine_mode="replace",
overlay_alpha=0.5,
overlay_color="green",
input_mask=None,
)
center_pixel = preview_image[0, 10, 10]
self.assertGreater(float(center_pixel[1].item()), float(center_pixel[0].item()))
self.assertGreater(float(center_pixel[1].item()), float(center_pixel[2].item()))
def test_union_subtract_and_intersect_with_input_mask(self):
image = make_image(width=32, height=32, color=0.2)
input_mask = make_mask(width=32, height=32, value=1.0)
_image_out, union_mask, _preview, union_info = self.node.create_mask(
image=image,
union_mask, _mask_image, union_info = self.node.create_mask(
image_width=32,
image_height=32,
shape_type="rectangle",
coordinate_mode="pixel",
center_x=16.0,
@@ -206,15 +176,14 @@ class TestMaskCreateNode(unittest.TestCase):
blur=0.0,
invert_mask=False,
combine_mode="union",
overlay_alpha=0.4,
overlay_color="red",
input_mask=input_mask,
)
self.assertEqual(union_info["binary_area_px"], 32 * 32)
self.assertTrue(torch.all(union_mask == 1.0))
_image_out, subtract_mask, _preview, subtract_info = self.node.create_mask(
image=image,
subtract_mask, _mask_image, subtract_info = self.node.create_mask(
image_width=32,
image_height=32,
shape_type="rectangle",
coordinate_mode="pixel",
center_x=16.0,
@@ -226,15 +195,14 @@ class TestMaskCreateNode(unittest.TestCase):
blur=0.0,
invert_mask=False,
combine_mode="subtract",
overlay_alpha=0.4,
overlay_color="red",
input_mask=input_mask,
)
self.assertLess(subtract_info["binary_area_px"], 32 * 32)
self.assertEqual(float(subtract_mask[0, 16, 16].item()), 0.0)
_image_out, intersect_mask, _preview, intersect_info = self.node.create_mask(
image=image,
intersect_mask, _mask_image, intersect_info = self.node.create_mask(
image_width=32,
image_height=32,
shape_type="rectangle",
coordinate_mode="pixel",
center_x=16.0,
@@ -246,8 +214,6 @@ class TestMaskCreateNode(unittest.TestCase):
blur=0.0,
invert_mask=False,
combine_mode="intersect",
overlay_alpha=0.4,
overlay_color="red",
input_mask=input_mask,
)
self.assertEqual(intersect_info["binary_area_px"], 100)
@@ -255,14 +221,14 @@ class TestMaskCreateNode(unittest.TestCase):
self.assertEqual(float(intersect_mask[0, 1, 1].item()), 0.0)
def test_invert_mask_flips_final_mask(self):
_image_out, mask, _preview, area_info = self._create(invert_mask=True)
mask, _mask_image, area_info = self._create(invert_mask=True)
self.assertEqual(area_info["invert_mask"], True)
self.assertEqual(float(mask[0, 40, 50].item()), 0.0)
self.assertEqual(float(mask[0, 5, 5].item()), 1.0)
def test_feather_and_blur_create_soft_mask(self):
_image_out, mask, _preview, area_info = self._create(
mask, _mask_image, area_info = self._create(
shape_type="rectangle",
coordinate_mode="pixel",
center_x=50.0,
@@ -278,7 +244,7 @@ class TestMaskCreateNode(unittest.TestCase):
self.assertTrue(bool(((mask > 0.0) & (mask < 1.0)).any().item()))
def test_shape_outside_image_returns_black_mask_without_crashing(self):
_image_out, mask, _preview, area_info = self._create(
mask, _mask_image, area_info = self._create(
shape_type="circle",
coordinate_mode="percent",
center_x=2.0,
@@ -294,8 +260,9 @@ class TestMaskCreateNode(unittest.TestCase):
image = make_image(width=64, height=64, color=0.25)
input_mask = make_mask(width=16, height=16, value=1.0)
image_out, mask_out, preview_out, area_info = self.node.create_mask(
image=image,
mask_out, _mask_image, area_info = self.node.create_mask(
image_width=64,
image_height=64,
shape_type="rectangle",
coordinate_mode="pixel",
center_x=32.0,
@@ -307,13 +274,11 @@ class TestMaskCreateNode(unittest.TestCase):
blur=0.0,
invert_mask=False,
combine_mode="union",
overlay_alpha=0.4,
overlay_color="red",
input_mask=input_mask,
)
latent, work_image, work_mask, repair_info, recommended_denoise, positive, negative = self.repair_prepare_node.prepare(
image=image_out,
image=image,
mask=mask_out,
vae=_FakeVAE(),
repair_scope="region",
@@ -339,7 +304,6 @@ class TestMaskCreateNode(unittest.TestCase):
negative=None,
)
self.assertEqual(tuple(preview_out.shape), (1, 64, 64, 3))
self.assertEqual(tuple(mask_out.shape), (1, 64, 64))
self.assertEqual(tuple(work_image.shape), (1, 64, 64, 3))
self.assertEqual(tuple(work_mask.shape), (1, 64, 64))
tests/test_mask_create_registration.py
+32 -4
View File
@@ -23,13 +23,16 @@ class TestMaskCreateRegistration(unittest.TestCase):
self.assertEqual(node_cls.CATEGORY, "LLS/Mask")
self.assertEqual(node_cls.FUNCTION, "create_mask")
self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE", "MASK", "IMAGE", "LLS_MASK_INFO"))
self.assertEqual(node_cls.RETURN_NAMES, ("image", "mask", "preview_image", "area_info"))
self.assertEqual(node_cls.RETURN_TYPES, ("MASK", "IMAGE", "LLS_MASK_INFO"))
self.assertEqual(node_cls.RETURN_NAMES, ("mask", "mask_image", "area_info"))
required = schema["required"]
optional = schema["optional"]
self.assertEqual(required["image"], ("IMAGE",))
self.assertEqual(required["image_width"][0], "INT")
self.assertEqual(required["image_height"][0], "INT")
self.assertEqual(required["image_width"][1]["default"], 1024)
self.assertEqual(required["image_height"][1]["default"], 1024)
self.assertEqual(optional["input_mask"], ("MASK",))
self.assertEqual(required["shape_type"][0], ["rectangle", "square", "circle", "ellipse"])
self.assertEqual(required["coordinate_mode"][0], ["pixel", "percent"])
@@ -46,10 +49,35 @@ class TestMaskCreateRegistration(unittest.TestCase):
self.assertEqual(required["invert_mask"][0], "BOOLEAN")
self.assertEqual(required["combine_mode"][0], ["replace", "union", "subtract", "intersect"])
self.assertEqual(required["combine_mode"][1]["default"], "replace")
self.assertTrue(callable(getattr(node_cls, "create_mask", None)))
def test_plugin_registers_mask_preview_node(self):
plugin = load_plugin_package()
self.assertIn("LLSSimpleMaskPreview", plugin.NODE_CLASS_MAPPINGS)
self.assertEqual(
plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSSimpleMaskPreview"],
"LLS Simple Mask Preview",
)
def test_mask_preview_schema_matches_contract(self):
plugin = load_plugin_package()
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskPreview"]
schema = node_cls.INPUT_TYPES()
self.assertEqual(node_cls.CATEGORY, "LLS/Mask")
self.assertEqual(node_cls.FUNCTION, "preview_mask")
self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE",))
self.assertEqual(node_cls.RETURN_NAMES, ("preview_image",))
required = schema["required"]
self.assertEqual(required["image"], ("IMAGE",))
self.assertEqual(required["mask"], ("MASK",))
self.assertEqual(required["overlay_alpha"][0], "FLOAT")
self.assertEqual(required["overlay_alpha"][1]["default"], 0.4)
self.assertEqual(required["overlay_color"][0], ["red", "green", "blue"])
self.assertTrue(callable(getattr(node_cls, "create_mask", None)))
self.assertTrue(callable(getattr(node_cls, "preview_mask", None)))
if __name__ == "__main__":
tests/test_mask_preview_node.py
+73
View File
@@ -0,0 +1,73 @@
import unittest
try:
from .test_mask_draw_helpers import (
load_plugin_package,
torch,
)
except ImportError:
from test_mask_draw_helpers import (
load_plugin_package,
torch,
)
HAS_MASK_PREVIEW_RUNTIME_DEPS = torch is not None
MASK_PREVIEW_RUNTIME_DEPS_MESSAGE = "mask preview node tests require torch"
def _require_mask_preview_runtime_deps():
if not HAS_MASK_PREVIEW_RUNTIME_DEPS:
raise unittest.SkipTest(MASK_PREVIEW_RUNTIME_DEPS_MESSAGE)
def make_image(width=8, height=8, color=0.25):
_require_mask_preview_runtime_deps()
return torch.full((1, height, width, 3), float(color), dtype=torch.float32)
def make_mask(width=8, height=8, value=0.0):
_require_mask_preview_runtime_deps()
return torch.full((1, height, width), float(value), dtype=torch.float32)
@unittest.skipUnless(HAS_MASK_PREVIEW_RUNTIME_DEPS, MASK_PREVIEW_RUNTIME_DEPS_MESSAGE)
class TestMaskPreviewNode(unittest.TestCase):
@classmethod
def setUpClass(cls):
plugin = load_plugin_package()
cls.node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskPreview"]()
def test_preview_overlays_requested_color(self):
image = make_image(width=20, height=20, color=0.2)
mask = make_mask(width=20, height=20, value=1.0)
(preview_image,) = self.node.preview_mask(
image=image,
mask=mask,
overlay_alpha=0.5,
overlay_color="green",
)
center_pixel = preview_image[0, 10, 10]
self.assertEqual(tuple(preview_image.shape), (1, 20, 20, 3))
self.assertGreater(float(center_pixel[1].item()), float(center_pixel[0].item()))
self.assertGreater(float(center_pixel[1].item()), float(center_pixel[2].item()))
def test_preview_resizes_mask_to_match_image(self):
image = make_image(width=20, height=20, color=0.2)
mask = make_mask(width=10, height=10, value=1.0)
(preview_image,) = self.node.preview_mask(
image=image,
mask=mask,
overlay_alpha=0.4,
overlay_color="red",
)
self.assertEqual(tuple(preview_image.shape), (1, 20, 20, 3))
self.assertGreater(float(preview_image[..., 0].mean().item()), float(image[..., 0].mean().item()))
if __name__ == "__main__":
unittest.main()