Add files via upload
This commit is contained in:
@@ -0,0 +1,54 @@
|
|||||||
|
import torch
|
||||||
|
|
||||||
|
class AlphaChanelAddByMask:
|
||||||
|
def __init__(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"images": ("IMAGE",),
|
||||||
|
"mask": ("MASK",),
|
||||||
|
"method": (["default", "invert"],),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("IMAGE",)
|
||||||
|
FUNCTION = "node"
|
||||||
|
CATEGORY = "Zho模块组/image"
|
||||||
|
|
||||||
|
def node(self, images, mask, method):
|
||||||
|
img_height, img_width = images[0, :, :, 0].shape
|
||||||
|
mask_height, mask_width = mask.shape
|
||||||
|
|
||||||
|
if img_height != mask_height or img_width != mask_width:
|
||||||
|
raise ValueError(
|
||||||
|
"[AlphaChanelByMask]: Size of images not equals size of mask. " +
|
||||||
|
"Images: [" + str(img_width) + ", " + str(img_height) + "] - " +
|
||||||
|
"Mask: [" + str(mask_width) + ", " + str(mask_height) + "]."
|
||||||
|
)
|
||||||
|
|
||||||
|
if method == "default":
|
||||||
|
return (torch.stack([
|
||||||
|
torch.stack((
|
||||||
|
images[i, :, :, 0],
|
||||||
|
images[i, :, :, 1],
|
||||||
|
images[i, :, :, 2],
|
||||||
|
1. - mask
|
||||||
|
), dim=-1) for i in range(len(images))
|
||||||
|
]),)
|
||||||
|
else:
|
||||||
|
return (torch.stack([
|
||||||
|
torch.stack((
|
||||||
|
images[i, :, :, 0],
|
||||||
|
images[i, :, :, 1],
|
||||||
|
images[i, :, :, 2],
|
||||||
|
mask
|
||||||
|
), dim=-1) for i in range(len(images))
|
||||||
|
]),)
|
||||||
|
|
||||||
|
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"AlphaChanelAddByMask": AlphaChanelAddByMask,
|
||||||
|
}
|
||||||
@@ -0,0 +1,417 @@
|
|||||||
|
import torchvision.transforms as t
|
||||||
|
import torch
|
||||||
|
from PIL import Image as ImageF
|
||||||
|
from PIL.Image import Image as ImageB
|
||||||
|
from torch import Tensor, dtype
|
||||||
|
|
||||||
|
def tensor_to_image(self):
|
||||||
|
return t.ToPILImage()(self.permute(2, 0, 1))
|
||||||
|
|
||||||
|
def image_to_tensor(self):
|
||||||
|
return t.ToTensor()(self).permute(1, 2, 0)
|
||||||
|
|
||||||
|
Tensor.tensor_to_image = tensor_to_image
|
||||||
|
ImageB.image_to_tensor = image_to_tensor
|
||||||
|
|
||||||
|
#--------------------------------------------------------
|
||||||
|
class ImageComposite_Zho:
|
||||||
|
def __init__(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"images_a": ("IMAGE",),
|
||||||
|
"images_b": ("IMAGE",),
|
||||||
|
"alpha_a": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # 添加控制图像a透明度的参数
|
||||||
|
"alpha_b": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # 添加控制图像b透明度的参数
|
||||||
|
"images_a_x": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"images_a_y": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"images_b_x": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"images_b_y": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"container_width": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"container_height": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"background": (["images_a", "images_b"],),
|
||||||
|
"method": (["pair", "matrix"],),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("IMAGE",)
|
||||||
|
FUNCTION = "node"
|
||||||
|
CATEGORY = "Zho模块组/image"
|
||||||
|
|
||||||
|
def node(
|
||||||
|
self,
|
||||||
|
images_a,
|
||||||
|
images_b,
|
||||||
|
images_a_x,
|
||||||
|
images_a_y,
|
||||||
|
images_b_x,
|
||||||
|
images_b_y,
|
||||||
|
container_width,
|
||||||
|
container_height,
|
||||||
|
background,
|
||||||
|
method,
|
||||||
|
alpha_a=1.0,
|
||||||
|
alpha_b=1.0,
|
||||||
|
):
|
||||||
|
def clip(value: float):
|
||||||
|
return value if value >= 0 else 0
|
||||||
|
|
||||||
|
# noinspection PyUnresolvedReferences
|
||||||
|
def composite(image_a, image_b):
|
||||||
|
img_a_height, img_a_width, img_a_dim = image_a.shape
|
||||||
|
img_b_height, img_b_width, img_b_dim = image_b.shape
|
||||||
|
|
||||||
|
if img_a_dim == 3:
|
||||||
|
image_a = torch.stack([
|
||||||
|
image_a[:, :, 0],
|
||||||
|
image_a[:, :, 1],
|
||||||
|
image_a[:, :, 2],
|
||||||
|
torch.ones((img_a_height, img_a_width)) * alpha_a
|
||||||
|
], dim=2)
|
||||||
|
|
||||||
|
if img_b_dim == 3:
|
||||||
|
image_b = torch.stack([
|
||||||
|
image_b[:, :, 0],
|
||||||
|
image_b[:, :, 1],
|
||||||
|
image_b[:, :, 2],
|
||||||
|
torch.ones((img_b_height, img_b_width)) * alpha_b
|
||||||
|
], dim=2)
|
||||||
|
|
||||||
|
container_x = max(img_a_width, img_b_width) if container_width == 0 else container_width
|
||||||
|
container_y = max(img_a_height, img_b_height) if container_height == 0 else container_height
|
||||||
|
|
||||||
|
container_a = torch.zeros((container_y, container_x, 4))
|
||||||
|
container_b = torch.zeros((container_y, container_x, 4))
|
||||||
|
|
||||||
|
img_a_height_c, img_a_width_c = [
|
||||||
|
clip((images_a_y + img_a_height) - container_y),
|
||||||
|
clip((images_a_x + img_a_width) - container_x)
|
||||||
|
]
|
||||||
|
|
||||||
|
img_b_height_c, img_b_width_c = [
|
||||||
|
clip((images_b_y + img_b_height) - container_y),
|
||||||
|
clip((images_b_x + img_b_width) - container_x)
|
||||||
|
]
|
||||||
|
|
||||||
|
if img_a_height_c <= img_a_height and img_a_width_c <= img_a_width:
|
||||||
|
container_a[
|
||||||
|
images_a_y:img_a_height + images_a_y - img_a_height_c,
|
||||||
|
images_a_x:img_a_width + images_a_x - img_a_width_c
|
||||||
|
] = image_a[
|
||||||
|
:img_a_height - img_a_height_c,
|
||||||
|
:img_a_width - img_a_width_c
|
||||||
|
]
|
||||||
|
|
||||||
|
if img_b_height_c <= img_b_height and img_b_width_c <= img_b_width:
|
||||||
|
container_b[
|
||||||
|
images_b_y:img_b_height + images_b_y - img_b_height_c,
|
||||||
|
images_b_x:img_b_width + images_b_x - img_b_width_c
|
||||||
|
] = image_b[
|
||||||
|
:img_b_height - img_b_height_c,
|
||||||
|
:img_b_width - img_b_width_c
|
||||||
|
]
|
||||||
|
|
||||||
|
if background == "images_a":
|
||||||
|
return ImageF.alpha_composite(
|
||||||
|
container_a.tensor_to_image(),
|
||||||
|
container_b.tensor_to_image()
|
||||||
|
).image_to_tensor()
|
||||||
|
else:
|
||||||
|
return ImageF.alpha_composite(
|
||||||
|
container_b.tensor_to_image(),
|
||||||
|
container_a.tensor_to_image()
|
||||||
|
).image_to_tensor()
|
||||||
|
|
||||||
|
if method == "pair":
|
||||||
|
if len(images_a) != len(images_b):
|
||||||
|
raise ValueError("Size of image_a and image_b not equals for pair batch type.")
|
||||||
|
|
||||||
|
return (torch.stack([
|
||||||
|
composite(images_a[i], images_b[i]) for i in range(len(images_a))
|
||||||
|
]),)
|
||||||
|
elif method == "matrix":
|
||||||
|
return (torch.stack([
|
||||||
|
composite(images_a[i], images_b[j]) for i in range(len(images_a)) for j in range(len(images_b))
|
||||||
|
]),)
|
||||||
|
|
||||||
|
return None
|
||||||
|
|
||||||
|
#--------------------------------------------------------
|
||||||
|
class ImageComposite_BG_Zho:
|
||||||
|
def __init__(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"container": ("IMAGE",),
|
||||||
|
"images_a": ("IMAGE",),
|
||||||
|
"images_b": ("IMAGE",),
|
||||||
|
"alpha_a": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # 添加控制图像a透明度的参数
|
||||||
|
"alpha_b": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # 添加控制图像b透明度的参数
|
||||||
|
"images_a_x": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"images_a_y": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"images_b_x": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"images_b_y": ("INT", {
|
||||||
|
"default": 0,
|
||||||
|
"step": 1
|
||||||
|
}),
|
||||||
|
"background": (["images_a", "images_b"],),
|
||||||
|
"method": (["pair", "matrix"],),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("IMAGE",)
|
||||||
|
FUNCTION = "node"
|
||||||
|
CATEGORY = "Zho模块组/image"
|
||||||
|
|
||||||
|
def node(
|
||||||
|
self,
|
||||||
|
container,
|
||||||
|
images_a,
|
||||||
|
images_b,
|
||||||
|
images_a_x,
|
||||||
|
images_a_y,
|
||||||
|
images_b_x,
|
||||||
|
images_b_y,
|
||||||
|
background,
|
||||||
|
method,
|
||||||
|
alpha_a=1.0,
|
||||||
|
alpha_b=1.0,
|
||||||
|
):
|
||||||
|
return ImageComposite_Zho().node(
|
||||||
|
images_a,
|
||||||
|
images_b,
|
||||||
|
images_a_x,
|
||||||
|
images_a_y,
|
||||||
|
images_b_x,
|
||||||
|
images_b_y,
|
||||||
|
container[0, :, :, 0].shape[1],
|
||||||
|
container[0, :, :, 0].shape[0],
|
||||||
|
background,
|
||||||
|
method,
|
||||||
|
alpha_a=alpha_a,
|
||||||
|
alpha_b=alpha_b,
|
||||||
|
)
|
||||||
|
|
||||||
|
#--------------------------------------------------------
|
||||||
|
class ImageCompositeBy_Zho:
|
||||||
|
def __init__(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"images_a": ("IMAGE",),
|
||||||
|
"images_b": ("IMAGE",),
|
||||||
|
"alpha_a": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # 添加控制图像a透明度的参数
|
||||||
|
"alpha_b": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # 添加控制图像b透明度的参数
|
||||||
|
"images_a_x": ("FLOAT", {
|
||||||
|
"default": 0.0,
|
||||||
|
"max": 1.0,
|
||||||
|
"step": 0.01
|
||||||
|
}),
|
||||||
|
"images_a_y": ("FLOAT", {
|
||||||
|
"default": 0.0,
|
||||||
|
"max": 1.0,
|
||||||
|
"step": 0.01
|
||||||
|
}),
|
||||||
|
"images_b_x": ("FLOAT", {
|
||||||
|
"default": 0.0,
|
||||||
|
"max": 1.0,
|
||||||
|
"step": 0.01
|
||||||
|
}),
|
||||||
|
"images_b_y": ("FLOAT", {
|
||||||
|
"default": 0.0,
|
||||||
|
"max": 1.0,
|
||||||
|
"step": 0.01
|
||||||
|
}),
|
||||||
|
"background": (["images_a", "images_b"],),
|
||||||
|
"container_size_type": (["max", "sum", "sum_width", "sum_height"],),
|
||||||
|
"method": (["pair", "matrix"],),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("IMAGE",)
|
||||||
|
FUNCTION = "node"
|
||||||
|
CATEGORY = "Zho模块组/image"
|
||||||
|
|
||||||
|
def node(
|
||||||
|
self,
|
||||||
|
images_a,
|
||||||
|
images_b,
|
||||||
|
images_a_x,
|
||||||
|
images_a_y,
|
||||||
|
images_b_x,
|
||||||
|
images_b_y,
|
||||||
|
background,
|
||||||
|
container_size_type,
|
||||||
|
method,
|
||||||
|
alpha_a=1.0,
|
||||||
|
alpha_b=1.0,
|
||||||
|
):
|
||||||
|
def offset_by_percent(container_size: int, image_size: int, percent: float):
|
||||||
|
return int((container_size - image_size) * percent)
|
||||||
|
|
||||||
|
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 == "max":
|
||||||
|
container_width = max(img_a_width, img_b_width)
|
||||||
|
container_height = max(img_a_height, img_b_height)
|
||||||
|
elif container_size_type == "sum":
|
||||||
|
container_width = img_a_width + img_b_width
|
||||||
|
container_height = img_a_height + img_b_height
|
||||||
|
elif container_size_type == "sum_width":
|
||||||
|
if img_a_height != img_b_height:
|
||||||
|
raise ValueError()
|
||||||
|
|
||||||
|
container_width = img_a_width + img_b_width
|
||||||
|
container_height = img_a_height
|
||||||
|
elif container_size_type == "sum_height":
|
||||||
|
if img_b_width != img_b_width:
|
||||||
|
raise ValueError()
|
||||||
|
|
||||||
|
container_width = img_a_width
|
||||||
|
container_height = img_a_height + img_a_height
|
||||||
|
else:
|
||||||
|
raise ValueError()
|
||||||
|
|
||||||
|
return ImageComposite_Zho().node(
|
||||||
|
images_a,
|
||||||
|
images_b,
|
||||||
|
offset_by_percent(container_width, img_a_width, images_a_x),
|
||||||
|
offset_by_percent(container_height, img_a_height, images_a_y),
|
||||||
|
offset_by_percent(container_width, img_b_width, images_b_x),
|
||||||
|
offset_by_percent(container_height, img_b_height, images_b_y),
|
||||||
|
container_width,
|
||||||
|
container_height,
|
||||||
|
background,
|
||||||
|
method,
|
||||||
|
alpha_a=alpha_a,
|
||||||
|
alpha_b=alpha_b
|
||||||
|
)
|
||||||
|
|
||||||
|
#--------------------------------------------------------
|
||||||
|
class ImageCompositeBy_BG_Zho:
|
||||||
|
def __init__(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"container": ("IMAGE",),
|
||||||
|
"images_a": ("IMAGE",),
|
||||||
|
"images_b": ("IMAGE",),
|
||||||
|
"alpha_a": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # 添加控制图像a透明度的参数
|
||||||
|
"alpha_b": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # 添加控制图像b透明度的参数
|
||||||
|
"images_a_x": ("FLOAT", {
|
||||||
|
"default": 0.0,
|
||||||
|
"max": 1.0,
|
||||||
|
"step": 0.01
|
||||||
|
}),
|
||||||
|
"images_a_y": ("FLOAT", {
|
||||||
|
"default": 0.0,
|
||||||
|
"max": 1.0,
|
||||||
|
"step": 0.01
|
||||||
|
}),
|
||||||
|
"images_b_x": ("FLOAT", {
|
||||||
|
"default": 0.0,
|
||||||
|
"max": 1.0,
|
||||||
|
"step": 0.01
|
||||||
|
}),
|
||||||
|
"images_b_y": ("FLOAT", {
|
||||||
|
"default": 0.0,
|
||||||
|
"max": 1.0,
|
||||||
|
"step": 0.01
|
||||||
|
}),
|
||||||
|
"background": (["images_a", "images_b"],),
|
||||||
|
"method": (["pair", "matrix"],),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("IMAGE",)
|
||||||
|
FUNCTION = "node"
|
||||||
|
CATEGORY = "Zho模块组/image"
|
||||||
|
|
||||||
|
def node(
|
||||||
|
self,
|
||||||
|
container,
|
||||||
|
images_a,
|
||||||
|
images_b,
|
||||||
|
images_a_x,
|
||||||
|
images_a_y,
|
||||||
|
images_b_x,
|
||||||
|
images_b_y,
|
||||||
|
background,
|
||||||
|
method,
|
||||||
|
alpha_a=1.0,
|
||||||
|
alpha_b=1.0,
|
||||||
|
):
|
||||||
|
def offset_by_percent(container_size: int, image_size: int, percent: float):
|
||||||
|
return int((container_size - image_size) * percent)
|
||||||
|
|
||||||
|
img_a_height, img_a_width = images_a[0, :, :, 0].shape
|
||||||
|
img_b_height, img_b_width = images_b[0, :, :, 0].shape
|
||||||
|
|
||||||
|
container_width = container[0, :, :, 0].shape[1]
|
||||||
|
container_height = container[0, :, :, 0].shape[0]
|
||||||
|
|
||||||
|
if container_width < max(img_a_width, img_b_width) or container_height < max(img_a_height, img_b_height):
|
||||||
|
raise ValueError("Container can't be smaller then max width or height of images.")
|
||||||
|
|
||||||
|
return ImageComposite_Zho().node(
|
||||||
|
images_a,
|
||||||
|
images_b,
|
||||||
|
offset_by_percent(container_width, img_a_width, images_a_x),
|
||||||
|
offset_by_percent(container_height, img_a_height, images_a_y),
|
||||||
|
offset_by_percent(container_width, img_b_width, images_b_x),
|
||||||
|
offset_by_percent(container_height, img_b_height, images_b_y),
|
||||||
|
container_width,
|
||||||
|
container_height,
|
||||||
|
background,
|
||||||
|
method,
|
||||||
|
alpha_a=alpha_a,
|
||||||
|
alpha_b=alpha_b
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"ImageComposite_Zho": ImageComposite_Zho,
|
||||||
|
"ImageComposite_BG_Zho": ImageComposite_BG_Zho,
|
||||||
|
"ImageCompositeBy_Zho": ImageCompositeBy_Zho,
|
||||||
|
"ImageCompositeBy_BG_Zho": ImageCompositeBy_BG_Zho
|
||||||
|
}
|
||||||
@@ -0,0 +1,52 @@
|
|||||||
|
import torch
|
||||||
|
from PIL import Image
|
||||||
|
from typing import List, Optional, Union
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
def pil2tensor(image: Union[Image.Image, List[Image.Image]]) -> torch.Tensor:
|
||||||
|
if isinstance(image, list):
|
||||||
|
return torch.cat([pil2tensor(img) for img in image], dim=0)
|
||||||
|
|
||||||
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
||||||
|
|
||||||
|
# 添加一个辅助函数,用于交换宽度和高度
|
||||||
|
def swap_width_height(width, height):
|
||||||
|
return height, width
|
||||||
|
|
||||||
|
class RGB_Image_Zho:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"width": ("INT", {"default": 512, "min": 16, "max": 8160}),
|
||||||
|
"height": ("INT", {"default": 512, "min": 16, "max": 8160}),
|
||||||
|
"swap": ("BOOLEAN", {"default": False}), # 添加交换宽度和高度的按钮
|
||||||
|
"color": ("COLOR",),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("IMAGE",)
|
||||||
|
FUNCTION = "rgb_image"
|
||||||
|
CATEGORY = "Zho模块组/image"
|
||||||
|
|
||||||
|
def rgb_image(self, color, width, height, swap=False):
|
||||||
|
# 如果用户选择交换宽度和高度,则调用交换函数
|
||||||
|
if swap:
|
||||||
|
width, height = swap_width_height(width, height)
|
||||||
|
|
||||||
|
# 创建RGBA图像
|
||||||
|
image = Image.new("RGB", (width, height), color=color)
|
||||||
|
|
||||||
|
# 转换为张量
|
||||||
|
image = pil2tensor(image)
|
||||||
|
|
||||||
|
return (image,)
|
||||||
|
|
||||||
|
|
||||||
|
#----------------------------------------------------------------------------
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"RGB_Image_Zho": RGB_Image_Zho,
|
||||||
|
}
|
||||||
@@ -0,0 +1,426 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
from typing import cast
|
||||||
|
from typing import List, Optional, Union
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
#----------------------------------------------------------------------------
|
||||||
|
here = Path(__file__).parent.absolute()
|
||||||
|
comfy_dir = here.parent.parent
|
||||||
|
|
||||||
|
#----------------------------------------------------------------------------
|
||||||
|
def pil2tensor(image: Union[Image.Image, List[Image.Image]]) -> torch.Tensor:
|
||||||
|
if isinstance(image, list):
|
||||||
|
return torch.cat([pil2tensor(img) for img in image], dim=0)
|
||||||
|
|
||||||
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
||||||
|
|
||||||
|
#----------------------------------------------------------------------------
|
||||||
|
# 添加一个辅助函数,用于交换宽度和高度
|
||||||
|
def swap_width_height(width, height):
|
||||||
|
return height, width
|
||||||
|
|
||||||
|
#----------------------------------------------------------------------------
|
||||||
|
import logging
|
||||||
|
import re
|
||||||
|
import os
|
||||||
|
|
||||||
|
base_log_level = logging.DEBUG if os.environ.get("MTB_DEBUG") else logging.INFO
|
||||||
|
|
||||||
|
|
||||||
|
# Custom object that discards the output
|
||||||
|
class NullWriter:
|
||||||
|
def write(self, text):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class Formatter(logging.Formatter):
|
||||||
|
grey = "\x1b[38;20m"
|
||||||
|
cyan = "\x1b[36;20m"
|
||||||
|
purple = "\x1b[35;20m"
|
||||||
|
yellow = "\x1b[33;20m"
|
||||||
|
red = "\x1b[31;20m"
|
||||||
|
bold_red = "\x1b[31;1m"
|
||||||
|
reset = "\x1b[0m"
|
||||||
|
# format = "%(asctime)s - [%(name)s] - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
|
||||||
|
format = "[%(name)s] | %(levelname)s -> %(message)s"
|
||||||
|
|
||||||
|
FORMATS = {
|
||||||
|
logging.DEBUG: purple + format + reset,
|
||||||
|
logging.INFO: cyan + format + reset,
|
||||||
|
logging.WARNING: yellow + format + reset,
|
||||||
|
logging.ERROR: red + format + reset,
|
||||||
|
logging.CRITICAL: bold_red + format + reset,
|
||||||
|
}
|
||||||
|
|
||||||
|
def format(self, record):
|
||||||
|
log_fmt = self.FORMATS.get(record.levelno)
|
||||||
|
formatter = logging.Formatter(log_fmt)
|
||||||
|
return formatter.format(record)
|
||||||
|
|
||||||
|
|
||||||
|
def mklog(name, level=base_log_level):
|
||||||
|
logger = logging.getLogger(name)
|
||||||
|
logger.setLevel(level)
|
||||||
|
|
||||||
|
for handler in logger.handlers:
|
||||||
|
logger.removeHandler(handler)
|
||||||
|
|
||||||
|
ch = logging.StreamHandler()
|
||||||
|
ch.setLevel(level)
|
||||||
|
ch.setFormatter(Formatter())
|
||||||
|
logger.addHandler(ch)
|
||||||
|
|
||||||
|
# Disable log propagation
|
||||||
|
logger.propagate = False
|
||||||
|
|
||||||
|
return logger
|
||||||
|
|
||||||
|
|
||||||
|
# - The main app logger
|
||||||
|
log = mklog(__package__, base_log_level)
|
||||||
|
|
||||||
|
|
||||||
|
def log_user(arg):
|
||||||
|
print("\033[34mComfy MTB Utils:\033[0m {arg}")
|
||||||
|
|
||||||
|
|
||||||
|
def get_summary(docstring):
|
||||||
|
return docstring.strip().split("\n\n", 1)[0]
|
||||||
|
|
||||||
|
|
||||||
|
def blue_text(text):
|
||||||
|
return f"\033[94m{text}\033[0m"
|
||||||
|
|
||||||
|
|
||||||
|
def cyan_text(text):
|
||||||
|
return f"\033[96m{text}\033[0m"
|
||||||
|
|
||||||
|
|
||||||
|
def get_label(label):
|
||||||
|
words = re.findall(r"(?:^|[A-Z])[a-z]*", label)
|
||||||
|
return " ".join(words).strip()
|
||||||
|
|
||||||
|
# 禁用 aiohttp 的访问日志记录器
|
||||||
|
logging.getLogger('aiohttp.access').disabled = True
|
||||||
|
#----------------------------------------------------------------------------
|
||||||
|
def bbox_dim(bbox):
|
||||||
|
left, upper, right, lower = bbox
|
||||||
|
width = right - left
|
||||||
|
height = lower - upper
|
||||||
|
return width, height
|
||||||
|
|
||||||
|
#----------------------------------------------------------------------------
|
||||||
|
class Text_Image_Zho:
|
||||||
|
|
||||||
|
fonts = {}
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
# - This is executed when the graph is executed, we could conditionaly reload fonts there
|
||||||
|
pass
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def CACHE_FONTS(cls):
|
||||||
|
font_extensions = ["*.ttf", "*.otf", "*.woff", "*.woff2", "*.eot"]
|
||||||
|
fonts = []
|
||||||
|
|
||||||
|
for extension in font_extensions:
|
||||||
|
fonts.extend(comfy_dir.glob(f"**/{extension}"))
|
||||||
|
|
||||||
|
if not fonts:
|
||||||
|
log.warn(
|
||||||
|
"> No fonts found in the comfy folder, place at least one font file somewhere in ComfyUI's hierarchy"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
log.debug(f"> Found {len(fonts)} fonts")
|
||||||
|
|
||||||
|
for font in fonts:
|
||||||
|
log.debug(f"Adding font {font}")
|
||||||
|
cls.fonts[font.stem] = font.as_posix()
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
if not cls.fonts:
|
||||||
|
cls.CACHE_FONTS()
|
||||||
|
else:
|
||||||
|
log.debug(f"Using cached fonts (count: {len(cls.fonts)})")
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"text": (
|
||||||
|
"STRING",
|
||||||
|
{"default": "ZHOZHOZHO"},
|
||||||
|
),
|
||||||
|
"selected_font": ((sorted(cls.fonts.keys())),),
|
||||||
|
"align": (["left", "center", "right"],
|
||||||
|
),
|
||||||
|
"wrap": (
|
||||||
|
"INT",
|
||||||
|
{"default": 0, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"font_size": (
|
||||||
|
"INT",
|
||||||
|
{"default": 12, "min": 1, "max": 2500, "step": 1},
|
||||||
|
),
|
||||||
|
"color": (
|
||||||
|
"COLOR",
|
||||||
|
{"default": "red"},
|
||||||
|
),
|
||||||
|
"outline_size": (
|
||||||
|
"INT",
|
||||||
|
{"default": 0, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"outline_color": (
|
||||||
|
"COLOR",
|
||||||
|
{"default": "blue"}, # 设置默认的描边颜色
|
||||||
|
),
|
||||||
|
"margin_x": (
|
||||||
|
"INT",
|
||||||
|
{"default": 0, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"margin_y": (
|
||||||
|
"INT",
|
||||||
|
{"default": 0, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"width": (
|
||||||
|
"INT",
|
||||||
|
{"default": 512, "min": 1, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"height": (
|
||||||
|
"INT",
|
||||||
|
{"default": 512, "min": 1, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"swap": ("BOOLEAN", {"default": False}), # 添加交换宽度和高度的按钮
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("IMAGE",)
|
||||||
|
RETURN_NAMES = ("image",)
|
||||||
|
FUNCTION = "text_to_image"
|
||||||
|
CATEGORY = "Zho模块组/text"
|
||||||
|
|
||||||
|
def text_to_image(
|
||||||
|
self, text, selected_font, align, wrap, font_size, width, height, color, outline_size, outline_color, margin_x, margin_y, swap=False
|
||||||
|
):
|
||||||
|
from PIL import Image, ImageDraw, ImageFont
|
||||||
|
import textwrap
|
||||||
|
|
||||||
|
# 如果用户选择交换宽度和高度,则调用交换函数
|
||||||
|
if swap:
|
||||||
|
width, height = swap_width_height(width, height)
|
||||||
|
|
||||||
|
font_path = self.fonts[selected_font]
|
||||||
|
(_, top, _, _) = ImageFont.truetype(font_path, font_size).getbbox(text)
|
||||||
|
font = cast(ImageFont.FreeTypeFont, ImageFont.truetype(font_path, font_size))
|
||||||
|
if wrap == 0:
|
||||||
|
wrap = width / font_size
|
||||||
|
wrap = int(wrap)
|
||||||
|
lines = textwrap.wrap(text, width=wrap)
|
||||||
|
log.debug(f"Lines: {lines}")
|
||||||
|
line_height = bbox_dim(font.getbbox("hg"))[1]
|
||||||
|
img_height = height # line_height * len(lines)
|
||||||
|
img_width = width # max(font.getsize(line)[0] for line in lines)
|
||||||
|
|
||||||
|
img = Image.new("RGBA", (img_width, img_height), (0, 0, 0, 0))
|
||||||
|
draw = ImageDraw.Draw(img)
|
||||||
|
|
||||||
|
# 初始化 y_text
|
||||||
|
y_text = margin_y + outline_size - top
|
||||||
|
|
||||||
|
for line in lines:
|
||||||
|
width, height = bbox_dim(font.getbbox(line))
|
||||||
|
|
||||||
|
# 根据 align 参数计算文本的 x 坐标
|
||||||
|
if align == "left":
|
||||||
|
x_text = margin_x
|
||||||
|
elif align == "center":
|
||||||
|
x_text = (img_width - width) // 2
|
||||||
|
elif align == "right":
|
||||||
|
x_text = img_width - width - margin_x
|
||||||
|
else:
|
||||||
|
x_text = margin_x # 默认为左对齐
|
||||||
|
|
||||||
|
draw.text(
|
||||||
|
(x_text, y_text),
|
||||||
|
text=line,
|
||||||
|
fill=color,
|
||||||
|
stroke_fill=outline_color,
|
||||||
|
stroke_width=outline_size,
|
||||||
|
font=font,
|
||||||
|
)
|
||||||
|
y_text += height
|
||||||
|
|
||||||
|
return (pil2tensor(img),)
|
||||||
|
|
||||||
|
#----------------------------------------------------------------------------
|
||||||
|
class Text_Image_Multiline_Zho:
|
||||||
|
|
||||||
|
fonts = {}
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
# - This is executed when the graph is executed, we could conditionaly reload fonts there
|
||||||
|
pass
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def CACHE_FONTS(cls):
|
||||||
|
font_extensions = ["*.ttf", "*.otf", "*.woff", "*.woff2", "*.eot"]
|
||||||
|
fonts = []
|
||||||
|
|
||||||
|
for extension in font_extensions:
|
||||||
|
fonts.extend(comfy_dir.glob(f"**/{extension}"))
|
||||||
|
|
||||||
|
if not fonts:
|
||||||
|
log.warn(
|
||||||
|
"> No fonts found in the comfy folder, place at least one font file somewhere in ComfyUI's hierarchy"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
log.debug(f"> Found {len(fonts)} fonts")
|
||||||
|
|
||||||
|
for font in fonts:
|
||||||
|
log.debug(f"Adding font {font}")
|
||||||
|
cls.fonts[font.stem] = font.as_posix()
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def INPUT_TYPES(cls):
|
||||||
|
if not cls.fonts:
|
||||||
|
cls.CACHE_FONTS()
|
||||||
|
else:
|
||||||
|
log.debug(f"Using cached fonts (count: {len(cls.fonts)})")
|
||||||
|
return {
|
||||||
|
"required": {
|
||||||
|
"text": (
|
||||||
|
"STRING",
|
||||||
|
{"default": "ZHOZHOZHO", "multiline": True},
|
||||||
|
),
|
||||||
|
"selected_font": ((sorted(cls.fonts.keys())),),
|
||||||
|
"align": (["left", "center", "right"],
|
||||||
|
),
|
||||||
|
"wrap": (
|
||||||
|
"INT",
|
||||||
|
{"default": 120, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"graphspace": (
|
||||||
|
"INT",
|
||||||
|
{"default": 10, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"linespace": (
|
||||||
|
"INT",
|
||||||
|
{"default": 2, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"font_size": (
|
||||||
|
"INT",
|
||||||
|
{"default": 12, "min": 1, "max": 2500, "step": 1},
|
||||||
|
),
|
||||||
|
"color": (
|
||||||
|
"COLOR",
|
||||||
|
{"default": "red"},
|
||||||
|
),
|
||||||
|
"outline_size": (
|
||||||
|
"INT",
|
||||||
|
{"default": 0, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"outline_color": (
|
||||||
|
"COLOR",
|
||||||
|
{"default": "blue"}, # 设置默认的描边颜色
|
||||||
|
),
|
||||||
|
"margin_x": (
|
||||||
|
"INT",
|
||||||
|
{"default": 0, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"margin_y": (
|
||||||
|
"INT",
|
||||||
|
{"default": 0, "min": 0, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"width": (
|
||||||
|
"INT",
|
||||||
|
{"default": 512, "min": 1, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"height": (
|
||||||
|
"INT",
|
||||||
|
{"default": 512, "min": 1, "max": 8096, "step": 1},
|
||||||
|
),
|
||||||
|
"swap": ("BOOLEAN", {"default": False}), # 添加交换宽度和高度的按钮
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
RETURN_TYPES = ("IMAGE",)
|
||||||
|
RETURN_NAMES = ("image",)
|
||||||
|
FUNCTION = "text_to_image_multiline"
|
||||||
|
CATEGORY = "Zho模块组/text"
|
||||||
|
|
||||||
|
def text_to_image_multiline(
|
||||||
|
self, text, selected_font, align, wrap, graphspace, linespace, font_size, width, height, color, outline_size, outline_color, margin_x, margin_y, swap=False
|
||||||
|
):
|
||||||
|
from PIL import Image, ImageDraw, ImageFont
|
||||||
|
import textwrap
|
||||||
|
|
||||||
|
# 如果用户选择交换宽度和高度,则调用交换函数
|
||||||
|
if swap:
|
||||||
|
width, height = swap_width_height(width, height)
|
||||||
|
|
||||||
|
font_path = self.fonts[selected_font]
|
||||||
|
(_, top, _, _) = ImageFont.truetype(font_path, font_size).getbbox(text)
|
||||||
|
font = cast(ImageFont.FreeTypeFont, ImageFont.truetype(font_path, font_size))
|
||||||
|
if wrap == 0:
|
||||||
|
wrap = width / font_size
|
||||||
|
wrap = int(wrap)
|
||||||
|
|
||||||
|
paragraphs = text.split('\n')
|
||||||
|
|
||||||
|
log.debug(f"Paragraphs: {paragraphs}")
|
||||||
|
|
||||||
|
img_height = height # line_height * len(lines)
|
||||||
|
img_width = width # max(font.getsize(line)[0] for line in lines)
|
||||||
|
|
||||||
|
img = Image.new("RGBA", (img_width, img_height), (0, 0, 0, 0))
|
||||||
|
draw = ImageDraw.Draw(img)
|
||||||
|
|
||||||
|
# 初始化 y_text
|
||||||
|
y_text = margin_y + outline_size
|
||||||
|
|
||||||
|
for paragraph in paragraphs:
|
||||||
|
lines = textwrap.wrap(paragraph, width=wrap, expand_tabs=False, replace_whitespace=False)
|
||||||
|
|
||||||
|
for line in lines:
|
||||||
|
width, height = bbox_dim(font.getbbox(line))
|
||||||
|
|
||||||
|
# 根据 align 参数重新计算 x 坐标
|
||||||
|
if align == "left":
|
||||||
|
x_text = margin_x
|
||||||
|
elif align == "center":
|
||||||
|
x_text = (img_width - width) // 2
|
||||||
|
elif align == "right":
|
||||||
|
x_text = img_width - width - margin_x
|
||||||
|
else:
|
||||||
|
x_text = margin_x # 默认为左对齐
|
||||||
|
|
||||||
|
draw.text(
|
||||||
|
(x_text, y_text),
|
||||||
|
text=line,
|
||||||
|
fill=color,
|
||||||
|
stroke_fill=outline_color,
|
||||||
|
stroke_width=outline_size,
|
||||||
|
font=font,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 更新 y 坐标,加上当前行的高度和一些额外的间距
|
||||||
|
y_text += height + linespace # linespace 是行之间的额外间距
|
||||||
|
|
||||||
|
# 段落之间添加一些额外的间距
|
||||||
|
y_text += graphspace # 可以根据需要调整
|
||||||
|
|
||||||
|
return (pil2tensor(img),)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
#----------------------------------------------------------------------------
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"Text_Image_Zho": Text_Image_Zho,
|
||||||
|
"Text_Image_Multiline_Zho": Text_Image_Multiline_Zho,
|
||||||
|
}
|
||||||
+46
@@ -0,0 +1,46 @@
|
|||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import filecmp
|
||||||
|
import shutil
|
||||||
|
import __main__
|
||||||
|
|
||||||
|
|
||||||
|
python = sys.executable
|
||||||
|
|
||||||
|
|
||||||
|
extentions_folder = os.path.join(os.path.dirname(os.path.realpath(__main__.__file__)),
|
||||||
|
"web" + os.sep + "extensions" + os.sep + "ZHO")
|
||||||
|
javascript_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "mtb")
|
||||||
|
|
||||||
|
if not os.path.exists(extentions_folder):
|
||||||
|
print('Making the "web\extensions\ZHO" folder')
|
||||||
|
os.mkdir(extentions_folder)
|
||||||
|
|
||||||
|
result = filecmp.dircmp(javascript_folder, extentions_folder)
|
||||||
|
|
||||||
|
if result.left_only or result.diff_files:
|
||||||
|
print('Update to javascripts files detected')
|
||||||
|
file_list = list(result.left_only)
|
||||||
|
file_list.extend(x for x in result.diff_files if x not in file_list)
|
||||||
|
|
||||||
|
for file in file_list:
|
||||||
|
print(f'Copying {file} to extensions folder')
|
||||||
|
src_file = os.path.join(javascript_folder, file)
|
||||||
|
dst_file = os.path.join(extentions_folder, file)
|
||||||
|
if os.path.exists(dst_file):
|
||||||
|
os.remove(dst_file)
|
||||||
|
#print("disabled")
|
||||||
|
shutil.copy(src_file, dst_file)
|
||||||
|
|
||||||
|
|
||||||
|
from .Zho_TextImage import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_TI
|
||||||
|
from .Zho_RGB_Image import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_RGB
|
||||||
|
from .Zho_ImageComposite import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_IC
|
||||||
|
from .Zho_AlphaChanel import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_AC
|
||||||
|
|
||||||
|
|
||||||
|
# Combine the dictionaries
|
||||||
|
NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS_TI, **NODE_CLASS_MAPPINGS_RGB, **NODE_CLASS_MAPPINGS_IC, **NODE_CLASS_MAPPINGS_AC}
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ['NODE_CLASS_MAPPINGS']
|
||||||
Reference in New Issue
Block a user