From 1ae1f8bf8a7969aea658246d73036be89189c6c8 Mon Sep 17 00:00:00 2001 From: fengwk <759543714@qq.com> Date: Sat, 4 Jul 2026 13:38:50 +0800 Subject: [PATCH] first commit --- .gitignore | 3 + README.md | 288 ++++++ __init__.py | 8 + ...al Remove - Interactive Region Editor.json | 723 +++++++++++++++ nodes.py | 843 ++++++++++++++++++ pytest.ini | 2 + requirements-dev.txt | 1 + tests/test_universal_remove.py | 244 +++++ web/region_editor.js | 535 +++++++++++ 9 files changed, 2647 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 example_workflows/Universal Remove - Interactive Region Editor.json create mode 100644 nodes.py create mode 100644 pytest.ini create mode 100644 requirements-dev.txt create mode 100644 tests/test_universal_remove.py create mode 100644 web/region_editor.js diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a160f49 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..2eebab9 --- /dev/null +++ b/README.md @@ -0,0 +1,288 @@ +# ComfyUI-UniversalRemove + +一个自包含的 ComfyUI remove 核心扩展。 + +## 设计目标 + +- 不硬编码任何本地路径 +- 不依赖其他自定义节点 +- 只使用 ComfyUI 常见核心类型：`IMAGE` / `MASK` / `STRING` / `INT` / `FLOAT` +- 面向 **图像批 / 视频帧批** 工作，适合作为通用 remove 流程的中间层 + +## 当前提供的节点 + +### 0. `UR Demo Image` + +生成一个无需外部图片文件的测试图像，主要用于示例工作流和安装验证。 + +### 1. `UR Remove Prompt Builder` + +输入目标描述，输出三路提示词： + +- `segmentation_prompt` +- `remove_positive_prompt` +- `remove_negative_prompt` + +用途： + +- `segmentation_prompt` 给任意文本选目标模块 +- `remove_positive_prompt` 给关键帧修复模型 +- `remove_negative_prompt` 给关键帧修复模型的负向提示 + +### 2. `UR Region Editor` + +交互式前端 region 编辑器，输出标准 `regions_json`。 + +支持： + +- 加载本地图片作为编辑背景 +- 加载本地视频作为编辑背景并拖动时间轴 +- 使用当前视频时间写入关键帧编号 +- 绘制矩形区域 +- 绘制圆形区域 +- 绘制多边形区域 +- 删除选中区域 +- 清空全部区域 + +注意：背景图片 / 视频只用于前端交互预览，不会写进 workflow。真正可移植、可复用的是输出的 `regions_json`。 + +### 3. `UR Boxes To Mask Batch` + +把 JSON 形式的矩形框转换成 `MASK` 批次。 + +支持： + +- `xyxy` +- `xywh` +- 归一化坐标 +- 三种广播模式： + - `repeat_to_batch` + - `single_frame_only` + - `hold_after_frame` + +用途： + +- 手工框选移除区域 +- 从外部工具导入 bbox + +### 4. `UR Regions To Mask Batch` + +把关键帧 region JSON 转换成整段图像批 / 视频帧批的 `MASK`。 + +支持： + +- `rect` +- `circle` +- `polygon` +- 单关键帧 +- 多关键帧 +- 四种广播模式： + - `repeat_to_batch` + - `single_frame_only` + - `hold_after_frame` + - `hold_until_next_keyframe` + +用途： + +- 用 JSON 描述“关键帧上的圈选区域” +- 为视频帧批生成对应 mask +- 后续接 `UR Mask Processor` / `UR Keyframe Prep` / `UR Soft Composite` + +### 5. `UR Mask Processor` + +对 mask 做 remove 常用处理： + +- 扩张 +- 高斯羽化 +- 预览叠加 + +### 6. `UR Keyframe Prep` + +从图像批中选出关键帧，并准备 remove 输入： + +- 选关键帧 +- 对应 mask 对齐 +- 按最长边缩放 +- pad 到指定倍数 +- 用填充色生成 `filled_frame` + +其中 `filled_frame` 可以直接送进任意关键帧修复 / inpaint / edit 模型。 + +### 7. `UR Soft Composite` + +一个不依赖 BrushNet 的软融合节点，用于把修复结果融合回原图 / 原视频帧批。 + +## 推荐工作流 + +> 这个扩展刻意只负责 remove 核心流程，不强绑某个检测器、分割器、视频模型或视频 IO 节点。 + +### A. 手工框选 remove + +```text +Frames/Image Batch + -> UR Region Editor + -> UR Regions To Mask Batch + -> UR Mask Processor + -> UR Keyframe Prep + -> 关键帧修复模型 + -> (可选) 视频传播模型 + -> UR Soft Composite + -> Output +``` + +### B. 提示词驱动 remove + +```text +Frames/Image Batch + -> UR Remove Prompt Builder + -> 任意文本选目标模块（输出 MASK） + -> UR Mask Processor + -> UR Keyframe Prep + -> 关键帧修复模型 + -> (可选) 视频传播模型 + -> UR Soft Composite + -> Output +``` + +## 为什么没有内置“文本直接找目标” + +这是刻意的设计边界： + +- remove 核心流程应该通用、可复用 +- 文本找目标属于检测 / 分割 / 跟踪能力 +- 这部分通常意味着额外模型、额外权重、额外依赖 + +因此当前版本先把 remove 主流程做成稳定、自包含、无外部节点硬依赖的核心层。 + +如果后续要继续扩展，最自然的下一步是在这个扩展内部新增： + +- `UR Prompt Mask Locator` +- `UR Video Track Mask` +- `UR Model Bundle` + +但这些应当作为后续迭代，而不是在 V1 就把依赖链做死。 + +## 安装方式 + +把整个目录放进： + +- `ComfyUI/custom_nodes/ComfyUI-UniversalRemove` + +然后重启 ComfyUI。 + +## 示例工作流 + +当前已提供一个可直接导入的最小工作流： + +- `example_workflows/Universal Remove - Interactive Region Editor.json` + +这个 workflow 的特点： + +- 只依赖 ComfyUI core + 本扩展 +- 可直接导入并运行，不需要额外图片文件 +- 可用 `UR Region Editor` 交互绘制 `regions_json` +- 也可以手工编辑 `regions_json` +- 可预览： + - mask preview + - filled keyframe + - soft composite result + +这个示例是 **remove 核心流程示例**，不内置外部检测模型或修复模型，因此更适合作为： + +- remove 管线起点 +- 接入你自己的 inpaint / edit / video propagation 模型前的可运行 scaffold + +## 测试 + +仓库内已提供基础单元测试： + +- `tests/test_universal_remove.py` + +运行方式： + +```bash +uv pip install --python /path/to/ComfyUI/.venv/bin/python -r requirements-dev.txt +/path/to/ComfyUI/.venv/bin/python -m pytest tests/test_universal_remove.py +``` + +## 视频输入与关键帧圈选 + +当前扩展不内置视频读取器，但已经内置 `UR Region Editor` 前端交互式圈选 UI。 + +推荐接法是： + +```text +任意视频读取节点 / 你的脚本 + -> IMAGE frame batch + -> UR Region Editor 输出 regions_json + -> UR Regions To Mask Batch +``` + +也就是说，本扩展支持的是： + +- 输入视频拆帧后的 `IMAGE` 批次 +- 通过 `frame_index` 指定关键帧 +- 通过 `UR Region Editor` 鼠标绘制或手写 `regions_json` 描述关键帧上的圈选区域 +- 输出整段帧批可用的 `MASK` + +`UR Region Editor` 的本地视频背景只用于交互定位，不替代实际视频读取节点；真正进入后端计算的是上游 `IMAGE` frame batch 和编辑器输出的 `regions_json`。 + +## region JSON 示例 + +### 单关键帧矩形 + 圆形 + +```json +[ + {"type": "rect", "frame_index": 0, "x": 160, "y": 120, "w": 260, "h": 140}, + {"type": "circle", "frame_index": 0, "cx": 520, "cy": 250, "r": 70} +] +``` + +### 多关键帧多边形 + +```json +[ + { + "type": "polygon", + "frame_index": 0, + "points": [[120, 160], [240, 150], [280, 260], [150, 290]] + }, + { + "type": "polygon", + "frame_index": 48, + "points": [[180, 170], [310, 160], [350, 280], [210, 310]] + } +] +``` + +配合 `broadcast_mode=hold_until_next_keyframe` 时，第一个关键帧 mask 会使用到下一个关键帧之前，第二个关键帧 mask 会使用到结尾。 + +## bbox JSON 示例 + +### `xyxy` + +```json +[ + [100, 120, 420, 220], + [520, 50, 700, 160] +] +``` + +### `xywh` + +```json +[ + {"x": 100, "y": 120, "w": 320, "h": 100}, + {"x": 520, "y": 50, "w": 180, "h": 110} +] +``` + +## 备注 + +这个扩展当前最适合做： + +- 通用 remove 流程骨架 +- 去水印 / 去字幕 / 去 logo +- 有上游 mask 时的角色移除 / 物体移除 + +如果你下一步要做“文本直接选目标并跟踪整段视频”，建议继续在这个目录里加新节点，而不是让核心 remove 节点反向依赖别的 custom nodes。 diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..78046ee --- /dev/null +++ b/__init__.py @@ -0,0 +1,8 @@ +try: + from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +except ImportError: + from nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +WEB_DIRECTORY = "./web" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/example_workflows/Universal Remove - Interactive Region Editor.json b/example_workflows/Universal Remove - Interactive Region Editor.json new file mode 100644 index 0000000..8b12d1d --- /dev/null +++ b/example_workflows/Universal Remove - Interactive Region Editor.json @@ -0,0 +1,723 @@ +{ + "revision": 0, + "last_node_id": 9, + "last_link_id": 12, + "nodes": [ + { + "id": 1, + "type": "UR Demo Image", + "pos": [ + -980, + 260 + ], + "size": [ + 400, + 480 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [ + { + "name": "width", + "type": "INT", + "widget": { + "name": "width" + }, + "link": null + }, + { + "name": "height", + "type": "INT", + "widget": { + "name": "height" + }, + "link": null + }, + { + "name": "batch_size", + "type": "INT", + "widget": { + "name": "batch_size" + }, + "link": null + } + ], + "outputs": [ + { + "name": "image", + "type": "IMAGE", + "links": [ + 1, + 4 + ] + } + ], + "properties": { + "Node name for S&R": "UR Demo Image" + }, + "widgets_values": [ + 768, + 512, + 1 + ] + }, + { + "id": 2, + "type": "UR Regions To Mask Batch", + "pos": [ + -460, + 230 + ], + "size": [ + 360, + 320 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 1 + }, + { + "name": "regions_json", + "type": "STRING", + "widget": { + "name": "regions_json" + }, + "link": 12 + }, + { + "name": "default_frame_index", + "type": "INT", + "widget": { + "name": "default_frame_index" + }, + "link": null + }, + { + "name": "broadcast_mode", + "type": "COMBO", + "widget": { + "name": "broadcast_mode" + }, + "link": null + }, + { + "name": "normalized", + "type": "BOOLEAN", + "widget": { + "name": "normalized" + }, + "link": null + }, + { + "name": "preview_frame", + "type": "INT", + "widget": { + "name": "preview_frame" + }, + "link": null + }, + { + "name": "preview_color", + "type": "STRING", + "widget": { + "name": "preview_color" + }, + "link": null + }, + { + "name": "preview_opacity", + "type": "FLOAT", + "widget": { + "name": "preview_opacity" + }, + "link": null + } + ], + "outputs": [ + { + "name": "mask", + "type": "MASK", + "links": [ + 2 + ] + }, + { + "name": "selected_frame", + "type": "IMAGE", + "links": [ + 3 + ] + }, + { + "name": "preview", + "type": "IMAGE", + "links": null + } + ], + "properties": { + "Node name for S&R": "UR Regions To Mask Batch" + }, + "widgets_values": [ + "[{\"type\":\"rect\",\"frame_index\":0,\"x\":160,\"y\":120,\"w\":260,\"h\":140},{\"type\":\"circle\",\"frame_index\":0,\"cx\":520,\"cy\":250,\"r\":70}]", + 0, + "repeat_to_batch", + false, + 0, + "#FF3B30", + 0.35 + ] + }, + { + "id": 3, + "type": "UR Mask Processor", + "pos": [ + -10, + 230 + ], + "size": [ + 330, + 260 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 2 + }, + { + "name": "expand_pixels", + "type": "INT", + "widget": { + "name": "expand_pixels" + }, + "link": null + }, + { + "name": "blur_sigma", + "type": "FLOAT", + "widget": { + "name": "blur_sigma" + }, + "link": null + }, + { + "name": "blur_kernel", + "type": "INT", + "widget": { + "name": "blur_kernel" + }, + "link": null + }, + { + "name": "preview_color", + "type": "STRING", + "widget": { + "name": "preview_color" + }, + "link": null + }, + { + "name": "preview_opacity", + "type": "FLOAT", + "widget": { + "name": "preview_opacity" + }, + "link": null + }, + { + "name": "image", + "type": "IMAGE", + "link": 3 + } + ], + "outputs": [ + { + "name": "mask", + "type": "MASK", + "links": [ + 5, + 7 + ] + }, + { + "name": "mask_image", + "type": "IMAGE", + "links": null + }, + { + "name": "preview", + "type": "IMAGE", + "links": [ + 11 + ] + } + ], + "properties": { + "Node name for S&R": "UR Mask Processor" + }, + "widgets_values": [ + 16, + 4.0, + 0, + "#00FF00", + 0.35 + ] + }, + { + "id": 4, + "type": "UR Keyframe Prep", + "pos": [ + 400, + 230 + ], + "size": [ + 340, + 250 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 4 + }, + { + "name": "mask", + "type": "MASK", + "link": 5 + }, + { + "name": "frame_index", + "type": "INT", + "widget": { + "name": "frame_index" + }, + "link": null + }, + { + "name": "fill_color", + "type": "STRING", + "widget": { + "name": "fill_color" + }, + "link": null + }, + { + "name": "resize_longest_side", + "type": "INT", + "widget": { + "name": "resize_longest_side" + }, + "link": null + }, + { + "name": "pad_to_multiple", + "type": "INT", + "widget": { + "name": "pad_to_multiple" + }, + "link": null + } + ], + "outputs": [ + { + "name": "original_frame", + "type": "IMAGE", + "links": [ + 6 + ] + }, + { + "name": "prepared_frame", + "type": "IMAGE", + "links": null + }, + { + "name": "prepared_mask", + "type": "MASK", + "links": null + }, + { + "name": "filled_frame", + "type": "IMAGE", + "links": [ + 8, + 9 + ] + }, + { + "name": "prepared_width", + "type": "INT", + "links": null + }, + { + "name": "prepared_height", + "type": "INT", + "links": null + }, + { + "name": "frame_count", + "type": "INT", + "links": null + } + ], + "properties": { + "Node name for S&R": "UR Keyframe Prep" + }, + "widgets_values": [ + 0, + "#00FF00", + 0, + 16 + ] + }, + { + "id": 5, + "type": "UR Soft Composite", + "pos": [ + 840, + 230 + ], + "size": [ + 300, + 160 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "inpaint", + "type": "IMAGE", + "link": 8 + }, + { + "name": "original", + "type": "IMAGE", + "link": 6 + }, + { + "name": "mask", + "type": "MASK", + "link": 7 + }, + { + "name": "kernel", + "type": "INT", + "widget": { + "name": "kernel" + }, + "link": null + }, + { + "name": "sigma", + "type": "FLOAT", + "widget": { + "name": "sigma" + }, + "link": null + } + ], + "outputs": [ + { + "name": "image", + "type": "IMAGE", + "links": [ + 10 + ] + }, + { + "name": "soft_mask", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "UR Soft Composite" + }, + "widgets_values": [ + 25, + 10.0 + ] + }, + { + "id": 6, + "type": "PreviewImage", + "pos": [ + 300, + 620 + ], + "size": [ + 320, + 260 + ], + "flags": {}, + "order": 6, + "mode": 4, + "inputs": [ + { + "localized_name": "images", + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 7, + "type": "PreviewImage", + "pos": [ + 700, + 620 + ], + "size": [ + 320, + 260 + ], + "flags": {}, + "order": 7, + "mode": 4, + "inputs": [ + { + "localized_name": "images", + "name": "images", + "type": "IMAGE", + "link": 10 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 8, + "type": "PreviewImage", + "pos": [ + -40, + 620 + ], + "size": [ + 320, + 260 + ], + "flags": {}, + "order": 8, + "mode": 4, + "inputs": [ + { + "localized_name": "images", + "name": "images", + "type": "IMAGE", + "link": 11 + } + ], + "outputs": [], + "properties": { + "cnr_id": "comfy-core", + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 9, + "type": "UR Region Editor", + "pos": [ + -900, + -250 + ], + "size": [ + 680, + 680 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "regions_json", + "type": "STRING", + "widget": { + "name": "regions_json" + }, + "link": null + }, + { + "name": "canvas_width", + "type": "INT", + "widget": { + "name": "canvas_width" + }, + "link": null + }, + { + "name": "canvas_height", + "type": "INT", + "widget": { + "name": "canvas_height" + }, + "link": null + }, + { + "name": "default_frame_index", + "type": "INT", + "widget": { + "name": "default_frame_index" + }, + "link": null + } + ], + "outputs": [ + { + "name": "regions_json", + "type": "STRING", + "links": [ + 12 + ] + }, + { + "name": "canvas_width", + "type": "INT", + "links": null + }, + { + "name": "canvas_height", + "type": "INT", + "links": null + }, + { + "name": "region_count", + "type": "INT", + "links": null + } + ], + "properties": { + "Node name for S&R": "UR Region Editor" + }, + "widgets_values": [ + "[{\"type\":\"rect\",\"frame_index\":0,\"x\":160,\"y\":120,\"w\":260,\"h\":140},{\"type\":\"circle\",\"frame_index\":0,\"cx\":520,\"cy\":250,\"r\":70}]", + 768, + 512, + 0 + ] + } + ], + "links": [ + [ + 1, + 1, + 0, + 2, + 0, + "IMAGE" + ], + [ + 2, + 2, + 0, + 3, + 0, + "MASK" + ], + [ + 3, + 2, + 1, + 3, + 6, + "IMAGE" + ], + [ + 4, + 1, + 0, + 4, + 0, + "IMAGE" + ], + [ + 5, + 3, + 0, + 4, + 1, + "MASK" + ], + [ + 6, + 4, + 0, + 5, + 1, + "IMAGE" + ], + [ + 7, + 3, + 0, + 5, + 2, + "MASK" + ], + [ + 8, + 4, + 3, + 5, + 0, + "IMAGE" + ], + [ + 9, + 4, + 3, + 6, + 0, + "IMAGE" + ], + [ + 10, + 5, + 0, + 7, + 0, + "IMAGE" + ], + [ + 11, + 3, + 2, + 8, + 0, + "IMAGE" + ], + [ + 12, + 9, + 0, + 2, + 1, + "STRING" + ] + ], + "version": 0.4 +} diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..80748b2 --- /dev/null +++ b/nodes.py @@ -0,0 +1,843 @@ +import json +import math +from typing import List, Sequence, Tuple + +import torch +import torch.nn.functional as F + + +CATEGORY = "Universal Remove" + + +def _ensure_image_batch(images: torch.Tensor) -> torch.Tensor: + if not isinstance(images, torch.Tensor): + raise TypeError("Expected IMAGE input as torch.Tensor") + if images.ndim == 3: + images = images.unsqueeze(0) + if images.ndim != 4: + raise ValueError(f"Expected IMAGE tensor with 4 dims, got {images.ndim}") + return images.clamp(0.0, 1.0) + + +def _ensure_mask_batch(mask: torch.Tensor) -> torch.Tensor: + if not isinstance(mask, torch.Tensor): + raise TypeError("Expected MASK input as torch.Tensor") + if mask.ndim == 2: + mask = mask.unsqueeze(0) + elif mask.ndim == 4 and mask.shape[1] == 1: + mask = mask[:, 0] + elif mask.ndim == 4 and mask.shape[-1] == 1: + mask = mask[..., 0] + if mask.ndim != 3: + raise ValueError(f"Expected MASK tensor with 3 dims, got {mask.ndim}") + return mask.float().clamp(0.0, 1.0) + + +def _tile_to_batch(tensor: torch.Tensor, target_batch: int) -> torch.Tensor: + if tensor.shape[0] == target_batch: + return tensor + if tensor.shape[0] <= 0: + raise ValueError("Cannot tile an empty batch") + repeats = math.ceil(target_batch / tensor.shape[0]) + repeat_shape = [repeats] + [1] * (tensor.ndim - 1) + return tensor.repeat(*repeat_shape)[:target_batch] + + +def _match_mask_to_images(mask: torch.Tensor, images: torch.Tensor) -> torch.Tensor: + mask = _ensure_mask_batch(mask) + images = _ensure_image_batch(images) + if mask.shape[0] != images.shape[0]: + mask = _tile_to_batch(mask, images.shape[0]) + if mask.shape[1] != images.shape[1] or mask.shape[2] != images.shape[2]: + mask = F.interpolate( + mask.unsqueeze(1), + size=(images.shape[1], images.shape[2]), + mode="nearest", + ).squeeze(1) + return mask.clamp(0.0, 1.0) + + +def _to_bchw(images: torch.Tensor) -> torch.Tensor: + return images.permute(0, 3, 1, 2) + + +def _to_bhwc(images: torch.Tensor) -> torch.Tensor: + return images.permute(0, 2, 3, 1) + + +def _resize_images(images: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor: + bchw = _to_bchw(images) + resized = F.interpolate(bchw, size=size, mode="bilinear", align_corners=False) + return _to_bhwc(resized).clamp(0.0, 1.0) + + +def _resize_masks(mask: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor: + resized = F.interpolate(mask.unsqueeze(1), size=size, mode="nearest") + return resized.squeeze(1).clamp(0.0, 1.0) + + +def _pad_images(images: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor: + pad_h = max(0, target_h - images.shape[1]) + pad_w = max(0, target_w - images.shape[2]) + if pad_h == 0 and pad_w == 0: + return images + bchw = _to_bchw(images) + padded = F.pad(bchw, (0, pad_w, 0, pad_h), mode="constant", value=0.0) + return _to_bhwc(padded).clamp(0.0, 1.0) + + +def _pad_masks(mask: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor: + pad_h = max(0, target_h - mask.shape[1]) + pad_w = max(0, target_w - mask.shape[2]) + if pad_h == 0 and pad_w == 0: + return mask + padded = F.pad(mask.unsqueeze(1), (0, pad_w, 0, pad_h), mode="constant", value=0.0) + return padded.squeeze(1).clamp(0.0, 1.0) + + +def _parse_hex_color(color: str) -> torch.Tensor: + text = color.strip() + if text.startswith("#"): + text = text[1:] + if len(text) != 6: + raise ValueError("Color must be a 6-digit hex string, e.g. #00FF00") + rgb = [int(text[i : i + 2], 16) / 255.0 for i in range(0, 6, 2)] + return torch.tensor(rgb, dtype=torch.float32) + + +def _overlay_mask(images: torch.Tensor, mask: torch.Tensor, color: torch.Tensor, opacity: float) -> torch.Tensor: + images = _ensure_image_batch(images) + mask = _match_mask_to_images(mask, images) + color = color.to(device=images.device, dtype=images.dtype).view(1, 1, 1, 3) + alpha = mask.unsqueeze(-1).clamp(0.0, 1.0) * float(opacity) + return (images * (1.0 - alpha) + color * alpha).clamp(0.0, 1.0) + + +def _gaussian_kernel_1d(kernel_size: int, sigma: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if kernel_size % 2 == 0: + kernel_size += 1 + radius = kernel_size // 2 + if sigma <= 0: + sigma = max(float(kernel_size) / 6.0, 1e-6) + coords = torch.arange(-radius, radius + 1, device=device, dtype=dtype) + kernel = torch.exp(-0.5 * (coords / sigma) ** 2) + kernel = kernel / kernel.sum() + return kernel + + +def _gaussian_blur_mask(mask: torch.Tensor, kernel_size: int, sigma: float) -> torch.Tensor: + mask = _ensure_mask_batch(mask) + if kernel_size <= 1 and sigma <= 0: + return mask + if kernel_size <= 1: + kernel_size = int(max(3, math.ceil(sigma * 6))) + if kernel_size % 2 == 0: + kernel_size += 1 + + kernel_1d = _gaussian_kernel_1d(kernel_size, sigma, mask.device, mask.dtype) + kernel_x = kernel_1d.view(1, 1, 1, -1) + kernel_y = kernel_1d.view(1, 1, -1, 1) + + blurred = mask.unsqueeze(1) + pad = kernel_size // 2 + blurred = F.pad(blurred, (pad, pad, 0, 0), mode="replicate") + blurred = F.conv2d(blurred, kernel_x) + blurred = F.pad(blurred, (0, 0, pad, pad), mode="replicate") + blurred = F.conv2d(blurred, kernel_y) + return blurred.squeeze(1).clamp(0.0, 1.0) + + +def _expand_mask(mask: torch.Tensor, expand_pixels: int) -> torch.Tensor: + mask = _ensure_mask_batch(mask) + if expand_pixels <= 0: + return mask + kernel = expand_pixels * 2 + 1 + expanded = F.max_pool2d(mask.unsqueeze(1), kernel_size=kernel, stride=1, padding=expand_pixels) + return expanded.squeeze(1).clamp(0.0, 1.0) + + +def _fill_masked_region(images: torch.Tensor, mask: torch.Tensor, color: torch.Tensor) -> torch.Tensor: + images = _ensure_image_batch(images) + mask = _match_mask_to_images(mask, images) + color = color.to(device=images.device, dtype=images.dtype).view(1, 1, 1, 3) + return (images * (1.0 - mask.unsqueeze(-1)) + color * mask.unsqueeze(-1)).clamp(0.0, 1.0) + + +def _normalize_box_values(values: Sequence[float], width: int, height: int, normalized: bool) -> Tuple[float, float, float, float]: + x1, y1, x2, y2 = values + if normalized: + x1 *= width + x2 *= width + y1 *= height + y2 *= height + return x1, y1, x2, y2 + + +def _parse_single_box(item, coordinate_format: str, width: int, height: int, normalized: bool) -> Tuple[int, int, int, int]: + if isinstance(item, dict): + if {"x", "y", "w", "h"}.issubset(item): + x1 = float(item["x"]) + y1 = float(item["y"]) + w = float(item["w"]) + h = float(item["h"]) + if normalized: + x1 *= width + y1 *= height + w *= width + h *= height + x2 = x1 + w + y2 = y1 + h + elif {"x1", "y1", "x2", "y2"}.issubset(item): + x1, y1, x2, y2 = _normalize_box_values( + [float(item["x1"]), float(item["y1"]), float(item["x2"]), float(item["y2"])], + width, + height, + normalized, + ) + elif {"left", "top", "right", "bottom"}.issubset(item): + x1, y1, x2, y2 = _normalize_box_values( + [float(item["left"]), float(item["top"]), float(item["right"]), float(item["bottom"])], + width, + height, + normalized, + ) + else: + raise ValueError(f"Unsupported box object: {item}") + elif isinstance(item, (list, tuple)) and len(item) == 4: + values = [float(v) for v in item] + if coordinate_format == "xywh": + x1, y1, w, h = values + if normalized: + x1 *= width + y1 *= height + w *= width + h *= height + x2 = x1 + w + y2 = y1 + h + else: + x1, y1, x2, y2 = _normalize_box_values(values, width, height, normalized) + else: + raise ValueError(f"Unsupported box entry: {item}") + + left = int(math.floor(min(x1, x2))) + right = int(math.ceil(max(x1, x2))) + top = int(math.floor(min(y1, y2))) + bottom = int(math.ceil(max(y1, y2))) + + left = max(0, min(width, left)) + right = max(0, min(width, right)) + top = max(0, min(height, top)) + bottom = max(0, min(height, bottom)) + return left, top, right, bottom + + +def _parse_boxes_json(boxes_json: str, coordinate_format: str, width: int, height: int, normalized: bool) -> List[Tuple[int, int, int, int]]: + text = boxes_json.strip() + if not text: + return [] + data = json.loads(text) + if isinstance(data, dict): + if "boxes" in data: + data = data["boxes"] + else: + data = [data] + if not isinstance(data, list): + raise ValueError("boxes_json must be a JSON array or object") + boxes = [] + for item in data: + box = _parse_single_box(item, coordinate_format, width, height, normalized) + if box[2] > box[0] and box[3] > box[1]: + boxes.append(box) + return boxes + + +def _scale_point(x: float, y: float, width: int, height: int, normalized: bool) -> Tuple[float, float]: + if normalized: + x *= width + y *= height + return x, y + + +def _region_frame_index(region: dict, default_frame_index: int, batch_size: int) -> int: + index = int(region.get("frame_index", default_frame_index)) + return max(0, min(index, batch_size - 1)) + + +def _region_type(region: dict) -> str: + value = str(region.get("type", region.get("shape", "rect"))).strip().lower() + aliases = { + "rectangle": "rect", + "box": "rect", + "bbox": "rect", + "ellipse": "circle", + "poly": "polygon", + } + return aliases.get(value, value) + + +def _polygon_to_mask(points: Sequence[Sequence[float]], width: int, height: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if len(points) < 3: + return torch.zeros((height, width), dtype=dtype, device=device) + + px = torch.arange(width, dtype=dtype, device=device).view(1, width) + 0.5 + py = torch.arange(height, dtype=dtype, device=device).view(height, 1) + 0.5 + inside = torch.zeros((height, width), dtype=torch.bool, device=device) + + wrapped = list(points) + [points[0]] + eps = torch.tensor(1e-6, dtype=dtype, device=device) + for current, nxt in zip(wrapped[:-1], wrapped[1:]): + x1 = torch.tensor(float(current[0]), dtype=dtype, device=device) + y1 = torch.tensor(float(current[1]), dtype=dtype, device=device) + x2 = torch.tensor(float(nxt[0]), dtype=dtype, device=device) + y2 = torch.tensor(float(nxt[1]), dtype=dtype, device=device) + crosses_y = torch.logical_xor(y1 > py, y2 > py) + x_intersection = (x2 - x1) * (py - y1) / (y2 - y1 + eps) + x1 + intersects = crosses_y & (px < x_intersection) + inside = torch.logical_xor(inside, intersects) + + return inside.to(dtype=dtype) + + +def _draw_region(mask: torch.Tensor, region: dict, width: int, height: int, normalized: bool) -> None: + shape = _region_type(region) + if shape == "rect": + box = _parse_single_box(region, "xyxy", width, height, normalized) + x1, y1, x2, y2 = box + if x2 > x1 and y2 > y1: + mask[y1:y2, x1:x2] = 1.0 + return + + if shape == "circle": + if "center" in region: + cx, cy = region["center"] + else: + cx = region.get("cx", region.get("x", 0.0)) + cy = region.get("cy", region.get("y", 0.0)) + cx, cy = _scale_point(float(cx), float(cy), width, height, normalized) + radius = float(region.get("r", region.get("radius", 0.0))) + if normalized: + radius *= min(width, height) + if radius <= 0: + return + yy, xx = torch.meshgrid( + torch.arange(height, dtype=mask.dtype, device=mask.device), + torch.arange(width, dtype=mask.dtype, device=mask.device), + indexing="ij", + ) + circle = (xx + 0.5 - cx) ** 2 + (yy + 0.5 - cy) ** 2 <= radius ** 2 + mask[circle] = 1.0 + return + + if shape == "polygon": + raw_points = region.get("points", []) + points = [] + for point in raw_points: + if isinstance(point, dict): + x = float(point.get("x", 0.0)) + y = float(point.get("y", 0.0)) + else: + x = float(point[0]) + y = float(point[1]) + points.append(_scale_point(x, y, width, height, normalized)) + mask.copy_(torch.maximum(mask, _polygon_to_mask(points, width, height, mask.device, mask.dtype))) + return + + raise ValueError(f"Unsupported region type: {shape}") + + +def _parse_regions_json(regions_json: str) -> List[dict]: + text = regions_json.strip() + if not text: + return [] + data = json.loads(text) + if isinstance(data, dict): + if "regions" in data: + data = data["regions"] + else: + data = [data] + if not isinstance(data, list): + raise ValueError("regions_json must be a JSON array or object") + for region in data: + if not isinstance(region, dict): + raise ValueError(f"Region entries must be objects, got: {region}") + return data + + +def _apply_keyframe_masks(keyframe_masks: dict, batch: int, height: int, width: int, dtype: torch.dtype, device: torch.device, broadcast_mode: str) -> torch.Tensor: + mask = torch.zeros((batch, height, width), dtype=dtype, device=device) + if not keyframe_masks: + return mask + + sorted_items = sorted(keyframe_masks.items()) + if broadcast_mode == "repeat_to_batch": + combined = torch.zeros((height, width), dtype=dtype, device=device) + for _, key_mask in sorted_items: + combined = torch.maximum(combined, key_mask) + return combined.unsqueeze(0).repeat(batch, 1, 1) + + if broadcast_mode == "hold_until_next_keyframe": + for idx, (start_frame, key_mask) in enumerate(sorted_items): + end_frame = sorted_items[idx + 1][0] if idx + 1 < len(sorted_items) else batch + if end_frame > start_frame: + mask[start_frame:end_frame] = key_mask + return mask + + for frame, key_mask in sorted_items: + if broadcast_mode == "hold_after_frame": + mask[frame:] = key_mask + else: + mask[frame] = key_mask + return mask + + +class URDemoImage: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 8}), + "height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 256}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "generate" + CATEGORY = CATEGORY + + def generate(self, width: int, height: int, batch_size: int): + x = torch.linspace(0.0, 1.0, int(width), dtype=torch.float32).view(1, 1, width) + y = torch.linspace(0.0, 1.0, int(height), dtype=torch.float32).view(1, height, 1) + red = 0.15 + 0.35 * x.expand(1, height, width) + green = 0.18 + 0.45 * y.expand(1, height, width) + blue = 0.25 + 0.25 * (1.0 - x.expand(1, height, width)) + image = torch.stack([red, green, blue], dim=-1).repeat(int(batch_size), 1, 1, 1) + + x1 = max(0, min(int(width), int(width * 0.22))) + x2 = max(0, min(int(width), int(width * 0.55))) + y1 = max(0, min(int(height), int(height * 0.24))) + y2 = max(0, min(int(height), int(height * 0.50))) + image[:, y1:y2, x1:x2, 0] = 0.92 + image[:, y1:y2, x1:x2, 1] = 0.16 + image[:, y1:y2, x1:x2, 2] = 0.12 + return (image.clamp(0.0, 1.0),) + + +class URRemovePromptBuilder: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "target_prompt": ("STRING", {"multiline": False, "default": "object to remove"}), + "scene_hint": ("STRING", {"multiline": True, "default": ""}), + "preserve_prompt": ("STRING", {"multiline": True, "default": ""}), + "extra_negative_prompt": ("STRING", {"multiline": True, "default": ""}), + "fill_mode": (["background_reconstruct", "clean_plate", "direct_replace"],), + "masked_edit_only": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("STRING", "STRING", "STRING") + RETURN_NAMES = ("segmentation_prompt", "remove_positive_prompt", "remove_negative_prompt") + FUNCTION = "build" + CATEGORY = CATEGORY + + def build( + self, + target_prompt: str, + scene_hint: str, + preserve_prompt: str, + extra_negative_prompt: str, + fill_mode: str, + masked_edit_only: bool, + ): + target = target_prompt.strip() or "target object" + segmentation_prompt = target + + positive_parts = [f"Remove {target} from the masked region."] + if fill_mode == "background_reconstruct": + positive_parts.append("Reconstruct a plausible background behind the removed content.") + elif fill_mode == "clean_plate": + positive_parts.append("Create a clean plate result that matches the surrounding scene.") + else: + positive_parts.append("Replace the masked region with visually coherent scene content.") + + if masked_edit_only: + positive_parts.append("Only modify the masked region and keep unmasked content unchanged.") + if scene_hint.strip(): + positive_parts.append(f"Scene context: {scene_hint.strip()}.") + if preserve_prompt.strip(): + positive_parts.append(f"Preserve: {preserve_prompt.strip()}.") + + negative_parts = [ + "artifacts", + "ghosting", + "duplicate subject", + "warped anatomy", + "broken perspective", + "extra text", + "watermark", + "subtitle", + "logo", + "blur", + "low quality", + ] + if extra_negative_prompt.strip(): + negative_parts.append(extra_negative_prompt.strip()) + + return ( + segmentation_prompt, + " ".join(positive_parts), + ", ".join(negative_parts), + ) + + +class URRegionEditor: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "regions_json": ( + "STRING", + { + "multiline": True, + "default": "[]", + }, + ), + "canvas_width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 8}), + "canvas_height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}), + "default_frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}), + } + } + + RETURN_TYPES = ("STRING", "INT", "INT", "INT") + RETURN_NAMES = ("regions_json", "canvas_width", "canvas_height", "region_count") + FUNCTION = "emit" + CATEGORY = CATEGORY + + def emit(self, regions_json: str, canvas_width: int, canvas_height: int, default_frame_index: int): + regions = _parse_regions_json(regions_json) + normalized_regions = [] + for region in regions: + copied = dict(region) + copied["frame_index"] = int(copied.get("frame_index", default_frame_index)) + normalized_regions.append(copied) + text = json.dumps(normalized_regions, ensure_ascii=False) + return (text, int(canvas_width), int(canvas_height), len(normalized_regions)) + + +class URBoxesToMaskBatch: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "boxes_json": ("STRING", {"multiline": True, "default": "[]"}), + "frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "coordinate_format": (["xyxy", "xywh"],), + "broadcast_mode": (["repeat_to_batch", "single_frame_only", "hold_after_frame"],), + "normalized": ("BOOLEAN", {"default": False}), + "preview_color": ("STRING", {"multiline": False, "default": "#FF3B30"}), + "preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("MASK", "IMAGE", "IMAGE") + RETURN_NAMES = ("mask", "selected_frame", "preview") + FUNCTION = "build_masks" + CATEGORY = CATEGORY + + def build_masks( + self, + images: torch.Tensor, + boxes_json: str, + frame_index: int, + coordinate_format: str, + broadcast_mode: str, + normalized: bool, + preview_color: str, + preview_opacity: float, + ): + images = _ensure_image_batch(images) + batch, height, width, _ = images.shape + index = max(0, min(int(frame_index), batch - 1)) + + boxes = _parse_boxes_json(boxes_json, coordinate_format, width, height, normalized) + selected_mask = torch.zeros((1, height, width), dtype=images.dtype, device=images.device) + for x1, y1, x2, y2 in boxes: + selected_mask[:, y1:y2, x1:x2] = 1.0 + + if broadcast_mode == "repeat_to_batch": + mask = selected_mask.repeat(batch, 1, 1) + elif broadcast_mode == "hold_after_frame": + mask = torch.zeros((batch, height, width), dtype=images.dtype, device=images.device) + mask[index:] = selected_mask + else: + mask = torch.zeros((batch, height, width), dtype=images.dtype, device=images.device) + mask[index] = selected_mask[0] + + preview = _overlay_mask( + images[index : index + 1], + selected_mask, + _parse_hex_color(preview_color), + preview_opacity, + ) + return (mask, images[index : index + 1], preview) + + +class URRegionsToMaskBatch: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "regions_json": ( + "STRING", + { + "multiline": True, + "default": '[{"type":"rect","frame_index":0,"x":160,"y":120,"w":260,"h":140}]', + }, + ), + "default_frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "broadcast_mode": (["repeat_to_batch", "single_frame_only", "hold_after_frame", "hold_until_next_keyframe"],), + "normalized": ("BOOLEAN", {"default": False}), + "preview_frame": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "preview_color": ("STRING", {"multiline": False, "default": "#FF3B30"}), + "preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("MASK", "IMAGE", "IMAGE") + RETURN_NAMES = ("mask", "selected_frame", "preview") + FUNCTION = "build_masks" + CATEGORY = CATEGORY + + def build_masks( + self, + images: torch.Tensor, + regions_json: str, + default_frame_index: int, + broadcast_mode: str, + normalized: bool, + preview_frame: int, + preview_color: str, + preview_opacity: float, + ): + images = _ensure_image_batch(images) + batch, height, width, _ = images.shape + default_index = max(0, min(int(default_frame_index), batch - 1)) + preview_index = max(0, min(int(preview_frame), batch - 1)) + + keyframe_masks = {} + for region in _parse_regions_json(regions_json): + frame = _region_frame_index(region, default_index, batch) + if frame not in keyframe_masks: + keyframe_masks[frame] = torch.zeros((height, width), dtype=images.dtype, device=images.device) + _draw_region(keyframe_masks[frame], region, width, height, normalized) + + mask = _apply_keyframe_masks( + keyframe_masks, + batch, + height, + width, + images.dtype, + images.device, + broadcast_mode, + ) + preview = _overlay_mask( + images[preview_index : preview_index + 1], + mask[preview_index : preview_index + 1], + _parse_hex_color(preview_color), + preview_opacity, + ) + return (mask, images[preview_index : preview_index + 1], preview) + + +class URMaskProcessor: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK",), + "expand_pixels": ("INT", {"default": 0, "min": 0, "max": 2048}), + "blur_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 512.0, "step": 0.01}), + "blur_kernel": ("INT", {"default": 0, "min": 0, "max": 2049}), + "preview_color": ("STRING", {"multiline": False, "default": "#00FF00"}), + "preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + "optional": { + "image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("MASK", "IMAGE", "IMAGE") + RETURN_NAMES = ("mask", "mask_image", "preview") + FUNCTION = "process" + CATEGORY = CATEGORY + + def process( + self, + mask: torch.Tensor, + expand_pixels: int, + blur_sigma: float, + blur_kernel: int, + preview_color: str, + preview_opacity: float, + image: torch.Tensor = None, + ): + mask = _ensure_mask_batch(mask) + if expand_pixels > 0: + mask = _expand_mask(mask, expand_pixels) + if blur_sigma > 0.0 or blur_kernel > 1: + mask = _gaussian_blur_mask(mask, blur_kernel, blur_sigma) + mask = mask.clamp(0.0, 1.0) + + mask_image = mask.unsqueeze(-1).repeat(1, 1, 1, 3) + if image is not None: + image = _ensure_image_batch(image) + preview = _overlay_mask(image, mask, _parse_hex_color(preview_color), preview_opacity) + else: + preview = _overlay_mask(mask_image, mask, _parse_hex_color(preview_color), preview_opacity) + return (mask, mask_image, preview) + + +class URKeyframePrep: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "mask": ("MASK",), + "frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}), + "fill_color": ("STRING", {"multiline": False, "default": "#00FF00"}), + "resize_longest_side": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 8}), + "pad_to_multiple": ("INT", {"default": 16, "min": 1, "max": 256}), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "IMAGE", "INT", "INT", "INT") + RETURN_NAMES = ( + "original_frame", + "prepared_frame", + "prepared_mask", + "filled_frame", + "prepared_width", + "prepared_height", + "frame_count", + ) + FUNCTION = "prepare" + CATEGORY = CATEGORY + + def prepare( + self, + images: torch.Tensor, + mask: torch.Tensor, + frame_index: int, + fill_color: str, + resize_longest_side: int, + pad_to_multiple: int, + ): + images = _ensure_image_batch(images) + mask = _match_mask_to_images(mask, images) + batch, height, width, _ = images.shape + index = max(0, min(int(frame_index), batch - 1)) + + original_frame = images[index : index + 1] + prepared_frame = original_frame.clone() + prepared_mask = mask[index : index + 1].clone() + + if resize_longest_side > 0: + longest = max(height, width) + scale = float(resize_longest_side) / float(longest) + target_h = max(1, int(round(height * scale))) + target_w = max(1, int(round(width * scale))) + prepared_frame = _resize_images(prepared_frame, (target_h, target_w)) + prepared_mask = _resize_masks(prepared_mask, (target_h, target_w)) + + prepared_h = prepared_frame.shape[1] + prepared_w = prepared_frame.shape[2] + if pad_to_multiple > 1: + target_h = int(math.ceil(prepared_h / pad_to_multiple) * pad_to_multiple) + target_w = int(math.ceil(prepared_w / pad_to_multiple) * pad_to_multiple) + prepared_frame = _pad_images(prepared_frame, target_h, target_w) + prepared_mask = _pad_masks(prepared_mask, target_h, target_w) + + color = _parse_hex_color(fill_color) + filled_frame = _fill_masked_region(prepared_frame, prepared_mask, color) + return ( + original_frame, + prepared_frame, + prepared_mask, + filled_frame, + prepared_frame.shape[2], + prepared_frame.shape[1], + batch, + ) + + +class URSoftComposite: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "inpaint": ("IMAGE",), + "original": ("IMAGE",), + "mask": ("MASK",), + "kernel": ("INT", {"default": 15, "min": 1, "max": 2049}), + "sigma": ("FLOAT", {"default": 7.5, "min": 0.0, "max": 512.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "soft_mask") + FUNCTION = "composite" + CATEGORY = CATEGORY + + def composite( + self, + inpaint: torch.Tensor, + original: torch.Tensor, + mask: torch.Tensor, + kernel: int, + sigma: float, + ): + original = _ensure_image_batch(original) + inpaint = _ensure_image_batch(inpaint) + target_batch = max(original.shape[0], inpaint.shape[0], _ensure_mask_batch(mask).shape[0]) + + original = _tile_to_batch(original, target_batch) + inpaint = _tile_to_batch(inpaint, target_batch) + mask = _tile_to_batch(_ensure_mask_batch(mask), target_batch) + mask = _match_mask_to_images(mask, original) + + if inpaint.shape[1] != original.shape[1] or inpaint.shape[2] != original.shape[2]: + inpaint = _resize_images(inpaint, (original.shape[1], original.shape[2])) + + soft_mask = _gaussian_blur_mask(mask, kernel, sigma).clamp(0.0, 1.0) + composited = original * (1.0 - soft_mask.unsqueeze(-1)) + inpaint * soft_mask.unsqueeze(-1) + return (composited.clamp(0.0, 1.0), soft_mask) + + +NODE_CLASS_MAPPINGS = { + "UR Demo Image": URDemoImage, + "UR Remove Prompt Builder": URRemovePromptBuilder, + "UR Region Editor": URRegionEditor, + "UR Boxes To Mask Batch": URBoxesToMaskBatch, + "UR Regions To Mask Batch": URRegionsToMaskBatch, + "UR Mask Processor": URMaskProcessor, + "UR Keyframe Prep": URKeyframePrep, + "UR Soft Composite": URSoftComposite, +} + + +NODE_DISPLAY_NAME_MAPPINGS = { + "UR Demo Image": "UR Demo Image", + "UR Remove Prompt Builder": "UR Remove Prompt Builder", + "UR Region Editor": "UR Region Editor", + "UR Boxes To Mask Batch": "UR Boxes To Mask Batch", + "UR Regions To Mask Batch": "UR Regions To Mask Batch", + "UR Mask Processor": "UR Mask Processor", + "UR Keyframe Prep": "UR Keyframe Prep", + "UR Soft Composite": "UR Soft Composite", +} diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..5ee6477 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +testpaths = tests diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..0587502 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1 @@ +pytest>=7.8.0 diff --git a/tests/test_universal_remove.py b/tests/test_universal_remove.py new file mode 100644 index 0000000..8378d09 --- /dev/null +++ b/tests/test_universal_remove.py @@ -0,0 +1,244 @@ +import importlib.util +import sys +from pathlib import Path + +import torch + + +MODULE_PATH = Path(__file__).resolve().parents[1] / "nodes.py" +SPEC = importlib.util.spec_from_file_location("universal_remove_nodes", MODULE_PATH) +MODULE = importlib.util.module_from_spec(SPEC) +assert SPEC is not None and SPEC.loader is not None +SPEC.loader.exec_module(MODULE) + + +def test_extension_entrypoint_exports_nodes_and_web_directory(): + package_path = MODULE_PATH.parent / "__init__.py" + spec = importlib.util.spec_from_file_location( + "ComfyUI_UniversalRemove_test", + package_path, + submodule_search_locations=[str(MODULE_PATH.parent)], + ) + package = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = package + assert spec is not None and spec.loader is not None + spec.loader.exec_module(package) + + assert package.WEB_DIRECTORY == "./web" + assert "UR Region Editor" in package.NODE_CLASS_MAPPINGS + assert "UR Regions To Mask Batch" in package.NODE_CLASS_MAPPINGS + + +def test_prompt_builder_outputs_expected_strings(): + node = MODULE.URRemovePromptBuilder() + segmentation_prompt, positive_prompt, negative_prompt = node.build( + target_prompt="watermark", + scene_hint="night city street", + preserve_prompt="keep the lighting and camera framing", + extra_negative_prompt="text remnants", + fill_mode="background_reconstruct", + masked_edit_only=True, + ) + + assert segmentation_prompt == "watermark" + assert "Remove watermark" in positive_prompt + assert "night city street" in positive_prompt + assert "keep the lighting and camera framing" in positive_prompt + assert "text remnants" in negative_prompt + + +def test_demo_image_is_self_contained_and_batched(): + node = MODULE.URDemoImage() + (image,) = node.generate(width=64, height=32, batch_size=2) + + assert image.shape == (2, 32, 64, 3) + assert image.min().item() >= 0.0 + assert image.max().item() <= 1.0 + + +def test_region_editor_normalizes_regions_and_counts_them(): + node = MODULE.URRegionEditor() + regions_json, canvas_width, canvas_height, region_count = node.emit( + regions_json='[{"type":"rect","x":1,"y":2,"w":3,"h":4}]', + canvas_width=64, + canvas_height=32, + default_frame_index=5, + ) + + assert '"frame_index": 5' in regions_json + assert canvas_width == 64 + assert canvas_height == 32 + assert region_count == 1 + + +def test_boxes_to_mask_batch_repeats_mask_to_batch(): + node = MODULE.URBoxesToMaskBatch() + images = torch.zeros((2, 4, 4, 3), dtype=torch.float32) + + mask, selected_frame, preview = node.build_masks( + images=images, + boxes_json="[[1, 1, 3, 3]]", + frame_index=0, + coordinate_format="xyxy", + broadcast_mode="repeat_to_batch", + normalized=False, + preview_color="#FF0000", + preview_opacity=0.5, + ) + + expected = torch.tensor( + [ + [0, 0, 0, 0], + [0, 1, 1, 0], + [0, 1, 1, 0], + [0, 0, 0, 0], + ], + dtype=torch.float32, + ) + assert mask.shape == (2, 4, 4) + assert torch.equal(mask[0], expected) + assert torch.equal(mask[1], expected) + assert selected_frame.shape == (1, 4, 4, 3) + assert preview.shape == (1, 4, 4, 3) + + +def test_boxes_to_mask_batch_handles_reversed_out_of_bounds_box(): + node = MODULE.URBoxesToMaskBatch() + images = torch.zeros((1, 4, 4, 3), dtype=torch.float32) + + mask, _, _ = node.build_masks( + images=images, + boxes_json="[[5, 5, -1, -1]]", + frame_index=0, + coordinate_format="xyxy", + broadcast_mode="single_frame_only", + normalized=False, + preview_color="#FF0000", + preview_opacity=0.5, + ) + + assert mask.shape == (1, 4, 4) + assert mask.sum().item() == 16.0 + + +def test_regions_to_mask_batch_supports_rect_circle_and_polygon(): + node = MODULE.URRegionsToMaskBatch() + images = torch.zeros((1, 10, 10, 3), dtype=torch.float32) + regions_json = """ + [ + {"type": "rect", "frame_index": 0, "x": 1, "y": 1, "w": 2, "h": 2}, + {"type": "circle", "frame_index": 0, "cx": 7, "cy": 7, "r": 1.5}, + {"type": "polygon", "frame_index": 0, "points": [[5, 1], [8, 1], [8, 3], [5, 3]]} + ] + """ + + mask, selected_frame, preview = node.build_masks( + images=images, + regions_json=regions_json, + default_frame_index=0, + broadcast_mode="single_frame_only", + normalized=False, + preview_frame=0, + preview_color="#FF0000", + preview_opacity=0.5, + ) + + assert mask.shape == (1, 10, 10) + assert mask[0, 1:3, 1:3].sum().item() == 4.0 + assert mask[0, 7, 7].item() == 1.0 + assert mask[0, 1:3, 5:8].sum().item() > 0.0 + assert selected_frame.shape == (1, 10, 10, 3) + assert preview.shape == (1, 10, 10, 3) + + +def test_regions_to_mask_batch_holds_between_keyframes(): + node = MODULE.URRegionsToMaskBatch() + images = torch.zeros((4, 6, 6, 3), dtype=torch.float32) + regions_json = """ + [ + {"type": "rect", "frame_index": 0, "x": 0, "y": 0, "w": 2, "h": 2}, + {"type": "rect", "frame_index": 2, "x": 4, "y": 4, "w": 2, "h": 2} + ] + """ + + mask, _, _ = node.build_masks( + images=images, + regions_json=regions_json, + default_frame_index=0, + broadcast_mode="hold_until_next_keyframe", + normalized=False, + preview_frame=0, + preview_color="#FF0000", + preview_opacity=0.5, + ) + + assert mask.shape == (4, 6, 6) + assert mask[0, 0:2, 0:2].sum().item() == 4.0 + assert mask[1, 0:2, 0:2].sum().item() == 4.0 + assert mask[2, 4:6, 4:6].sum().item() == 4.0 + assert mask[3, 4:6, 4:6].sum().item() == 4.0 + assert mask[1, 4:6, 4:6].sum().item() == 0.0 + + +def test_mask_processor_can_expand_mask_without_image(): + node = MODULE.URMaskProcessor() + mask = torch.zeros((1, 5, 5), dtype=torch.float32) + mask[:, 2, 2] = 1.0 + + processed_mask, mask_image, preview = node.process( + mask=mask, + expand_pixels=1, + blur_sigma=0.0, + blur_kernel=0, + preview_color="#00FF00", + preview_opacity=0.5, + image=None, + ) + + assert processed_mask.sum().item() == 9.0 + assert mask_image.shape == (1, 5, 5, 3) + assert preview.shape == (1, 5, 5, 3) + + +def test_keyframe_prep_fills_masked_region_and_pads(): + node = MODULE.URKeyframePrep() + images = torch.zeros((1, 3, 5, 3), dtype=torch.float32) + mask = torch.zeros((1, 3, 5), dtype=torch.float32) + mask[:, 1, 2] = 1.0 + + original_frame, prepared_frame, prepared_mask, filled_frame, prepared_width, prepared_height, frame_count = node.prepare( + images=images, + mask=mask, + frame_index=0, + fill_color="#00FF00", + resize_longest_side=0, + pad_to_multiple=4, + ) + + assert original_frame.shape == (1, 3, 5, 3) + assert prepared_frame.shape == (1, 4, 8, 3) + assert prepared_mask.shape == (1, 4, 8) + assert prepared_width == 8 + assert prepared_height == 4 + assert frame_count == 1 + assert torch.allclose(filled_frame[0, 1, 2], torch.tensor([0.0, 1.0, 0.0])) + + +def test_soft_composite_uses_mask_for_direct_blend_when_no_blur(): + node = MODULE.URSoftComposite() + original = torch.zeros((1, 4, 4, 3), dtype=torch.float32) + inpaint = torch.ones((1, 4, 4, 3), dtype=torch.float32) + mask = torch.zeros((1, 4, 4), dtype=torch.float32) + mask[:, 1:3, 1:3] = 1.0 + + composited, soft_mask = node.composite( + inpaint=inpaint, + original=original, + mask=mask, + kernel=1, + sigma=0.0, + ) + + assert torch.equal(soft_mask, mask) + assert torch.allclose(composited[0, 0, 0], torch.tensor([0.0, 0.0, 0.0])) + assert torch.allclose(composited[0, 1, 1], torch.tensor([1.0, 1.0, 1.0])) diff --git a/web/region_editor.js b/web/region_editor.js new file mode 100644 index 0000000..379c806 --- /dev/null +++ b/web/region_editor.js @@ -0,0 +1,535 @@ +import { app } from "../../scripts/app.js"; + +function chainCallback(object, property, callback) { + const original = object[property]; + if (original) { + object[property] = function (...args) { + const result = original.apply(this, args); + callback.apply(this, args); + return result; + }; + } else { + object[property] = callback; + } +} + +function findWidget(node, name) { + return node.widgets?.find((widget) => widget.name === name); +} + +function clamp(value, min, max) { + return Math.max(min, Math.min(max, value)); +} + +function parseRegions(value) { + try { + const parsed = JSON.parse(value || "[]"); + if (Array.isArray(parsed)) return parsed; + if (parsed && Array.isArray(parsed.regions)) return parsed.regions; + if (parsed && typeof parsed === "object") return [parsed]; + } catch (error) { + console.warn("UR Region Editor: failed to parse regions_json", error); + } + return []; +} + +function normalizeRect(region) { + const x = Number(region.x ?? region.x1 ?? 0); + const y = Number(region.y ?? region.y1 ?? 0); + if (region.w != null || region.h != null) { + return { + x, + y, + w: Number(region.w ?? 0), + h: Number(region.h ?? 0), + }; + } + const x2 = Number(region.x2 ?? region.right ?? x); + const y2 = Number(region.y2 ?? region.bottom ?? y); + return { + x: Math.min(x, x2), + y: Math.min(y, y2), + w: Math.abs(x2 - x), + h: Math.abs(y2 - y), + }; +} + +function createButton(label, title) { + const button = document.createElement("button"); + button.textContent = label; + button.title = title || label; + button.style.cssText = "height:24px;padding:0 8px;border:1px solid #555;background:#2b2b2b;color:#ddd;border-radius:4px;cursor:pointer;font-size:12px;"; + button.addEventListener("mouseenter", () => { + button.style.background = "#3a3a3a"; + }); + button.addEventListener("mouseleave", () => { + button.style.background = button.dataset.active === "true" ? "#4b5f7a" : "#2b2b2b"; + }); + return button; +} + +function createNumberInput(value, min, title) { + const input = document.createElement("input"); + input.type = "number"; + input.value = String(value); + input.min = String(min ?? 0); + input.title = title || ""; + input.style.cssText = "height:22px;width:76px;background:#111;color:#ddd;border:1px solid #555;border-radius:4px;padding:0 4px;font-size:12px;"; + return input; +} + +class RegionEditorWidget { + constructor(node) { + this.node = node; + this.regionsWidget = findWidget(node, "regions_json"); + this.widthWidget = findWidget(node, "canvas_width"); + this.heightWidget = findWidget(node, "canvas_height"); + this.defaultFrameWidget = findWidget(node, "default_frame_index"); + + this.regions = parseRegions(this.regionsWidget?.value); + this.mode = "rect"; + this.activeFrame = Number(this.defaultFrameWidget?.value ?? 0); + this.selectedIndex = -1; + this.dragStart = null; + this.currentDraft = null; + this.polygonDraft = []; + this.background = null; + this.video = null; + this.videoFps = 24; + + this.root = document.createElement("div"); + this.root.style.cssText = "display:flex;flex-direction:column;gap:6px;width:100%;height:100%;box-sizing:border-box;background:#121212;border:1px solid #303030;border-radius:6px;padding:6px;overflow:hidden;"; + + this.toolbar = document.createElement("div"); + this.toolbar.style.cssText = "display:flex;align-items:center;gap:6px;flex-wrap:wrap;"; + + this.rectBtn = createButton("Rect", "Draw rectangle by dragging"); + this.circleBtn = createButton("Circle", "Draw circle by dragging"); + this.polyBtn = createButton("Poly", "Click points, then Finish Poly"); + this.finishPolyBtn = createButton("Finish Poly", "Finish current polygon"); + this.undoPointBtn = createButton("Undo Point", "Remove last polygon point"); + this.deleteBtn = createButton("Delete", "Delete selected region"); + this.clearBtn = createButton("Clear", "Clear all regions"); + this.loadBgBtn = createButton("Load BG", "Load local image/video as editor background only"); + + this.frameInput = createNumberInput(this.activeFrame, 0, "Current keyframe index"); + this.frameLabel = document.createElement("span"); + this.frameLabel.textContent = "Frame"; + this.frameLabel.style.cssText = "color:#aaa;font-size:12px;"; + + this.fileInput = document.createElement("input"); + this.fileInput.type = "file"; + this.fileInput.accept = "image/*,video/*"; + this.fileInput.style.display = "none"; + + this.videoControls = document.createElement("div"); + this.videoControls.style.cssText = "display:none;align-items:center;gap:6px;width:100%;"; + this.playBtn = createButton("Play", "Play/pause background video"); + this.videoSlider = document.createElement("input"); + this.videoSlider.type = "range"; + this.videoSlider.min = "0"; + this.videoSlider.max = "0"; + this.videoSlider.step = "0.001"; + this.videoSlider.value = "0"; + this.videoSlider.style.cssText = "flex:1;"; + this.fpsInput = createNumberInput(this.videoFps, 1, "FPS used to convert video time to frame index"); + this.useVideoFrameBtn = createButton("Use Time", "Set keyframe from current video time"); + + this.canvas = document.createElement("canvas"); + this.canvas.width = Number(this.widthWidget?.value ?? 768); + this.canvas.height = Number(this.heightWidget?.value ?? 512); + this.canvas.style.cssText = "width:100%;height:auto;max-height:520px;object-fit:contain;background:#181818;border:1px solid #333;border-radius:4px;cursor:crosshair;"; + this.ctx = this.canvas.getContext("2d"); + + this.status = document.createElement("div"); + this.status.style.cssText = "color:#aaa;font-size:12px;white-space:pre-wrap;line-height:1.3;"; + + this.toolbar.append( + this.rectBtn, + this.circleBtn, + this.polyBtn, + this.finishPolyBtn, + this.undoPointBtn, + this.deleteBtn, + this.clearBtn, + this.loadBgBtn, + this.frameLabel, + this.frameInput, + this.fileInput, + ); + this.videoControls.append(this.playBtn, this.videoSlider, document.createTextNode("FPS"), this.fpsInput, this.useVideoFrameBtn); + this.root.append(this.toolbar, this.videoControls, this.canvas, this.status); + + this.bindEvents(); + this.setMode("rect"); + this.syncWidgets(false); + this.render(); + } + + bindEvents() { + this.rectBtn.onclick = () => this.setMode("rect"); + this.circleBtn.onclick = () => this.setMode("circle"); + this.polyBtn.onclick = () => this.setMode("polygon"); + this.finishPolyBtn.onclick = () => this.finishPolygon(); + this.undoPointBtn.onclick = () => { + this.polygonDraft.pop(); + this.render(); + }; + this.deleteBtn.onclick = () => this.deleteSelected(); + this.clearBtn.onclick = () => { + this.regions = []; + this.selectedIndex = -1; + this.polygonDraft = []; + this.syncWidgets(true); + this.render(); + }; + this.loadBgBtn.onclick = () => this.fileInput.click(); + + this.frameInput.addEventListener("change", () => { + this.activeFrame = Math.max(0, Number(this.frameInput.value || 0)); + if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame; + this.render(); + }); + this.fpsInput.addEventListener("change", () => { + this.videoFps = Math.max(1, Number(this.fpsInput.value || 24)); + }); + + this.fileInput.addEventListener("change", () => this.loadBackgroundFile()); + this.playBtn.onclick = () => this.toggleVideo(); + this.videoSlider.addEventListener("input", () => { + if (!this.video) return; + this.video.currentTime = Number(this.videoSlider.value || 0); + this.drawVideoFrame(); + }); + this.useVideoFrameBtn.onclick = () => { + if (!this.video) return; + this.activeFrame = Math.max(0, Math.round(this.video.currentTime * this.videoFps)); + this.frameInput.value = String(this.activeFrame); + if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame; + this.render(); + }; + + this.canvas.addEventListener("mousedown", (event) => this.onMouseDown(event)); + this.canvas.addEventListener("mousemove", (event) => this.onMouseMove(event)); + this.canvas.addEventListener("mouseup", (event) => this.onMouseUp(event)); + this.canvas.addEventListener("dblclick", () => this.finishPolygon()); + this.canvas.addEventListener("contextmenu", (event) => { + event.preventDefault(); + const point = this.toCanvasPoint(event); + const index = this.hitTest(point.x, point.y); + if (index >= 0) { + this.selectedIndex = index; + this.deleteSelected(); + } + }); + document.addEventListener("keydown", (event) => { + if (!this.root.isConnected) return; + if (event.key === "Delete" || event.key === "Backspace") this.deleteSelected(); + if (event.key === "Enter" && this.mode === "polygon") this.finishPolygon(); + if (event.key === "Escape") { + this.currentDraft = null; + this.polygonDraft = []; + this.render(); + } + }); + + if (this.regionsWidget) { + const previous = this.regionsWidget.callback; + this.regionsWidget.callback = (...args) => { + if (previous) previous.apply(this.regionsWidget, args); + this.regions = parseRegions(this.regionsWidget.value); + this.render(); + }; + } + } + + setMode(mode) { + this.mode = mode; + for (const [button, name] of [[this.rectBtn, "rect"], [this.circleBtn, "circle"], [this.polyBtn, "polygon"]]) { + button.dataset.active = name === mode ? "true" : "false"; + button.style.background = name === mode ? "#4b5f7a" : "#2b2b2b"; + } + this.currentDraft = null; + this.render(); + } + + syncWidgets(markDirty) { + const text = JSON.stringify(this.regions); + if (this.regionsWidget) this.regionsWidget.value = text; + if (this.widthWidget) this.widthWidget.value = this.canvas.width; + if (this.heightWidget) this.heightWidget.value = this.canvas.height; + if (this.defaultFrameWidget) this.defaultFrameWidget.value = this.activeFrame; + if (markDirty) app.graph?.setDirtyCanvas(true, true); + } + + loadBackgroundFile() { + const file = this.fileInput.files?.[0]; + if (!file) return; + const url = URL.createObjectURL(file); + this.background = null; + this.video = null; + this.videoControls.style.display = "none"; + if (file.type.startsWith("video/")) { + const video = document.createElement("video"); + video.src = url; + video.muted = true; + video.playsInline = true; + video.addEventListener("loadedmetadata", () => { + this.video = video; + this.canvas.width = video.videoWidth || this.canvas.width; + this.canvas.height = video.videoHeight || this.canvas.height; + this.videoSlider.max = String(video.duration || 0); + this.videoControls.style.display = "flex"; + this.drawVideoFrame(); + this.syncWidgets(true); + }); + video.addEventListener("timeupdate", () => { + this.videoSlider.value = String(video.currentTime); + this.drawVideoFrame(); + }); + return; + } + + const image = new Image(); + image.onload = () => { + this.background = image; + this.canvas.width = image.width; + this.canvas.height = image.height; + this.syncWidgets(true); + this.render(); + }; + image.src = url; + } + + toggleVideo() { + if (!this.video) return; + if (this.video.paused) { + this.video.play(); + this.playBtn.textContent = "Pause"; + } else { + this.video.pause(); + this.playBtn.textContent = "Play"; + } + } + + drawVideoFrame() { + this.render(); + } + + toCanvasPoint(event) { + const rect = this.canvas.getBoundingClientRect(); + return { + x: clamp((event.clientX - rect.left) * (this.canvas.width / rect.width), 0, this.canvas.width), + y: clamp((event.clientY - rect.top) * (this.canvas.height / rect.height), 0, this.canvas.height), + }; + } + + onMouseDown(event) { + event.preventDefault(); + const point = this.toCanvasPoint(event); + const hit = this.hitTest(point.x, point.y); + if (hit >= 0 && event.shiftKey) { + this.selectedIndex = hit; + this.render(); + return; + } + if (this.mode === "polygon") { + this.polygonDraft.push(point); + this.render(); + return; + } + this.dragStart = point; + this.currentDraft = { type: this.mode, frame_index: this.activeFrame, start: point, end: point }; + } + + onMouseMove(event) { + if (!this.currentDraft) return; + this.currentDraft.end = this.toCanvasPoint(event); + this.render(); + } + + onMouseUp(event) { + if (!this.currentDraft) return; + this.currentDraft.end = this.toCanvasPoint(event); + const region = this.draftToRegion(this.currentDraft); + this.currentDraft = null; + if (region) { + this.regions.push(region); + this.selectedIndex = this.regions.length - 1; + this.syncWidgets(true); + } + this.render(); + } + + draftToRegion(draft) { + const x1 = draft.start.x; + const y1 = draft.start.y; + const x2 = draft.end.x; + const y2 = draft.end.y; + if (draft.type === "rect") { + const x = Math.min(x1, x2); + const y = Math.min(y1, y2); + const w = Math.abs(x2 - x1); + const h = Math.abs(y2 - y1); + if (w < 3 || h < 3) return null; + return { type: "rect", frame_index: this.activeFrame, x: Math.round(x), y: Math.round(y), w: Math.round(w), h: Math.round(h) }; + } + if (draft.type === "circle") { + const r = Math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2); + if (r < 3) return null; + return { type: "circle", frame_index: this.activeFrame, cx: Math.round(x1), cy: Math.round(y1), r: Math.round(r) }; + } + return null; + } + + finishPolygon() { + if (this.polygonDraft.length < 3) return; + this.regions.push({ + type: "polygon", + frame_index: this.activeFrame, + points: this.polygonDraft.map((point) => [Math.round(point.x), Math.round(point.y)]), + }); + this.selectedIndex = this.regions.length - 1; + this.polygonDraft = []; + this.syncWidgets(true); + this.render(); + } + + deleteSelected() { + if (this.selectedIndex < 0 || this.selectedIndex >= this.regions.length) return; + this.regions.splice(this.selectedIndex, 1); + this.selectedIndex = -1; + this.syncWidgets(true); + this.render(); + } + + hitTest(x, y) { + for (let i = this.regions.length - 1; i >= 0; i--) { + const region = this.regions[i]; + if (Number(region.frame_index ?? 0) !== this.activeFrame) continue; + const type = String(region.type || "rect").toLowerCase(); + if (type === "rect") { + const rect = normalizeRect(region); + if (x >= rect.x && x <= rect.x + rect.w && y >= rect.y && y <= rect.y + rect.h) return i; + } else if (type === "circle") { + const dx = x - Number(region.cx ?? region.x ?? 0); + const dy = y - Number(region.cy ?? region.y ?? 0); + if (dx * dx + dy * dy <= Number(region.r ?? region.radius ?? 0) ** 2) return i; + } + } + return -1; + } + + render() { + const ctx = this.ctx; + ctx.clearRect(0, 0, this.canvas.width, this.canvas.height); + ctx.fillStyle = "#181818"; + ctx.fillRect(0, 0, this.canvas.width, this.canvas.height); + if (this.video) ctx.drawImage(this.video, 0, 0, this.canvas.width, this.canvas.height); + else if (this.background) ctx.drawImage(this.background, 0, 0, this.canvas.width, this.canvas.height); + else this.drawGrid(); + + this.regions.forEach((region, index) => this.drawRegion(region, index === this.selectedIndex)); + if (this.currentDraft) this.drawDraft(this.currentDraft); + if (this.polygonDraft.length) this.drawPolygonDraft(); + this.status.textContent = `Mode: ${this.mode} | Frame: ${this.activeFrame} | Regions: ${this.regions.length}\nLoad BG is editor-only; regions_json is portable and saved in the workflow.`; + } + + drawGrid() { + const ctx = this.ctx; + ctx.strokeStyle = "#2a2a2a"; + ctx.lineWidth = 1; + for (let x = 0; x <= this.canvas.width; x += 64) { + ctx.beginPath(); + ctx.moveTo(x, 0); + ctx.lineTo(x, this.canvas.height); + ctx.stroke(); + } + for (let y = 0; y <= this.canvas.height; y += 64) { + ctx.beginPath(); + ctx.moveTo(0, y); + ctx.lineTo(this.canvas.width, y); + ctx.stroke(); + } + ctx.fillStyle = "#777"; + ctx.font = "16px sans-serif"; + ctx.textAlign = "center"; + ctx.fillText("Load image/video background, then draw regions", this.canvas.width / 2, this.canvas.height / 2); + } + + drawRegion(region, selected) { + const ctx = this.ctx; + const type = String(region.type || "rect").toLowerCase(); + ctx.save(); + ctx.strokeStyle = selected ? "#ffcc00" : "#35c759"; + ctx.fillStyle = selected ? "rgba(255,204,0,0.20)" : "rgba(53,199,89,0.20)"; + ctx.lineWidth = selected ? 3 : 2; + if (type === "rect") { + const rect = normalizeRect(region); + ctx.fillRect(rect.x, rect.y, rect.w, rect.h); + ctx.strokeRect(rect.x, rect.y, rect.w, rect.h); + } else if (type === "circle") { + ctx.beginPath(); + ctx.arc(Number(region.cx ?? region.x ?? 0), Number(region.cy ?? region.y ?? 0), Number(region.r ?? region.radius ?? 0), 0, Math.PI * 2); + ctx.fill(); + ctx.stroke(); + } else if (type === "polygon") { + this.drawPolygonPath(region.points || []); + ctx.fill(); + ctx.stroke(); + } + ctx.restore(); + } + + drawPolygonPath(points) { + const ctx = this.ctx; + ctx.beginPath(); + points.forEach((point, index) => { + const x = Array.isArray(point) ? point[0] : point.x; + const y = Array.isArray(point) ? point[1] : point.y; + if (index === 0) ctx.moveTo(x, y); + else ctx.lineTo(x, y); + }); + ctx.closePath(); + } + + drawDraft(draft) { + const region = this.draftToRegion(draft); + if (!region) return; + this.drawRegion(region, true); + } + + drawPolygonDraft() { + const ctx = this.ctx; + ctx.save(); + ctx.strokeStyle = "#ffcc00"; + ctx.fillStyle = "#ffcc00"; + ctx.lineWidth = 2; + ctx.beginPath(); + this.polygonDraft.forEach((point, index) => { + if (index === 0) ctx.moveTo(point.x, point.y); + else ctx.lineTo(point.x, point.y); + ctx.fillRect(point.x - 3, point.y - 3, 6, 6); + }); + ctx.stroke(); + ctx.restore(); + } +} + +app.registerExtension({ + name: "ComfyUI.UniversalRemove.RegionEditor", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== "UR Region Editor") return; + chainCallback(nodeType.prototype, "onNodeCreated", function () { + this.serialize_widgets = true; + const editor = new RegionEditorWidget(this); + this.urRegionEditor = editor; + const widget = this.addDOMWidget("region_editor", "region_editor", editor.root); + widget.computeSize = (width) => [width, Math.max(420, this.size?.[1] ? this.size[1] - 160 : 420)]; + this.setSize([Math.max(this.size?.[0] || 680, 680), Math.max(this.size?.[1] || 680, 680)]); + chainCallback(this, "onResize", function () { + editor.render(); + }); + }); + }, +});