Add mask-only batch processing nodes

README.md

@@ -24,6 +24,10 @@ RefineAnything LoRA:
 
 You can try different options in the `Edit Model Reference Method` node, but it is best to avoid `index_timestep_zero` because it can introduce a noticeable color shift.
 
+For workflows with several painted regions, insert the mask-only nodes before `RefineNode Preprocess Mask`:
+
+`Load Image mask` -> `RefineNode Mask Batch Process` -> optional `RefineNode Combine Nearby Masks` -> `RefineNode Preprocess Mask`
+
 `example_workflows/Reference-based Logo Refinement.json`
 
 - Reference-based workflow for refining logos, text, or product details with a clean reference image.
@@ -39,6 +43,36 @@ You can try different options in the `Edit Model Reference Method` node, but it
 
 ## Node Details
 
+### RefineNode Mask Batch Process
+
+Prepares a `MASK` tensor before it is sent into `RefineNode Preprocess Mask`.
+
+Inputs:
+
+- `mask`: Input mask. This can be a normal mask, a batched mask, or a ComfyUI mask list.
+- `combined_mask`: When disabled, each disconnected painted island is split into a separate mask in the output batch. When enabled, all incoming masks are unioned into one mask.
+
+Outputs:
+
+- `mask`: A processed `MASK` batch.
+
+Use this node when a single `Load Image` mask contains multiple separated painted regions and you want either one refinement job per region or one combined refinement job.
+
+### RefineNode Combine Nearby Masks
+
+Groups close masks together without touching the image or paste-back metadata.
+
+Inputs:
+
+- `mask`: Input mask batch or mask list. All masks must have the same spatial size.
+- `max_gap`: Maximum bounding-box gap, in pixels, for two masks to be grouped. The default is `32`.
+
+Outputs:
+
+- `mask`: A `MASK` batch where each output mask is the union of one nearby group.
+
+Grouping is deterministic and transitive: if mask A is near B, and B is near C, A/B/C become one output mask. Distant groups remain separate.
+
 ### RefineNode Preprocess Mask
 
 Prepares the target image and optional mask before model inference, and creates the `info` data required for accurate paste-back.
@@ -50,7 +84,6 @@ Inputs:
 - `focus_crop`: When enabled, crops around the masked region to reduce unrelated image area before generation.
 - `focus_crop_margin`: Extra context retained around the mask crop.
 - `spatial_prompt_source`: Controls the `spatial_mask_image` output. `mask` keeps the original mask shape, while `bbox` uses the mask bounding box.
-- `combined_mask`: When disabled, a multi-mask `MASK` input is expanded into separate image-list refinement jobs. This also splits disconnected painted mask islands from a single `Load Image` mask. When enabled, all masks for each source image are combined into one union mask before preprocessing.
 
 Outputs:
 
@@ -59,10 +92,11 @@ Outputs:
 - `mask`: Mask aligned with the output `image`.
 - `info`: Metadata containing the original image, crop position, mask, and paste-back coordinates. Pass this to downstream paste-back nodes.
 
-Multi-mask behavior:
+Mask behavior:
 
-- With `combined_mask` disabled, one input image plus multiple masks becomes an image list: one refinement job per mask or disconnected painted mask island.
-- With `combined_mask` enabled, multiple masks are merged first and only one refinement job is created per source image.
+- `RefineNode Preprocess Mask` no longer decides whether masks should be split or combined.
+- A batched `MASK` input is treated as the exact refinement job list: one output image-list item per incoming mask, using the existing image/mask batch assignment rules.
+- Use `RefineNode Mask Batch Process` and `RefineNode Combine Nearby Masks` before this node when you need to split painted islands, combine all masks, or group nearby mask regions.
 
 ### RefineNode Reference Image Process
 

example_workflows/Reference-based Logo Refinement.json

