281 lines
13 KiB
Python
281 lines
13 KiB
Python
|
import os
|
|||
|
|
import folder_paths
|
|||
|
|
import re
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from .utils import is_module_imported, pil2tensor, get_device_by_name, comfy_tensor_Image2np_Image
|
|||
|
|
|
|||
|
|
comfy_temp_dir = folder_paths.get_temp_directory()
|
|||
|
|
Random_Gen_Mask_path = os.path.join(comfy_temp_dir, "AnyText_random_mask_pos_img.png")
|
|||
|
|
tmp_pose_img_path = os.path.join(comfy_temp_dir, "AnyText_manual_mask_pos_img.png")
|
|||
|
|
tmp_ori_img_path = os.path.join(comfy_temp_dir, "AnyText_ori_img.png")
|
|||
|
|
|
|||
|
|
class AnyText:
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self.model = None
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def INPUT_TYPES(cls):
|
|||
|
|
return {
|
|||
|
|
"required": {
|
|||
|
|
"AnyText_Loader": ("AnyText_Loader", {"forceInput": True}),
|
|||
|
|
"prompt": ("STRING", {"default": "A raccoon stands in front of the blackboard with the words \"你好呀~Hello!\" written on it.", "multiline": True}),
|
|||
|
|
"a_prompt": ("STRING", {"default": "best quality, extremely detailed,4k, HD, supper legible text, clear text edges, clear strokes, neat writing, no watermarks", "multiline": True}),
|
|||
|
|
"n_prompt": ("STRING", {"default": "low-res, bad anatomy, extra digit, fewer digits, cropped, worst quality, low quality, watermark, unreadable text, messy words, distorted text, disorganized writing, advertising picture", "multiline": True}),
|
|||
|
|
"mode": (['text-generation', 'text-editing'],{"default": 'text-generation'}),
|
|||
|
|
"sort_radio": (["↕", "↔"],{"default": "↔"}),
|
|||
|
|
"revise_pos": ("BOOLEAN", {"default": False}),
|
|||
|
|
"img_count": ("INT", {"default": 1, "min": 1, "max": 10}),
|
|||
|
|
"ddim_steps": ("INT", {"default": 20, "min": 2, "max": 100}),
|
|||
|
|
"seed": ("INT", {"default": 9999, "min": -1, "max": 99999999}),
|
|||
|
|
"nonEdit_random_gen_width": ("INT", {"default": 512, "min": 128, "max": 1920, "step": 64}),
|
|||
|
|
"nonEdit_random_gen_height": ("INT", {"default": 512, "min": 128, "max": 1920, "step": 64}),
|
|||
|
|
# "width": ("INT", {"forceInput": True}),
|
|||
|
|
# "height": ("INT", {"forceInput": True}),
|
|||
|
|
"Random_Gen": ("BOOLEAN", {"default": False}),
|
|||
|
|
"strength": ("FLOAT", {
|
|||
|
|
"default": 1.00,
|
|||
|
|
"min": -999999,
|
|||
|
|
"max": 9999999,
|
|||
|
|
"step": 0.01
|
|||
|
|
}),
|
|||
|
|
"cfg_scale": ("FLOAT", {
|
|||
|
|
"default": 9,
|
|||
|
|
"min": 1,
|
|||
|
|
"max": 99,
|
|||
|
|
"step": 0.1
|
|||
|
|
}),
|
|||
|
|
"eta": ("FLOAT", {
|
|||
|
|
"default": 0,
|
|||
|
|
"min": 0,
|
|||
|
|
"max": 1,
|
|||
|
|
"step": 0.1
|
|||
|
|
}),
|
|||
|
|
"device": (["auto", "cuda", "cpu", "mps", "xpu"],{"default": "auto"}),
|
|||
|
|
"fp16": ("BOOLEAN", {"default": True}),
|
|||
|
|
"cpu_offload": ("BOOLEAN", {"default": False, "label_on": "model_to_cpu", "label_off": "unload_model"}),
|
|||
|
|
"all_to_device": ("BOOLEAN", {"default": False}),
|
|||
|
|
},
|
|||
|
|
"optional": {
|
|||
|
|
"ori_image": ("IMAGE", {"forceInput": True}),
|
|||
|
|
"pos_image": ("IMAGE", {"forceInput": True}),
|
|||
|
|
# "show_debug": ("BOOLEAN", {"default": False}),
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
RETURN_TYPES = ("IMAGE",)
|
|||
|
|
CATEGORY = "ExtraModels/AnyText"
|
|||
|
|
FUNCTION = "anytext_process"
|
|||
|
|
TITLE = "AnyText Geneation"
|
|||
|
|
|
|||
|
|
def anytext_process(self,
|
|||
|
|
mode,
|
|||
|
|
AnyText_Loader,
|
|||
|
|
ori_image,
|
|||
|
|
pos_image,
|
|||
|
|
sort_radio,
|
|||
|
|
revise_pos,
|
|||
|
|
Random_Gen,
|
|||
|
|
prompt,
|
|||
|
|
cpu_offload,
|
|||
|
|
# show_debug,
|
|||
|
|
img_count,
|
|||
|
|
fp16,
|
|||
|
|
device,
|
|||
|
|
all_to_device,
|
|||
|
|
ddim_steps=20,
|
|||
|
|
strength=1,
|
|||
|
|
cfg_scale=9,
|
|||
|
|
seed="",
|
|||
|
|
eta=0.0,
|
|||
|
|
a_prompt="",
|
|||
|
|
n_prompt="",
|
|||
|
|
nonEdit_random_gen_width=512,
|
|||
|
|
nonEdit_random_gen_height=512,
|
|||
|
|
):
|
|||
|
|
def prompt_replace(prompt):
|
|||
|
|
prompt = prompt.replace('“', '"')
|
|||
|
|
prompt = prompt.replace('”', '"')
|
|||
|
|
p = '"(.*?)"'
|
|||
|
|
strs = re.findall(p, prompt)
|
|||
|
|
if len(strs) == 0:
|
|||
|
|
strs = [' ']
|
|||
|
|
else:
|
|||
|
|
for s in strs:
|
|||
|
|
prompt = prompt.replace(f'"{s}"', f' * ', 1)
|
|||
|
|
return prompt
|
|||
|
|
|
|||
|
|
def check_overlap_polygon(rect_pts1, rect_pts2):
|
|||
|
|
poly1 = cv2.convexHull(rect_pts1)
|
|||
|
|
poly2 = cv2.convexHull(rect_pts2)
|
|||
|
|
rect1 = cv2.boundingRect(poly1)
|
|||
|
|
rect2 = cv2.boundingRect(poly2)
|
|||
|
|
if rect1[0] + rect1[2] >= rect2[0] and rect2[0] + rect2[2] >= rect1[0] and rect1[1] + rect1[3] >= rect2[1] and rect2[1] + rect2[3] >= rect1[1]:
|
|||
|
|
return True
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def count_lines(prompt):
|
|||
|
|
prompt = prompt.replace('“', '"')
|
|||
|
|
prompt = prompt.replace('”', '"')
|
|||
|
|
p = '"(.*?)"'
|
|||
|
|
strs = re.findall(p, prompt)
|
|||
|
|
if len(strs) == 0:
|
|||
|
|
strs = [' ']
|
|||
|
|
return len(strs)
|
|||
|
|
|
|||
|
|
def generate_rectangles(w, h, n, max_trys=200):
|
|||
|
|
img = np.zeros((h, w, 1), dtype=np.uint8)
|
|||
|
|
rectangles = []
|
|||
|
|
attempts = 0
|
|||
|
|
n_pass = 0
|
|||
|
|
low_edge = int(max(w, h)*0.3 if n <= 3 else max(w, h)*0.2) # ~150, ~100
|
|||
|
|
while attempts < max_trys:
|
|||
|
|
rect_w = min(np.random.randint(max((w*0.5)//n, low_edge), w), int(w*0.8))
|
|||
|
|
ratio = np.random.uniform(4, 10)
|
|||
|
|
rect_h = max(low_edge, int(rect_w/ratio))
|
|||
|
|
rect_h = min(rect_h, int(h*0.8))
|
|||
|
|
# gen rotate angle
|
|||
|
|
rotation_angle = 0
|
|||
|
|
rand_value = np.random.rand()
|
|||
|
|
if rand_value < 0.7:
|
|||
|
|
pass
|
|||
|
|
elif rand_value < 0.8:
|
|||
|
|
rotation_angle = np.random.randint(0, 40)
|
|||
|
|
elif rand_value < 0.9:
|
|||
|
|
rotation_angle = np.random.randint(140, 180)
|
|||
|
|
else:
|
|||
|
|
rotation_angle = np.random.randint(85, 95)
|
|||
|
|
# rand position
|
|||
|
|
x = np.random.randint(0, w - rect_w)
|
|||
|
|
y = np.random.randint(0, h - rect_h)
|
|||
|
|
# get vertex
|
|||
|
|
rect_pts = cv2.boxPoints(((rect_w/2, rect_h/2), (rect_w, rect_h), rotation_angle))
|
|||
|
|
rect_pts = np.int32(rect_pts)
|
|||
|
|
# move
|
|||
|
|
rect_pts += (x, y)
|
|||
|
|
# check boarder
|
|||
|
|
if np.any(rect_pts < 0) or np.any(rect_pts[:, 0] >= w) or np.any(rect_pts[:, 1] >= h):
|
|||
|
|
attempts += 1
|
|||
|
|
continue
|
|||
|
|
# check overlap
|
|||
|
|
if any(check_overlap_polygon(rect_pts, rp) for rp in rectangles): # type: ignore
|
|||
|
|
attempts += 1
|
|||
|
|
continue
|
|||
|
|
n_pass += 1
|
|||
|
|
img = cv2.fillPoly(img, [rect_pts], 255)
|
|||
|
|
cv2.imwrite(Random_Gen_Mask_path, 255-img[..., ::-1])
|
|||
|
|
rectangles.append(rect_pts)
|
|||
|
|
if n_pass == n:
|
|||
|
|
break
|
|||
|
|
print("attempts:", attempts)
|
|||
|
|
if len(rectangles) != n:
|
|||
|
|
raise Exception(f'Failed in auto generate positions after {attempts} attempts, try again!')
|
|||
|
|
return img
|
|||
|
|
|
|||
|
|
if not is_module_imported('AnyText_Pipeline'):
|
|||
|
|
from .AnyText_scripts.AnyText_pipeline import AnyText_Pipeline
|
|||
|
|
|
|||
|
|
#check if prompt is chinese to decide whether to load translator，检测是否为中文提示词，否则不适用翻译。
|
|||
|
|
prompt_modify = prompt_replace(prompt)
|
|||
|
|
bool_is_chinese = AnyText_Pipeline.is_chinese(self, prompt_modify)
|
|||
|
|
|
|||
|
|
device = get_device_by_name(device)
|
|||
|
|
loader_out = AnyText_Loader.split("|")
|
|||
|
|
|
|||
|
|
if bool_is_chinese == False:
|
|||
|
|
use_translator = False
|
|||
|
|
else:
|
|||
|
|
use_translator = True
|
|||
|
|
if 'damo/nlp_csanmt_translation_zh2en' in loader_out[3]:
|
|||
|
|
if not os.access(os.path.join(folder_paths.models_dir, "prompt_generator", "nlp_csanmt_translation_zh2en", "tf_ckpts", "ckpt-0.data-00000-of-00001"), os.F_OK):
|
|||
|
|
if not is_module_imported('snapshot_download'):
|
|||
|
|
from modelscope.hub.snapshot_download import snapshot_download
|
|||
|
|
snapshot_download('damo/nlp_csanmt_translation_zh2en')
|
|||
|
|
else:
|
|||
|
|
if not os.access(os.path.join(folder_paths.models_dir, "prompt_generator", "models--utrobinmv--t5_translate_en_ru_zh_small_1024", "model.safetensors"), os.F_OK):
|
|||
|
|
if not is_module_imported('hg_snapshot_download'):
|
|||
|
|
from huggingface_hub import snapshot_download as hg_snapshot_download
|
|||
|
|
hg_snapshot_download(repo_id="utrobinmv/t5_translate_en_ru_zh_small_1024")
|
|||
|
|
|
|||
|
|
pipe = AnyText_Pipeline(ckpt_path=loader_out[1], clip_path=loader_out[2], translator_path=loader_out[3], cfg_path=loader_out[4], use_translator=use_translator, device=device, use_fp16=fp16, all_to_device=all_to_device, loaded_model_tensor=self.model)
|
|||
|
|
|
|||
|
|
# tensor图片转换为numpy图片
|
|||
|
|
pos_image = comfy_tensor_Image2np_Image(self, pos_image)
|
|||
|
|
ori_image = comfy_tensor_Image2np_Image(self, ori_image)
|
|||
|
|
# 保存转换后的numpy图片到ComfyUI临时文件夹
|
|||
|
|
pos_image.save(tmp_pose_img_path)
|
|||
|
|
ori_image.save(tmp_ori_img_path)
|
|||
|
|
|
|||
|
|
ori = tmp_ori_img_path
|
|||
|
|
pos = tmp_pose_img_path
|
|||
|
|
|
|||
|
|
if mode == "text-generation":
|
|||
|
|
ori_image = None
|
|||
|
|
revise_pos = revise_pos
|
|||
|
|
else:
|
|||
|
|
revise_pos = False
|
|||
|
|
ori_image = ori
|
|||
|
|
|
|||
|
|
n_lines = count_lines(prompt)
|
|||
|
|
if Random_Gen == True:
|
|||
|
|
generate_rectangles(nonEdit_random_gen_width, nonEdit_random_gen_height, n_lines, max_trys=500)
|
|||
|
|
pos_img = Random_Gen_Mask_path
|
|||
|
|
else:
|
|||
|
|
pos_img = pos
|
|||
|
|
|
|||
|
|
# lora_path = r"D:\AI\ComfyUI_windows_portable\ComfyUI\models\loras\ys艺术\sd15_mw_bpch_扁平风格插画v1d1.safetensors"
|
|||
|
|
# lora_ratio = 1
|
|||
|
|
# lora_path_ratio = str(lora_path)+ " " + str(lora_ratio)
|
|||
|
|
# print("\033[93m", lora_path_ratio, "\033[0m")
|
|||
|
|
|
|||
|
|
params = {
|
|||
|
|
"mode": mode,
|
|||
|
|
"use_fp16": fp16,
|
|||
|
|
"Random_Gen": Random_Gen,
|
|||
|
|
"sort_priority": sort_radio,
|
|||
|
|
"revise_pos": revise_pos,
|
|||
|
|
# "show_debug": show_debug,
|
|||
|
|
"image_count": img_count,
|
|||
|
|
"ddim_steps": ddim_steps - 1,
|
|||
|
|
"image_width": nonEdit_random_gen_width,
|
|||
|
|
"image_height": nonEdit_random_gen_height,
|
|||
|
|
"strength": strength,
|
|||
|
|
"cfg_scale": cfg_scale,
|
|||
|
|
"eta": eta,
|
|||
|
|
"a_prompt": a_prompt,
|
|||
|
|
"n_prompt": n_prompt,
|
|||
|
|
# "lora_path_ratio": lora_path_ratio,
|
|||
|
|
}
|
|||
|
|
input_data = {
|
|||
|
|
"prompt": prompt,
|
|||
|
|
"seed": seed,
|
|||
|
|
"draw_pos": pos_img,
|
|||
|
|
"ori_image": ori_image,
|
|||
|
|
}
|
|||
|
|
# if show_debug ==True:
|
|||
|
|
# print(f'\033[93mloader from .util(从.util输入的loader): {AnyText_Loader}, \033[0m\n \
|
|||
|
|
# \033[93mloader_out split form loader(分割loader得到4个参数): {loader_out}, \033[0m\n \
|
|||
|
|
# \033[93mFont(字体)--loader_out[0]: {loader_out[0]}, \033[0m\n \
|
|||
|
|
# \033[93mAnyText Model(AnyText模型)--loader_out[1]: {loader_out[1]}, \033[0m\n \
|
|||
|
|
# \033[93mclip model(clip模型)--loader_out[2]: {loader_out[2]}, \033[0m\n \
|
|||
|
|
# \033[93mTranslator(翻译模型)--loader_out[3]: {loader_out[3]}, \033[0m\n \
|
|||
|
|
# \033[93myaml_file(yaml配置文件): {loader_out[4]}, \033[0m\n) \
|
|||
|
|
# \033[93mIs Chinese Input(是否中文输入): {use_translator}, \033[0m\n \
|
|||
|
|
# \033[93mNumber of text-content to generate(需要生成的文本数量): {n_lines}, \033[0m\n \
|
|||
|
|
# \033[93mpos_image location(遮罩图位置): {pos}, \033[0m\n \
|
|||
|
|
# \033[93mori_image location(原图位置): {ori}, \033[0m\n \
|
|||
|
|
# \033[93mSort Position(文本生成位置排序): {sort_radio}, \033[0m\n \
|
|||
|
|
# \033[93mEnable revise_pos(启用位置修正): {revise_pos}, \033[0m')
|
|||
|
|
x_samples, results, rtn_code, rtn_warning, debug_info, self.model = pipe(input_data, font_path=loader_out[0], cpu_offload=cpu_offload, **params)
|
|||
|
|
if rtn_code < 0:
|
|||
|
|
raise Exception(f"Error in AnyText pipeline: {rtn_warning}")
|
|||
|
|
output = pil2tensor(x_samples)
|
|||
|
|
print("\n", debug_info)
|
|||
|
|
return(output)
|
|||
|
|
|
|||
|
|
# Node class and display name mappings
|
|||
|
|
NODE_CLASS_MAPPINGS = {
|
|||
|
|
"AnyText": AnyText,
|
|||
|
|
}
|