Delete AILab_InpaintTools.py
This commit is contained in:
AI Lab
@@ -1,146 +0,0 @@
|
||||
# ComfyUI-RMBG
|
||||
#
|
||||
# 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/🎭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) 🖼️🎭",
|
||||
}
|
||||
Reference in New Issue
Block a user