@@ -846,8 +846,7 @@
       "widgets_values": [
         true,
         32,
-        "mask",
-        false
+        "mask"
       ]
     },
     {
@@ -911,8 +910,7 @@
       "widgets_values": [
         true,
         32,
-        "mask",
-        false
+        "mask"
       ]
     },
     {

example_workflows/Reference-free Text Refinement.json

@@ -879,8 +879,7 @@
       "widgets_values": [
         true,
         64,
-        "mask",
-        false
+        "mask"
       ]
     },
     {

nodes.py

@@ -285,6 +285,21 @@ def combine_masks(masks: torch.Tensor, indices: list[int], size: tuple[int, int]
     return Image.fromarray(combined, mode="L")
 
 
+def union_mask_images(mask_images: list[Image.Image], size: tuple[int, int] | None = None) -> Image.Image:
+    if not mask_images:
+        if size is None:
+            raise ValueError("Cannot union an empty mask list without a target size.")
+        return Image.new("L", size, 0)
+    target_size = size or mask_images[0].size
+    combined = np.zeros((target_size[1], target_size[0]), dtype=np.uint8)
+    for mask_l in mask_images:
+        current = mask_l.convert("L")
+        if current.size != target_size:
+            current = current.resize(target_size, _NEAREST)
+        combined = np.maximum(combined, np.asarray(current, dtype=np.uint8))
+    return Image.fromarray(combined, mode="L")
+
+
 def connected_component_labels(binary: np.ndarray) -> tuple[np.ndarray, int]:
     binary_u8 = binary.astype(np.uint8, copy=False)
     try:
@@ -352,6 +367,56 @@ def split_mask_components(mask_l: Image.Image) -> list[Image.Image]:
     return [component for _, _, component in components]
 
 
+def flatten_mask_input(masks: torch.Tensor | list[torch.Tensor]) -> list[Image.Image]:
+    values = masks if isinstance(masks, list) else [masks]
+    mask_images = []
+    expected_size = None
+    for value in values:
+        if value is None:
+            continue
+        count = mask_batch_size(value)
+        if count <= 0:
+            continue
+        if value.ndim == 2:
+            size = (int(value.shape[1]), int(value.shape[0]))
+        elif value.ndim == 3 and value.shape[-1] == 1 and value.shape[0] > 4 and value.shape[1] > 4:
+            size = (int(value.shape[1]), int(value.shape[0]))
+        elif value.ndim == 4:
+            size = (int(value.shape[2]), int(value.shape[1]))
+        else:
+            size = (int(value.shape[-1]), int(value.shape[-2]))
+        for index in range(count):
+            mask_l = tensor_mask_to_pil(value, index, size)
+            if expected_size is None:
+                expected_size = mask_l.size
+            elif mask_l.size != expected_size:
+                raise ValueError("All input masks must have the same size.")
+            mask_images.append(mask_l)
+    if not mask_images:
+        raise ValueError("Missing input mask.")
+    return mask_images
+
+
+def stack_mask_images(mask_images: list[Image.Image]) -> torch.Tensor:
+    if not mask_images:
+        raise ValueError("Cannot output an empty mask batch.")
+    size = mask_images[0].size
+    normalized = []
+    for mask_l in mask_images:
+        if mask_l.size != size:
+            raise ValueError("All output masks must have the same size.")
+        normalized.append(pil_to_tensor_mask(mask_l))
+    return torch.stack(normalized, dim=0)
+
+
+def bbox_gap(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> int:
+    ax1, ay1, ax2, ay2 = a
+    bx1, by1, bx2, by2 = b
+    dx = max(bx1 - ax2, ax1 - bx2, 0)
+    dy = max(by1 - ay2, ay1 - by2, 0)
+    return max(int(dx), int(dy))
+
+
 def pil_to_tensor_image(image: Image.Image) -> torch.Tensor:
     arr = np.asarray(image.convert("RGB")).astype(np.float32) / 255.0
     return torch.from_numpy(arr)
@@ -606,6 +671,108 @@ def restore_generated_to_model_space(
     return generated.resize(model_image.size, _BICUBIC)
 
 
+class RefineNodeMaskBatchProcess:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "combined_mask": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("MASK",)
+    RETURN_NAMES = ("mask",)
+    INPUT_IS_LIST = True
+    FUNCTION = "process"
+    CATEGORY = "RefineNode/Mask"
+
+    def process(self, mask: torch.Tensor | list[torch.Tensor], combined_mask: bool | list[bool]):
+        combined_mask = bool(combined_mask[0]) if isinstance(combined_mask, list) else bool(combined_mask)
+        mask_images = flatten_mask_input(mask)
+        size = mask_images[0].size
+
+        if combined_mask:
+            return (stack_mask_images([union_mask_images(mask_images, size)]),)
+
+        output_masks = []
+        for mask_l in mask_images:
+            components = split_mask_components(mask_l)
+            if components:
+                output_masks.extend(components)
+
+        if not output_masks:
+            output_masks = [Image.new("L", size, 0)]
+
+        return (stack_mask_images(output_masks),)
+
+
+class RefineNodeCombineNearbyMasks:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "max_gap": ("INT", {"default": 32, "min": 0, "max": 4096}),
+            },
+        }
+
+    RETURN_TYPES = ("MASK",)
+    RETURN_NAMES = ("mask",)
+    INPUT_IS_LIST = True
+    FUNCTION = "combine"
+    CATEGORY = "RefineNode/Mask"
+
+    def combine(self, mask: torch.Tensor | list[torch.Tensor], max_gap: int | list[int]):
+        max_gap = int(max_gap[0]) if isinstance(max_gap, list) else int(max_gap)
+        mask_images = flatten_mask_input(mask)
+        size = mask_images[0].size
+        entries = []
+        for index, mask_l in enumerate(mask_images):
+            bbox = bbox_from_mask_or_none(mask_l)
+            if bbox is not None:
+                entries.append((index, mask_l, bbox))
+
+        if not entries:
+            return (stack_mask_images([Image.new("L", size, 0)]),)
+
+        parents = list(range(len(entries)))
+
+        def find(index: int) -> int:
+            while parents[index] != index:
+                parents[index] = parents[parents[index]]
+                index = parents[index]
+            return index
+
+        def union(a: int, b: int) -> None:
+            root_a = find(a)
+            root_b = find(b)
+            if root_a == root_b:
+                return
+            if root_a < root_b:
+                parents[root_b] = root_a
+            else:
+                parents[root_a] = root_b
+
+        for left_index in range(len(entries)):
+            for right_index in range(left_index + 1, len(entries)):
+                if bbox_gap(entries[left_index][2], entries[right_index][2]) <= max_gap:
+                    union(left_index, right_index)
+
+        groups: dict[int, list[Image.Image]] = {}
+        first_indices: dict[int, int] = {}
+        for entry_index, (original_index, mask_l, _) in enumerate(entries):
+            root = find(entry_index)
+            groups.setdefault(root, []).append(mask_l)
+            first_indices[root] = min(first_indices.get(root, original_index), original_index)
+
+        output_masks = [
+            union_mask_images(groups[root], size)
+            for root in sorted(groups, key=lambda value: first_indices[value])
+        ]
+        return (stack_mask_images(output_masks),)
+
+
 class RefineNodePreprocessMask:
     @classmethod
     def INPUT_TYPES(cls):
@@ -615,7 +782,6 @@ class RefineNodePreprocessMask:
                 "focus_crop": ("BOOLEAN", {"default": True}),
                 "focus_crop_margin": ("INT", {"default": 64, "min": 0, "max": 2048}),
                 "spatial_prompt_source": (["mask", "bbox"], {"default": "mask"}),
-                "combined_mask": ("BOOLEAN", {"default": False}),
             },
             "optional": {
                 "mask": ("MASK",),
@@ -634,7 +800,6 @@ class RefineNodePreprocessMask:
         focus_crop: bool,
         focus_crop_margin: int,
         spatial_prompt_source: str,
-        combined_mask: bool = False,
         mask: torch.Tensor | None = None,
     ):
         image_count = image_batch_size(image)
@@ -650,9 +815,6 @@ class RefineNodePreprocessMask:
             source_image_index: int,
             mask_index: int | None,
             mask_indices: list[int],
-            component_index: int | None,
-            component_count: int,
-            is_combined: bool,
         ) -> None:
             bbox_raw = bbox_from_mask_or_none(mask_l)
             has_region = bbox_raw is not None
