316 lines
9.2 KiB
Python
316 lines
9.2 KiB
Python
import torch
|
|
|
|
from .Utils import create_rgba_image
|
|
|
|
|
|
class ImageContainer:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"width": ("INT", {
|
|
"default": 512,
|
|
"min": 1,
|
|
"step": 1
|
|
}),
|
|
"height": ("INT", {
|
|
"default": 512,
|
|
"min": 1,
|
|
"step": 1
|
|
}),
|
|
"red": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"green": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"blue": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"alpha": ("FLOAT", {
|
|
"default": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "image/container"
|
|
|
|
def node(self, width, height, red, green, blue, alpha):
|
|
return (create_rgba_image(width, height, (red, green, blue, int(alpha * 255))).image_to_tensor().unsqueeze(0),)
|
|
|
|
|
|
class ImageContainerInheritanceAdd:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"add_width": ("INT", {
|
|
"default": 0,
|
|
"step": 1
|
|
}),
|
|
"add_height": ("INT", {
|
|
"default": 0,
|
|
"step": 1
|
|
}),
|
|
"red": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"green": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"blue": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"alpha": ("FLOAT", {
|
|
"default": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"method": (["single", "for_each"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "image/container"
|
|
|
|
def node(self, images, add_width, add_height, red, green, blue, alpha, method):
|
|
width, height = images[0, :, :, 0].shape
|
|
|
|
width = width + add_width
|
|
height = height + add_height
|
|
|
|
image = create_rgba_image(width, height, (red, green, blue, int(alpha * 255))).image_to_tensor()
|
|
|
|
if method == "single":
|
|
return (image.unsqueeze(0),)
|
|
else:
|
|
length = len(images)
|
|
|
|
images = torch.zeros(length, height, width, 4)
|
|
images[:, :, :] = image
|
|
return (images,)
|
|
|
|
|
|
class ImageContainerInheritanceScale:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"scale_width": ("FLOAT", {
|
|
"default": 1.0,
|
|
"step": 0.1
|
|
}),
|
|
"scale_height": ("FLOAT", {
|
|
"default": 1.0,
|
|
"step": 0.1
|
|
}),
|
|
"red": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"green": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"blue": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"alpha": ("FLOAT", {
|
|
"default": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"method": (["single", "for_each"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "image/container"
|
|
|
|
def node(self, images, scale_width, scale_height, red, green, blue, alpha, method):
|
|
height, width = images[0, :, :, 0].shape
|
|
|
|
width = int((width * scale_width) - width)
|
|
height = int((height * scale_height) - height)
|
|
|
|
return ImageContainerInheritanceAdd() \
|
|
.node(images, width, height, red, green, blue, alpha, method)
|
|
|
|
|
|
class ImageContainerInheritanceMax:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images_a": ("IMAGE",),
|
|
"images_b": ("IMAGE",),
|
|
"red": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"green": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"blue": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"alpha": ("FLOAT", {
|
|
"default": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"method": (["single", "for_each_pair", "for_each_matrix"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "image/container"
|
|
|
|
def node(self, images_a, images_b, red, green, blue, alpha, method):
|
|
img_a_height, img_a_width = images_a[0, :, :, 0].shape
|
|
img_b_height, img_b_width = images_b[0, :, :, 0].shape
|
|
|
|
width = max(img_a_width, img_b_width)
|
|
height = max(img_a_height, img_b_height)
|
|
|
|
image = create_rgba_image(width, height, (red, green, blue, int(alpha * 255))).image_to_tensor()
|
|
|
|
if method == "single":
|
|
return (image.unsqueeze(0),)
|
|
elif method == "for_each_pair":
|
|
length = len(images_a)
|
|
images = torch.zeros(length, height, width, 4)
|
|
else:
|
|
length = len(images_a) * len(images_b)
|
|
images = torch.zeros(length, height, width, 4)
|
|
|
|
images[:, :, :] = image
|
|
return (images,)
|
|
|
|
|
|
class ImageContainerInheritanceSum:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images_a": ("IMAGE",),
|
|
"images_b": ("IMAGE",),
|
|
"red": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"green": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"blue": ("INT", {
|
|
"default": 255,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"alpha": ("FLOAT", {
|
|
"default": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"container_size_type": (["sum", "sum_width", "sum_height"],),
|
|
"method": (["single", "for_each_pair", "for_each_matrix"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "node"
|
|
CATEGORY = "image/container"
|
|
|
|
def node(self, images_a, images_b, red, green, blue, alpha, container_size_type, method):
|
|
img_a_height, img_a_width = images_a[0, :, :, 0].shape
|
|
img_b_height, img_b_width = images_b[0, :, :, 0].shape
|
|
|
|
if container_size_type == "sum":
|
|
width = img_a_width + img_b_width
|
|
height = img_a_height + img_b_height
|
|
elif container_size_type == "sum_width":
|
|
if img_a_height != img_b_height:
|
|
raise ValueError()
|
|
|
|
width = img_a_width + img_b_width
|
|
height = img_a_height
|
|
elif container_size_type == "sum_height":
|
|
if img_a_width != img_b_width:
|
|
raise ValueError()
|
|
|
|
width = img_a_width
|
|
height = img_a_height + img_b_height
|
|
else:
|
|
raise ValueError()
|
|
|
|
image = create_rgba_image(width, height, (red, green, blue, int(alpha * 255))).image_to_tensor()
|
|
|
|
if method == "single":
|
|
return (image.unsqueeze(0),)
|
|
elif method == "for_each_pair":
|
|
length = len(images_a)
|
|
images = torch.zeros(length, height, width, 4)
|
|
else:
|
|
length = len(images_a) * len(images_b)
|
|
images = torch.zeros(length, height, width, 4)
|
|
|
|
images[:, :, :] = image
|
|
return (images,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageContainer": ImageContainer,
|
|
"ImageContainerInheritanceAdd": ImageContainerInheritanceAdd,
|
|
"ImageContainerInheritanceScale": ImageContainerInheritanceScale,
|
|
"ImageContainerInheritanceMax": ImageContainerInheritanceMax,
|
|
"ImageContainerInheritanceSum": ImageContainerInheritanceSum
|
|
}
|