146 lines
5.0 KiB
Python
146 lines
5.0 KiB
Python
# ComfyUI-RMBG v2.6.0
|
|
#
|
|
# AILab Inpaint Tools
|
|
# A collection of specialized nodes for inpainting tasks in ComfyUI.
|
|
# Features a set of utilities for mask processing, latent conditioning, and inpainting workflows.
|
|
#
|
|
# 1. Inpaint Nodes:
|
|
# - AILab_ReferenceLatentMask: A node for inpainting tasks with the Flux Kontext model, using a reference latent and mask for precise region conditioning
|
|
#
|
|
# License: GPL-3.0
|
|
# These nodes are crafted to streamline common image and mask operations within ComfyUI workflows.
|
|
|
|
import torch
|
|
import node_helpers
|
|
|
|
def expand_mask(mask, expand_amount):
|
|
if expand_amount == 0:
|
|
return mask
|
|
|
|
import torch.nn.functional as F
|
|
|
|
binary_mask = (mask > 0.5).float()
|
|
kernel_size = abs(expand_amount) * 2 + 1
|
|
kernel_size = max(3, kernel_size)
|
|
|
|
kernel = torch.ones(1, 1, kernel_size, kernel_size, device=mask.device)
|
|
|
|
if expand_amount > 0:
|
|
expanded = F.conv2d(
|
|
binary_mask.reshape(-1, 1, mask.shape[-2], mask.shape[-1]),
|
|
kernel,
|
|
padding=kernel_size // 2
|
|
)
|
|
result = (expanded > 0).float()
|
|
else:
|
|
eroded = F.conv2d(
|
|
binary_mask.reshape(-1, 1, mask.shape[-2], mask.shape[-1]),
|
|
kernel,
|
|
padding=kernel_size // 2
|
|
)
|
|
result = (eroded >= kernel_size * kernel_size).float()
|
|
|
|
if len(mask.shape) == 3:
|
|
result = result.squeeze(1)
|
|
|
|
return result
|
|
|
|
|
|
def blur_mask(mask, blur_amount):
|
|
if blur_amount == 0:
|
|
return mask
|
|
|
|
import torch.nn.functional as F
|
|
import math
|
|
|
|
x = mask.reshape(-1, 1, mask.shape[-2], mask.shape[-1])
|
|
kernel_size = max(3, math.ceil(blur_amount * 3) * 2 + 1)
|
|
|
|
sigma = blur_amount
|
|
half_kernel = kernel_size // 2
|
|
grid = torch.arange(-half_kernel, half_kernel + 1, device=mask.device).float()
|
|
|
|
gaussian = torch.exp(-0.5 * (grid / sigma) ** 2)
|
|
gaussian = gaussian / gaussian.sum()
|
|
|
|
gaussian_x = gaussian.view(1, 1, 1, kernel_size)
|
|
gaussian_y = gaussian.view(1, 1, kernel_size, 1)
|
|
|
|
blurred = F.conv2d(x, gaussian_x, padding=(0, half_kernel))
|
|
blurred = F.conv2d(blurred, gaussian_y, padding=(half_kernel, 0))
|
|
|
|
if len(mask.shape) == 3:
|
|
blurred = blurred.squeeze(1)
|
|
|
|
return blurred
|
|
|
|
class AILab_ReferenceLatentMask:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
tooltips = {
|
|
"conditioning": "Base conditioning input for inpainting task",
|
|
"latent": "Encoded latent from VAE",
|
|
"mask": "Area to inpaint (white regions)",
|
|
"expand": "Grow mask (+) or shrink mask (-)",
|
|
"blur": "Soften mask edges",
|
|
"mask_only": "Only generate content in masked area"
|
|
}
|
|
|
|
return {
|
|
"required": {
|
|
"conditioning": ("CONDITIONING", {"tooltip": tooltips["conditioning"]}),
|
|
"latent": ("LATENT", {"tooltip": tooltips["latent"]}),
|
|
"mask": ("MASK", {"tooltip": tooltips["mask"]}),
|
|
"expand": ("INT", {"default": 5, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["expand"]}),
|
|
"blur": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 64.0, "step": 0.1, "tooltip": tooltips["blur"]}),
|
|
"mask_only": ("BOOLEAN", {"default": True, "tooltip": tooltips["mask_only"]}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "LATENT", "MASK")
|
|
RETURN_NAMES = ("CONDITIONING", "LATENT", "MASK")
|
|
FUNCTION = "prepare_inpaint_conditioning"
|
|
CATEGORY = "🧪AILab/🧽RMBG/🎭Inpaint"
|
|
|
|
def add_latent_to_conditioning(self, conditioning, latent=None):
|
|
if latent is not None:
|
|
return node_helpers.conditioning_set_values(
|
|
conditioning,
|
|
{"reference_latents": [latent["samples"]]},
|
|
append=True
|
|
)
|
|
return conditioning
|
|
|
|
def prepare_inpaint_conditioning(self, conditioning, latent, mask, expand=5, blur=3.0, mask_only=True):
|
|
processed_mask = mask
|
|
|
|
if expand != 0:
|
|
processed_mask = expand_mask(processed_mask, expand)
|
|
|
|
if blur > 0:
|
|
processed_mask = blur_mask(processed_mask, blur)
|
|
|
|
modified_cond = node_helpers.conditioning_set_values(
|
|
conditioning,
|
|
{
|
|
"concat_latent_image": latent["samples"],
|
|
"concat_mask": processed_mask
|
|
}
|
|
)
|
|
|
|
final_cond = self.add_latent_to_conditioning(modified_cond, latent)
|
|
|
|
output_latent = {"samples": latent["samples"]}
|
|
if mask_only:
|
|
output_latent["noise_mask"] = processed_mask
|
|
|
|
return (final_cond, output_latent, processed_mask)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AILab_ReferenceLatentMask": AILab_ReferenceLatentMask,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AILab_ReferenceLatentMask": "Reference Latent Mask (RMBG) 🖼️🎭",
|
|
} |