214 lines
9.6 KiB
Python
214 lines
9.6 KiB
Python
import os
|
||
import folder_paths
|
||
import torch
|
||
import numpy as np
|
||
import time
|
||
from PIL import Image
|
||
|
||
current_directory = os.path.dirname(os.path.abspath(__file__))
|
||
comfyui_models_dir = folder_paths.models_dir
|
||
comfy_temp_dir = folder_paths.get_temp_directory()
|
||
temp_txt_path = os.path.join(comfy_temp_dir, "AnyText_temp.txt")
|
||
|
||
class AnyText_loader:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
font_list = os.listdir(os.path.join(comfyui_models_dir, "fonts"))
|
||
checkpoints_list = folder_paths.get_filename_list("checkpoints")
|
||
clip_list = os.listdir(os.path.join(comfyui_models_dir, "clip"))
|
||
font_list.insert(0, "Auto_DownLoad")
|
||
checkpoints_list.insert(0, "Auto_DownLoad")
|
||
clip_list.insert(0, "Auto_DownLoad")
|
||
|
||
return {
|
||
"required": {
|
||
"font": (font_list, ),
|
||
"ckpt_name": (checkpoints_list, ),
|
||
"clip": (clip_list, ),
|
||
"translator": (["utrobinmv/t5_translate_en_ru_zh_small_1024", "damo/nlp_csanmt_translation_zh2en"],{"default": "utrobinmv/t5_translate_en_ru_zh_small_1024"}),
|
||
# "show_debug": ("BOOLEAN", {"default": False}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("AnyText_Loader", )
|
||
RETURN_NAMES = ("AnyText_Loader", )
|
||
FUNCTION = "AnyText_loader_fn"
|
||
CATEGORY = "ExtraModels/AnyText"
|
||
TITLE = "AnyText Loader"
|
||
|
||
def AnyText_loader_fn(self,
|
||
font,
|
||
ckpt_name,
|
||
clip,
|
||
translator,
|
||
# show_debug
|
||
):
|
||
font_path = os.path.join(comfyui_models_dir, "fonts", font)
|
||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||
cfg_path = os.path.join(current_directory, 'models_yaml', 'anytext_sd15.yaml')
|
||
if clip != 'Auto_DownLoad':
|
||
clip_path = os.path.join(comfyui_models_dir, "clip", clip)
|
||
else:
|
||
clip_path = clip
|
||
if translator != 'Auto_DownLoad':
|
||
translator_path = os.path.join(comfyui_models_dir, "prompt_generator", translator)
|
||
else:
|
||
translator_path = translator
|
||
|
||
#将输入参数合并到一个参数里面传递到.nodes
|
||
loader = (font_path + "|" + str(ckpt_path) + "|" + clip_path + "|" + translator_path + "|" + cfg_path)
|
||
|
||
# if show_debug == True:
|
||
# print(f'\033[93mloader(合并后的4个输入参数，传递给nodes): {loader} \033[0m\n \
|
||
# \033[93mfont_path(字体): {font_path} \033[0m\n \
|
||
# \033[93mckpt_path(AnyText模型): {ckpt_path} \033[0m\n \
|
||
# \033[93mclip_path(clip模型): {clip_path} \033[0m\n \
|
||
# \033[93mtranslator_path(翻译模型): {translator_path} \033[0m\n \
|
||
# \033[93myaml_file(yaml配置文件): {cfg_path} \033[0m\n')
|
||
return (loader, )
|
||
|
||
class AnyText_translator:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model": (["utrobinmv/t5_translate_en_ru_zh_small_1024", "damo/nlp_csanmt_translation_zh2en"],{"default": "utrobinmv/t5_translate_en_ru_zh_small_1024"}),
|
||
"prompt": ("STRING", {"default": "这里是单批次翻译文本输入。\n声明补充说，沃伦的同事都深感震惊，并且希望他能够投案自首。\n尽量输入单句文本，如果是多句长文本建议人工分句，否则可能出现漏译或未译等情况！！！\n使用换行，效果可能更佳。", "multiline": True}),
|
||
"Batch_prompt": ("STRING", {"default": "这里是多批次翻译文本输入，使用换行进行分割。\n天上掉馅饼啦，快去看超人！！！\n飞流直下三千尺，疑似银河落九天。\n启用Batch_Newline表示输出的翻译会按换行输入进行二次换行,否则是用空格合并起来的整篇文本。", "multiline": True}),
|
||
"t5_Target_Language": (["en", "zh", "ru", ],{"default": "en"}),
|
||
"if_Batch": ("BOOLEAN", {"default": False}),
|
||
"Batch_Newline" :("BOOLEAN", {"default": True}),
|
||
"device": (["auto", "cuda", "cpu", "mps", "xpu"],{"default": "auto"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("text",)
|
||
CATEGORY = "ExtraModels/AnyText"
|
||
FUNCTION = "AnyText_translator"
|
||
TITLE = "AnyText Translator"
|
||
|
||
def AnyText_translator(self, prompt, model, Batch_prompt, if_Batch, device, Batch_Newline, t5_Target_Language):
|
||
device = get_device_by_name(device)
|
||
# 使用换行(\n)作为分隔符
|
||
Batch_prompt = Batch_prompt.split("\n")
|
||
input_sequence = prompt
|
||
if model == 'damo/nlp_csanmt_translation_zh2en':
|
||
sttime = time.time()
|
||
if if_Batch == True:
|
||
input_sequence = Batch_prompt
|
||
# 用特定的连接符<SENT_SPLIT>，将多个句子进行串联
|
||
input_sequence = '<SENT_SPLIT>'.join(input_sequence)
|
||
if os.access(os.path.join(comfyui_models_dir, "prompt_generator", "nlp_csanmt_translation_zh2en", "tf_ckpts", "ckpt-0.data-00000-of-00001"), os.F_OK):
|
||
zh2en_path = os.path.join(comfyui_models_dir, 'prompt_generator', 'nlp_csanmt_translation_zh2en')
|
||
else:
|
||
zh2en_path = "damo/nlp_csanmt_translation_zh2en"
|
||
|
||
if not is_module_imported('pipeline'):
|
||
from modelscope.pipelines import pipeline
|
||
if not is_module_imported('Tasks'):
|
||
from modelscope.utils.constant import Tasks
|
||
if device == 'cuda':
|
||
pipeline_ins = pipeline(task=Tasks.translation, model=zh2en_path, device='gpu')
|
||
outputs = pipeline_ins(input=input_sequence)
|
||
if if_Batch == True:
|
||
results = outputs['translation'].split('<SENT_SPLIT>')
|
||
if Batch_Newline == True:
|
||
results = '\n\n'.join(results)
|
||
else:
|
||
results = ' '.join(results)
|
||
else:
|
||
results = outputs['translation']
|
||
endtime = time.time()
|
||
print("\033[93mTime for translating(翻译耗时): ", endtime - sttime, "\033[0m")
|
||
del pipeline_ins
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
else:
|
||
if if_Batch == True:
|
||
input_sequence = Batch_prompt
|
||
# 用特定的连接符<SENT_SPLIT>，将多个句子进行串联
|
||
input_sequence = '|'.join(input_sequence)
|
||
self.zh2en_path = os.path.join(folder_paths.models_dir, "prompt_generator", "models--utrobinmv--t5_translate_en_ru_zh_small_1024")
|
||
if not os.access(os.path.join(self.zh2en_path, "model.safetensors"), os.F_OK):
|
||
self.zh2en_path = "utrobinmv/t5_translate_en_ru_zh_small_1024"
|
||
outputs = t5_translate_en_ru_zh(t5_Target_Language, input_sequence, self.zh2en_path, device)[0]
|
||
if if_Batch == True:
|
||
results = outputs.split('| ')
|
||
if Batch_Newline == True:
|
||
results = '\n\n'.join(results)
|
||
else:
|
||
results = ' '.join(results)
|
||
else:
|
||
results = outputs
|
||
|
||
with open(temp_txt_path, "w", encoding="UTF-8") as text_file:
|
||
text_file.write(results)
|
||
return (results, )
|
||
|
||
def is_module_imported(module_name):
|
||
try:
|
||
__import__(module_name)
|
||
except ImportError:
|
||
return False
|
||
else:
|
||
return True
|
||
|
||
def pil2tensor(image):
|
||
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
||
|
||
def is_folder_exist(folder_path):
|
||
result = os.path.exists(folder_path)
|
||
return result
|
||
|
||
def get_device_by_name(device):
|
||
if device == 'auto':
|
||
try:
|
||
device = "cpu"
|
||
if torch.cuda.is_available():
|
||
device = "cuda"
|
||
elif torch.backends.mps.is_available():
|
||
device = "mps"
|
||
elif torch.xpu.is_available():
|
||
device = "xpu"
|
||
except:
|
||
raise AttributeError("What's your device(到底用什么设备跑的)？")
|
||
print("\033[93mUse Device(使用设备):", device, "\033[0m")
|
||
return device
|
||
|
||
# Node class and display name mappings
|
||
NODE_CLASS_MAPPINGS = {
|
||
"AnyText_loader": AnyText_loader,
|
||
"AnyText_translator": AnyText_translator,
|
||
}
|
||
|
||
def t5_translate_en_ru_zh(Target_Language, prompt, model_path, device):
|
||
|
||
# prefix = 'translate to en: '
|
||
sttime = time.time()
|
||
if not is_module_imported('T5ForConditionalGeneration'):
|
||
from transformers import T5ForConditionalGeneration
|
||
if not is_module_imported('T5Tokenizer'):
|
||
from transformers import T5Tokenizer
|
||
model = T5ForConditionalGeneration.from_pretrained(model_path,)
|
||
tokenizer = T5Tokenizer.from_pretrained(model_path)
|
||
if Target_Language == 'zh':
|
||
prefix = 'translate to zh: '
|
||
elif Target_Language == 'en':
|
||
prefix = 'translate to en: '
|
||
else:
|
||
prefix = 'translate to ru: '
|
||
src_text = prefix + prompt
|
||
input_ids = tokenizer(src_text, return_tensors="pt")
|
||
generated_tokens = model.generate(**input_ids).to(device, torch.float32)
|
||
result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
|
||
model.to('cpu')
|
||
endtime = time.time()
|
||
print("\033[93mTime for translating(翻译耗时): ", endtime - sttime, "\033[0m")
|
||
return result
|
||
|
||
def comfy_tensor_Image2np_Image(self,comfy_tensor_image):
|
||
comfyimage = comfy_tensor_image.numpy()[0] * 255
|
||
image_np = comfyimage.astype(np.uint8)
|
||
image = Image.fromarray(image_np)
|
||
return image |