@@ -691,10 +853,10 @@ class RefineNodePreprocessMask:
                     "source_image_index": int(source_image_index),
                     "mask_index": None if mask_index is None else int(mask_index),
                     "mask_indices": [int(value) for value in mask_indices],
-                    "component_index": None if component_index is None else int(component_index),
-                    "component_count": int(component_count),
+                    "component_index": None,
+                    "component_count": 0,
                     "group_id": group_id,
-                    "combined_mask": bool(is_combined),
+                    "combined_mask": False,
                 }
             )
 
@@ -707,9 +869,6 @@ class RefineNodePreprocessMask:
                     image_index,
                     None,
                     [],
-                    None,
-                    0,
-                    bool(combined_mask),
                 )
                 continue
 
@@ -721,51 +880,18 @@ class RefineNodePreprocessMask:
                     image_index,
                     None,
                     [],
-                    None,
-                    0,
-                    bool(combined_mask),
-                )
-                continue
-            if combined_mask:
-                append_job(
-                    original,
-                    combine_masks(mask, indices, original.size),
-                    image_index,
-                    None,
-                    indices,
-                    None,
-                    1,
-                    True,
                 )
                 continue
 
             for mask_index in indices:
                 mask_l = tensor_mask_to_pil(mask, mask_index, original.size)
-                components = split_mask_components(mask_l)
-                if not components:
-                    append_job(
-                        original,
-                        mask_l,
-                        image_index,
-                        mask_index,
-                        [mask_index],
-                        None,
-                        0,
-                        False,
-                    )
-                    continue
-                component_count = len(components)
-                for component_index, component_mask in enumerate(components):
-                    append_job(
-                        original,
-                        component_mask,
-                        image_index,
-                        mask_index,
-                        [mask_index],
-                        component_index,
-                        component_count,
-                        False,
-                    )
+                append_job(
+                    original,
+                    mask_l,
+                    image_index,
+                    mask_index,
+                    [mask_index],
+                )
 
         return (
             model_images,
@@ -1042,12 +1168,16 @@ def focus_crop_region(
 
 
 NODE_CLASS_MAPPINGS = {
+    "RefineNodeMaskBatchProcess": RefineNodeMaskBatchProcess,
+    "RefineNodeCombineNearbyMasks": RefineNodeCombineNearbyMasks,
     "RefineNodePreprocessMask": RefineNodePreprocessMask,
     "RefineNodeReferenceImageProcess": RefineNodeReferenceImageProcess,
     "RefineNodePasteBack": RefineNodePasteBack,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
+    "RefineNodeMaskBatchProcess": "RefineNode Mask Batch Process",
+    "RefineNodeCombineNearbyMasks": "RefineNode Combine Nearby Masks",
     "RefineNodePreprocessMask": "RefineNode Preprocess Mask",
     "RefineNodeReferenceImageProcess": "RefineNode Reference Image Process",
     "RefineNodePasteBack": "RefineNode Paste Back",