diff --git a/README.md b/README.md index b1981d4..102503a 100644 --- a/README.md +++ b/README.md @@ -80,16 +80,87 @@ node/ `LLS-node` 现在提供一套面向本地工作流的修复链路，由以下节点组成： +- `LLS Simple Mask Draw` - `LLS Simple Repair Prepare` - `LLS Simple KSampler` - `LLS Simple Repair Finish` 修复工作流不是从 `LLS Simple Empty Latent` 开始，而是从 `image + mask + vae` 开始。`LLS Simple Repair Prepare` 会把这些输入转换为带有 `repair_info` 的修复态 latent，再交给同一个 `LLS Simple KSampler` 继续采样。 +### `LLS Simple Mask Draw` + +`LLS Simple Mask Draw` 是一个面向局部重绘的交互式遮罩输入节点，用来直接在输入图像上手动画 `mask`，并把结果输出给后续修复链路。 + +**推荐工作流：** + +- `Load Image -> LLS Simple Mask Draw -> Preview Image` +- `Load Image -> LLS Simple Mask Draw -> LLS Simple Repair Prepare` + +**推荐连接方式：** + +- `Load Image.image -> LLS Simple Mask Draw.image` +- `LLS Simple Mask Draw.image -> LLS Simple Repair Prepare.image` +- `LLS Simple Mask Draw.mask -> LLS Simple Repair Prepare.mask` + +**输入：** + +| 名称 | 类型 | 说明 | +|------|------|------| +| `image` | `IMAGE` | 必填，原图输入 | +| `input_mask` | `MASK` | 可选，已有遮罩；节点会在这个遮罩基础上继续编辑 | + +**输出：** + +| 名称 | 类型 | 说明 | +|------|------|------| +| `image` | `IMAGE` | 原图透传 | +| `mask` | `MASK` | 最终遮罩，白色 / `1` 表示要重绘，黑色 / `0` 表示不重绘 | +| `preview_image` | `IMAGE` | 原图与半透明红色遮罩叠加后的预览图 | + +**当前支持的交互功能：** + +- 直接在图像上绘制白色遮罩 +- `brush` / `erase` 两种模式 +- 调整 `brush_size` +- 调整 `brush_softness` +- 调整 `overlay_alpha` +- 预览半透明红色遮罩叠加 +- `Clear` +- `Undo` +- `Redo` +- `Invert` +- 在已有 `input_mask` 基础上继续补画或擦除 + +**如何手动画 mask：** + +1. 连接 `Load Image` 到 `LLS Simple Mask Draw.image` +2. 如果有已有遮罩，可再连接到 `input_mask` +3. 在节点预览区域直接涂抹需要修复的区域 +4. 用 `erase` 擦除不需要重绘的部分 +5. 用 `Clear / Undo / Redo / Invert` 调整最终结果 +6. 将 `mask` 输出接到 `LLS Simple Repair Prepare.mask` + +**典型用途：** + +- 手动指定删除区域 +- 手动指定修复区域 +- 手动指定去阴影区域 +- 手动指定局部增强区域 + +**当前版本限制：** + +- 第一版只支持 `brush` 和 `erase` +- 还没有实现 `polygon`、`rectangle`、`ellipse` +- 还没有实现魔棒、自动分割、智能抠图 + ### Minimal Repair Workflow `Load Image -> Load Mask -> LLS Simple Checkpoint Loader -> LLS Simple Prompt Encode -> LLS Simple Repair Prepare -> LLS Simple KSampler -> VAE Decode -> LLS Simple Repair Finish -> Preview Image` +### Minimal Manual Mask Workflow + +`Load Image -> LLS Simple Mask Draw -> LLS Simple Repair Prepare -> LLS Simple KSampler -> VAE Decode -> LLS Simple Repair Finish -> Preview Image` + ### `repair_scope` - `region`：在原图边界内直接修复 diff --git a/__init__.py b/__init__.py index f428a7c..9dfbacf 100644 --- a/__init__.py +++ b/__init__.py @@ -36,6 +36,8 @@ from __future__ import annotations import importlib import types +WEB_DIRECTORY = "./web" + # ---------- 子包白名单（控制加载顺序，同时作为文档） ---------- # 按功能分类总览章节顺序排列；新增子包追加到列表末尾即可。 _SUBPACKAGES: list[str] = [ @@ -46,6 +48,7 @@ _SUBPACKAGES: list[str] = [ "latent", # 第 4 节：Latent 空间操作 "image", # 第 5 节：图像处理与后处理 "repair", # 第 5 节（修复子域）：图像修复 + "mask_draw", # 第 5 节（修复子域）：交互式遮罩绘制 "upscale", # 第 5 节（超分子域）：图像超分 "mask", # 第 6 节：遮罩操作 "controlnet", # 第 7 节：ControlNet / 控制引导 @@ -121,4 +124,4 @@ _merge_root_nodes(NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS) for _pkg in _SUBPACKAGES: _merge_subpackage(_pkg, NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS) -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/docs/superpowers/plans/2026-05-21-lls-simple-mask-draw.md b/docs/superpowers/plans/2026-05-21-lls-simple-mask-draw.md new file mode 100644 index 0000000..8935aff --- /dev/null +++ b/docs/superpowers/plans/2026-05-21-lls-simple-mask-draw.md @@ -0,0 +1,1411 @@ +# LLS Simple Mask Draw Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a single interactive `LLS Simple Mask Draw` node that lets the user paint, erase, clear, preview, persist, and output a repair mask directly over an input image for downstream local repair workflows. + +**Architecture:** Add a dedicated `mask_draw/` backend package plus one `WEB_DIRECTORY` frontend extension in `web/js/lls_mask_draw.js`. Keep the backend responsible for mask-state decoding, output mask resolution, and preview-image generation; keep the frontend responsible for interactive drawing, session-only undo/redo, file-backed source-image recovery, and workflow-persisted mask raster state. + +**Tech Stack:** Python, unittest, JSON/base64, Pillow, torch, ComfyUI node definitions, ComfyUI frontend JavaScript extension hooks, HTML Canvas, `node --check` + +--- + +## File Structure + +- Create: `mask_draw/__init__.py` + - Expose the new node registration maps. +- Create: `mask_draw/node.py` + - Define `LLSSimpleMaskDraw`, its schema, and runtime orchestration. +- Create: `mask_draw/utils.py` + - Decode persisted state, resolve final mask precedence, resize masks, and generate `preview_image`. +- Create: `web/js/lls_mask_draw.js` + - Register the frontend extension, hide the persistence widget, render the editor, manage strokes, and persist `mask_state_json`. +- Create: `tests/test_mask_draw_helpers.py` + - Shared plugin loader, tiny-image builders, and state-payload helpers. +- Create: `tests/test_mask_draw_registration.py` + - Verify plugin registration, node category, and node schema. +- Create: `tests/test_mask_draw_utils.py` + - Verify persisted-state parsing, fallback precedence, inversion, resize, and preview-image shape. +- Create: `tests/test_mask_draw_node.py` + - Verify node orchestration, clear semantics, and repair-prepare compatibility. +- Create: `tests/test_mask_draw_frontend.py` + - Verify `WEB_DIRECTORY` export and presence of the JS extension entrypoints. +- Create: `tests/test_mask_draw_docs.py` + - Verify the README documents the node and connection pattern. +- Modify: `__init__.py` + - Append `mask_draw` to `_SUBPACKAGES` and export `WEB_DIRECTORY = "./web"`. +- Modify: `README.md` + - Document the new node, usage, wiring, and limitations. + +### Task 1: Add Package Scaffolding and Registration Contract + +**Files:** +- Create: `tests/test_mask_draw_helpers.py` +- Create: `tests/test_mask_draw_registration.py` +- Create: `mask_draw/__init__.py` +- Create: `mask_draw/node.py` +- Create: `mask_draw/utils.py` +- Modify: `__init__.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_mask_draw_helpers.py +import importlib.util +import pathlib +import sys + + +ROOT = pathlib.Path(__file__).resolve().parents[1] +MODULE_NAME = "lls_node_test_mask_draw" + + +def load_plugin_package(): + for name in list(sys.modules): + if name == MODULE_NAME or name.startswith(f"{MODULE_NAME}."): + sys.modules.pop(name) + + spec = importlib.util.spec_from_file_location( + MODULE_NAME, + ROOT / "__init__.py", + submodule_search_locations=[str(ROOT)], + ) + module = importlib.util.module_from_spec(spec) + sys.modules[MODULE_NAME] = module + spec.loader.exec_module(module) + return module +``` + +```python +# tests/test_mask_draw_registration.py +import unittest + +try: + from .test_mask_draw_helpers import load_plugin_package +except ImportError: + from test_mask_draw_helpers import load_plugin_package + + +class TestMaskDrawRegistration(unittest.TestCase): + def test_plugin_registers_mask_draw_node(self): + plugin = load_plugin_package() + + self.assertIn("LLSSimpleMaskDraw", plugin.NODE_CLASS_MAPPINGS) + self.assertEqual( + plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSSimpleMaskDraw"], + "LLS Simple Mask Draw", + ) + + def test_mask_draw_schema_matches_contract(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskDraw"] + schema = node_cls.INPUT_TYPES() + + self.assertEqual(node_cls.CATEGORY, "LLS/Image Repair") + self.assertEqual(node_cls.FUNCTION, "draw_mask") + self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE", "MASK", "IMAGE")) + self.assertEqual(node_cls.RETURN_NAMES, ("image", "mask", "preview_image")) + + required = schema["required"] + optional = schema["optional"] + hidden = schema["hidden"] + + self.assertEqual(required["image"], ("IMAGE",)) + self.assertEqual(optional["input_mask"], ("MASK",)) + self.assertEqual(required["draw_mode"][0], ["brush", "erase"]) + self.assertEqual(required["brush_size"][0], "INT") + self.assertEqual(required["brush_softness"][0], "FLOAT") + self.assertEqual(required["overlay_alpha"][0], "FLOAT") + self.assertEqual(required["invert_mask"][0], "BOOLEAN") + self.assertEqual(required["mask_state_json"][0], "STRING") + self.assertEqual(hidden["node_id"], "UNIQUE_ID") + + +if __name__ == "__main__": + unittest.main() +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_registration.py' -v` +Expected: FAIL because `mask_draw` is not registered and the node class does not exist yet. + +- [ ] **Step 3: Implement the minimal scaffolding** + +```python +# __init__.py +from __future__ import annotations + +import importlib +import types + + +WEB_DIRECTORY = "./web" + +_SUBPACKAGES: list[str] = [ + "model_loader", + "conditioning", + "sampling", + "qwen", + "latent", + "image", + "repair", + "upscale", + "mask", + "mask_draw", + "controlnet", + "lora", + "video", + "audio", + "utils", +] + + +def _merge_subpackage(pkg_name: str, classes: dict, display_names: dict) -> None: + full_name = f"{__name__}.{pkg_name}" + try: + mod: types.ModuleType = importlib.import_module(full_name) + except Exception as exc: # pragma: no cover + print(f"[LLS] WARNING: Failed to import subpackage '{pkg_name}': {exc}") + return + + sub_classes: dict = getattr(mod, "NODE_CLASS_MAPPINGS", {}) + sub_names: dict = getattr(mod, "NODE_DISPLAY_NAME_MAPPINGS", {}) + + for key, cls in sub_classes.items(): + if key in classes: + continue + classes[key] = cls + if key in sub_names: + display_names[key] = sub_names[key] + + +def _merge_root_nodes(classes: dict, display_names: dict) -> None: + try: + mod: types.ModuleType = importlib.import_module(f"{__name__}.nodes") + except Exception as exc: # pragma: no cover + print(f"[LLS] WARNING: Failed to import root nodes module: {exc}") + return + + root_classes: dict = getattr(mod, "NODE_CLASS_MAPPINGS", {}) + root_names: dict = getattr(mod, "NODE_DISPLAY_NAME_MAPPINGS", {}) + + for key, cls in root_classes.items(): + if key in classes: + continue + classes[key] = cls + if key in root_names: + display_names[key] = root_names[key] + + +NODE_CLASS_MAPPINGS: dict[str, type] = {} +NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = {} + +_merge_root_nodes(NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS) + +for _pkg in _SUBPACKAGES: + _merge_subpackage(_pkg, NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS) + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] +``` + +```python +# mask_draw/__init__.py +from .node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +``` + +```python +# mask_draw/utils.py +def parse_mask_state(mask_state_json): + return { + "version": 1, + "mask_png_base64": "", + "touched": False, + "editor": {}, + } +``` + +```python +# mask_draw/node.py +class LLSSimpleMaskDraw: + CATEGORY = "LLS/Image Repair" + FUNCTION = "draw_mask" + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("image", "mask", "preview_image") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "draw_mode": (["brush", "erase"], {"default": "brush"}), + "brush_size": ("INT", {"default": 32, "min": 1, "max": 512, "step": 1}), + "brush_softness": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "overlay_alpha": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}), + "invert_mask": ("BOOLEAN", {"default": False}), + "mask_state_json": ("STRING", {"default": "{}", "multiline": False, "advanced": True}), + }, + "optional": { + "input_mask": ("MASK",), + }, + "hidden": { + "node_id": "UNIQUE_ID", + }, + } + + def draw_mask( + self, + image, + draw_mode, + brush_size, + brush_softness, + overlay_alpha, + invert_mask, + mask_state_json, + input_mask=None, + node_id=None, + ): + raise RuntimeError("[LLS] LLS Simple Mask Draw is not implemented yet.") + + +NODE_CLASS_MAPPINGS = {"LLSSimpleMaskDraw": LLSSimpleMaskDraw} +NODE_DISPLAY_NAME_MAPPINGS = {"LLSSimpleMaskDraw": "LLS Simple Mask Draw"} +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_registration.py' -v` +Expected: PASS with the registration and schema contract green. + +- [ ] **Step 5: Commit** + +```bash +git add __init__.py mask_draw/__init__.py mask_draw/node.py mask_draw/utils.py tests/test_mask_draw_helpers.py tests/test_mask_draw_registration.py +git commit -m "feat: scaffold LLS simple mask draw node" +``` + +### Task 2: Implement Backend State Decoding and Mask Resolution Helpers + +**Files:** +- Modify: `tests/test_mask_draw_helpers.py` +- Create: `tests/test_mask_draw_utils.py` +- Modify: `mask_draw/utils.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_mask_draw_helpers.py +import base64 +import io +import json + +from PIL import Image +import torch + + +def make_image(width=8, height=8, color=0.25): + return torch.full((1, height, width, 3), float(color), dtype=torch.float32) + + +def make_mask(width=8, height=8, value=0.0): + return torch.full((1, height, width), float(value), dtype=torch.float32) + + +def make_mask_state_json(width=8, height=8, value=1.0, touched=True): + pixels = Image.new("L", (width, height), color=int(round(float(value) * 255.0))) + buffer = io.BytesIO() + pixels.save(buffer, format="PNG") + payload = { + "version": 1, + "mask_png_base64": base64.b64encode(buffer.getvalue()).decode("ascii"), + "touched": bool(touched), + "editor": { + "draw_mode": "brush", + "brush_size": 32, + "brush_softness": 0.5, + "overlay_alpha": 0.4, + }, + } + return json.dumps(payload) +``` + +```python +# tests/test_mask_draw_utils.py +import unittest + +import torch + +try: + from .test_mask_draw_helpers import make_image, make_mask, make_mask_state_json + from ..mask_draw.utils import build_preview_image, parse_mask_state, resolve_output_mask +except ImportError: + from test_mask_draw_helpers import make_image, make_mask, make_mask_state_json + from lls_node_test_mask_draw.mask_draw.utils import build_preview_image, parse_mask_state, resolve_output_mask + + +class TestMaskDrawUtils(unittest.TestCase): + def test_parse_mask_state_defaults_invalid_json(self): + state = parse_mask_state("{invalid") + + self.assertFalse(state["touched"]) + self.assertEqual(state["mask_png_base64"], "") + self.assertEqual(state["version"], 1) + + def test_resolve_output_mask_prefers_saved_mask_when_touched(self): + image = make_image(width=6, height=4) + input_mask = make_mask(width=6, height=4, value=0.0) + state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True) + + mask = resolve_output_mask(image=image, input_mask=input_mask, mask_state_json=state_json, invert_mask=False) + + self.assertEqual(tuple(mask.shape), (1, 4, 6)) + self.assertTrue(torch.all(mask == 1.0)) + + def test_resolve_output_mask_uses_input_mask_when_untouched(self): + image = make_image(width=6, height=4) + input_mask = make_mask(width=3, height=2, value=1.0) + + mask = resolve_output_mask(image=image, input_mask=input_mask, mask_state_json="{}", invert_mask=False) + + self.assertEqual(tuple(mask.shape), (1, 4, 6)) + self.assertTrue(torch.all(mask == 1.0)) + + def test_resolve_output_mask_returns_black_after_clear_state(self): + image = make_image(width=6, height=4) + input_mask = make_mask(width=6, height=4, value=1.0) + state_json = make_mask_state_json(width=6, height=4, value=0.0, touched=True) + + mask = resolve_output_mask(image=image, input_mask=input_mask, mask_state_json=state_json, invert_mask=False) + + self.assertTrue(torch.all(mask == 0.0)) + + def test_resolve_output_mask_applies_invert(self): + image = make_image(width=6, height=4) + state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True) + + mask = resolve_output_mask(image=image, input_mask=None, mask_state_json=state_json, invert_mask=True) + + self.assertTrue(torch.all(mask == 0.0)) + + def test_build_preview_image_preserves_shape(self): + image = make_image(width=6, height=4, color=0.2) + mask = make_mask(width=6, height=4, value=1.0) + + preview = build_preview_image(image=image, mask=mask, overlay_alpha=0.4) + + self.assertEqual(tuple(preview.shape), (1, 4, 6, 3)) + self.assertGreater(float(preview[..., 0].mean()), float(image[..., 0].mean())) + + +if __name__ == "__main__": + unittest.main() +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_utils.py' -v` +Expected: FAIL because `resolve_output_mask` and `build_preview_image` do not exist yet and `parse_mask_state` is incomplete. + +- [ ] **Step 3: Implement the minimal helper logic** + +```python +# mask_draw/utils.py +from __future__ import annotations + +import base64 +import io +import json + +import numpy as np +from PIL import Image + +try: + import torch + import torch.nn.functional as F +except Exception as exc: # pragma: no cover + torch = None + F = None + _TORCH_ERR = exc +else: # pragma: no cover + _TORCH_ERR = None + + +def get_image_size(image) -> tuple[int, int]: + shape = tuple(getattr(image, "shape", ())) + if len(shape) != 4: + raise RuntimeError("[LLS] image must have shape [batch, height, width, channels].") + return int(shape[2]), int(shape[1]) + + +def parse_mask_state(mask_state_json): + default = { + "version": 1, + "mask_png_base64": "", + "touched": False, + "editor": {}, + } + if not mask_state_json: + return default + try: + raw = json.loads(mask_state_json) + except Exception: + return default + if not isinstance(raw, dict): + return default + return { + "version": int(raw.get("version", 1)), + "mask_png_base64": str(raw.get("mask_png_base64") or ""), + "touched": bool(raw.get("touched", False)), + "editor": dict(raw.get("editor") or {}), + } + + +def make_black_mask(image): + if torch is None: + raise RuntimeError("[LLS] torch is required for mask generation.") from _TORCH_ERR + width, height = get_image_size(image) + return torch.zeros((image.shape[0], height, width), dtype=image.dtype, device=image.device) + + +def resize_mask_to_image(mask, image): + width, height = get_image_size(image) + if mask is None: + return None + if tuple(mask.shape[1:]) == (height, width): + return mask.clamp(0.0, 1.0) + resized = F.interpolate(mask.unsqueeze(1), size=(height, width), mode="bilinear", align_corners=False) + return resized.squeeze(1).clamp(0.0, 1.0) + + +def decode_mask_png(mask_png_base64, image): + if not mask_png_base64: + return None + width, height = get_image_size(image) + try: + payload = base64.b64decode(mask_png_base64) + loaded = Image.open(io.BytesIO(payload)).convert("L") + except Exception: + return None + if loaded.size != (width, height): + loaded = loaded.resize((width, height), Image.BILINEAR) + array = np.asarray(loaded).astype(np.float32) / 255.0 + tensor = torch.from_numpy(array).to(device=image.device, dtype=image.dtype) + return tensor.unsqueeze(0).clamp(0.0, 1.0) + + +def resolve_output_mask(image, input_mask, mask_state_json, invert_mask): + state = parse_mask_state(mask_state_json) + resolved = None + if state["touched"] and state["mask_png_base64"]: + resolved = decode_mask_png(state["mask_png_base64"], image) + if resolved is None and input_mask is not None: + resolved = resize_mask_to_image(input_mask, image) + if resolved is None: + resolved = make_black_mask(image) + if invert_mask: + resolved = 1.0 - resolved + return resolved.clamp(0.0, 1.0) + + +def build_preview_image(image, mask, overlay_alpha): + alpha = max(0.0, min(1.0, float(overlay_alpha))) + expanded = mask.unsqueeze(-1) + overlay = torch.zeros_like(image) + overlay[..., 0] = 1.0 + return (image * (1.0 - expanded * alpha) + overlay * (expanded * alpha)).clamp(0.0, 1.0) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_utils.py' -v` +Expected: PASS with green coverage for saved-state precedence, fallback, invert, resize, and preview shape. + +- [ ] **Step 5: Commit** + +```bash +git add mask_draw/utils.py tests/test_mask_draw_helpers.py tests/test_mask_draw_utils.py +git commit -m "feat: add mask draw backend resolution helpers" +``` + +### Task 3: Implement Node Orchestration and Repair-Prepare Compatibility + +**Files:** +- Create: `tests/test_mask_draw_node.py` +- Modify: `mask_draw/node.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_mask_draw_node.py +import unittest + +import torch + +try: + from .test_mask_draw_helpers import load_plugin_package, make_image, make_mask, make_mask_state_json +except ImportError: + from test_mask_draw_helpers import load_plugin_package, make_image, make_mask, make_mask_state_json + + +class FakeVAE: + def encode(self, image): + batch, height, width, _channels = image.shape + return torch.zeros((batch, 4, max(1, height // 8), max(1, width // 8)), dtype=image.dtype, device=image.device) + + +class TestMaskDrawNode(unittest.TestCase): + def test_node_uses_saved_mask_when_touched(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskDraw"]() + image = make_image(width=6, height=4, color=0.2) + state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True) + + image_out, mask_out, preview_out = node.draw_mask( + image=image, + draw_mode="brush", + brush_size=32, + brush_softness=0.5, + overlay_alpha=0.4, + invert_mask=False, + mask_state_json=state_json, + input_mask=None, + node_id="123", + ) + + self.assertIs(image_out, image) + self.assertEqual(tuple(mask_out.shape), (1, 4, 6)) + self.assertEqual(tuple(preview_out.shape), (1, 4, 6, 3)) + self.assertTrue(torch.all(mask_out == 1.0)) + + def test_node_falls_back_to_input_mask_when_untouched(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskDraw"]() + image = make_image(width=6, height=4) + input_mask = make_mask(width=3, height=2, value=1.0) + + _image_out, mask_out, _preview_out = node.draw_mask( + image=image, + draw_mode="brush", + brush_size=32, + brush_softness=0.5, + overlay_alpha=0.4, + invert_mask=False, + mask_state_json="{}", + input_mask=input_mask, + node_id="123", + ) + + self.assertEqual(tuple(mask_out.shape), (1, 4, 6)) + self.assertTrue(torch.all(mask_out == 1.0)) + + def test_node_returns_black_after_clear_state(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskDraw"]() + image = make_image(width=6, height=4) + input_mask = make_mask(width=6, height=4, value=1.0) + cleared_state = make_mask_state_json(width=6, height=4, value=0.0, touched=True) + + _image_out, mask_out, _preview_out = node.draw_mask( + image=image, + draw_mode="erase", + brush_size=32, + brush_softness=0.5, + overlay_alpha=0.4, + invert_mask=False, + mask_state_json=cleared_state, + input_mask=input_mask, + node_id="123", + ) + + self.assertTrue(torch.all(mask_out == 0.0)) + + def test_node_output_connects_to_repair_prepare(self): + plugin = load_plugin_package() + draw_node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskDraw"]() + prepare_node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]() + image = make_image(width=64, height=64, color=0.1) + state_json = make_mask_state_json(width=64, height=64, value=1.0, touched=True) + + image_out, mask_out, _preview_out = draw_node.draw_mask( + image=image, + draw_mode="brush", + brush_size=32, + brush_softness=0.5, + overlay_alpha=0.4, + invert_mask=False, + mask_state_json=state_json, + input_mask=None, + node_id="123", + ) + + latent, work_image, work_mask, repair_info, recommended, *_rest = prepare_node.prepare( + image=image_out, + mask=mask_out, + 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(tuple(work_image.shape), (1, 64, 64, 3)) + self.assertEqual(tuple(work_mask.shape), (1, 64, 64)) + self.assertEqual(repair_info["repair_scope"], "region") + self.assertIn("noise_mask", latent) + self.assertEqual(recommended, 0.45) + + +if __name__ == "__main__": + unittest.main() +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_node.py' -v` +Expected: FAIL because the node method still raises `RuntimeError`. + +- [ ] **Step 3: Implement the node orchestration** + +```python +# mask_draw/node.py +from __future__ import annotations + +from .utils import build_preview_image, resolve_output_mask + + +class LLSSimpleMaskDraw: + CATEGORY = "LLS/Image Repair" + FUNCTION = "draw_mask" + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("image", "mask", "preview_image") + DESCRIPTION = "Interactively draw and refine a repair mask directly on an input image." + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "draw_mode": (["brush", "erase"], {"default": "brush"}), + "brush_size": ("INT", {"default": 32, "min": 1, "max": 512, "step": 1}), + "brush_softness": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "overlay_alpha": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}), + "invert_mask": ("BOOLEAN", {"default": False}), + "mask_state_json": ("STRING", {"default": "{}", "multiline": False, "advanced": True}), + }, + "optional": { + "input_mask": ("MASK",), + }, + "hidden": { + "node_id": "UNIQUE_ID", + }, + } + + def draw_mask( + self, + image, + draw_mode, + brush_size, + brush_softness, + overlay_alpha, + invert_mask, + mask_state_json, + input_mask=None, + node_id=None, + ): + if image is None: + raise RuntimeError("[LLS] image input is required for LLS Simple Mask Draw.") + + mask = resolve_output_mask( + image=image, + input_mask=input_mask, + mask_state_json=mask_state_json, + invert_mask=bool(invert_mask), + ) + preview_image = build_preview_image(image=image, mask=mask, overlay_alpha=overlay_alpha) + return (image, mask, preview_image) + + +NODE_CLASS_MAPPINGS = {"LLSSimpleMaskDraw": LLSSimpleMaskDraw} +NODE_DISPLAY_NAME_MAPPINGS = {"LLSSimpleMaskDraw": "LLS Simple Mask Draw"} +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_node.py' -v` +Expected: PASS with green coverage for touched-state precedence, untouched fallback, clear semantics, and repair-prepare compatibility. + +- [ ] **Step 5: Commit** + +```bash +git add mask_draw/node.py tests/test_mask_draw_node.py +git commit -m "feat: implement LLS simple mask draw backend node" +``` + +### Task 4: Add the Frontend Extension and JS Asset Registration + +**Files:** +- Create: `tests/test_mask_draw_frontend.py` +- Create: `web/js/lls_mask_draw.js` +- Modify: `__init__.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_mask_draw_frontend.py +import pathlib +import unittest + +try: + from .test_mask_draw_helpers import ROOT, load_plugin_package +except ImportError: + from test_mask_draw_helpers import ROOT, load_plugin_package + + +class TestMaskDrawFrontend(unittest.TestCase): + def test_plugin_exports_web_directory(self): + plugin = load_plugin_package() + + self.assertEqual(plugin.WEB_DIRECTORY, "./web") + + def test_frontend_asset_exists(self): + asset = ROOT / "web" / "js" / "lls_mask_draw.js" + self.assertTrue(asset.exists(), msg=f"Missing frontend asset: {asset}") + + def test_frontend_asset_registers_target_extension(self): + asset = (ROOT / "web" / "js" / "lls_mask_draw.js").read_text(encoding="utf-8") + + self.assertIn("app.registerExtension", asset) + self.assertIn("LLS Simple Mask Draw", asset) + self.assertIn("mask_state_json", asset) + self.assertIn("beforeRegisterNodeDef", asset) + + +if __name__ == "__main__": + unittest.main() +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_frontend.py' -v` +Expected: FAIL because the `web/js/lls_mask_draw.js` asset does not exist yet. + +- [ ] **Step 3: Implement the frontend extension** + +```javascript +// web/js/lls_mask_draw.js +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +const TARGET_NODE_NAME = "LLS Simple Mask Draw"; +const HISTORY_LIMIT = 20; + +function findWidget(node, name) { + return node.widgets?.find((widget) => widget.name === name) ?? null; +} + +function hideWidget(widget) { + if (!widget) return; + widget.type = "hidden"; + widget.computeSize = () => [0, -4]; +} + +function parseState(value) { + try { + const parsed = JSON.parse(value || "{}"); + if (!parsed || typeof parsed !== "object") return {}; + return parsed; + } catch { + return {}; + } +} + +function serializeState(stateWidget, state) { + stateWidget.value = JSON.stringify({ + version: 1, + mask_png_base64: state.maskPngBase64 || "", + touched: !!state.touched, + editor: { + draw_mode: state.drawMode || "brush", + brush_size: state.brushSize || 32, + brush_softness: state.brushSoftness ?? 0.5, + overlay_alpha: state.overlayAlpha ?? 0.4, + }, + }); + app.graph.setDirtyCanvas(true, true); +} + +function dataUrlToBase64(dataUrl) { + const commaIndex = dataUrl.indexOf(","); + return commaIndex >= 0 ? dataUrl.slice(commaIndex + 1) : ""; +} + +function getNodeInputLink(node, inputName) { + const input = node.inputs?.find((item) => item.name === inputName); + if (!input || input.link == null) return null; + return node.graph?.links?.[input.link] ?? null; +} + +function getUpstreamNode(node) { + const link = getNodeInputLink(node, "image"); + if (!link) return null; + return node.graph?.getNodeById?.(link.origin_id) ?? null; +} + +function getUpstreamNodeForInput(node, inputName) { + const link = getNodeInputLink(node, inputName); + if (!link) return null; + return node.graph?.getNodeById?.(link.origin_id) ?? null; +} + +function resolveViewUrl(widgetValue, folderType) { + if (!widgetValue) return null; + const params = new URLSearchParams(); + params.set("filename", widgetValue); + if (!String(widgetValue).startsWith("blake3:") && folderType) { + params.set("type", folderType); + } + return `/view?${params.toString()}`; +} + +async function tryLoadFileBackedSource(node) { + const upstream = getUpstreamNode(node); + if (!upstream) return null; + + const widget = upstream.widgets?.find((item) => item.name === "image"); + if (!widget || !widget.value) return null; + + const typeName = String(upstream.type || upstream.comfyClass || upstream.title || ""); + if (typeName.includes("LoadImageOutput")) { + return resolveViewUrl(widget.value, "output"); + } + if (typeName.includes("LoadImage")) { + return resolveViewUrl(widget.value, "input"); + } + return null; +} + +async function tryLoadFileBackedMask(node) { + const upstream = getUpstreamNodeForInput(node, "input_mask"); + if (!upstream) return null; + + const imageWidget = upstream.widgets?.find((item) => item.name === "image"); + if (!imageWidget || !imageWidget.value) return null; + + const typeName = String(upstream.type || upstream.comfyClass || upstream.title || ""); + if (typeName.includes("LoadImageMask")) { + const params = new URLSearchParams(); + params.set("filename", imageWidget.value); + params.set("type", "input"); + params.set("channel", "a"); + return `/view?${params.toString()}`; + } + if (typeName.includes("LoadImage")) { + const params = new URLSearchParams(); + params.set("filename", imageWidget.value); + params.set("type", "input"); + params.set("channel", "a"); + return `/view?${params.toString()}`; + } + return null; +} + +function pushHistory(history, dataUrl) { + history.undo.push(dataUrl); + if (history.undo.length > HISTORY_LIMIT) { + history.undo.shift(); + } + history.redo.length = 0; +} + +function drawOverlay(displayCtx, maskCanvas, overlayAlpha, invertMask) { + const width = displayCtx.canvas.width; + const height = displayCtx.canvas.height; + const scratch = document.createElement("canvas"); + scratch.width = width; + scratch.height = height; + const scratchCtx = scratch.getContext("2d"); + scratchCtx.drawImage(maskCanvas, 0, 0, width, height); + const imageData = scratchCtx.getImageData(0, 0, width, height); + const data = imageData.data; + for (let i = 0; i < data.length; i += 4) { + const value = invertMask ? 255 - data[i] : data[i]; + data[i] = 255; + data[i + 1] = 0; + data[i + 2] = 0; + data[i + 3] = Math.round((value / 255) * overlayAlpha * 255); + } + scratchCtx.putImageData(imageData, 0, 0); + displayCtx.drawImage(scratch, 0, 0); +} + +function buildEditor(node) { + const stateWidget = findWidget(node, "mask_state_json"); + const drawModeWidget = findWidget(node, "draw_mode"); + const brushSizeWidget = findWidget(node, "brush_size"); + const brushSoftnessWidget = findWidget(node, "brush_softness"); + const overlayAlphaWidget = findWidget(node, "overlay_alpha"); + const invertWidget = findWidget(node, "invert_mask"); + + hideWidget(stateWidget); + + const container = document.createElement("div"); + container.className = "lls-mask-draw"; + container.style.display = "flex"; + container.style.flexDirection = "column"; + container.style.gap = "8px"; + container.style.minWidth = "320px"; + + const toolbar = document.createElement("div"); + toolbar.style.display = "flex"; + toolbar.style.gap = "6px"; + + const stage = document.createElement("canvas"); + stage.width = 512; + stage.height = 512; + stage.style.width = "100%"; + stage.style.border = "1px solid #555"; + stage.style.borderRadius = "8px"; + stage.style.background = "#111"; + + const maskCanvas = document.createElement("canvas"); + const history = { undo: [], redo: [] }; + const stageCtx = stage.getContext("2d"); + const maskCtx = maskCanvas.getContext("2d"); + const state = { + touched: false, + drawMode: drawModeWidget?.value || "brush", + brushSize: brushSizeWidget?.value || 32, + brushSoftness: brushSoftnessWidget?.value ?? 0.5, + overlayAlpha: overlayAlphaWidget?.value ?? 0.4, + baseImage: null, + drawing: false, + maskPngBase64: "", + }; + + function snapshotMask() { + return maskCanvas.toDataURL("image/png"); + } + + function restoreMask(dataUrl) { + return new Promise((resolve) => { + const img = new Image(); + img.onload = () => { + maskCtx.clearRect(0, 0, maskCanvas.width, maskCanvas.height); + maskCtx.drawImage(img, 0, 0, maskCanvas.width, maskCanvas.height); + resolve(); + }; + img.src = dataUrl; + }); + } + + function renderStage() { + stageCtx.clearRect(0, 0, stage.width, stage.height); + if (state.baseImage) { + stageCtx.drawImage(state.baseImage, 0, 0, stage.width, stage.height); + } else { + stageCtx.fillStyle = "#1a1a1a"; + stageCtx.fillRect(0, 0, stage.width, stage.height); + stageCtx.fillStyle = "#888"; + stageCtx.fillText("Run once or connect Load Image to initialize preview", 16, 24); + } + drawOverlay(stageCtx, maskCanvas, state.overlayAlpha, !!invertWidget?.value); + } + + function persistMask() { + const pngData = snapshotMask(); + state.maskPngBase64 = dataUrlToBase64(pngData); + serializeState(stateWidget, state); + renderStage(); + } + + function makeButton(label, handler) { + const button = document.createElement("button"); + button.type = "button"; + button.textContent = label; + button.onclick = handler; + return button; + } + + async function loadInitialImage() { + const url = await tryLoadFileBackedSource(node); + if (!url) return; + const image = new Image(); + image.onload = () => { + state.baseImage = image; + if (!maskCanvas.width || !maskCanvas.height) { + maskCanvas.width = image.naturalWidth; + maskCanvas.height = image.naturalHeight; + maskCtx.fillStyle = "black"; + maskCtx.fillRect(0, 0, maskCanvas.width, maskCanvas.height); + } + renderStage(); + + if (!state.maskPngBase64) { + tryLoadFileBackedMask(node).then((maskUrl) => { + if (!maskUrl) return; + const maskImage = new Image(); + maskImage.onload = () => { + maskCtx.clearRect(0, 0, maskCanvas.width, maskCanvas.height); + maskCtx.drawImage(maskImage, 0, 0, maskCanvas.width, maskCanvas.height); + renderStage(); + }; + maskImage.src = maskUrl; + }); + } + }; + image.src = url; + } + + function syncWidgetState() { + state.drawMode = drawModeWidget?.value || "brush"; + state.brushSize = brushSizeWidget?.value || 32; + state.brushSoftness = brushSoftnessWidget?.value ?? 0.5; + state.overlayAlpha = overlayAlphaWidget?.value ?? 0.4; + renderStage(); + } + + const savedState = parseState(stateWidget?.value); + if (savedState.editor) { + state.drawMode = savedState.editor.draw_mode || state.drawMode; + state.brushSize = savedState.editor.brush_size || state.brushSize; + state.brushSoftness = savedState.editor.brush_softness ?? state.brushSoftness; + state.overlayAlpha = savedState.editor.overlay_alpha ?? state.overlayAlpha; + } + state.touched = !!savedState.touched; + state.maskPngBase64 = savedState.mask_png_base64 || ""; + + toolbar.appendChild( + makeButton("Clear", () => { + if (!maskCanvas.width || !maskCanvas.height) return; + pushHistory(history, snapshotMask()); + maskCtx.clearRect(0, 0, maskCanvas.width, maskCanvas.height); + maskCtx.fillStyle = "black"; + maskCtx.fillRect(0, 0, maskCanvas.width, maskCanvas.height); + state.touched = true; + persistMask(); + }), + ); + toolbar.appendChild( + makeButton("Undo", async () => { + if (!history.undo.length) return; + history.redo.push(snapshotMask()); + const previous = history.undo.pop(); + await restoreMask(previous); + state.touched = true; + persistMask(); + }), + ); + toolbar.appendChild( + makeButton("Redo", async () => { + if (!history.redo.length) return; + history.undo.push(snapshotMask()); + const next = history.redo.pop(); + await restoreMask(next); + state.touched = true; + persistMask(); + }), + ); + toolbar.appendChild( + makeButton("Invert", () => { + if (invertWidget) { + invertWidget.value = !invertWidget.value; + invertWidget.callback?.(invertWidget.value); + } + renderStage(); + }), + ); + + stage.addEventListener("pointerdown", (event) => { + if (!maskCanvas.width || !maskCanvas.height) return; + state.drawing = true; + pushHistory(history, snapshotMask()); + const rect = stage.getBoundingClientRect(); + const scaleX = maskCanvas.width / rect.width; + const scaleY = maskCanvas.height / rect.height; + const x = (event.clientX - rect.left) * scaleX; + const y = (event.clientY - rect.top) * scaleY; + const radius = Math.max(1, Number(state.brushSize || 32)) / 2; + const gradient = maskCtx.createRadialGradient(x, y, 0, x, y, radius); + const softness = Math.max(0.0, Math.min(1.0, Number(state.brushSoftness ?? 0.5))); + const edge = 1.0 - softness; + const color = state.drawMode === "erase" ? "0,0,0" : "255,255,255"; + gradient.addColorStop(0.0, `rgba(${color},1)`); + gradient.addColorStop(edge, `rgba(${color},1)`); + gradient.addColorStop(1.0, `rgba(${color},0)`); + maskCtx.fillStyle = gradient; + maskCtx.beginPath(); + maskCtx.arc(x, y, radius, 0, Math.PI * 2); + maskCtx.fill(); + state.touched = true; + persistMask(); + }); + + stage.addEventListener("pointermove", (event) => { + if (!state.drawing || !maskCanvas.width || !maskCanvas.height) return; + const rect = stage.getBoundingClientRect(); + const scaleX = maskCanvas.width / rect.width; + const scaleY = maskCanvas.height / rect.height; + const x = (event.clientX - rect.left) * scaleX; + const y = (event.clientY - rect.top) * scaleY; + const radius = Math.max(1, Number(state.brushSize || 32)) / 2; + const gradient = maskCtx.createRadialGradient(x, y, 0, x, y, radius); + const softness = Math.max(0.0, Math.min(1.0, Number(state.brushSoftness ?? 0.5))); + const edge = 1.0 - softness; + const color = state.drawMode === "erase" ? "0,0,0" : "255,255,255"; + gradient.addColorStop(0.0, `rgba(${color},1)`); + gradient.addColorStop(edge, `rgba(${color},1)`); + gradient.addColorStop(1.0, `rgba(${color},0)`); + maskCtx.fillStyle = gradient; + maskCtx.beginPath(); + maskCtx.arc(x, y, radius, 0, Math.PI * 2); + maskCtx.fill(); + state.touched = true; + persistMask(); + }); + + stage.addEventListener("pointerup", () => { + state.drawing = false; + }); + stage.addEventListener("pointerleave", () => { + state.drawing = false; + }); + + container.appendChild(toolbar); + container.appendChild(stage); + + node.addDOMWidget("lls_mask_editor", "lls_mask_editor", container, { + serialize: false, + hideOnZoom: false, + getValue: () => stateWidget?.value, + setValue: (value) => { + if (stateWidget) { + stateWidget.value = value; + } + }, + }); + + if (state.maskPngBase64) { + const restored = new Image(); + restored.onload = () => { + if (!maskCanvas.width || !maskCanvas.height) { + maskCanvas.width = restored.naturalWidth; + maskCanvas.height = restored.naturalHeight; + } + maskCtx.drawImage(restored, 0, 0, maskCanvas.width, maskCanvas.height); + renderStage(); + }; + restored.src = `data:image/png;base64,${state.maskPngBase64}`; + } + + drawModeWidget && (drawModeWidget.callback = syncWidgetState); + brushSizeWidget && (brushSizeWidget.callback = syncWidgetState); + brushSoftnessWidget && (brushSoftnessWidget.callback = syncWidgetState); + overlayAlphaWidget && (overlayAlphaWidget.callback = syncWidgetState); + invertWidget && (invertWidget.callback = syncWidgetState); + + loadInitialImage(); + syncWidgetState(); +} + +app.registerExtension({ + name: "LLS.MaskDraw", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== TARGET_NODE_NAME) { + return; + } + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function onMaskDrawNodeCreated() { + const result = onNodeCreated?.apply(this, arguments); + if (!this.__llsMaskDrawEditorAttached) { + this.__llsMaskDrawEditorAttached = true; + buildEditor(this); + } + return result; + }; + }, +}); +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_frontend.py' -v` +Expected: PASS with green coverage for `WEB_DIRECTORY` export and asset presence. + +- [ ] **Step 5: Run a syntax check for the asset** + +Run: `node --check web/js/lls_mask_draw.js` +Expected: PASS with no output. + +- [ ] **Step 6: Commit** + +```bash +git add __init__.py web/js/lls_mask_draw.js tests/test_mask_draw_frontend.py +git commit -m "feat: add LLS simple mask draw frontend editor" +``` + +### Task 5: Update README and Run Full Verification + +**Files:** +- Create: `tests/test_mask_draw_docs.py` +- Modify: `README.md` + +- [ ] **Step 1: Write the failing README contract test** + +```python +# tests/test_mask_draw_docs.py +import pathlib +import unittest + + +ROOT = pathlib.Path(__file__).resolve().parents[1] + + +class TestMaskDrawDocs(unittest.TestCase): + def test_readme_documents_mask_draw_node_and_repair_wiring(self): + readme = (ROOT / "README.md").read_text(encoding="utf-8") + + self.assertIn("LLS Simple Mask Draw", readme) + self.assertIn("Load Image", readme) + self.assertIn("LLS Simple Repair Prepare", readme) + self.assertIn("preview_image", readme) + self.assertIn("手动指定删除区域", readme) + self.assertIn("brush", readme) + self.assertIn("erase", readme) + + +if __name__ == "__main__": + unittest.main() +``` + +- [ ] **Step 2: Run the README test to verify it fails** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_docs.py' -v` +Expected: FAIL because the README does not document the new node yet. + +- [ ] **Step 3: Update the README** + +```markdown +## Image Repair + +`LLS-node` now includes a direct manual mask-drawing helper for local repair workflows: + +- `LLS Simple Mask Draw` +- `LLS Simple Repair Prepare` +- `LLS Simple KSampler` +- `LLS Simple Repair Finish` + +### `LLS Simple Mask Draw` + +Draw or erase a repair mask directly on the input image inside the ComfyUI node. + +**Inputs** + +- `image: IMAGE` +- `input_mask: MASK` (optional) + +**Outputs** + +- `image: IMAGE` +- `mask: MASK` +- `preview_image: IMAGE` + +**Supported first-version interactions** + +- `brush` +- `erase` +- `Clear` +- `Undo` +- `Redo` +- overlay alpha preview +- optional invert mask + +### Minimal manual workflow + +`Load Image -> LLS Simple Mask Draw -> Preview Image` + +### Repair workflow wiring + +- `Load Image.image -> LLS Simple Mask Draw.image` +- `LLS Simple Mask Draw.image -> LLS Simple Repair Prepare.image` +- `LLS Simple Mask Draw.mask -> LLS Simple Repair Prepare.mask` + +### Typical uses + +- 手动指定删除区域 +- 手动指定修复区域 +- 手动指定去阴影区域 +- 手动指定局部增强区域 + +### First-version limits + +- no polygon tool +- no rectangle tool +- no ellipse tool +- no magic wand +- no automatic segmentation +- some non-file upstream images may need one execution before the editor can show the base image +``` + +- [ ] **Step 4: Run the README test to verify it passes** + +Run: `python3 -m unittest discover -s tests -p 'test_mask_draw_docs.py' -v` +Expected: PASS with the new node documented. + +- [ ] **Step 5: Run the focused automated verification suite** + +Run: `python3 -m unittest tests.test_mask_draw_registration tests.test_mask_draw_utils tests.test_mask_draw_node tests.test_mask_draw_frontend tests.test_mask_draw_docs -v` +Expected: PASS with all new mask-draw tests green. + +- [ ] **Step 6: Run regression checks for adjacent repair behavior** + +Run: `python3 -m unittest tests.test_repair_registration tests.test_repair_prepare -v` +Expected: PASS with no regression in repair registration or prepare behavior. + +- [ ] **Step 7: Run a lightweight compile check** + +Run: `python3 -m compileall __init__.py mask_draw tests` +Expected: PASS with no syntax errors. + +- [ ] **Step 8: Perform the manual frontend verification checklist** + +Run these checks in ComfyUI: + +```text +1. Add `LLS Simple Mask Draw` to the graph and confirm it appears under `LLS/Image Repair`. +2. Connect `Load Image -> LLS Simple Mask Draw -> Preview Image`. +3. Confirm the image preview appears immediately for a file-backed `Load Image`. +4. Paint with `brush` and verify `preview_image` shows a red overlay. +5. Switch to `erase` and verify mask removal. +6. Click `Clear` and verify the output mask becomes black. +7. Use `Undo` and `Redo` repeatedly and verify no crash or stale state. +8. Save the workflow, close it, reopen it, and verify the current mask raster is restored. +9. Connect `Load Image Mask` into `LLS Simple Mask Draw.input_mask` and verify the loaded mask appears as the editable starting mask when the source is file-backed. +10. Connect `LLS Simple Mask Draw.image` and `.mask` into `LLS Simple Repair Prepare`. +11. Queue the repair-prepare path and verify it accepts the mask output without errors. +``` + +Expected: All manual checks succeed, with only the documented first-version limitation for non-file upstream image previews. + +- [ ] **Step 9: Commit** + +```bash +git add README.md tests/test_mask_draw_docs.py +git commit -m "docs: document LLS simple mask draw" +``` diff --git a/mask_draw/__init__.py b/mask_draw/__init__.py new file mode 100644 index 0000000..051786f --- /dev/null +++ b/mask_draw/__init__.py @@ -0,0 +1,8 @@ +""" +node.mask_draw +============== +Mask draw node package exports. +""" +from .node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/mask_draw/node.py b/mask_draw/node.py new file mode 100644 index 0000000..e87f3e0 --- /dev/null +++ b/mask_draw/node.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from .utils import build_preview_image, resolve_output_mask + + +_DRAW_MODE_CHOICES = ["brush", "erase"] + + +class LLSSimpleMaskDraw: + CATEGORY = "LLS/Image Repair" + FUNCTION = "draw_mask" + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("image", "mask", "preview_image") + DESCRIPTION = "Resolve the edited mask output and build a preview overlay." + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "draw_mode": (_DRAW_MODE_CHOICES, {"default": "brush"}), + "brush_size": ("INT", {"default": 32, "min": 1, "max": 512, "step": 1}), + "brush_softness": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "overlay_alpha": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}), + "invert_mask": ("BOOLEAN", {"default": False}), + "mask_state_json": ("STRING", {"default": "{}", "multiline": False, "advanced": True}), + }, + "optional": { + "input_mask": ("MASK",), + }, + "hidden": { + "node_id": "UNIQUE_ID", + }, + } + + def draw_mask( + self, + image, + draw_mode, + brush_size, + brush_softness, + overlay_alpha, + invert_mask, + mask_state_json, + node_id, + input_mask=None, + ): + del draw_mode, brush_size, brush_softness, node_id + + mask = resolve_output_mask( + image=image, + input_mask=input_mask, + mask_state_json=mask_state_json, + invert_mask=invert_mask, + ) + preview_image = build_preview_image( + image=image, + mask=mask, + overlay_alpha=overlay_alpha, + ) + return image, mask, preview_image + + +NODE_CLASS_MAPPINGS = {"LLSSimpleMaskDraw": LLSSimpleMaskDraw} +NODE_DISPLAY_NAME_MAPPINGS = {"LLSSimpleMaskDraw": "LLS Simple Mask Draw"} diff --git a/mask_draw/utils.py b/mask_draw/utils.py new file mode 100644 index 0000000..aca61a6 --- /dev/null +++ b/mask_draw/utils.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +import base64 +import io +import json + +try: + import numpy as np +except Exception as exc: # pragma: no cover - optional runtime dependency + np = None + _NUMPY_ERR = exc +else: # pragma: no cover - exercised only in ComfyUI-like runtime + _NUMPY_ERR = None + +try: + from PIL import Image +except Exception as exc: # pragma: no cover - optional runtime dependency + Image = None + _PIL_ERR = exc +else: # pragma: no cover - exercised only in ComfyUI-like runtime + _PIL_ERR = None + +try: + import torch + import torch.nn.functional as torch_nn_functional +except Exception as exc: # pragma: no cover - optional runtime dependency + torch = None + torch_nn_functional = None + _TORCH_ERR = exc +else: # pragma: no cover - exercised only in ComfyUI-like runtime + _TORCH_ERR = None + + +_DEFAULT_MASK_STATE = { + "version": 1, + "mask_png_base64": "", + "touched": False, + "editor": {}, +} + + +def parse_mask_state(mask_state_json): + if not mask_state_json: + return dict(_DEFAULT_MASK_STATE) + + try: + payload = json.loads(mask_state_json) + except (TypeError, ValueError): + return dict(_DEFAULT_MASK_STATE) + + if not isinstance(payload, dict): + return dict(_DEFAULT_MASK_STATE) + + state = dict(_DEFAULT_MASK_STATE) + state["version"] = _coerce_int(payload.get("version"), 1) + state["mask_png_base64"] = str(payload.get("mask_png_base64") or "") + state["touched"] = _coerce_bool(payload.get("touched"), False) + state["editor"] = _coerce_dict(payload.get("editor")) + return state + + +def resolve_output_mask(image, input_mask, mask_state_json, invert_mask): + state = parse_mask_state(mask_state_json) + resolved = None + + if state["touched"]: + resolved = _decode_mask_png_base64(state["mask_png_base64"], image) + + if resolved is None and input_mask is not None: + resolved = _resize_mask_to_image(input_mask, image) + + if resolved is None: + resolved = _make_black_mask(image) + + if invert_mask: + resolved = 1.0 - resolved + + return resolved.clamp(0.0, 1.0) + + +def build_preview_image(image, mask, overlay_alpha): + _require_torch("[LLS] torch is required for preview image generation.") + + alpha = max(0.0, min(1.0, _coerce_float(overlay_alpha, 0.4))) + normalized_mask = _resize_mask_to_image(mask, image).unsqueeze(-1) + overlay_strength = normalized_mask * alpha + overlay = torch.zeros_like(image) + overlay[..., 0] = 1.0 + return (image * (1.0 - overlay_strength) + (overlay * overlay_strength)).clamp(0.0, 1.0) + + +def _get_image_size(image): + 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, _ = tuple(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_image_batch_size(image): + 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].") + + batch = _coerce_int(shape[0], 0) + if batch <= 0: + raise RuntimeError("[LLS] image batch size must be positive.") + return batch + + +def _make_black_mask(image): + _require_torch("[LLS] torch is required for mask generation.") + + width, height = _get_image_size(image) + return torch.zeros((image.shape[0], height, width), dtype=image.dtype, device=image.device) + + +def _resize_mask_to_image(mask, image): + _require_torch("[LLS] torch is required for mask resizing.") + + 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].") + + image_batch = _get_image_batch_size(image) + mask_batch = _coerce_int(shape[0], 0) + if mask_batch <= 0: + raise RuntimeError("[LLS] mask batch size must be positive.") + + width, height = _get_image_size(image) + normalized = mask.to(device=image.device, dtype=image.dtype).clamp(0.0, 1.0) + if mask_batch != image_batch: + if mask_batch == 1 and image_batch > 1: + normalized = normalized.repeat(image_batch, 1, 1) + else: + raise RuntimeError( + f"[LLS] mask batch size must match image batch size or be 1; got mask batch {mask_batch} and image batch {image_batch}." + ) + if tuple(normalized.shape[1:]) == (height, width): + return normalized + + resized = torch_nn_functional.interpolate( + normalized.unsqueeze(1), + size=(height, width), + mode="bilinear", + align_corners=False, + ) + return resized.squeeze(1).clamp(0.0, 1.0) + + +def _decode_mask_png_base64(mask_png_base64, image): + _require_torch("[LLS] torch is required for mask decoding.") + _require_mask_decode_deps() + + if not mask_png_base64: + return None + + width, height = _get_image_size(image) + try: + payload = base64.b64decode(mask_png_base64) + decoded = Image.open(io.BytesIO(payload)).convert("L") + except Exception: + return None + + if decoded.size != (width, height): + decoded = decoded.resize((width, height), Image.BILINEAR) + + array = np.asarray(decoded, dtype=np.float32) / 255.0 + tensor = torch.from_numpy(array).to(device=image.device, dtype=image.dtype) + tensor = tensor.unsqueeze(0) + if image.shape[0] > 1: + tensor = tensor.repeat(image.shape[0], 1, 1) + return tensor.clamp(0.0, 1.0) + + +def _require_torch(message): + if torch is None or torch_nn_functional is None: + raise RuntimeError(message) from _TORCH_ERR + + +def _require_mask_decode_deps(): + if np is None: + raise RuntimeError("[LLS] numpy is required for mask decoding.") from _NUMPY_ERR + if Image is None: + raise RuntimeError("[LLS] Pillow is required for mask decoding.") from _PIL_ERR + + +def _coerce_bool(value, default): + if isinstance(value, bool): + return value + return default + + +def _coerce_dict(value): + if isinstance(value, dict): + return dict(value) + return {} + + +def _coerce_int(value, default): + try: + return int(value) + except (TypeError, ValueError): + return default + + +def _coerce_float(value, default): + try: + return float(value) + except (TypeError, ValueError): + return default diff --git a/repair/repair_prepare.py b/repair/repair_prepare.py index b220920..c96b4d3 100644 --- a/repair/repair_prepare.py +++ b/repair/repair_prepare.py @@ -46,8 +46,8 @@ _CANVAS_FILL_CHOICES = ["edge", "blur", "black", "white", "neutral"] class LLSSimpleRepairPrepare: CATEGORY = "LLS/Image Repair" FUNCTION = "prepare" - RETURN_TYPES = ("LATENT", "IMAGE", "MASK", "LLS_REPAIR_INFO", "FLOAT") - RETURN_NAMES = ("latent", "work_image", "work_mask", "repair_info", "recommended_denoise") + RETURN_TYPES = ("LATENT", "IMAGE", "MASK", "LLS_REPAIR_INFO", "FLOAT", "CONDITIONING", "CONDITIONING") + RETURN_NAMES = ("latent", "work_image", "work_mask", "repair_info", "recommended_denoise", "positive", "negative") DESCRIPTION = "Prepare repair inputs, work area, and repair metadata." @classmethod @@ -110,8 +110,6 @@ class LLSSimpleRepairPrepare: positive=None, negative=None, ): - del positive, negative - if vae is None: raise RuntimeError("[LLS] Missing VAE for LLS Simple Repair Prepare.") @@ -190,7 +188,7 @@ class LLSSimpleRepairPrepare: 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 + return latent, work_image, work_mask, repair_info, recommended, positive, negative if effective_scope == "crop": if mask_bbox is None or mask_area_ratio <= 0.0: @@ -221,7 +219,7 @@ class LLSSimpleRepairPrepare: 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 + return latent, work_image, work_mask, repair_info, recommended, positive, negative if effective_scope == "canvas": canvas_info = build_canvas_info( @@ -264,7 +262,7 @@ class LLSSimpleRepairPrepare: 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 + return latent, work_image, work_mask, repair_info, recommended, positive, negative raise RuntimeError(f"[LLS] Unsupported repair scope '{effective_scope}'.") diff --git a/sampling/nodes.py b/sampling/nodes.py index 39f9af3..cbcb933 100644 --- a/sampling/nodes.py +++ b/sampling/nodes.py @@ -240,13 +240,13 @@ class LLSSimpleKSampler: "sampler_name": (_get_samplers(), {"default": "euler_ancestral"}), "scheduler": (_get_schedulers(), {"default": "karras"}), "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - "denoise_mode": (_DENOISE_MODE_CHOICES, {"default": "manual"}), - "adapter_mode": (_ADAPTER_MODE_CHOICES, {"default": "auto"}), "flux_guidance": ( _PRIMITIVE_NUMBER_INPUT, {"default": 3.5, "widgetType": "FLOAT", "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.1}, ), "model_family": (MODEL_FAMILY_CHOICES, {"default": "Auto"}), + "denoise_mode": (_DENOISE_MODE_CHOICES, {"default": "manual"}), + "adapter_mode": (_ADAPTER_MODE_CHOICES, {"default": "auto"}), }, "optional": { "repair_info": ("LLS_REPAIR_INFO",), diff --git a/tests/test_loader_prompt_refactor.py b/tests/test_loader_prompt_refactor.py index 0cfd088..f968a7f 100644 --- a/tests/test_loader_prompt_refactor.py +++ b/tests/test_loader_prompt_refactor.py @@ -593,7 +593,7 @@ class TestLoaderPromptRefactor(unittest.TestCase): self.assertEqual(sample_data["sampler_name"], "euler") self.assertEqual(latent["samples"].shape, (1, 4, 64, 64)) - def test_ksampler_schema_appends_model_family_after_legacy_widgets(self): + def test_ksampler_schema_appends_repair_widgets_after_legacy_widgets(self): plugin = load_plugin_package() node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleKSampler"] required = node_cls.INPUT_TYPES()["required"] @@ -612,10 +612,10 @@ class TestLoaderPromptRefactor(unittest.TestCase): "sampler_name", "scheduler", "denoise", - "denoise_mode", - "adapter_mode", "flux_guidance", "model_family", + "denoise_mode", + "adapter_mode", ), ) self.assertEqual(required["model_family"][1]["default"], "Auto") diff --git a/tests/test_mask_draw_docs.py b/tests/test_mask_draw_docs.py new file mode 100644 index 0000000..07e899c --- /dev/null +++ b/tests/test_mask_draw_docs.py @@ -0,0 +1,30 @@ +import unittest + +try: + from .test_mask_draw_helpers import ROOT +except ImportError: + from test_mask_draw_helpers import ROOT + + +class TestMaskDrawDocs(unittest.TestCase): + def test_readme_documents_mask_draw_node(self): + readme = (ROOT / "README.md").read_text(encoding="utf-8") + + self.assertIn("LLS Simple Mask Draw", readme) + self.assertIn("Load Image", readme) + self.assertIn("LLS Simple Repair Prepare", readme) + self.assertIn("preview_image", readme) + self.assertIn("brush", readme) + self.assertIn("erase", readme) + + def test_readme_lists_recommended_use_cases(self): + readme = (ROOT / "README.md").read_text(encoding="utf-8") + + self.assertIn("手动指定删除区域", readme) + self.assertIn("手动指定修复区域", readme) + self.assertIn("手动指定去阴影区域", readme) + self.assertIn("手动指定局部增强区域", readme) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mask_draw_frontend.py b/tests/test_mask_draw_frontend.py new file mode 100644 index 0000000..ea79b5d --- /dev/null +++ b/tests/test_mask_draw_frontend.py @@ -0,0 +1,36 @@ +import unittest + +try: + from .test_mask_draw_helpers import ROOT, load_plugin_package +except ImportError: + from test_mask_draw_helpers import ROOT, load_plugin_package + + +class TestMaskDrawFrontend(unittest.TestCase): + def test_plugin_exports_web_directory(self): + plugin = load_plugin_package() + + self.assertEqual(plugin.WEB_DIRECTORY, "./web") + + def test_frontend_asset_exists(self): + asset = ROOT / "web" / "js" / "lls_mask_draw.js" + + self.assertTrue(asset.exists(), msg=f"Missing frontend asset: {asset}") + + def test_frontend_asset_registers_mask_draw_extension(self): + asset = (ROOT / "web" / "js" / "lls_mask_draw.js").read_text(encoding="utf-8") + + self.assertIn("app.registerExtension", asset) + self.assertIn("LLSSimpleMaskDraw", asset) + self.assertIn("LLS Simple Mask Draw", asset) + self.assertIn("mask_state_json", asset) + self.assertIn("beforeRegisterNodeDef", asset) + self.assertIn("addDOMWidget", asset) + self.assertIn("Clear", asset) + self.assertIn("Undo", asset) + self.assertIn("Redo", asset) + self.assertIn("Invert", asset) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mask_draw_helpers.py b/tests/test_mask_draw_helpers.py new file mode 100644 index 0000000..7501c5a --- /dev/null +++ b/tests/test_mask_draw_helpers.py @@ -0,0 +1,82 @@ +import base64 +import importlib.util +import pathlib +import sys +import importlib +import io +import json +import unittest + + +ROOT = pathlib.Path(__file__).resolve().parents[1] +MODULE_NAME = "lls_node_test_mask_draw" + + +def _import_optional(module_name): + try: + return importlib.import_module(module_name) + except Exception: + return None + + +torch = _import_optional("torch") +np = _import_optional("numpy") +_pil_image_module = _import_optional("PIL.Image") +Image = _pil_image_module +HAS_MASK_DRAW_RUNTIME_DEPS = torch is not None and np is not None and Image is not None +MASK_DRAW_RUNTIME_DEPS_MESSAGE = "mask draw helper tests require torch, numpy, and Pillow" + + +def load_plugin_package(): + for name in list(sys.modules): + if name == MODULE_NAME or name.startswith(f"{MODULE_NAME}."): + sys.modules.pop(name) + + spec = importlib.util.spec_from_file_location( + MODULE_NAME, + ROOT / "__init__.py", + submodule_search_locations=[str(ROOT)], + ) + module = importlib.util.module_from_spec(spec) + sys.modules[MODULE_NAME] = module + spec.loader.exec_module(module) + return module + + +def load_mask_draw_utils(): + load_plugin_package() + return importlib.import_module(f"{MODULE_NAME}.mask_draw.utils") + + +def make_image(width=8, height=8, color=0.25): + require_mask_draw_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_draw_runtime_deps() + return torch.full((1, height, width), float(value), dtype=torch.float32) + + +def make_mask_state_json(width=8, height=8, value=1.0, touched=True): + require_mask_draw_runtime_deps() + pixels = Image.new("L", (width, height), color=int(round(float(value) * 255.0))) + buffer = io.BytesIO() + pixels.save(buffer, format="PNG") + payload = { + "version": 1, + "mask_png_base64": base64.b64encode(buffer.getvalue()).decode("ascii"), + "touched": bool(touched), + "editor": { + "draw_mode": "brush", + "brush_size": 32, + "brush_softness": 0.5, + "overlay_alpha": 0.4, + }, + } + return json.dumps(payload) + + +def require_mask_draw_runtime_deps(): + if not HAS_MASK_DRAW_RUNTIME_DEPS: + raise unittest.SkipTest(MASK_DRAW_RUNTIME_DEPS_MESSAGE) diff --git a/tests/test_mask_draw_node.py b/tests/test_mask_draw_node.py new file mode 100644 index 0000000..e99b868 --- /dev/null +++ b/tests/test_mask_draw_node.py @@ -0,0 +1,159 @@ +import unittest + +try: + from .test_mask_draw_helpers import ( + HAS_MASK_DRAW_RUNTIME_DEPS, + MASK_DRAW_RUNTIME_DEPS_MESSAGE, + load_plugin_package, + make_image, + make_mask, + make_mask_state_json, + torch, + ) +except ImportError: + from test_mask_draw_helpers import ( + HAS_MASK_DRAW_RUNTIME_DEPS, + MASK_DRAW_RUNTIME_DEPS_MESSAGE, + load_plugin_package, + make_image, + make_mask, + make_mask_state_json, + torch, + ) + + +class _FakeVAE: + def encode(self, image): + batch, height, width, _channels = tuple(image.shape) + latent_height = max(1, int(height) // 8) + latent_width = max(1, int(width) // 8) + return torch.zeros((batch, 4, latent_height, latent_width), dtype=image.dtype, device=image.device) + + +@unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) +class TestMaskDrawNode(unittest.TestCase): + @classmethod + def setUpClass(cls): + plugin = load_plugin_package() + cls.mask_draw_node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskDraw"]() + cls.repair_prepare_node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]() + + def test_draw_mask_uses_saved_mask_when_state_was_touched(self): + image = make_image(width=6, height=4, color=0.2) + input_mask = make_mask(width=6, height=4, value=0.0) + state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True) + + image_out, mask_out, preview_out = self.mask_draw_node.draw_mask( + image=image, + draw_mode="brush", + brush_size=32, + brush_softness=0.5, + overlay_alpha=0.4, + invert_mask=False, + mask_state_json=state_json, + node_id="node-1", + input_mask=input_mask, + ) + + self.assertIs(image_out, image) + self.assertEqual(tuple(mask_out.shape), (1, 4, 6)) + self.assertTrue(torch.all(mask_out == 1.0)) + self.assertEqual(tuple(preview_out.shape), (1, 4, 6, 3)) + self.assertGreater(float(preview_out[..., 0].mean()), float(image[..., 0].mean())) + + def test_draw_mask_falls_back_to_input_mask_when_state_is_untouched(self): + image = make_image(width=6, height=4, color=0.2) + input_mask = make_mask(width=3, height=2, value=1.0) + + _image_out, mask_out, preview_out = self.mask_draw_node.draw_mask( + image=image, + draw_mode="brush", + brush_size=32, + brush_softness=0.5, + overlay_alpha=0.4, + invert_mask=False, + mask_state_json="{}", + node_id="node-2", + input_mask=input_mask, + ) + + self.assertEqual(tuple(mask_out.shape), (1, 4, 6)) + self.assertTrue(torch.all(mask_out == 1.0)) + self.assertEqual(tuple(preview_out.shape), (1, 4, 6, 3)) + + def test_draw_mask_returns_black_mask_for_clear_state(self): + image = make_image(width=6, height=4, color=0.2) + input_mask = make_mask(width=6, height=4, value=1.0) + state_json = make_mask_state_json(width=6, height=4, value=0.0, touched=True) + + _image_out, mask_out, preview_out = self.mask_draw_node.draw_mask( + image=image, + draw_mode="erase", + brush_size=32, + brush_softness=0.5, + overlay_alpha=0.4, + invert_mask=False, + mask_state_json=state_json, + node_id="node-3", + input_mask=input_mask, + ) + + self.assertTrue(torch.all(mask_out == 0.0)) + self.assertTrue(torch.allclose(preview_out, image)) + + def test_draw_mask_output_connects_to_repair_prepare_region_path(self): + image = make_image(width=64, height=64, color=0.25) + state_json = make_mask_state_json(width=64, height=64, value=1.0, touched=True) + + image_out, mask_out, preview_out = self.mask_draw_node.draw_mask( + image=image, + draw_mode="brush", + brush_size=32, + brush_softness=0.5, + overlay_alpha=0.4, + invert_mask=False, + mask_state_json=state_json, + node_id="node-4", + input_mask=None, + ) + + latent, work_image, work_mask, repair_info, recommended_denoise, positive, negative = self.repair_prepare_node.prepare( + image=image_out, + mask=mask_out, + vae=_FakeVAE(), + repair_scope="region", + repair_kernel="vae_inpaint", + task_hint="repair", + mask_grow=0, + mask_blur=0.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=None, + positive=None, + negative=None, + ) + + self.assertEqual(tuple(preview_out.shape), (1, 64, 64, 3)) + self.assertEqual(tuple(work_image.shape), (1, 64, 64, 3)) + self.assertEqual(tuple(work_mask.shape), (1, 64, 64)) + self.assertEqual(latent["source"], "repair_prepare_region") + self.assertTrue(repair_info["has_mask"]) + self.assertGreater(repair_info["mask_area_ratio"], 0.0) + self.assertEqual(recommended_denoise, 0.45) + self.assertIsNone(positive) + self.assertIsNone(negative) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mask_draw_registration.py b/tests/test_mask_draw_registration.py new file mode 100644 index 0000000..93624ac --- /dev/null +++ b/tests/test_mask_draw_registration.py @@ -0,0 +1,52 @@ +import unittest + +try: + from .test_mask_draw_helpers import load_plugin_package +except ImportError: + from test_mask_draw_helpers import load_plugin_package + + +class TestMaskDrawRegistration(unittest.TestCase): + def test_plugin_registers_mask_draw_node(self): + plugin = load_plugin_package() + + self.assertIn("LLSSimpleMaskDraw", plugin.NODE_CLASS_MAPPINGS) + self.assertEqual( + plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSSimpleMaskDraw"], + "LLS Simple Mask Draw", + ) + + def test_mask_draw_schema_matches_contract(self): + plugin = load_plugin_package() + node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleMaskDraw"] + schema = node_cls.INPUT_TYPES() + + self.assertEqual(node_cls.CATEGORY, "LLS/Image Repair") + self.assertEqual(node_cls.FUNCTION, "draw_mask") + self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE", "MASK", "IMAGE")) + self.assertEqual(node_cls.RETURN_NAMES, ("image", "mask", "preview_image")) + + required = schema["required"] + optional = schema["optional"] + hidden = schema["hidden"] + + self.assertEqual(required["image"], ("IMAGE",)) + self.assertEqual(optional["input_mask"], ("MASK",)) + self.assertEqual(required["draw_mode"][0], ["brush", "erase"]) + self.assertEqual(required["brush_size"][0], "INT") + self.assertEqual(required["brush_size"][1]["max"], 512) + self.assertEqual(required["brush_softness"][0], "FLOAT") + self.assertEqual(required["overlay_alpha"][0], "FLOAT") + self.assertEqual(required["overlay_alpha"][1]["default"], 0.4) + self.assertEqual(required["invert_mask"][0], "BOOLEAN") + self.assertEqual(required["mask_state_json"][0], "STRING") + self.assertEqual( + required["mask_state_json"][1], + {"default": "{}", "multiline": False, "advanced": True}, + ) + self.assertEqual(hidden["node_id"], "UNIQUE_ID") + self.assertTrue(callable(getattr(node_cls, "draw_mask", None))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mask_draw_utils.py b/tests/test_mask_draw_utils.py new file mode 100644 index 0000000..a050683 --- /dev/null +++ b/tests/test_mask_draw_utils.py @@ -0,0 +1,182 @@ +import json +import unittest + +try: + from .test_mask_draw_helpers import ( + HAS_MASK_DRAW_RUNTIME_DEPS, + MASK_DRAW_RUNTIME_DEPS_MESSAGE, + load_mask_draw_utils, + make_image, + make_mask, + make_mask_state_json, + torch, + ) +except ImportError: + from test_mask_draw_helpers import ( + HAS_MASK_DRAW_RUNTIME_DEPS, + MASK_DRAW_RUNTIME_DEPS_MESSAGE, + load_mask_draw_utils, + make_image, + make_mask, + make_mask_state_json, + torch, + ) + + +class TestMaskDrawUtils(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.utils = load_mask_draw_utils() + + def test_parse_mask_state_defaults_invalid_json(self): + state = self.utils.parse_mask_state("{invalid") + + self.assertFalse(state["touched"]) + self.assertEqual(state["mask_png_base64"], "") + self.assertEqual(state["version"], 1) + + def test_parse_mask_state_defaults_invalid_editor_and_dirty_touched_values(self): + for payload in ( + {"editor": [1], "touched": "false"}, + {"editor": "bad", "touched": "true"}, + ): + with self.subTest(payload=payload): + state = self.utils.parse_mask_state(json.dumps(payload)) + + self.assertEqual(state["editor"], {}) + self.assertFalse(state["touched"]) + + @unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) + def test_resolve_output_mask_prefers_saved_mask_when_touched(self): + image = make_image(width=6, height=4) + input_mask = make_mask(width=6, height=4, value=0.0) + state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True) + + mask = self.utils.resolve_output_mask( + image=image, + input_mask=input_mask, + mask_state_json=state_json, + invert_mask=False, + ) + + self.assertEqual(tuple(mask.shape), (1, 4, 6)) + self.assertTrue(torch.all(mask == 1.0)) + + @unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) + def test_resolve_output_mask_repeats_saved_mask_to_image_batch(self): + image = torch.full((3, 4, 6, 3), 0.25, dtype=torch.float32) + state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True) + + mask = self.utils.resolve_output_mask( + image=image, + input_mask=None, + mask_state_json=state_json, + invert_mask=False, + ) + + self.assertEqual(tuple(mask.shape), (3, 4, 6)) + self.assertTrue(torch.all(mask == 1.0)) + + @unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) + def test_resolve_output_mask_uses_input_mask_when_untouched(self): + image = make_image(width=6, height=4) + input_mask = make_mask(width=3, height=2, value=1.0) + + mask = self.utils.resolve_output_mask( + image=image, + input_mask=input_mask, + mask_state_json="{}", + invert_mask=False, + ) + + self.assertEqual(tuple(mask.shape), (1, 4, 6)) + self.assertTrue(torch.all(mask == 1.0)) + + @unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) + def test_resolve_output_mask_repeats_single_input_mask_to_image_batch(self): + image = torch.full((3, 4, 6, 3), 0.25, dtype=torch.float32) + input_mask = make_mask(width=6, height=4, value=1.0) + + mask = self.utils.resolve_output_mask( + image=image, + input_mask=input_mask, + mask_state_json="{}", + invert_mask=False, + ) + + self.assertEqual(tuple(mask.shape), (3, 4, 6)) + self.assertTrue(torch.all(mask == 1.0)) + + @unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) + def test_resolve_output_mask_rejects_input_mask_batch_mismatch(self): + image = torch.full((3, 4, 6, 3), 0.25, dtype=torch.float32) + input_mask = torch.full((2, 4, 6), 1.0, dtype=torch.float32) + + with self.assertRaisesRegex(RuntimeError, "mask batch size must match image batch size"): + self.utils.resolve_output_mask( + image=image, + input_mask=input_mask, + mask_state_json="{}", + invert_mask=False, + ) + + @unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) + def test_resolve_output_mask_returns_black_after_clear_state(self): + image = make_image(width=6, height=4) + input_mask = make_mask(width=6, height=4, value=1.0) + state_json = make_mask_state_json(width=6, height=4, value=0.0, touched=True) + + mask = self.utils.resolve_output_mask( + image=image, + input_mask=input_mask, + mask_state_json=state_json, + invert_mask=False, + ) + + self.assertTrue(torch.all(mask == 0.0)) + + @unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) + def test_resolve_output_mask_falls_back_to_input_mask_when_saved_mask_is_corrupt(self): + image = make_image(width=6, height=4) + input_mask = make_mask(width=6, height=4, value=1.0) + state_json = '{"version": 1, "mask_png_base64": "not-base64", "touched": true, "editor": {}}' + + mask = self.utils.resolve_output_mask( + image=image, + input_mask=input_mask, + mask_state_json=state_json, + invert_mask=False, + ) + + self.assertEqual(tuple(mask.shape), (1, 4, 6)) + self.assertTrue(torch.all(mask == 1.0)) + + @unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) + def test_resolve_output_mask_applies_invert(self): + image = make_image(width=6, height=4) + state_json = make_mask_state_json(width=6, height=4, value=1.0, touched=True) + + mask = self.utils.resolve_output_mask( + image=image, + input_mask=None, + mask_state_json=state_json, + invert_mask=True, + ) + + self.assertTrue(torch.all(mask == 0.0)) + + @unittest.skipUnless(HAS_MASK_DRAW_RUNTIME_DEPS, MASK_DRAW_RUNTIME_DEPS_MESSAGE) + def test_build_preview_image_preserves_shape_and_adds_red_overlay(self): + image = make_image(width=6, height=4, color=0.2) + mask = make_mask(width=6, height=4, value=1.0) + + preview = self.utils.build_preview_image(image=image, mask=mask, overlay_alpha=0.4) + + self.assertEqual(tuple(preview.shape), (1, 4, 6, 3)) + self.assertGreater(float(preview[..., 0].mean()), float(image[..., 0].mean())) + self.assertLess(float(preview[..., 1].mean()), float(image[..., 1].mean())) + self.assertLess(float(preview[..., 2].mean()), float(image[..., 2].mean())) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_repair_prepare.py b/tests/test_repair_prepare.py index 24c5734..4b0eafe 100644 --- a/tests/test_repair_prepare.py +++ b/tests/test_repair_prepare.py @@ -19,7 +19,7 @@ class TestRepairPrepare(unittest.TestCase): label="region-mask", ) - latent, work_image, work_mask, repair_info, recommended_denoise = node.prepare( + latent, work_image, work_mask, repair_info, recommended_denoise, positive, negative = node.prepare( image=image, mask=mask, vae=FakeVAE(), @@ -50,6 +50,8 @@ class TestRepairPrepare(unittest.TestCase): self.assertEqual(repair_info["repair_scope"], "region") self.assertEqual(repair_info["repair_kernel"], "latent_mask") self.assertEqual(recommended_denoise, 0.45) + self.assertIsNone(positive) + self.assertIsNone(negative) def test_prepare_crop_with_vae_inpaint_populates_crop_metadata(self): plugin = load_plugin_package() @@ -63,7 +65,7 @@ class TestRepairPrepare(unittest.TestCase): label="crop-mask", ) - latent, work_image, work_mask, repair_info, recommended_denoise = node.prepare( + latent, work_image, work_mask, repair_info, recommended_denoise, positive, negative = node.prepare( image=image, mask=mask, vae=FakeVAE(), @@ -99,6 +101,8 @@ class TestRepairPrepare(unittest.TestCase): self.assertIn(":crop[", work_mask.label) self.assertEqual(work_image.crop_box, tuple(repair_info["crop_box"])) self.assertEqual(work_mask.crop_box, tuple(repair_info["crop_box"])) + self.assertIsNone(positive) + self.assertIsNone(negative) def test_prepare_canvas_with_vae_inpaint_supports_empty_mask(self): plugin = load_plugin_package() @@ -107,7 +111,7 @@ class TestRepairPrepare(unittest.TestCase): 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( + latent, work_image, work_mask, repair_info, recommended_denoise, positive, negative = node.prepare( image=image, mask=mask, vae=FakeVAE(), @@ -142,6 +146,8 @@ class TestRepairPrepare(unittest.TestCase): self.assertTrue(repair_info["has_mask"]) self.assertIsNotNone(repair_info["mask_bbox"]) self.assertEqual(work_image.fill_mode, "edge") + self.assertIsNone(positive) + self.assertIsNone(negative) def test_prepare_region_applies_mask_preprocessing_before_metrics(self): plugin = load_plugin_package() @@ -155,7 +161,7 @@ class TestRepairPrepare(unittest.TestCase): label="preprocess-mask", ) - _latent, _work_image, work_mask, repair_info, _recommended_denoise = node.prepare( + _latent, _work_image, work_mask, repair_info, _recommended_denoise, _positive, _negative = node.prepare( image=image, mask=mask, vae=FakeVAE(), @@ -189,6 +195,45 @@ class TestRepairPrepare(unittest.TestCase): 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_passthroughs_positive_and_negative_outputs(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]() + + image = FakeTensor((1, 512, 512, 3), label="passthrough-image") + mask = FakeMask((1, 512, 512), mask_bbox=(64, 64, 128, 128), mask_area_ratio=0.02) + positive = [["pos", {"strength": 1.0}]] + negative = [["neg", {"strength": 1.0}]] + + _latent, _work_image, _work_mask, _repair_info, _recommended_denoise, positive_out, negative_out = 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"}, + positive=positive, + negative=negative, + ) + + self.assertIs(positive_out, positive) + self.assertIs(negative_out, negative) + def test_prepare_requires_vae(self): plugin = load_plugin_package() node = plugin.NODE_CLASS_MAPPINGS["LLSSimpleRepairPrepare"]() diff --git a/tests/test_repair_registration.py b/tests/test_repair_registration.py index d547aa3..5273bb9 100644 --- a/tests/test_repair_registration.py +++ b/tests/test_repair_registration.py @@ -30,11 +30,11 @@ class TestRepairRegistration(unittest.TestCase): self.assertEqual(node_cls.FUNCTION, "prepare") self.assertEqual( node_cls.RETURN_TYPES, - ("LATENT", "IMAGE", "MASK", "LLS_REPAIR_INFO", "FLOAT"), + ("LATENT", "IMAGE", "MASK", "LLS_REPAIR_INFO", "FLOAT", "CONDITIONING", "CONDITIONING"), ) self.assertEqual( node_cls.RETURN_NAMES, - ("latent", "work_image", "work_mask", "repair_info", "recommended_denoise"), + ("latent", "work_image", "work_mask", "repair_info", "recommended_denoise", "positive", "negative"), ) required = schema["required"] diff --git a/web/js/lls_mask_draw.js b/web/js/lls_mask_draw.js new file mode 100644 index 0000000..328a187 --- /dev/null +++ b/web/js/lls_mask_draw.js @@ -0,0 +1,918 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +const EXTENSION_NAME = "LLS.MaskDraw"; +const TARGET_NODE_CLASS = "LLSSimpleMaskDraw"; +const TARGET_NODE_DISPLAY_NAME = "LLS Simple Mask Draw"; +const HISTORY_LIMIT = 20; +const MAX_PREVIEW_EDGE = 420; +const MIN_PREVIEW_EDGE = 220; + +function clamp(value, min, max) { + return Math.min(max, Math.max(min, value)); +} + +function asNumber(value, fallback) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : fallback; +} + +function asBoolean(value, fallback = false) { + return typeof value === "boolean" ? value : fallback; +} + +function findWidget(node, name) { + return node.widgets?.find((widget) => widget.name === name) ?? null; +} + +function getNodeClassName(node) { + return String(node?.comfyClass || node?.constructor?.comfyClass || node?.type || node?.title || ""); +} + +function hideWidget(widget) { + if (!widget) { + return; + } + widget.type = "hidden"; + widget.computeSize = () => [0, -4]; +} + +function chainWidgetCallback(widget, handler) { + if (!widget || widget.__llsMaskDrawChained) { + return; + } + const previous = widget.callback; + widget.callback = function chainedMaskDrawCallback() { + previous?.apply(this, arguments); + handler?.apply(this, arguments); + }; + widget.__llsMaskDrawChained = true; +} + +function parseState(rawValue) { + try { + const parsed = JSON.parse(rawValue || "{}"); + if (!parsed || typeof parsed !== "object") { + throw new Error("Invalid state payload"); + } + return { + version: asNumber(parsed.version, 1), + mask_png_base64: String(parsed.mask_png_base64 || ""), + touched: asBoolean(parsed.touched, false), + editor: parsed.editor && typeof parsed.editor === "object" ? parsed.editor : {}, + }; + } catch { + return { + version: 1, + mask_png_base64: "", + touched: false, + editor: {}, + }; + } +} + +function dataUrlToBase64(dataUrl) { + const commaIndex = String(dataUrl || "").indexOf(","); + return commaIndex >= 0 ? dataUrl.slice(commaIndex + 1) : ""; +} + +function getInputLink(node, inputName) { + const input = node.inputs?.find((item) => item.name === inputName); + if (!input || input.link == null) { + return null; + } + return node.graph?.links?.[input.link] ?? null; +} + +function getUpstreamNode(node, inputName) { + const link = getInputLink(node, inputName); + if (!link) { + return null; + } + return node.graph?.getNodeById?.(link.origin_id) ?? null; +} + +function normalizePath(rawValue) { + return String(rawValue || "").trim().replace(/\\/g, "/"); +} + +function splitAnnotatedPath(rawValue) { + const normalized = normalizePath(rawValue); + const match = normalized.match(/^(.*?)(?:\s+\[([^\]]+)\])?$/); + const path = match?.[1] ?? normalized; + const annotatedType = match?.[2] ?? ""; + const slashIndex = path.lastIndexOf("/"); + return { + path, + annotatedType, + subfolder: slashIndex >= 0 ? path.slice(0, slashIndex) : "", + filename: slashIndex >= 0 ? path.slice(slashIndex + 1) : path, + }; +} + +function buildViewUrl(rawValue, fallbackType, extraParams = {}) { + const normalized = normalizePath(rawValue); + if (!normalized) { + return null; + } + + const params = new URLSearchParams(); + if (normalized.startsWith("blake3:")) { + params.set("filename", normalized); + } else { + const parsed = splitAnnotatedPath(normalized); + if (!parsed.filename) { + return null; + } + params.set("filename", parsed.filename); + if (parsed.subfolder) { + params.set("subfolder", parsed.subfolder); + } + params.set("type", parsed.annotatedType || fallbackType || "input"); + } + + for (const [key, value] of Object.entries(extraParams)) { + if (value != null && value !== "") { + params.set(key, String(value)); + } + } + + const relativePath = `/view?${params.toString()}`; + return typeof api?.apiURL === "function" ? api.apiURL(relativePath) : relativePath; +} + +function resolveImageSource(node) { + const upstream = getUpstreamNode(node, "image"); + const className = getNodeClassName(upstream); + const imageWidget = findWidget(upstream, "image"); + if (!imageWidget?.value) { + return null; + } + + if (className === "LoadImage") { + return { + url: buildViewUrl(imageWidget.value, "input"), + label: "LoadImage", + }; + } + + if (className === "LoadImageOutput") { + return { + url: buildViewUrl(imageWidget.value, "output"), + label: "LoadImageOutput", + }; + } + + return null; +} + +function resolveMaskSource(node) { + const upstream = getUpstreamNode(node, "input_mask"); + const className = getNodeClassName(upstream); + const imageWidget = findWidget(upstream, "image"); + if (!imageWidget?.value) { + return null; + } + + if (className === "LoadImageMask") { + const channelName = String(findWidget(upstream, "channel")?.value || "alpha").toLowerCase(); + if (channelName === "alpha") { + return { + url: buildViewUrl(imageWidget.value, "input", { channel: "a" }), + channel: "a", + invert: true, + label: "LoadImageMask.alpha", + }; + } + + const channelMap = { + red: "r", + green: "g", + blue: "b", + }; + const mappedChannel = channelMap[channelName]; + if (!mappedChannel) { + return null; + } + + return { + url: buildViewUrl(imageWidget.value, "input"), + channel: mappedChannel, + invert: false, + label: `LoadImageMask.${channelName}`, + }; + } + + if (className === "LoadImage") { + return { + url: buildViewUrl(imageWidget.value, "input", { channel: "a" }), + channel: "a", + invert: true, + label: "LoadImage.alpha", + }; + } + + return null; +} + +function loadImageElement(url) { + return new Promise((resolve, reject) => { + if (!url) { + resolve(null); + return; + } + + const image = new Image(); + image.onload = () => resolve(image); + image.onerror = () => reject(new Error(`Failed to load image: ${url}`)); + image.src = url; + }); +} + +function computeDisplaySize(width, height) { + const safeWidth = Math.max(1, Math.round(width)); + const safeHeight = Math.max(1, Math.round(height)); + const maxEdge = Math.max(safeWidth, safeHeight); + let scale = 1; + + if (maxEdge > MAX_PREVIEW_EDGE) { + scale = MAX_PREVIEW_EDGE / maxEdge; + } else if (maxEdge < MIN_PREVIEW_EDGE) { + scale = MIN_PREVIEW_EDGE / maxEdge; + } + + return { + width: Math.max(1, Math.round(safeWidth * scale)), + height: Math.max(1, Math.round(safeHeight * scale)), + }; +} + +function drawHintText(context, lines) { + context.fillStyle = "#171717"; + context.fillRect(0, 0, context.canvas.width, context.canvas.height); + context.fillStyle = "#9a9a9a"; + context.font = "12px sans-serif"; + lines.forEach((line, index) => { + context.fillText(line, 14, 24 + index * 18); + }); +} + +function drawMaskOverlay(stageContext, overlayContext, overlayCanvas, maskCanvas, overlayAlpha, invertMask) { + if (!maskCanvas.width || !maskCanvas.height) { + return; + } + + overlayContext.clearRect(0, 0, overlayCanvas.width, overlayCanvas.height); + overlayContext.drawImage(maskCanvas, 0, 0, overlayCanvas.width, overlayCanvas.height); + const imageData = overlayContext.getImageData(0, 0, overlayCanvas.width, overlayCanvas.height); + const pixels = imageData.data; + + for (let index = 0; index < pixels.length; index += 4) { + let value = pixels[index]; + if (invertMask) { + value = 255 - value; + } + pixels[index] = 255; + pixels[index + 1] = 0; + pixels[index + 2] = 0; + pixels[index + 3] = Math.round((value / 255) * overlayAlpha * 255); + } + + overlayContext.putImageData(imageData, 0, 0); + stageContext.drawImage(overlayCanvas, 0, 0); +} + +function attachMaskEditor(node) { + if (node.__llsMaskDrawEditor || typeof node.addDOMWidget !== "function") { + return node.__llsMaskDrawEditor ?? null; + } + + const stateWidget = findWidget(node, "mask_state_json"); + const drawModeWidget = findWidget(node, "draw_mode"); + const brushSizeWidget = findWidget(node, "brush_size"); + const brushSoftnessWidget = findWidget(node, "brush_softness"); + const overlayAlphaWidget = findWidget(node, "overlay_alpha"); + const invertWidget = findWidget(node, "invert_mask"); + + hideWidget(stateWidget); + + const savedState = parseState(stateWidget?.value); + + const container = document.createElement("div"); + container.className = "lls-mask-draw-editor"; + container.style.display = "flex"; + container.style.flexDirection = "column"; + container.style.alignItems = "flex-start"; + container.style.gap = "8px"; + container.style.width = "100%"; + container.style.boxSizing = "border-box"; + + const toolbar = document.createElement("div"); + toolbar.style.display = "flex"; + toolbar.style.flexWrap = "wrap"; + toolbar.style.gap = "6px"; + toolbar.style.width = "100%"; + + const frame = document.createElement("div"); + frame.style.display = "inline-flex"; + frame.style.border = "1px solid #4e4e4e"; + frame.style.borderRadius = "8px"; + frame.style.overflow = "hidden"; + frame.style.background = "#111"; + + const stage = document.createElement("canvas"); + stage.style.display = "block"; + stage.style.cursor = "crosshair"; + stage.style.touchAction = "none"; + stage.style.userSelect = "none"; + frame.appendChild(stage); + + const status = document.createElement("div"); + status.style.fontSize = "12px"; + status.style.lineHeight = "1.45"; + status.style.color = "#b8b8b8"; + status.style.maxWidth = "420px"; + status.style.whiteSpace = "normal"; + + container.append(toolbar, frame, status); + + const domWidget = node.addDOMWidget("lls_mask_editor", "lls_mask_editor", container); + domWidget.serialize = false; + if (domWidget.options) { + domWidget.options.hideOnZoom = false; + } + + const stageContext = stage.getContext("2d"); + const overlayCanvas = document.createElement("canvas"); + const overlayContext = overlayCanvas.getContext("2d", { willReadFrequently: true }); + const maskCanvas = document.createElement("canvas"); + let maskContext = maskCanvas.getContext("2d", { willReadFrequently: true }); + + const history = { + undo: [], + redo: [], + }; + + const state = { + baseImage: null, + localMaskActive: Boolean(savedState.touched && savedState.mask_png_base64), + drawMode: String(savedState.editor.draw_mode || "brush"), + brushSize: clamp(asNumber(savedState.editor.brush_size, 32), 1, 512), + brushSoftness: clamp(asNumber(savedState.editor.brush_softness, 0.5), 0, 1), + overlayAlpha: clamp(asNumber(savedState.editor.overlay_alpha, 0.4), 0, 1), + drawing: false, + dirtyStroke: false, + lastPoint: null, + loadToken: 0, + }; + + let canPersist = false; + + function updateStatus(message) { + status.textContent = message; + } + + function resetHistory() { + history.undo.length = 0; + history.redo.length = 0; + } + + function fillMaskBlack() { + if (!maskCanvas.width || !maskCanvas.height || !maskContext) { + return; + } + maskContext.save(); + maskContext.globalCompositeOperation = "source-over"; + maskContext.clearRect(0, 0, maskCanvas.width, maskCanvas.height); + maskContext.fillStyle = "rgb(0, 0, 0)"; + maskContext.fillRect(0, 0, maskCanvas.width, maskCanvas.height); + maskContext.restore(); + } + + function ensureMaskCanvasSize(width, height, options = {}) { + const preserve = options.preserve !== false; + const nextWidth = Math.max(1, Math.round(width)); + const nextHeight = Math.max(1, Math.round(height)); + if (maskCanvas.width === nextWidth && maskCanvas.height === nextHeight) { + return; + } + + let previousCanvas = null; + if (preserve && maskCanvas.width > 0 && maskCanvas.height > 0) { + previousCanvas = document.createElement("canvas"); + previousCanvas.width = maskCanvas.width; + previousCanvas.height = maskCanvas.height; + previousCanvas.getContext("2d").drawImage(maskCanvas, 0, 0); + } + + maskCanvas.width = nextWidth; + maskCanvas.height = nextHeight; + maskContext = maskCanvas.getContext("2d", { willReadFrequently: true }); + fillMaskBlack(); + + if (previousCanvas) { + maskContext.drawImage(previousCanvas, 0, 0, nextWidth, nextHeight); + } + } + + function snapshotMask() { + if (!maskCanvas.width || !maskCanvas.height) { + return ""; + } + return maskCanvas.toDataURL("image/png"); + } + + function persistState() { + if (!canPersist || !stateWidget) { + return; + } + + stateWidget.value = JSON.stringify({ + version: 1, + mask_png_base64: state.localMaskActive ? dataUrlToBase64(snapshotMask()) : "", + touched: state.localMaskActive, + editor: { + draw_mode: state.drawMode, + brush_size: state.brushSize, + brush_softness: state.brushSoftness, + overlay_alpha: state.overlayAlpha, + }, + }); + + app.graph?.setDirtyCanvas(true, true); + } + + function readWidgetValues(options = {}) { + const persist = options.persist !== false; + state.drawMode = String(drawModeWidget?.value ?? state.drawMode || "brush"); + state.brushSize = clamp(asNumber(brushSizeWidget?.value, state.brushSize || 32), 1, 512); + state.brushSoftness = clamp(asNumber(brushSoftnessWidget?.value, state.brushSoftness ?? 0.5), 0, 1); + state.overlayAlpha = clamp(asNumber(overlayAlphaWidget?.value, state.overlayAlpha ?? 0.4), 0, 1); + if (persist) { + persistState(); + } + } + + function updateNodeSize() { + const sourceWidth = state.baseImage?.naturalWidth || maskCanvas.width || 320; + const sourceHeight = state.baseImage?.naturalHeight || maskCanvas.height || 240; + const displaySize = computeDisplaySize(sourceWidth, sourceHeight); + + stage.width = displaySize.width; + stage.height = displaySize.height; + stage.style.width = `${displaySize.width}px`; + stage.style.height = `${displaySize.height}px`; + overlayCanvas.width = displaySize.width; + overlayCanvas.height = displaySize.height; + frame.style.width = `${displaySize.width}px`; + container.style.width = `${displaySize.width}px`; + + if (typeof node.setSize === "function") { + requestAnimationFrame(() => { + node.setSize([ + Math.max(360, displaySize.width + 28), + Math.max(250, displaySize.height + 150), + ]); + app.graph?.setDirtyCanvas(true, true); + }); + } + } + + function renderStage() { + updateNodeSize(); + stageContext.clearRect(0, 0, stage.width, stage.height); + + if (state.baseImage) { + stageContext.drawImage(state.baseImage, 0, 0, stage.width, stage.height); + } else { + drawHintText(stageContext, state.localMaskActive + ? [ + "Source preview unavailable.", + "Connect Load Image / Load Image Output", + "to show the edited mask over the image.", + ] + : [ + "Connect Load Image / Load Image Output", + "to start drawing before execution.", + ]); + } + + drawMaskOverlay( + stageContext, + overlayContext, + overlayCanvas, + maskCanvas, + state.overlayAlpha, + asBoolean(invertWidget?.value, false), + ); + + app.graph?.setDirtyCanvas(true, true); + } + + function renderSourceMask(image, channel, invert) { + if (!maskCanvas.width || !maskCanvas.height) { + return; + } + + const scratchCanvas = document.createElement("canvas"); + scratchCanvas.width = maskCanvas.width; + scratchCanvas.height = maskCanvas.height; + const scratchContext = scratchCanvas.getContext("2d", { willReadFrequently: true }); + scratchContext.drawImage(image, 0, 0, scratchCanvas.width, scratchCanvas.height); + + const imageData = scratchContext.getImageData(0, 0, scratchCanvas.width, scratchCanvas.height); + const pixels = imageData.data; + + for (let index = 0; index < pixels.length; index += 4) { + let value = pixels[index]; + if (channel === "g") { + value = pixels[index + 1]; + } else if (channel === "b") { + value = pixels[index + 2]; + } else if (channel === "a") { + value = pixels[index + 3]; + } + if (invert) { + value = 255 - value; + } + pixels[index] = value; + pixels[index + 1] = value; + pixels[index + 2] = value; + pixels[index + 3] = 255; + } + + scratchContext.putImageData(imageData, 0, 0); + fillMaskBlack(); + maskContext.drawImage(scratchCanvas, 0, 0); + } + + async function restoreMaskFromDataUrl(dataUrl) { + const image = await loadImageElement(dataUrl); + if (!image) { + return false; + } + + ensureMaskCanvasSize(image.naturalWidth || image.width, image.naturalHeight || image.height, { preserve: false }); + fillMaskBlack(); + maskContext.drawImage(image, 0, 0, maskCanvas.width, maskCanvas.height); + return true; + } + + async function restoreSerializedLocalMask(maskBase64) { + if (!maskBase64) { + return false; + } + + try { + const restored = await restoreMaskFromDataUrl(`data:image/png;base64,${maskBase64}`); + if (restored) { + resetHistory(); + } + return restored; + } catch { + return false; + } + } + + function pushUndoSnapshot() { + const snapshot = snapshotMask(); + if (!snapshot) { + return; + } + history.undo.push(snapshot); + if (history.undo.length > HISTORY_LIMIT) { + history.undo.shift(); + } + history.redo.length = 0; + } + + async function restoreHistorySnapshot(snapshot) { + if (!snapshot) { + return false; + } + const restored = await restoreMaskFromDataUrl(snapshot); + renderStage(); + return restored; + } + + function getMaskPoint(event) { + if (!maskCanvas.width || !maskCanvas.height) { + return null; + } + + const bounds = stage.getBoundingClientRect(); + if (!bounds.width || !bounds.height) { + return null; + } + + return { + x: clamp((event.clientX - bounds.left) * (maskCanvas.width / bounds.width), 0, maskCanvas.width), + y: clamp((event.clientY - bounds.top) * (maskCanvas.height / bounds.height), 0, maskCanvas.height), + }; + } + + function applyBrushStamp(x, y) { + const radius = Math.max(1, state.brushSize) / 2; + const edge = clamp(1 - state.brushSoftness, 0, 1); + const channelValue = state.drawMode === "erase" ? 0 : 255; + const gradient = maskContext.createRadialGradient(x, y, 0, x, y, radius); + gradient.addColorStop(0, `rgba(${channelValue}, ${channelValue}, ${channelValue}, 1)`); + gradient.addColorStop(edge, `rgba(${channelValue}, ${channelValue}, ${channelValue}, 1)`); + gradient.addColorStop(1, `rgba(${channelValue}, ${channelValue}, ${channelValue}, 0)`); + + maskContext.save(); + maskContext.globalCompositeOperation = "source-over"; + maskContext.fillStyle = gradient; + maskContext.beginPath(); + maskContext.arc(x, y, radius, 0, Math.PI * 2); + maskContext.fill(); + maskContext.restore(); + } + + function applyStroke(fromPoint, toPoint) { + if (!toPoint) { + return; + } + + if (!fromPoint) { + applyBrushStamp(toPoint.x, toPoint.y); + return; + } + + const deltaX = toPoint.x - fromPoint.x; + const deltaY = toPoint.y - fromPoint.y; + const distance = Math.hypot(deltaX, deltaY); + const spacing = Math.max(1, state.brushSize * 0.2); + const steps = Math.max(1, Math.ceil(distance / spacing)); + + for (let step = 1; step <= steps; step += 1) { + const progress = step / steps; + applyBrushStamp( + fromPoint.x + deltaX * progress, + fromPoint.y + deltaY * progress, + ); + } + } + + async function reloadFromInputs() { + const requestToken = ++state.loadToken; + const imageSource = resolveImageSource(node); + const maskSource = state.localMaskActive ? null : resolveMaskSource(node); + + const baseImagePromise = imageSource?.url + ? loadImageElement(imageSource.url).catch(() => null) + : Promise.resolve(null); + const maskImagePromise = maskSource?.url + ? loadImageElement(maskSource.url).catch(() => null) + : Promise.resolve(null); + + const [baseImage, maskImage] = await Promise.all([baseImagePromise, maskImagePromise]); + if (requestToken !== state.loadToken) { + return; + } + + state.baseImage = baseImage; + + const targetWidth = baseImage?.naturalWidth || maskCanvas.width || maskImage?.naturalWidth || maskImage?.width || 0; + const targetHeight = baseImage?.naturalHeight || maskCanvas.height || maskImage?.naturalHeight || maskImage?.height || 0; + + if (targetWidth && targetHeight) { + if (state.localMaskActive) { + ensureMaskCanvasSize(targetWidth, targetHeight, { preserve: true }); + } else { + resetHistory(); + ensureMaskCanvasSize(targetWidth, targetHeight, { preserve: false }); + if (maskImage) { + renderSourceMask(maskImage, maskSource.channel, maskSource.invert); + } + } + } + + if (!imageSource) { + updateStatus(state.localMaskActive + ? "Local mask restored. Connect Load Image or Load Image Output for source preview." + : "Connect Load Image or Load Image Output to paint before execution."); + } else if (!baseImage) { + updateStatus("Source preview could not be loaded from the upstream file-backed node."); + } else if (state.localMaskActive) { + updateStatus("Editing the local mask raster. input_mask is ignored until this local mask is replaced."); + } else if (maskSource && maskImage) { + updateStatus("input_mask loaded as the editable base layer."); + } else if (maskSource && !maskImage) { + updateStatus("Source image loaded. input_mask preview is unavailable for this upstream node."); + } else { + updateStatus("Source image loaded. Drawing starts from an empty mask."); + } + + renderStage(); + } + + function createButton(label, handler) { + const button = document.createElement("button"); + button.type = "button"; + button.textContent = label; + button.style.border = "1px solid #5d5d5d"; + button.style.borderRadius = "6px"; + button.style.background = "#202020"; + button.style.color = "#f0f0f0"; + button.style.padding = "4px 10px"; + button.style.cursor = "pointer"; + button.addEventListener("click", (event) => { + event.preventDefault(); + event.stopPropagation(); + handler(event); + }); + return button; + } + + toolbar.append( + createButton("Clear", () => { + if (!maskCanvas.width || !maskCanvas.height) { + return; + } + pushUndoSnapshot(); + fillMaskBlack(); + state.localMaskActive = true; + persistState(); + updateStatus("Local mask cleared to black."); + renderStage(); + }), + createButton("Undo", async () => { + if (!history.undo.length) { + return; + } + const previous = history.undo.pop(); + const current = snapshotMask(); + if (current) { + history.redo.push(current); + } + if (await restoreHistorySnapshot(previous)) { + state.localMaskActive = true; + persistState(); + updateStatus("Restored the previous local mask state."); + } + }), + createButton("Redo", async () => { + if (!history.redo.length) { + return; + } + const next = history.redo.pop(); + const current = snapshotMask(); + if (current) { + history.undo.push(current); + } + if (await restoreHistorySnapshot(next)) { + state.localMaskActive = true; + persistState(); + updateStatus("Reapplied the next local mask state."); + } + }), + createButton("Invert", () => { + if (!invertWidget) { + return; + } + invertWidget.value = !Boolean(invertWidget.value); + invertWidget.callback?.(invertWidget.value); + renderStage(); + }), + ); + + stage.addEventListener("pointerdown", (event) => { + event.preventDefault(); + event.stopPropagation(); + + const point = getMaskPoint(event); + if (!point) { + updateStatus("Source image is not ready. Connect Load Image or Load Image Output first."); + renderStage(); + return; + } + + stage.setPointerCapture?.(event.pointerId); + pushUndoSnapshot(); + state.drawing = true; + state.dirtyStroke = true; + state.localMaskActive = true; + applyStroke(null, point); + state.lastPoint = point; + renderStage(); + }); + + stage.addEventListener("pointermove", (event) => { + if (!state.drawing) { + return; + } + event.preventDefault(); + event.stopPropagation(); + + const point = getMaskPoint(event); + if (!point) { + return; + } + + applyStroke(state.lastPoint, point); + state.lastPoint = point; + renderStage(); + }); + + function endStroke(event) { + if (!state.drawing) { + return; + } + event?.preventDefault?.(); + event?.stopPropagation?.(); + + state.drawing = false; + state.lastPoint = null; + if (state.dirtyStroke) { + persistState(); + state.dirtyStroke = false; + updateStatus("Local mask updated."); + } + renderStage(); + } + + stage.addEventListener("pointerup", endStroke); + stage.addEventListener("pointercancel", endStroke); + stage.addEventListener("lostpointercapture", endStroke); + + chainWidgetCallback(drawModeWidget, () => { + readWidgetValues({ persist: true }); + renderStage(); + }); + chainWidgetCallback(brushSizeWidget, () => { + readWidgetValues({ persist: true }); + renderStage(); + }); + chainWidgetCallback(brushSoftnessWidget, () => { + readWidgetValues({ persist: true }); + renderStage(); + }); + chainWidgetCallback(overlayAlphaWidget, () => { + readWidgetValues({ persist: true }); + renderStage(); + }); + chainWidgetCallback(invertWidget, () => { + renderStage(); + }); + + async function initializeFromState(syncSerializedState) { + const parsedState = syncSerializedState ? parseState(stateWidget?.value) : savedState; + readWidgetValues({ persist: false }); + + if (parsedState.touched && parsedState.mask_png_base64) { + state.localMaskActive = await restoreSerializedLocalMask(parsedState.mask_png_base64); + } else if (!state.localMaskActive) { + state.localMaskActive = false; + } + + await reloadFromInputs(); + canPersist = true; + renderStage(); + } + + const editorApi = { + handleGraphConfigured() { + void initializeFromState(true); + }, + handleConnectionsChange() { + void reloadFromInputs(); + }, + }; + + node.__llsMaskDrawEditor = editorApi; + updateStatus("Connect Load Image or Load Image Output to start drawing before execution."); + renderStage(); + void initializeFromState(false); + return editorApi; +} + +app.registerExtension({ + name: EXTENSION_NAME, + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== TARGET_NODE_CLASS && nodeData.display_name !== TARGET_NODE_DISPLAY_NAME) { + return; + } + + const previousOnNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function onMaskDrawNodeCreated() { + const result = previousOnNodeCreated?.apply(this, arguments); + attachMaskEditor(this); + return result; + }; + + const previousOnGraphConfigured = nodeType.prototype.onGraphConfigured; + nodeType.prototype.onGraphConfigured = function onMaskDrawGraphConfigured() { + const result = previousOnGraphConfigured?.apply(this, arguments); + attachMaskEditor(this)?.handleGraphConfigured?.(); + return result; + }; + + const previousOnConnectionsChange = nodeType.prototype.onConnectionsChange; + nodeType.prototype.onConnectionsChange = function onMaskDrawConnectionsChange() { + const result = previousOnConnectionsChange?.apply(this, arguments); + attachMaskEditor(this)?.handleConnectionsChange?.(); + return result; + }; + }, +});