diff --git a/nodes/image_processing.py b/nodes/image_processing.py index 8ee5667..29ac47e 100644 --- a/nodes/image_processing.py +++ b/nodes/image_processing.py @@ -354,7 +354,11 @@ class ColoredImage: "color": ("COLOR",), "width": ("INT", {"default": 512, "min": 16, "max": 8160}), "height": ("INT", {"default": 512, "min": 16, "max": 8160}), - } + }, + "optional": { + "foreground_image": ("IMAGE",), + "foreground_mask": ("MASK",), + }, } CATEGORY = "mtb/generate" @@ -363,12 +367,46 @@ class ColoredImage: FUNCTION = "render_img" - def render_img(self, color, width, height): - image = Image.new("RGB", (width, height), color=color) + def render_img( + self, color, width, height, foreground_image=None, foreground_mask=None + ): + image = Image.new("RGBA", (width, height), color=color) + output = [] + if foreground_image is not None: + if foreground_mask is None: + fg_images = tensor2pil(foreground_image) + for img in fg_images: + if image.size != img.size: + raise ValueError( + f"Dimension mismatch: image {image.size}, img {img.size}" + ) - image = pil2tensor(image) + if img.mode != "RGBA": + raise ValueError( + f"Foreground image must be in 'RGBA' mode when no mask is provided, got {img.mode}" + ) - return (image,) + output.append(Image.alpha_composite(image, img).convert("RGB")) + + elif foreground_image.size[0] != foreground_mask.size[0]: + raise ValueError("Foreground image and mask must have same batch size") + else: + fg_images = tensor2pil(foreground_image) + fg_masks = tensor2pil(foreground_mask) + output.extend( + Image.composite( + fg_image.convert("RGBA"), + image, + fg_mask, + ).convert("RGB") + for fg_image, fg_mask in zip(fg_images, fg_masks) + ) + elif foreground_mask is not None: + log.warn("Mask ignored because no foreground image is given") + + output = pil2tensor(output) + + return (output,) class ImagePremultiply: