376 lines
21 KiB
Python
376 lines
21 KiB
Python
# import os
|
|
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
|
|
# import torch
|
|
# import re
|
|
import numpy as np
|
|
import cv2
|
|
# import einops
|
|
# import time
|
|
|
|
# from .cldm.ddim_hacked import DDIMSampler
|
|
from .AnyText_t3_dataset import draw_glyph, draw_glyph2
|
|
# from .AnyText_pipeline_util import check_channels, resize_image
|
|
# from .AnyText_bert_tokenizer import BasicTokenizer
|
|
# import folder_paths
|
|
# from comfy.utils import ProgressBar
|
|
# from comfy.model_management import soft_empty_cache
|
|
# from ...UL_common.common import clean_up
|
|
|
|
# checker = BasicTokenizer()
|
|
# BBOX_MAX_NUM = 20
|
|
# PLACE_HOLDER = '*'
|
|
# max_chars = 50
|
|
|
|
# comfyui_models_dir = folder_paths.models_dir
|
|
# class AnyText_Pipeline():
|
|
# def __init__(self, anytext_model, loaded_prompt, processed_prompt, texts, translator, use_translator, model_device, model_dtype, translator_device, keep_translator_loaded, keep_translator_device, Auto_Download_Path):
|
|
# self.device = model_device
|
|
# self.dtype = model_dtype
|
|
# self.use_translator = use_translator
|
|
# self.translator_path = translator
|
|
# self.translator_device = translator_device
|
|
# self.keep_translator_loaded = keep_translator_loaded
|
|
# self.keep_translator_device = keep_translator_device
|
|
# self.Auto_Download_Path = Auto_Download_Path
|
|
# self.anytext_model = anytext_model
|
|
|
|
# self.loaded_t5_translate_en_ru_zh_model_name = None
|
|
|
|
# if self.use_translator == True:
|
|
# if translator == 'damo/nlp_csanmt_translation_zh2en':
|
|
# self.trans_pipe = 'damo/nlp_csanmt_translation_zh2en'
|
|
# elif translator == "utrobinmv/t5_translate_en_ru_zh_base_200":
|
|
# self.trans_pipe = 'utrobinmv/t5_translate_en_ru_zh_base_200'
|
|
# elif translator == "utrobinmv/t5_translate_en_ru_zh_large_1024":
|
|
# self.trans_pipe = 'utrobinmv/t5_translate_en_ru_zh_large_1024'
|
|
# elif translator == "utrobinmv/t5_translate_en_ru_zh_small_1024":
|
|
# self.trans_pipe = 'utrobinmv/t5_translate_en_ru_zh_small_1024'
|
|
# else:
|
|
# self.trans_pipe = None
|
|
|
|
# self.processed_prompt = processed_prompt
|
|
# self.texts = texts
|
|
# self.loaded_prompt = loaded_prompt
|
|
|
|
# def __call__(self, font, input_tensor, **forward_params):
|
|
# tic = time.time()
|
|
|
|
# # if 'cpu' in self.anytext_model.device.type:
|
|
# # self.anytext_model.to(self.device)
|
|
|
|
# ddim_sampler = DDIMSampler(self.anytext_model, device=self.device)
|
|
|
|
# str_warning = ''
|
|
# # get inputs
|
|
# seed = input_tensor.get('seed', -1)
|
|
# self.prompt = input_tensor.get('prompt')
|
|
# draw_pos = input_tensor.get('draw_pos')
|
|
# ori_image = input_tensor.get('ori_image')
|
|
|
|
# mode = forward_params.get('mode')
|
|
# translator_device = forward_params.get('translator_device')
|
|
# Random_Gen = forward_params.get('Random_Gen')
|
|
# sort_priority = forward_params.get('sort_priority', '↕')
|
|
# show_debug = forward_params.get('show_debug', False)
|
|
# revise_pos = forward_params.get('revise_pos', False)
|
|
# img_count = forward_params.get('image_count', 1)
|
|
# ddim_steps = forward_params.get('ddim_steps', 20)
|
|
# w = forward_params.get('image_width', 512)
|
|
# h = forward_params.get('image_height', 512)
|
|
# strength = forward_params.get('strength', 1.0)
|
|
# cfg_scale = forward_params.get('cfg_scale', 9.0)
|
|
# eta = forward_params.get('eta', 0.0)
|
|
# a_prompt = forward_params.get('a_prompt', 'best quality, extremely detailed,4k, HD, supper legible text, clear text edges, clear strokes, neat writing, no watermarks')
|
|
# n_prompt = forward_params.get('n_prompt', 'low-res, bad anatomy, extra digit, fewer digits, cropped, worst quality, low quality, watermark, unreadable text, messy words, distorted text, disorganized writing, advertising picture')
|
|
|
|
# if self.loaded_prompt == None or self.loaded_prompt != self.prompt or self.processed_prompt == None or self.texts == None: # 如果prompt未改变，则不进行prompt处理，防止反复调用翻译。
|
|
# anytext_prompt, texts = self.modify_prompt(self.prompt, translator_device)
|
|
# self.processed_prompt = anytext_prompt
|
|
# self.texts = texts
|
|
# anytext_prompt =self.processed_prompt
|
|
# texts = self.texts
|
|
# self.loaded_prompt = self.prompt
|
|
|
|
# if anytext_prompt is None and texts is None:
|
|
# return None, -1, "You have input Chinese prompt but the translator is not loaded!", ""
|
|
# n_lines = len(texts)
|
|
# if mode in ['text-generation', 'gen']:
|
|
# if Random_Gen == True:
|
|
# edit_image = np.ones((h, w, 3)) * 127.5 # empty mask image
|
|
# edit_image = resize_image(edit_image, max_length=max(h, w))
|
|
# h, w = edit_image.shape[:2]
|
|
# else:
|
|
# # edit_image = cv2.imread(draw_pos)[..., ::-1]
|
|
# edit_image = draw_pos[..., ::-1]
|
|
# edit_image = resize_image(edit_image, max_length=max(h, w))
|
|
# h, w = edit_image.shape[:2]
|
|
# edit_image = np.ones((h, w, 3)) * 127.5 # empty mask image
|
|
# elif mode in ['text-editing', 'edit']:
|
|
# if draw_pos is None or ori_image is None:
|
|
# return None, -1, "Reference image and position image are needed for text editing!", ""
|
|
# # if isinstance(ori_image, str):
|
|
# if isinstance(ori_image, np.ndarray):
|
|
# # ori_image = cv2.imread(ori_image)[..., ::-1]
|
|
# ori_image = ori_image[..., ::-1]
|
|
# assert ori_image is not None, f"Can't read ori_image image from{ori_image}!"
|
|
# elif isinstance(ori_image, torch.Tensor):
|
|
# ori_image = ori_image.cpu().numpy()
|
|
# # else:
|
|
# # assert isinstance(ori_image, np.ndarray), f'Unknown format of ori_image: {type(ori_image)}'
|
|
# edit_image = ori_image.clip(1, 255) # for mask reason
|
|
# edit_image = check_channels(edit_image)
|
|
# edit_image = resize_image(edit_image, max_length=max(h, w)) # make w h multiple of 64, resize if w or h > max_length
|
|
# h, w = edit_image.shape[:2] # change h, w by input ref_img
|
|
# # preprocess pos_imgs(if numpy, make sure it's white pos in black bg)
|
|
# if draw_pos is None:
|
|
# pos_imgs = np.zeros((w, h, 1))
|
|
# # if isinstance(draw_pos, str):
|
|
# if isinstance(draw_pos, np.ndarray):
|
|
# # draw_pos = cv2.imread(draw_pos)[..., ::-1]
|
|
# draw_pos = draw_pos[..., ::-1]
|
|
# draw_pos = resize_image(draw_pos, max_length=max(h, w))
|
|
# draw_pos = cv2.resize(draw_pos, (w, h))
|
|
# assert draw_pos is not None, f"Can't read draw_pos image from{draw_pos}!"
|
|
# pos_imgs = 255-draw_pos
|
|
# elif isinstance(draw_pos, torch.Tensor):
|
|
# pos_imgs = draw_pos.cpu().numpy()
|
|
# # else:
|
|
# # assert isinstance(draw_pos, np.ndarray), f'Unknown format of draw_pos: {type(draw_pos)}'
|
|
|
|
# pos_imgs = pos_imgs[..., 0:1]
|
|
# pos_imgs = cv2.convertScaleAbs(pos_imgs)
|
|
# _, pos_imgs = cv2.threshold(pos_imgs, 254, 255, cv2.THRESH_BINARY)
|
|
# # seprate pos_imgs
|
|
# pos_imgs = separate_pos_imgs(pos_imgs, sort_priority)
|
|
# if len(pos_imgs) == 0:
|
|
# pos_imgs = [np.zeros((h, w, 1))]
|
|
# if len(pos_imgs) < n_lines:
|
|
# if n_lines == 1 and texts[0] == ' ':
|
|
# # pass # text-to-image without text
|
|
# print('\033[93m', f'Warning: text-to-image without text.', '\033[0m')
|
|
# else:
|
|
# # return None, -1, f'Found {len(pos_imgs)} positions that < needed {n_lines} from prompt, check and try again!', ''
|
|
# raise ValueError(f'Found {len(pos_imgs)} positions that < needed {n_lines} from prompt, check and try again(手绘遮罩数少于要绘制的文本数，检查修改再重试)!')
|
|
# elif len(pos_imgs) > n_lines:
|
|
# # str_warning = f'Warning: found {len(pos_imgs)} positions that > needed {n_lines} from prompt.'
|
|
# print('\033[93m', f'Warning: found {len(pos_imgs)} positions that > needed {n_lines} from prompt.', '\033[0m')
|
|
# # get pre_pos, poly_list, hint that needed for anytext
|
|
# pre_pos = []
|
|
# poly_list = []
|
|
# for input_pos in pos_imgs:
|
|
# if input_pos.mean() != 0:
|
|
# input_pos = input_pos[..., np.newaxis] if len(input_pos.shape) == 2 else input_pos
|
|
# poly, pos_img = find_polygon(input_pos)
|
|
# pre_pos += [pos_img/255.]
|
|
# poly_list += [poly]
|
|
# else:
|
|
# pre_pos += [np.zeros((h, w, 1))]
|
|
# poly_list += [None]
|
|
# np_hint = np.sum(pre_pos, axis=0).clip(0, 1)
|
|
# # prepare info dict
|
|
# info = {}
|
|
# info['glyphs'] = []
|
|
# info['gly_line'] = []
|
|
# info['positions'] = []
|
|
# info['n_lines'] = [len(texts)]*img_count
|
|
# gly_pos_imgs = []
|
|
# for i in range(len(texts)):
|
|
# text = texts[i]
|
|
# if len(text) > max_chars:
|
|
# str_warning = f'"{text}" length > max_chars: {max_chars}, will be cut off...'
|
|
# text = text[:max_chars]
|
|
# gly_scale = 2
|
|
# if pre_pos[i].mean() != 0:
|
|
# gly_line = draw_glyph(font, text)
|
|
# glyphs, _ = draw_glyph2(font, text, poly_list[i], scale=gly_scale, width=w, height=h, add_space=False)
|
|
# gly_pos_img = cv2.drawContours(glyphs*255, [poly_list[i]*gly_scale], 0, (255, 255, 255), 1)
|
|
# if revise_pos:
|
|
# resize_gly = cv2.resize(glyphs, (pre_pos[i].shape[1], pre_pos[i].shape[0]))
|
|
# new_pos = cv2.morphologyEx((resize_gly*255).astype(np.uint8), cv2.MORPH_CLOSE, kernel=np.ones((resize_gly.shape[0]//10, resize_gly.shape[1]//10), dtype=np.uint8), iterations=1)
|
|
# new_pos = new_pos[..., np.newaxis] if len(new_pos.shape) == 2 else new_pos
|
|
# contours, _ = cv2.findContours(new_pos, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|
# if len(contours) != 1:
|
|
# str_warning = f'Fail to revise position {i} to bounding rect, remain position unchanged...'
|
|
# else:
|
|
# rect = cv2.minAreaRect(contours[0])
|
|
# poly = np.int0(cv2.boxPoints(rect))
|
|
# pre_pos[i] = cv2.drawContours(new_pos, [poly], -1, 255, -1) / 255.
|
|
# gly_pos_img = cv2.drawContours(glyphs*255, [poly*gly_scale], 0, (255, 255, 255), 1)
|
|
# gly_pos_imgs += [gly_pos_img] # for show
|
|
# else:
|
|
# glyphs = np.zeros((h*gly_scale, w*gly_scale, 1))
|
|
# gly_line = np.zeros((80, 512, 1))
|
|
# gly_pos_imgs += [np.zeros((h*gly_scale, w*gly_scale, 1))] # for show
|
|
# pos = pre_pos[i]
|
|
# info['glyphs'] += [self.arr2tensor(glyphs, img_count, self.dtype)]
|
|
# info['gly_line'] += [self.arr2tensor(gly_line, img_count, self.dtype)]
|
|
# info['positions'] += [self.arr2tensor(pos, img_count, self.dtype)]
|
|
# # get masked_x
|
|
# masked_img = ((edit_image.astype(np.float32) / 127.5) - 1.0)*(1-np_hint)
|
|
# masked_img = np.transpose(masked_img, (2, 0, 1))
|
|
# masked_img = torch.from_numpy(masked_img.copy()).float().to(self.device, self.dtype)
|
|
|
|
# encoder_posterior = self.anytext_model.encode_first_stage(masked_img[None, ...])
|
|
# masked_x = (self.anytext_model.get_first_stage_encoding(encoder_posterior).detach()).to(self.dtype) #masked latent
|
|
|
|
# info['masked_x'] = torch.cat([masked_x for _ in range(img_count)], dim=0)
|
|
|
|
# hint = self.arr2tensor(np_hint, img_count, self.dtype)
|
|
|
|
# cond = self.anytext_model.get_learned_conditioning(dict(c_concat=[hint], c_crossattn=[[anytext_prompt + ' , ' + a_prompt] * img_count], text_info=info))
|
|
# un_cond = self.anytext_model.get_learned_conditioning(dict(c_concat=[hint], c_crossattn=[[n_prompt] * img_count], text_info=info))
|
|
# shape = (4, h // 8, w // 8)
|
|
# self.anytext_model.control_scales = ([strength] * 13)
|
|
|
|
# self.anytext_model.embedding_manager.to('cpu')
|
|
# self.anytext_model.cond_stage_model.to('cpu')
|
|
# self.anytext_model.first_stage_model.to('cpu')
|
|
# soft_empty_cache(True)
|
|
# clean_up()
|
|
# self.anytext_model.model.diffusion_model.to(self.device)
|
|
# self.anytext_model.control_model.to(self.device)
|
|
|
|
# pbar = ProgressBar(int(ddim_steps))
|
|
# def callback(*_):
|
|
# pbar.update(1)
|
|
# from pytorch_lightning import seed_everything
|
|
# seed_everything(seed)
|
|
# samples, intermediates = ddim_sampler.sample(ddim_steps, img_count,
|
|
# shape, cond, verbose=False, eta=eta,
|
|
# unconditional_guidance_scale=cfg_scale,
|
|
# unconditional_conditioning=un_cond,
|
|
# callback=callback, # 后端代码有timesteps的callback
|
|
# )
|
|
|
|
# self.anytext_model.model.diffusion_model.to('cpu')
|
|
# self.anytext_model.control_model.to('cpu')
|
|
# soft_empty_cache(True)
|
|
# clean_up()
|
|
# self.anytext_model.first_stage_model.to(self.device)
|
|
|
|
# x_samples = self.anytext_model.decode_first_stage(samples.to(self.dtype))
|
|
# self.anytext_model.first_stage_model.to('cpu')
|
|
# soft_empty_cache(True)
|
|
# clean_up()
|
|
# x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
|
|
# results = [x_samples[i] for i in range(img_count)]
|
|
# # if mode == 'edit' and False: # replace backgound in text editing but not ideal yet
|
|
# # results = [r*np_hint+edit_image*(1-np_hint) for r in results]
|
|
# # results = [r.clip(0, 255).astype(np.uint8) for r in results]
|
|
# if len(gly_pos_imgs) > 0 and show_debug:
|
|
# glyph_bs = np.stack(gly_pos_imgs, axis=2)
|
|
# glyph_img = np.sum(glyph_bs, axis=2) * 255
|
|
# glyph_img = glyph_img.clip(0, 255).astype(np.uint8)
|
|
# results += [np.repeat(glyph_img, 3, axis=2)]
|
|
# input_prompt = anytext_prompt
|
|
# for t in texts:
|
|
# input_prompt = input_prompt.replace('*', f'"{t}"', 1)
|
|
# print(f'Prompt: {input_prompt}')
|
|
# # debug_info
|
|
# if not show_debug:
|
|
# debug_info = ''
|
|
# else:
|
|
# debug_info = f'\033[93mPrompt(提示词): {input_prompt}\n\033[0m \
|
|
# \033[93mSize(尺寸): {w}x{h}\n\033[0m \
|
|
# \033[93mImage Count(生成数量): {img_count}\n\033[0m \
|
|
# \033[93mSeed(种子): {seed}\n\033[0m \
|
|
# \033[93mModel Precision(模型精度): {self.dtype}\n\033[0m \
|
|
# \033[93mCost Time(生成耗时): {(time.time()-tic):.2f}s\033[0m'
|
|
# rst_code = 1 if str_warning else 0
|
|
|
|
# return (x_samples, results, rst_code, str_warning, debug_info, self.processed_prompt, self.texts, self.loaded_prompt, )
|
|
|
|
# def modify_prompt(self, prompt, device):
|
|
# 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' {PLACE_HOLDER} ', 1)
|
|
# if self.is_chinese(prompt):
|
|
# if self.trans_pipe is None:
|
|
# return None, None
|
|
# old_prompt = prompt
|
|
|
|
# from ...Data_Process.utils import t5_translate_en_ru_zh, nlp_csanmt_translation_zh2en
|
|
# self.t5_translate_tensor = None
|
|
|
|
# if self.trans_pipe == 'utrobinmv/t5_translate_en_ru_zh_small_1024':
|
|
# zh2en_path = os.path.join(folder_paths.models_dir, "prompt_generator", "models--utrobinmv--t5_translate_en_ru_zh_small_1024")
|
|
# if not os.path.exists(os.path.join(zh2en_path, "model.safetensors")) and not self.Auto_Download_Path:
|
|
# zh2en_path = 'utrobinmv/t5_translate_en_ru_zh_small_1024'
|
|
# prompt, self.t5_translate_tensor, self.loaded_t5_translate_en_ru_zh_mdoel_name = t5_translate_en_ru_zh('en', prompt + ' .', zh2en_path, self.translator_device, self.keep_translator_loaded, self.t5_translate_tensor, self.loaded_t5_translate_en_ru_zh_model_name, keep_model_device=self.keep_translator_device)
|
|
# prompt = prompt[0]
|
|
|
|
# elif self.trans_pipe == 'utrobinmv/t5_translate_en_ru_zh_base_200':
|
|
# zh2en_path = os.path.join(folder_paths.models_dir, "prompt_generator", "models--utrobinmv--t5_translate_en_ru_zh_base_200")
|
|
# if not os.path.exists(os.path.join(zh2en_path, "model.safetensors")) and not self.Auto_Download_Path:
|
|
# zh2en_path = 'utrobinmv/t5_translate_en_ru_zh_base_200'
|
|
# prompt, self.t5_translate_tensor, self.loaded_t5_translate_en_ru_zh_mdoel_name = t5_translate_en_ru_zh('en', prompt + ' .', zh2en_path, self.translator_device, self.keep_translator_loaded, self.t5_translate_tensor, self.loaded_t5_translate_en_ru_zh_model_name, keep_model_device=self.keep_translator_device)
|
|
# prompt = prompt[0]
|
|
|
|
# elif self.trans_pipe == 'utrobinmv/t5_translate_en_ru_zh_large_1024':
|
|
# zh2en_path = os.path.join(folder_paths.models_dir, "prompt_generator", "models--utrobinmv--t5_translate_en_ru_zh_large_1024")
|
|
# if not os.path.exists(os.path.join(zh2en_path, "model.safetensors")) and not self.Auto_Download_Path:
|
|
# zh2en_path = 'utrobinmv/t5_translate_en_ru_zh_large_1024'
|
|
# prompt, self.t5_translate_tensor, self.loaded_t5_translate_en_ru_zh_mdoel_name = t5_translate_en_ru_zh('en', prompt + ' .', zh2en_path, self.translator_device, self.keep_translator_loaded, self.t5_translate_tensor, self.loaded_t5_translate_en_ru_zh_model_name, keep_model_device=self.keep_translator_device)
|
|
# prompt = prompt[0]
|
|
|
|
# else:
|
|
# nlp_device = 'gpu'
|
|
# if 'cpu' in device.type:
|
|
# nlp_device = 'cpu'
|
|
# zh2en_path = os.path.join(comfyui_models_dir, 'prompt_generator', 'modelscope--damo--nlp_csanmt_translation_zh2en')
|
|
# if not os.path.exists(os.path.join(zh2en_path, "tf_ckpts", "ckpt-0.data-00000-of-00001")) and not self.Auto_Download_Path:
|
|
# zh2en_path = "damo/nlp_csanmt_translation_zh2en"
|
|
# prompt = nlp_csanmt_translation_zh2en(nlp_device, prompt + ' .', zh2en_path)['translation']
|
|
# print(f'Translate: {old_prompt} --> {prompt}')
|
|
# return prompt, strs
|
|
|
|
# def is_chinese(self, text):
|
|
# text = checker._clean_text(text)
|
|
# for char in text:
|
|
# cp = ord(char)
|
|
# if checker._is_chinese_char(cp):
|
|
# return True
|
|
# return False
|
|
|
|
# def arr2tensor(self, arr, bs, dtype):
|
|
# arr = np.transpose(arr, (2, 0, 1))
|
|
# _arr = torch.from_numpy(arr.copy()).float().to(self.device)
|
|
# _arr = (torch.stack([_arr for _ in range(bs)], dim=0)).to(dtype)
|
|
# return _arr
|
|
|
|
def separate_pos_imgs(img, sort_priority, gap=102):
|
|
num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img)
|
|
components = []
|
|
for label in range(1, num_labels):
|
|
component = np.zeros_like(img)
|
|
component[labels == label] = 255
|
|
components.append((component, centroids[label]))
|
|
if sort_priority == '↕':
|
|
fir, sec = 1, 0 # top-down first
|
|
elif sort_priority == '↔':
|
|
fir, sec = 0, 1 # left-right first
|
|
components.sort(key=lambda c: (c[1][fir]//gap, c[1][sec]//gap))
|
|
sorted_components = [c[0] for c in components]
|
|
return sorted_components
|
|
|
|
def find_polygon(image, min_rect=False):
|
|
contours, hierarchy = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|
max_contour = max(contours, key=cv2.contourArea) # get contour with max area
|
|
if min_rect:
|
|
# get minimum enclosing rectangle
|
|
rect = cv2.minAreaRect(max_contour)
|
|
poly = np.int0(cv2.boxPoints(rect))
|
|
else:
|
|
# get approximate polygon
|
|
epsilon = 0.01 * cv2.arcLength(max_contour, True)
|
|
poly = cv2.approxPolyDP(max_contour, epsilon, True)
|
|
n, _, xy = poly.shape
|
|
poly = poly.reshape(n, xy)
|
|
cv2.drawContours(image, [poly], -1, 255, -1)
|
|
return poly, image |