Skip to content

Commit

Permalink
feat: 🚀 add optional inputs to colored image
Browse files Browse the repository at this point in the history
  • Loading branch information
melMass committed Oct 20, 2023
1 parent 255ac03 commit 049983d
Showing 1 changed file with 43 additions and 5 deletions.
48 changes: 43 additions & 5 deletions nodes/image_processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -354,7 +354,11 @@ def INPUT_TYPES(cls):
"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"
Expand All @@ -363,12 +367,46 @@ def INPUT_TYPES(cls):

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}"
)

if img.mode != "RGBA":
raise ValueError(
f"Foreground image must be in 'RGBA' mode when no mask is provided, got {img.mode}"
)

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")

image = pil2tensor(image)
output = pil2tensor(output)

return (image,)
return (output,)


class ImagePremultiply:
Expand Down

0 comments on commit 049983d

Please sign in to comment.