From 5e0fdbe0a945521a1827c320a2c414e25dc013ce Mon Sep 17 00:00:00 2001 From: DreamFly Date: Thu, 1 Jan 2026 13:17:43 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9EHY=E6=8F=90=E7=A4=BA=E8=AF=8D?= =?UTF-8?q?=E7=BF=BB=E8=AF=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- __init__.py | 4 +- modules/hunyuan_translation.py | 72 +++++++++++++++++++++++++++ nodes/prompt_translation.py | 89 ++++++++++++++++++++++++++++++++++ 3 files changed, 164 insertions(+), 1 deletion(-) create mode 100644 modules/hunyuan_translation.py diff --git a/__init__.py b/__init__.py index 09f7f05..0e9f227 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/modules/hunyuan_translation.py b/modules/hunyuan_translation.py new file mode 100644 index 0000000..f119fcb --- /dev/null +++ b/modules/hunyuan_translation.py @@ -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) \ No newline at end of file diff --git a/nodes/prompt_translation.py b/nodes/prompt_translation.py index f22aa94..d681ce6 100644 --- a/nodes/prompt_translation.py +++ b/nodes/prompt_translation.py @@ -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,) \ No newline at end of file