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,
|
||
}
|