zisb-comfyui-texture-packer/nodes.py

import torch
import torch.nn.functional as F


class AlphaEdgeSolidify:
    """
    Solidifies/bleeds edge colours outward from an alpha mask to create
    a fully solid RGB texture. Prevents edge halos and mipmap bleeding
    in game engines.
    """

    DESCRIPTION = "NOTE: ComfyUI auto-inverts masks from Load Image. You may need to use an 'Invert Mask' node before connecting to this node."

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "mask": ("MASK",),
                "bleed_pixels": ("INT", {
                    "default": 64,
                    "min": 1,
                    "max": 256,
                    "step": 1,
                    "display": "number"
                }),
                "mask_threshold": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "number"
                }),
            },
            "optional": {
                "kernel_size": ("INT", {
                    "default": 3,
                    "min": 3,
                    "max": 7,
                    "step": 2,
                    "display": "number"
                }),
                "mode": (["basic", "smooth"],),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "solidify"
    CATEGORY = "image/postprocessing"

    def solidify(self, image, mask, bleed_pixels, mask_threshold, kernel_size=3, mode="basic"):
        """
        Solidify the RGB image by bleeding edge colors into transparent regions.

        Args:
            image: RGB tensor [B, H, W, C] in range [0, 1]
            mask: Alpha mask tensor [B, H, W] or [H, W] in range [0, 1]
            bleed_pixels: Number of pixels to bleed outward
            mask_threshold: Threshold for what counts as "inside" the mask
            kernel_size: Size of the dilation kernel (3, 5, or 7)
            mode: "basic" for fast dilation, "smooth" for weighted average

        Returns:
            Solid RGB image tensor [B, H, W, C]
        """
        device = image.device
        batch_size = image.shape[0]

        # Ensure mask has batch dimension
        if mask.dim() == 2:
            mask = mask.unsqueeze(0)

        # Handle mask batch size mismatch
        if mask.shape[0] == 1 and batch_size > 1:
            mask = mask.repeat(batch_size, 1, 1)

        results = []

        for b in range(batch_size):
            img = image[b]  # [H, W, C]
            msk = mask[b]   # [H, W]

            # Convert to [C, H, W] for processing
            img = img.permute(2, 0, 1)  # [C, H, W]

            # Create binary mask based on threshold
            binary_mask = (msk > mask_threshold).float()  # [H, W]

            # Solidify the image
            if mode == "smooth":
                result = self._solidify_smooth(img, binary_mask, bleed_pixels, kernel_size, device)
            else:
                result = self._solidify_basic(img, binary_mask, bleed_pixels, kernel_size, device)

            # Convert back to [H, W, C]
            result = result.permute(1, 2, 0)
            results.append(result)

        return (torch.stack(results),)

    def _solidify_basic(self, img, mask, iterations, kernel_size, device):
        """
        Basic dilation: iteratively expand edge colors outward.
        Each iteration, empty pixels adjacent to filled pixels take the
        average color of their filled neighbors.
        """
        C, H, W = img.shape

        # Work with copies
        result = img.clone()
        filled = mask.clone()

        # Create dilation kernel
        kernel = torch.ones(1, 1, kernel_size, kernel_size, device=device)
        pad = kernel_size // 2

        for _ in range(iterations):
            # Find the current edge: pixels that are empty but adjacent to filled
            filled_expanded = F.conv2d(
                filled.unsqueeze(0).unsqueeze(0),
                kernel,
                padding=pad
            ).squeeze(0).squeeze(0)

            # Edge pixels: not yet filled, but have filled neighbors
            edge = (filled == 0) & (filled_expanded > 0)

            if not edge.any():
                break  # No more pixels to fill

            # For each edge pixel, compute average of filled neighbors
            for c in range(C):
                channel = result[c]

                # Sum of filled neighbor values
                neighbor_sum = F.conv2d(
                    (channel * filled).unsqueeze(0).unsqueeze(0),
                    kernel,
                    padding=pad
                ).squeeze(0).squeeze(0)

                # Count of filled neighbors
                neighbor_count = F.conv2d(
                    filled.unsqueeze(0).unsqueeze(0),
                    kernel,
                    padding=pad
                ).squeeze(0).squeeze(0)

                # Average (avoid division by zero)
                avg = neighbor_sum / neighbor_count.clamp(min=1)

                # Update edge pixels
                result[c] = torch.where(edge, avg, channel)

            # Mark edge pixels as now filled
            filled = filled + edge.float()

        return result

    def _solidify_smooth(self, img, mask, iterations, kernel_size, device):
        """
        Smooth dilation: uses distance-weighted averaging for smoother results.
        Better for large bleed distances.
        """
        C, H, W = img.shape

        result = img.clone()
        filled = mask.clone()

        # Create weighted kernel (higher weight for closer pixels)
        k = kernel_size
        center = k // 2
        kernel = torch.zeros(k, k, device=device)
        for i in range(k):
            for j in range(k):
                dist = ((i - center) ** 2 + (j - center) ** 2) ** 0.5
                if dist == 0:
                    kernel[i, j] = 1.0
                else:
                    kernel[i, j] = 1.0 / (1.0 + dist)

        kernel = kernel.unsqueeze(0).unsqueeze(0)
        pad = k // 2

        for _ in range(iterations):
            # Find edge pixels
            filled_check = F.conv2d(
                filled.unsqueeze(0).unsqueeze(0),
                torch.ones(1, 1, k, k, device=device),
                padding=pad
            ).squeeze(0).squeeze(0)

            edge = (filled == 0) & (filled_check > 0)

            if not edge.any():
                break

            for c in range(C):
                channel = result[c]

                # Weighted sum of filled neighbors
                weighted_sum = F.conv2d(
                    (channel * filled).unsqueeze(0).unsqueeze(0),
                    kernel,
                    padding=pad
                ).squeeze(0).squeeze(0)

                # Sum of weights for filled neighbors
                weight_sum = F.conv2d(
                    filled.unsqueeze(0).unsqueeze(0),
                    kernel,
                    padding=pad
                ).squeeze(0).squeeze(0)

                avg = weighted_sum / weight_sum.clamp(min=1e-8)
                result[c] = torch.where(edge, avg, channel)

            filled = filled + edge.float()

        return result


# Node class mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
    "AlphaEdgeSolidify": AlphaEdgeSolidify
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "AlphaEdgeSolidify": "Alpha Edge Solidify -> RGB"
}