194 lines
11 KiB
Python
194 lines
11 KiB
Python
import os
|
|
import folder_paths
|
|
import time
|
|
import torch
|
|
from ..UL_common.common import get_dtype_by_name
|
|
from comfy.model_management import get_torch_device, text_encoder_offload_device, soft_empty_cache
|
|
from comfy.utils import ProgressBar
|
|
|
|
current_directory = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
class UL_TranslatorLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": (["utrobinmv/t5_translate_en_ru_zh_small_1024", "iic/nlp_csanmt_translation_zh2en", "iic/nlp_csanmt_translation_en2zh", "iic/nlp_csanmt_translation_en2zh_base", "botisan-ai/mt5-translate-zh-yue", "facebook/nllb-200-distilled-600M", "facebook/nllb-200-distilled-1.3B", "facebook/nllb-200-3.3B"], {"default": "utrobinmv/t5_translate_en_ru_zh_small_1024", "tooltip": ""}),
|
|
"weight_dtype": (["auto", "fp16", "bf16", "fp32", "fp8_e4m3fn", "fp8_e4m3fnuz", "fp8_e5m2", "fp8_e5m2fnuz"],{"default":"auto", "tooltip": "For t5、nllb."}),
|
|
"Auto_Download_Path": ("BOOLEAN", {"default": True, "label_on": "models_local本地", "label_off": ".cache缓存", "tooltip": "Download to `ComfyUI\models\prompt_generator` or huggingface cache dir.\n下载模型到`ComfyUI\models\prompt_generator`或者缓存到huggingface缓存路径。"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TRANSLATE_MODEL", "STRING", )
|
|
RETURN_NAMES = ("model", "model_dir", )
|
|
CATEGORY = "UL Group/Data Process"
|
|
FUNCTION = "main"
|
|
TITLE = "Translator Loader"
|
|
DESCRIPTION = ""
|
|
|
|
def main(self, model_name, weight_dtype, Auto_Download_Path):
|
|
dtype = get_dtype_by_name(weight_dtype)
|
|
prompt_generator_dir = os.path.join(folder_paths.models_dir, "prompt_generator")
|
|
if 'nlp_csanmt_translation' in model_name.lower():
|
|
from modelscope.pipelines import pipeline
|
|
from modelscope.utils.constant import Tasks
|
|
|
|
nlp_device = 'cuda' if 'cuda' in get_torch_device().type else "cpu"
|
|
|
|
model_dir = os.path.join(prompt_generator_dir, "modelscope--damo--nlp_csanmt_translation_zh2en") if "zh2en" in model_name.lower() else os.path.join(prompt_generator_dir, "modelscope--damo--nlp_csanmt_translation_en2zh_base") if "en2zh_base" in model_name.lower() else os.path.join(prompt_generator_dir, "modelscope--damo--nlp_csanmt_translation_en2zh")
|
|
|
|
if not os.path.exists(os.path.join(model_dir, "tf_ckpts", "ckpt-0.data-00000-of-00001")):
|
|
if Auto_Download_Path:
|
|
from modelscope.hub.snapshot_download import snapshot_download
|
|
snapshot_download(
|
|
model_id=model_name,
|
|
local_dir=model_dir,
|
|
)
|
|
else:
|
|
model_dir = model_name
|
|
|
|
model = pipeline(task=Tasks.translation, model=model_dir, device=nlp_device)
|
|
tokenizer = "None"
|
|
model_type = "nlp"
|
|
elif model_name == "botisan-ai/mt5-translate-zh-yue":
|
|
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
|
|
|
model_dir = os.path.join(prompt_generator_dir, "models--botisan-ai--mt5-translate-zh-yue")
|
|
|
|
if not os.path.exists(os.path.join(model_dir, "model.safetensors")):
|
|
if Auto_Download_Path:
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(
|
|
repo_id=model_name,
|
|
local_dir=model_dir,
|
|
ignore_patterns=["*.bin"]
|
|
)
|
|
else:
|
|
model_dir = model_name
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_dir)
|
|
model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
|
|
model_type = "mt5"
|
|
elif "nllb-200" in model_name.lower():
|
|
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
|
|
|
|
|
model_dir = os.path.join(prompt_generator_dir, f'models--{model_name.replace("/", "--")}')
|
|
if not os.path.exists(os.path.join(model_dir, "pytorch_model.bin")):
|
|
if Auto_Download_Path:
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(
|
|
repo_id=model_name,
|
|
local_dir=model_dir,
|
|
)
|
|
else:
|
|
model_dir = model_name
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_dir)
|
|
model = AutoModelForSeq2SeqLM.from_pretrained(model_dir).to(dtype)
|
|
model_type = "nllb"
|
|
elif "t5_translate_en_ru_zh" in model_name.lower():
|
|
from transformers import T5ForConditionalGeneration, T5Tokenizer
|
|
|
|
model_dir = os.path.join(prompt_generator_dir, f'models--{model_name.replace("/", "--")}')
|
|
if not os.path.exists(os.path.join(model_dir, "model.safetensors")):
|
|
if Auto_Download_Path:
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(
|
|
repo_id=model_name,
|
|
local_dir=model_dir,
|
|
)
|
|
else:
|
|
model_dir = model_name
|
|
|
|
tokenizer = T5Tokenizer.from_pretrained(model_dir)
|
|
model = T5ForConditionalGeneration.from_pretrained(model_dir).to(dtype)
|
|
model_type = "t5"
|
|
|
|
model = {
|
|
"model": model,
|
|
"tokenizer": tokenizer,
|
|
"model_type": model_type,
|
|
"dtype": dtype,
|
|
}
|
|
|
|
return (model, model_dir, )
|
|
|
|
class UL_Translator:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
from .Translate_Scripts.translate_backbone import language_list
|
|
return {
|
|
"required": {
|
|
"model": ("TRANSLATE_MODEL", ),
|
|
"string": ("STRING", {'default': '', "multiline": True, "dynamicPrompts": True}),
|
|
"src_language": (language_list, {"default": "", "tooltip": "For nllb."}),
|
|
"tgt_language": (language_list, {"default": "", "tooltip": "For nllb、t5."}),
|
|
"max_new_tokens": ("INT", {"default": 512, "min": 0, "max": 0xffffffffffffffff, "tooltip": "For mt5、nllb."}),
|
|
"use_fast": ("BOOLEAN", {"default": False, "label_on": "yes", "label_off": "no", "tooltip": "For nllb."}),
|
|
"keep_model_loaded": ("BOOLEAN", {"default": True, "label_on": "yes", "label_off": "no", "tooltip": "Warning: do not delete model unless this node no longer needed, it will try release device_memory and ram. if checked and want to continue node generation, use ComfyUI-Manager `Free model and node cache` to reset node state or change parameter in Loader node to activate.\n注意：仅在这个节点不再需要时删除模型，将尽量尝试释放系统内存和设备专用内存。如果删除后想继续使用此节点，使用ComfyUI-Manager插件的`Free model and node cache`重置节点状态或者更换模型加载节点的参数来激活。"}),
|
|
"keep_model_device": ("BOOLEAN", {"default": True, "label_on": "comfy", "label_off": "device", "tooltip": "Keep model in comfy_auto_unet_offload_device (HIGH_VRAM: device, Others: cpu) or device_memory after generation.\n生图完成后，模型转移到comfy自动选择设备(HIGH_VRAM: device, 其他: cpu)或者保留在设备专用内存上。"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("string",)
|
|
CATEGORY = "UL Group/Data Process"
|
|
FUNCTION = "main"
|
|
TITLE = "Translator"
|
|
DESCRIPTION = ""
|
|
|
|
def main(self, model, string, src_language, tgt_language, max_new_tokens, keep_model_loaded, keep_model_device, use_fast):
|
|
from .Translate_Scripts.translate_backbone import nlp_tranlate, mt5_translate, nllb_translate, t5_language_map, t5_translate
|
|
sttime = time.time()
|
|
input_text = [string]
|
|
output_text = ""
|
|
if "\n" in string:
|
|
input_text = []
|
|
for txt in string.split("\n"):
|
|
input_text.append(txt)
|
|
pbar = ProgressBar(len(input_text))
|
|
|
|
for input in input_text:
|
|
if model['model_type'] == "nlp":
|
|
result = nlp_tranlate(model['model'], input)
|
|
output_text+=f'{result}\n'
|
|
pbar.update(1)
|
|
elif model['model_type'] == "mt5":
|
|
result = mt5_translate(model['model'], model['tokenizer'], input, max_new_tokens)
|
|
output_text+=f'{result}\n'
|
|
pbar.update(1)
|
|
elif model['model_type'] == "nllb":
|
|
if src_language != tgt_language:
|
|
result = nllb_translate(model['model'], model['tokenizer'], input, src_language, tgt_language, use_fast, max_new_tokens)
|
|
output_text+=f'{result}\n'
|
|
pbar.update(1)
|
|
else:
|
|
output_text = string
|
|
elif model['model_type'] == "t5":
|
|
if tgt_language in t5_language_map.keys():
|
|
result = t5_translate(model['model'], model['tokenizer'], input, t5_language_map[tgt_language])
|
|
output_text+=f'{result}\n'
|
|
pbar.update(1)
|
|
else:
|
|
output_text = string
|
|
|
|
endtime = time.time()
|
|
print("\n\033[93m翻译耗时：", endtime - sttime, "\033[0m")
|
|
if keep_model_loaded and keep_model_device and torch.cuda.is_available():
|
|
if model['model_type'] != "nlp":
|
|
model['model'].to(text_encoder_offload_device())
|
|
soft_empty_cache(True)
|
|
elif not keep_model_loaded:
|
|
del model['model'], model['tokenizer']
|
|
soft_empty_cache(True)
|
|
|
|
return (output_text, )
|
|
|
|
# Node class and display name mappings
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UL_TranslatorLoader": UL_TranslatorLoader,
|
|
"UL_Translator": UL_Translator,
|
|
}
|
|
|
|
target_language_list = ["source", "zh", "yue", "en", "ru", "ja", "ko", "es", "fr", "it", "uk", "hi", "ar", "af", "am", "as", "az", "ba", "be", "bg", "bn", "bo", "br", "bs", "ca", "cs", "cy", "da", "de", "el", "et", "eu", "fa", "fi", "fo", "gl", "gu", "ha", "haw", "he", "hr", "ht", "hu", "hy", "id", "is", "jw", "ka", "kk", "km", "kn", "la", "lb", "ln", "lo", "lt", "lv", "mg", "mi", "mk", "ml", "mn", "mr", "ms", "mt", "my", "ne", "nl", "nn", "no", "oc", "pa", "pl", "ps", "pt", "ro", "sa", "sd", "si", "sk", "sl", "sn", "so", "sq", "sr", "su", "sv", "sw", "ta", "te", "tg", "th", "tk", "tl", "tr", "tt", "ur", "uz", "vi", "yi", "yo"] |