3d039591fe
Previously before this commit, credits are already in entry and licenses are already in root. This commit will make info clearer.
207 lines
7.9 KiB
Python
207 lines
7.9 KiB
Python
# Taken from https://github.com/comfyanonymous/ComfyUI
|
|
# This file is only for reference, and not used in the backend or runtime.
|
|
|
|
|
|
import numpy as np
|
|
import torch
|
|
import ldm_patched.modules.utils
|
|
from enum import Enum
|
|
|
|
def resize_mask(mask, shape):
|
|
return torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[0], shape[1]), mode="bilinear").squeeze(1)
|
|
|
|
class PorterDuffMode(Enum):
|
|
ADD = 0
|
|
CLEAR = 1
|
|
DARKEN = 2
|
|
DST = 3
|
|
DST_ATOP = 4
|
|
DST_IN = 5
|
|
DST_OUT = 6
|
|
DST_OVER = 7
|
|
LIGHTEN = 8
|
|
MULTIPLY = 9
|
|
OVERLAY = 10
|
|
SCREEN = 11
|
|
SRC = 12
|
|
SRC_ATOP = 13
|
|
SRC_IN = 14
|
|
SRC_OUT = 15
|
|
SRC_OVER = 16
|
|
XOR = 17
|
|
|
|
|
|
def porter_duff_composite(src_image: torch.Tensor, src_alpha: torch.Tensor, dst_image: torch.Tensor, dst_alpha: torch.Tensor, mode: PorterDuffMode):
|
|
if mode == PorterDuffMode.ADD:
|
|
out_alpha = torch.clamp(src_alpha + dst_alpha, 0, 1)
|
|
out_image = torch.clamp(src_image + dst_image, 0, 1)
|
|
elif mode == PorterDuffMode.CLEAR:
|
|
out_alpha = torch.zeros_like(dst_alpha)
|
|
out_image = torch.zeros_like(dst_image)
|
|
elif mode == PorterDuffMode.DARKEN:
|
|
out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha
|
|
out_image = (1 - dst_alpha) * src_image + (1 - src_alpha) * dst_image + torch.min(src_image, dst_image)
|
|
elif mode == PorterDuffMode.DST:
|
|
out_alpha = dst_alpha
|
|
out_image = dst_image
|
|
elif mode == PorterDuffMode.DST_ATOP:
|
|
out_alpha = src_alpha
|
|
out_image = src_alpha * dst_image + (1 - dst_alpha) * src_image
|
|
elif mode == PorterDuffMode.DST_IN:
|
|
out_alpha = src_alpha * dst_alpha
|
|
out_image = dst_image * src_alpha
|
|
elif mode == PorterDuffMode.DST_OUT:
|
|
out_alpha = (1 - src_alpha) * dst_alpha
|
|
out_image = (1 - src_alpha) * dst_image
|
|
elif mode == PorterDuffMode.DST_OVER:
|
|
out_alpha = dst_alpha + (1 - dst_alpha) * src_alpha
|
|
out_image = dst_image + (1 - dst_alpha) * src_image
|
|
elif mode == PorterDuffMode.LIGHTEN:
|
|
out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha
|
|
out_image = (1 - dst_alpha) * src_image + (1 - src_alpha) * dst_image + torch.max(src_image, dst_image)
|
|
elif mode == PorterDuffMode.MULTIPLY:
|
|
out_alpha = src_alpha * dst_alpha
|
|
out_image = src_image * dst_image
|
|
elif mode == PorterDuffMode.OVERLAY:
|
|
out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha
|
|
out_image = torch.where(2 * dst_image < dst_alpha, 2 * src_image * dst_image,
|
|
src_alpha * dst_alpha - 2 * (dst_alpha - src_image) * (src_alpha - dst_image))
|
|
elif mode == PorterDuffMode.SCREEN:
|
|
out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha
|
|
out_image = src_image + dst_image - src_image * dst_image
|
|
elif mode == PorterDuffMode.SRC:
|
|
out_alpha = src_alpha
|
|
out_image = src_image
|
|
elif mode == PorterDuffMode.SRC_ATOP:
|
|
out_alpha = dst_alpha
|
|
out_image = dst_alpha * src_image + (1 - src_alpha) * dst_image
|
|
elif mode == PorterDuffMode.SRC_IN:
|
|
out_alpha = src_alpha * dst_alpha
|
|
out_image = src_image * dst_alpha
|
|
elif mode == PorterDuffMode.SRC_OUT:
|
|
out_alpha = (1 - dst_alpha) * src_alpha
|
|
out_image = (1 - dst_alpha) * src_image
|
|
elif mode == PorterDuffMode.SRC_OVER:
|
|
out_alpha = src_alpha + (1 - src_alpha) * dst_alpha
|
|
out_image = src_image + (1 - src_alpha) * dst_image
|
|
elif mode == PorterDuffMode.XOR:
|
|
out_alpha = (1 - dst_alpha) * src_alpha + (1 - src_alpha) * dst_alpha
|
|
out_image = (1 - dst_alpha) * src_image + (1 - src_alpha) * dst_image
|
|
else:
|
|
out_alpha = None
|
|
out_image = None
|
|
return out_image, out_alpha
|
|
|
|
|
|
class PorterDuffImageComposite:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"source": ("IMAGE",),
|
|
"source_alpha": ("MASK",),
|
|
"destination": ("IMAGE",),
|
|
"destination_alpha": ("MASK",),
|
|
"mode": ([mode.name for mode in PorterDuffMode], {"default": PorterDuffMode.DST.name}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "composite"
|
|
CATEGORY = "mask/compositing"
|
|
|
|
def composite(self, source: torch.Tensor, source_alpha: torch.Tensor, destination: torch.Tensor, destination_alpha: torch.Tensor, mode):
|
|
batch_size = min(len(source), len(source_alpha), len(destination), len(destination_alpha))
|
|
out_images = []
|
|
out_alphas = []
|
|
|
|
for i in range(batch_size):
|
|
src_image = source[i]
|
|
dst_image = destination[i]
|
|
|
|
assert src_image.shape[2] == dst_image.shape[2] # inputs need to have same number of channels
|
|
|
|
src_alpha = source_alpha[i].unsqueeze(2)
|
|
dst_alpha = destination_alpha[i].unsqueeze(2)
|
|
|
|
if dst_alpha.shape[:2] != dst_image.shape[:2]:
|
|
upscale_input = dst_alpha.unsqueeze(0).permute(0, 3, 1, 2)
|
|
upscale_output = ldm_patched.modules.utils.common_upscale(upscale_input, dst_image.shape[1], dst_image.shape[0], upscale_method='bicubic', crop='center')
|
|
dst_alpha = upscale_output.permute(0, 2, 3, 1).squeeze(0)
|
|
if src_image.shape != dst_image.shape:
|
|
upscale_input = src_image.unsqueeze(0).permute(0, 3, 1, 2)
|
|
upscale_output = ldm_patched.modules.utils.common_upscale(upscale_input, dst_image.shape[1], dst_image.shape[0], upscale_method='bicubic', crop='center')
|
|
src_image = upscale_output.permute(0, 2, 3, 1).squeeze(0)
|
|
if src_alpha.shape != dst_alpha.shape:
|
|
upscale_input = src_alpha.unsqueeze(0).permute(0, 3, 1, 2)
|
|
upscale_output = ldm_patched.modules.utils.common_upscale(upscale_input, dst_alpha.shape[1], dst_alpha.shape[0], upscale_method='bicubic', crop='center')
|
|
src_alpha = upscale_output.permute(0, 2, 3, 1).squeeze(0)
|
|
|
|
out_image, out_alpha = porter_duff_composite(src_image, src_alpha, dst_image, dst_alpha, PorterDuffMode[mode])
|
|
|
|
out_images.append(out_image)
|
|
out_alphas.append(out_alpha.squeeze(2))
|
|
|
|
result = (torch.stack(out_images), torch.stack(out_alphas))
|
|
return result
|
|
|
|
|
|
class SplitImageWithAlpha:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "mask/compositing"
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "split_image_with_alpha"
|
|
|
|
def split_image_with_alpha(self, image: torch.Tensor):
|
|
out_images = [i[:,:,:3] for i in image]
|
|
out_alphas = [i[:,:,3] if i.shape[2] > 3 else torch.ones_like(i[:,:,0]) for i in image]
|
|
result = (torch.stack(out_images), 1.0 - torch.stack(out_alphas))
|
|
return result
|
|
|
|
|
|
class JoinImageWithAlpha:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"alpha": ("MASK",),
|
|
}
|
|
}
|
|
|
|
CATEGORY = "mask/compositing"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "join_image_with_alpha"
|
|
|
|
def join_image_with_alpha(self, image: torch.Tensor, alpha: torch.Tensor):
|
|
batch_size = min(len(image), len(alpha))
|
|
out_images = []
|
|
|
|
alpha = 1.0 - resize_mask(alpha, image.shape[1:])
|
|
for i in range(batch_size):
|
|
out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2))
|
|
|
|
result = (torch.stack(out_images),)
|
|
return result
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PorterDuffImageComposite": PorterDuffImageComposite,
|
|
"SplitImageWithAlpha": SplitImageWithAlpha,
|
|
"JoinImageWithAlpha": JoinImageWithAlpha,
|
|
}
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PorterDuffImageComposite": "Porter-Duff Image Composite",
|
|
"SplitImageWithAlpha": "Split Image with Alpha",
|
|
"JoinImageWithAlpha": "Join Image with Alpha",
|
|
}
|