新增HY提示词翻译
This commit is contained in:
DreamFly
+3
-1
@@ -1,4 +1,4 @@
|
||||
from .nodes.prompt_translation import PromptTranslation
|
||||
from .nodes.prompt_translation import PromptTranslation, PromptTranslationHY
|
||||
from .nodes.prompt_generate import GeneratePrompt
|
||||
from .nodes.prompt_scenery import SceneryPrompt
|
||||
from .nodes.prompt_negative import NegativePrompt
|
||||
@@ -21,6 +21,7 @@ from .nodes.prompt_randoms import RandomsPrompt, RandomsWeightPrompt
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PromptTranslation": PromptTranslation,
|
||||
"PromptTranslationHY": PromptTranslationHY,
|
||||
"GeneratePrompt": GeneratePrompt,
|
||||
"SceneryPrompt" : SceneryPrompt,
|
||||
"NegativePrompt" : NegativePrompt,
|
||||
@@ -46,6 +47,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PromptTranslation": "Prompt Translation",
|
||||
"PromptTranslationHY": "Prompt Translation HY",
|
||||
"GeneratePrompt": "Generate Prompt",
|
||||
"SceneryPrompt": "Generate Scenery Prompt",
|
||||
"NegativePrompt": "Generate Negative Prompt",
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
|
||||
|
||||
def translation(model_path, model_file, target_language, input_text):
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, gguf_file=model_file, clean_up_tokenization_spaces=True)
|
||||
model = AutoModelForCausalLM.from_pretrained(model_path, gguf_file=model_file, device_map="auto") # You may want to use bfloat16 and/or move to GPU here
|
||||
messages = [
|
||||
{"role": "user", "content": "Translate the following segment into English, without additional explanation.\n\nGet something off your chest"},
|
||||
]
|
||||
tokenized_chat = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=False,
|
||||
return_tensors="pt"
|
||||
)
|
||||
|
||||
outputs = model.generate(tokenized_chat.to(model.device), max_new_tokens=2048)
|
||||
output_text = tokenizer.decode(outputs[0])
|
||||
|
||||
return output_text
|
||||
|
||||
|
||||
def translation2(model_path, target_language, input_text):
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto") # You may want to use bfloat16 and/or move to GPU here
|
||||
messages = [
|
||||
{"role": "user", "content": "Translate the following segment into {}, without additional explanation.\n\n{}".format(target_language, input_text)},
|
||||
]
|
||||
tokenized_chat = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=False,
|
||||
return_tensors="pt"
|
||||
)
|
||||
|
||||
outputs = model.generate(tokenized_chat.to(model.device), max_new_tokens=2048)
|
||||
output_content = tokenizer.decode(outputs[0])
|
||||
|
||||
output_text = ""
|
||||
if output_content is not None:
|
||||
text_length = len(output_content)
|
||||
if text_length > 0:
|
||||
start_text = "<｜hy_place▁holder▁no▁8｜>"
|
||||
end_text = "<｜hy_place▁holder▁no▁2｜>"
|
||||
start_index = output_content.find(start_text)
|
||||
end_index = output_content.rfind(end_text)
|
||||
if start_index >= 0 and end_index >= 0:
|
||||
start_index += len(start_text)
|
||||
output_text = output_content[start_index : end_index]
|
||||
|
||||
return output_text
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
model_path = r"D:\WorkSpace\AIPaint\sd-comfyui-app\models\unet"
|
||||
model_file = "Hunyuan-MT-7B.Q4_K_S.gguf"
|
||||
target_language = "English"
|
||||
input_text = "牡秋天的枫叶，是诗人笔下的灵感源泉。古往今来，无数诗人为枫叶写下了优美的诗篇，赞美它的美丽、它的坚韧、它的浪漫。这些诗篇，如同枫叶一样，在岁月的长河中流传千古，成为了中华文化的瑰宝。"
|
||||
|
||||
#output_text = translation(model_path, model_file, target_language, input_text)
|
||||
|
||||
model_name_or_path = r"D:\WorkSpace\AIPaint\sd-comfyui-app\models\nlp\tencent\HY-MT1.5-1.8B"
|
||||
output_text = translation2(model_name_or_path, target_language, input_text)
|
||||
|
||||
print("output:")
|
||||
print(output_text)
|
||||
@@ -1,6 +1,7 @@
|
||||
import os
|
||||
import folder_paths
|
||||
from ..modules import opus_translation
|
||||
from ..modules import hunyuan_translation
|
||||
|
||||
|
||||
|
||||
@@ -35,3 +36,91 @@ class PromptTranslation:
|
||||
output_text = opus_translation.translation(model_path, prompt)
|
||||
|
||||
return (output_text,)
|
||||
|
||||
|
||||
|
||||
# Hunyuan Translation Languages
|
||||
SupportLanguages = [
|
||||
("Chinese", "zh", "中文"),
|
||||
("English", "en", "英语"),
|
||||
("French", "fr", "法语"),
|
||||
("Portuguese", "pt", "葡萄牙语"),
|
||||
("Spanish", "es", "西班牙语"),
|
||||
("Japanese", "ja", "日语"),
|
||||
("Turkish", "tr", "土耳其语"),
|
||||
("Russian", "ru", "俄语"),
|
||||
("Arabic", "ar", "阿拉伯语"),
|
||||
("Korean", "ko", "韩语"),
|
||||
("Thai", "th", "泰语"),
|
||||
("Italian", "it", "意大利语"),
|
||||
("German", "de", "德语"),
|
||||
("Vietnamese", "vi", "越南语"),
|
||||
("Malay", "ms", "马来语"),
|
||||
("Indonesian", "id", "印尼语"),
|
||||
("Filipino", "tl", "菲律宾语"),
|
||||
("Hindi", "hi", "印地语"),
|
||||
("Traditional Chinese", "zh-Hant", "繁体中文"),
|
||||
("Polish", "pl", "波兰语"),
|
||||
("Czech", "cs", "捷克语"),
|
||||
("Dutch", "nl", "荷兰语"),
|
||||
("Khmer", "km", "高棉语"),
|
||||
("Burmese", "my", "缅甸语"),
|
||||
("Persian", "fa", "波斯语"),
|
||||
("Gujarati", "gu", "古吉拉特语"),
|
||||
("Urdu", "ur", "乌尔都语"),
|
||||
("Telugu", "te", "泰卢固语"),
|
||||
("Marathi", "mr", "马拉地语"),
|
||||
("Hebrew", "he", "希伯来语"),
|
||||
("Bengali", "bn", "孟加拉语"),
|
||||
("Tamil", "ta", "泰米尔语"),
|
||||
("Ukrainian", "uk", "乌克兰语"),
|
||||
("Tibetan", "bo ", "藏语"),
|
||||
("Kazakh", "kk", "哈萨克语"),
|
||||
("Mongolian", "mn", "蒙古语"),
|
||||
("Uyghur", "ug", "维吾尔语"),
|
||||
("Cantonese", "yue", "粤语")
|
||||
]
|
||||
|
||||
Languages = [ Language[0] for Language in SupportLanguages ]
|
||||
ChineseNames = [ Language[2] for Language in SupportLanguages ]
|
||||
|
||||
# Prompt Translation HY
|
||||
class PromptTranslationHY:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model_name": (["tencent/HY-MT1.5-1.8B", "tencent/Hunyuan-MT-7B"],),
|
||||
"target_language": (ChineseNames,),
|
||||
"prompt": ("STRING", {
|
||||
"multiline": True,
|
||||
"default": "A beautiful photo of"
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
FUNCTION = "prompt_translation"
|
||||
CATEGORY = "PromptWrapper"
|
||||
|
||||
def prompt_translation(self, model_name, target_language, prompt):
|
||||
|
||||
# 获取ComfyUI的models目录下的模型路径
|
||||
model_path = os.path.join(folder_paths.models_dir, "nlp", model_name)
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
raise ValueError(f"Local model not found at {model_path}. Please download the model and place it in the ComfyUI/models/nlp folder.")
|
||||
|
||||
select_index = 0
|
||||
select_language = Languages[select_index]
|
||||
for language_name in ChineseNames:
|
||||
if target_language == language_name:
|
||||
select_language = Languages[select_index]
|
||||
break
|
||||
select_index += 1
|
||||
|
||||
output_text = hunyuan_translation.translation2(model_path, select_language, prompt)
|
||||
|
||||
return (output_text,)
|
||||
Reference in New Issue
Block a user