brunvelop/ComfyUI
2
1import numpy as np2import scipy.ndimage3import torch4import comfy.utils5 6from nodes import MAX_RESOLUTION7 8def composite(destination, source, x, y, mask = None, multiplier = 8, resize_source = False):9 if resize_source:10 source = torch.nn.functional.interpolate(source, size=(destination.shape[2], destination.shape[3]), mode="bilinear")11 12 source = comfy.utils.repeat_to_batch_size(source, destination.shape[0])13 14 x = max(-source.shape[3] * multiplier, min(x, destination.shape[3] * multiplier))15 y = max(-source.shape[2] * multiplier, min(y, destination.shape[2] * multiplier))16 17 left, top = (x // multiplier, y // multiplier)18 right, bottom = (left + source.shape[3], top + source.shape[2],)19 20 if mask is None:21 mask = torch.ones_like(source)22 else:23 mask = mask.clone()24 mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[2], source.shape[3]), mode="bilinear")25 mask = comfy.utils.repeat_to_batch_size(mask, source.shape[0])26 27 # calculate the bounds of the source that will be overlapping the destination28 # this prevents the source trying to overwrite latent pixels that are out of bounds29 # of the destination30 visible_width, visible_height = (destination.shape[3] - left + min(0, x), destination.shape[2] - top + min(0, y),)31 32 mask = mask[:, :, :visible_height, :visible_width]33 inverse_mask = torch.ones_like(mask) - mask34 35 source_portion = mask * source[:, :, :visible_height, :visible_width]36 destination_portion = inverse_mask * destination[:, :, top:bottom, left:right]37 38 destination[:, :, top:bottom, left:right] = source_portion + destination_portion39 return destination40 41class LatentCompositeMasked:42 @classmethod43 def INPUT_TYPES(s):44 return {45 "required": {46 "destination": ("LATENT",),47 "source": ("LATENT",),48 "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),49 "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),50 "resize_source": ("BOOLEAN", {"default": False}),51 },52 "optional": {53 "mask": ("MASK",),54 }55 }56 RETURN_TYPES = ("LATENT",)57 FUNCTION = "composite"58 59 CATEGORY = "latent"60 61 def composite(self, destination, source, x, y, resize_source, mask = None):62 output = destination.copy()63 destination = destination["samples"].clone()64 source = source["samples"]65 output["samples"] = composite(destination, source, x, y, mask, 8, resize_source)66 return (output,)67 68class ImageCompositeMasked:69 @classmethod70 def INPUT_TYPES(s):71 return {72 "required": {73 "destination": ("IMAGE",),74 "source": ("IMAGE",),75 "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),76 "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),77 "resize_source": ("BOOLEAN", {"default": False}),78 },79 "optional": {80 "mask": ("MASK",),81 }82 }83 RETURN_TYPES = ("IMAGE",)84 FUNCTION = "composite"85 86 CATEGORY = "image"87 88 def composite(self, destination, source, x, y, resize_source, mask = None):89 destination = destination.clone().movedim(-1, 1)90 output = composite(destination, source.movedim(-1, 1), x, y, mask, 1, resize_source).movedim(1, -1)91 return (output,)92 93class MaskToImage:94 @classmethod95 def INPUT_TYPES(s):96 return {97 "required": {98 "mask": ("MASK",),99 }100 }101 102 CATEGORY = "mask"103 104 RETURN_TYPES = ("IMAGE",)105 FUNCTION = "mask_to_image"106 107 def mask_to_image(self, mask):108 result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)109 return (result,)110 111class ImageToMask:112 @classmethod113 def INPUT_TYPES(s):114 return {115 "required": {116 "image": ("IMAGE",),117 "channel": (["red", "green", "blue", "alpha"],),118 }119 }120 121 CATEGORY = "mask"122 123 RETURN_TYPES = ("MASK",)124 FUNCTION = "image_to_mask"125 126 def image_to_mask(self, image, channel):127 channels = ["red", "green", "blue", "alpha"]128 mask = image[:, :, :, channels.index(channel)]129 return (mask,)130 131class ImageColorToMask:132 @classmethod133 def INPUT_TYPES(s):134 return {135 "required": {136 "image": ("IMAGE",),137 "color": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFF, "step": 1, "display": "color"}),138 }139 }140 141 CATEGORY = "mask"142 143 RETURN_TYPES = ("MASK",)144 FUNCTION = "image_to_mask"145 146 def image_to_mask(self, image, color):147 temp = (torch.clamp(image, 0, 1.0) * 255.0).round().to(torch.int)148 temp = torch.bitwise_left_shift(temp[:,:,:,0], 16) + torch.bitwise_left_shift(temp[:,:,:,1], 8) + temp[:,:,:,2]149 mask = torch.where(temp == color, 255, 0).float()150 return (mask,)151 152class SolidMask:153 @classmethod154 def INPUT_TYPES(cls):155 return {156 "required": {157 "value": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),158 "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),159 "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),160 }161 }162 163 CATEGORY = "mask"164 165 RETURN_TYPES = ("MASK",)166 167 FUNCTION = "solid"168 169 def solid(self, value, width, height):170 out = torch.full((1, height, width), value, dtype=torch.float32, device="cpu")171 return (out,)172 173class InvertMask:174 @classmethod175 def INPUT_TYPES(cls):176 return {177 "required": {178 "mask": ("MASK",),179 }180 }181 182 CATEGORY = "mask"183 184 RETURN_TYPES = ("MASK",)185 186 FUNCTION = "invert"187 188 def invert(self, mask):189 out = 1.0 - mask190 return (out,)191 192class CropMask:193 @classmethod194 def INPUT_TYPES(cls):195 return {196 "required": {197 "mask": ("MASK",),198 "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),199 "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),200 "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),201 "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),202 }203 }204 205 CATEGORY = "mask"206 207 RETURN_TYPES = ("MASK",)208 209 FUNCTION = "crop"210 211 def crop(self, mask, x, y, width, height):212 mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))213 out = mask[:, y:y + height, x:x + width]214 return (out,)215 216class MaskComposite:217 @classmethod218 def INPUT_TYPES(cls):219 return {220 "required": {221 "destination": ("MASK",),222 "source": ("MASK",),223 "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),224 "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),225 "operation": (["multiply", "add", "subtract", "and", "or", "xor"],),226 }227 }228 229 CATEGORY = "mask"230 231 RETURN_TYPES = ("MASK",)232 233 FUNCTION = "combine"234 235 def combine(self, destination, source, x, y, operation):236 output = destination.reshape((-1, destination.shape[-2], destination.shape[-1])).clone()237 source = source.reshape((-1, source.shape[-2], source.shape[-1]))238 239 left, top = (x, y,)240 right, bottom = (min(left + source.shape[-1], destination.shape[-1]), min(top + source.shape[-2], destination.shape[-2]))241 visible_width, visible_height = (right - left, bottom - top,)242 243 source_portion = source[:, :visible_height, :visible_width]244 destination_portion = destination[:, top:bottom, left:right]245 246 if operation == "multiply":247 output[:, top:bottom, left:right] = destination_portion * source_portion248 elif operation == "add":249 output[:, top:bottom, left:right] = destination_portion + source_portion250 elif operation == "subtract":251 output[:, top:bottom, left:right] = destination_portion - source_portion252 elif operation == "and":253 output[:, top:bottom, left:right] = torch.bitwise_and(destination_portion.round().bool(), source_portion.round().bool()).float()254 elif operation == "or":255 output[:, top:bottom, left:right] = torch.bitwise_or(destination_portion.round().bool(), source_portion.round().bool()).float()256 elif operation == "xor":257 output[:, top:bottom, left:right] = torch.bitwise_xor(destination_portion.round().bool(), source_portion.round().bool()).float()258 259 output = torch.clamp(output, 0.0, 1.0)260 261 return (output,)262 263class FeatherMask:264 @classmethod265 def INPUT_TYPES(cls):266 return {267 "required": {268 "mask": ("MASK",),269 "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),270 "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),271 "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),272 "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),273 }274 }275 276 CATEGORY = "mask"277 278 RETURN_TYPES = ("MASK",)279 280 FUNCTION = "feather"281 282 def feather(self, mask, left, top, right, bottom):283 output = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).clone()284 285 left = min(left, output.shape[-1])286 right = min(right, output.shape[-1])287 top = min(top, output.shape[-2])288 bottom = min(bottom, output.shape[-2])289 290 for x in range(left):291 feather_rate = (x + 1.0) / left292 output[:, :, x] *= feather_rate293 294 for x in range(right):295 feather_rate = (x + 1) / right296 output[:, :, -x] *= feather_rate297 298 for y in range(top):299 feather_rate = (y + 1) / top300 output[:, y, :] *= feather_rate301 302 for y in range(bottom):303 feather_rate = (y + 1) / bottom304 output[:, -y, :] *= feather_rate305 306 return (output,)307 308class GrowMask:309 @classmethod310 def INPUT_TYPES(cls):311 return {312 "required": {313 "mask": ("MASK",),314 "expand": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}),315 "tapered_corners": ("BOOLEAN", {"default": True}),316 },317 }318 319 CATEGORY = "mask"320 321 RETURN_TYPES = ("MASK",)322 323 FUNCTION = "expand_mask"324 325 def expand_mask(self, mask, expand, tapered_corners):326 c = 0 if tapered_corners else 1327 kernel = np.array([[c, 1, c],328 [1, 1, 1],329 [c, 1, c]])330 mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))331 out = []332 for m in mask:333 output = m.numpy()334 for _ in range(abs(expand)):335 if expand < 0:336 output = scipy.ndimage.grey_erosion(output, footprint=kernel)337 else:338 output = scipy.ndimage.grey_dilation(output, footprint=kernel)339 output = torch.from_numpy(output)340 out.append(output)341 return (torch.stack(out, dim=0),)342 343 344 345NODE_CLASS_MAPPINGS = {346 "LatentCompositeMasked": LatentCompositeMasked,347 "ImageCompositeMasked": ImageCompositeMasked,348 "MaskToImage": MaskToImage,349 "ImageToMask": ImageToMask,350 "ImageColorToMask": ImageColorToMask,351 "SolidMask": SolidMask,352 "InvertMask": InvertMask,353 "CropMask": CropMask,354 "MaskComposite": MaskComposite,355 "FeatherMask": FeatherMask,356 "GrowMask": GrowMask,357}358 359NODE_DISPLAY_NAME_MAPPINGS = {360 "ImageToMask": "Convert Image to Mask",361 "MaskToImage": "Convert Mask to Image",362}363