494 lines
17 KiB
Python
494 lines
17 KiB
Python
from pathlib import Path
|
|
from typing import cast
|
|
from typing import List, Optional, Union
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import math
|
|
|
|
#----------------------------------------------------------------------------
|
|
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}), # 添加交换宽度和高度的按钮
|
|
"arc_text": ("BOOLEAN", {"default": False}),
|
|
"arc_radius": (
|
|
"INT",
|
|
{"default": 100, "min": 1, "max": 2500, "step": 1},
|
|
),
|
|
"arc_start_angle": (
|
|
"INT",
|
|
{"default": 180, "min": 0, "max": 360, "step": 1},
|
|
),
|
|
"arc_end_angle": (
|
|
"INT",
|
|
{"default": 360, "min": 0, "max": 360, "step": 1},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "text_to_image"
|
|
CATEGORY = "Zho模块组/text"
|
|
|
|
def draw_text_in_arc(self, image, draw, text, font, font_path, font_size, center, radius, start_angle, end_angle, fill='black', stroke_fill='blue', stroke_width=0):
|
|
# 文字的总长度
|
|
text_width = sum(font.getsize(char)[0] for char in text[:-1])
|
|
|
|
# 根据扇形角度计算角度步长
|
|
angle_range = end_angle - start_angle
|
|
angle_step = (angle_range / (len(text) - 1) if len(text) > 1 else 1)
|
|
|
|
# 开始绘制文字
|
|
current_angle = start_angle
|
|
for char in text:
|
|
char_width, char_height = font.getsize(char)
|
|
angle = math.radians(current_angle)
|
|
print(current_angle, angle, angle_step)
|
|
|
|
# 创建单独的图像用于旋转字符
|
|
super_sampling_multiplier = 10
|
|
char_image = Image.new("RGBA", (char_width * super_sampling_multiplier, char_height * super_sampling_multiplier), (0, 0, 0, 0))
|
|
char_draw = ImageDraw.Draw(char_image)
|
|
super_sampling_font = cast(ImageFont.FreeTypeFont, ImageFont.truetype(font_path, font_size * super_sampling_multiplier))
|
|
char_draw.text((0, 0), char, font=super_sampling_font, fill=fill, stroke_fill=stroke_fill, stroke_width=stroke_width)
|
|
|
|
# 计算旋转角度和字符位置
|
|
rotate_angle = current_angle - 90 - (current_angle - 180) * 2 # 使字符面向圆心
|
|
rotated_char_image = char_image.rotate(rotate_angle, expand=1, resample=Image.BICUBIC)
|
|
# 缩小图像大小
|
|
new_size = (int(rotated_char_image.width / 10), int(rotated_char_image.height / 10))
|
|
rotated_char_image_resized = rotated_char_image.resize(new_size, Image.ANTIALIAS)
|
|
|
|
# 计算缩小后的图像放置位置
|
|
x = center[0] + radius * math.cos(angle) - rotated_char_image_resized.size[0] / 2
|
|
y = center[1] + radius * math.sin(angle) - rotated_char_image_resized.size[1] / 2
|
|
|
|
# 粘贴缩小后的图像
|
|
image.paste(rotated_char_image_resized, (int(x), int(y)), rotated_char_image_resized)
|
|
|
|
# 更新下一个字符的开始角度
|
|
current_angle += angle_step
|
|
|
|
|
|
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, arc_text=False, arc_radius=100, arc_start_angle=180, arc_end_angle=360
|
|
):
|
|
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)
|
|
|
|
# 曲线文字
|
|
if arc_text:
|
|
width, height = bbox_dim(font.getbbox(text))
|
|
|
|
center_x = (img_width) // 2
|
|
center_y = arc_radius + (height)
|
|
|
|
if align == "left":
|
|
center_x = arc_radius + (height) // 2
|
|
elif align == "right":
|
|
center_x = img_width - arc_radius - (height) // 2
|
|
center = (center_x + margin_x, center_y + margin_y)
|
|
self.draw_text_in_arc(img, draw, text, font, font_path, font_size, center, arc_radius, arc_start_angle, arc_end_angle,
|
|
fill=color, stroke_fill=outline_color, stroke_width=outline_size)
|
|
else:
|
|
# 初始化 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,
|
|
} |