From 6bd85b242a22dcc998675a1e7bc9aff0018dac6f Mon Sep 17 00:00:00 2001 From: DreamFly Date: Tue, 25 Mar 2025 00:18:39 +0800 Subject: [PATCH] =?UTF-8?q?=E6=8F=90=E4=BA=A4=E7=89=88=E6=9C=AC=E5=BA=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 24 +++++++++++++++++++ README.md | 2 ++ __init__.py | 12 ++++++++++ modules/opus_translation.py | 31 +++++++++++++++++++++++++ modules/t5_translation.py | 46 +++++++++++++++++++++++++++++++++++++ nodes/prompt_translation.py | 37 +++++++++++++++++++++++++++++ pyproject.toml | 14 +++++++++++ requirements.txt | 2 ++ 8 files changed, 168 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 modules/opus_translation.py create mode 100644 modules/t5_translation.py create mode 100644 nodes/prompt_translation.py create mode 100644 pyproject.toml create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e23306f --- /dev/null +++ b/.gitignore @@ -0,0 +1,24 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class + +# Distribution / packaging +dist/ +build/ +*.egg-info/ + +# Virtual Environment +venv/ +env/ +.env/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo + +# OS +.DS_Store +Thumbs.db \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..d60be25 --- /dev/null +++ b/README.md @@ -0,0 +1,2 @@ +# ComfyUI-PromptWrapper + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..9882677 --- /dev/null +++ b/__init__.py @@ -0,0 +1,12 @@ +from .nodes.prompt_translation import PromptTranslation + + +NODE_CLASS_MAPPINGS = { + "PromptTranslation": PromptTranslation, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PromptTranslation": "Prompt Translation", +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/modules/opus_translation.py b/modules/opus_translation.py new file mode 100644 index 0000000..4025e86 --- /dev/null +++ b/modules/opus_translation.py @@ -0,0 +1,31 @@ + +from transformers import AutoTokenizer, AutoModelForSeq2SeqLM +from transformers import pipeline +import torch + + +def translation(model_path, input_text): + + tokenizer = AutoTokenizer.from_pretrained(model_path) + + model = AutoModelForSeq2SeqLM.from_pretrained(model_path) + + device = "cuda:0" if torch.cuda.is_available() else "cpu" + + OutputPipeline = pipeline("translation", model=model, tokenizer=tokenizer, device=device) + + output_text = OutputPipeline(input_text) + + return output_text[0]["translation_text"] + + + +if __name__ == '__main__': + + model_path = r"D:\WorkSpace\AIPaint\sd-comfyui-app\models\nlp\Helsinki-NLP\opus-mt-zh-en" + input_text = "中文翻译英文" + + output_text = translation(model_path, input_text) + + print("output:") + print(output_text) \ No newline at end of file diff --git a/modules/t5_translation.py b/modules/t5_translation.py new file mode 100644 index 0000000..5e7224f --- /dev/null +++ b/modules/t5_translation.py @@ -0,0 +1,46 @@ + +from transformers import T5Tokenizer, T5ForConditionalGeneration +import torch + + +def translation(model_path, task_prefix, input_text, max_length): + + # + tokenizer = T5Tokenizer.from_pretrained(model_path) + # + model = T5ForConditionalGeneration.from_pretrained(model_path) + + # when generating, we will use the logits of right-most token to predict the next token + # so the padding should be on the left + tokenizer.padding_side = "left" + tokenizer.pad_token = tokenizer.eos_token # to avoid an error + + input_task_text = task_prefix + input_text + inputs = tokenizer(input_task_text, return_tensors="pt", padding=False) + + output_sequences = model.generate( + input_ids=inputs['input_ids'], + attention_mask=inputs['attention_mask'], + do_sample=False, # disable sampling to test if batching affects output + max_length=max_length, + num_beams=4, + early_stopping=True + ) + + output_text = tokenizer.batch_decode(output_sequences, skip_special_tokens=True) + + return output_text[0][len(" "):] + + + +if __name__ == '__main__': + + model_path = r"D:\WorkSpace\AIPaint\sd-comfyui-app\models\nlp\t5-small" + task_prefix = "translate English to German: " + input_text = "Create a new string object from the given object. " + max_length = 512 + + output_text = translation(model_path, task_prefix, input_text, max_length) + + print("output:") + print(output_text) \ No newline at end of file diff --git a/nodes/prompt_translation.py b/nodes/prompt_translation.py new file mode 100644 index 0000000..1438ebf --- /dev/null +++ b/nodes/prompt_translation.py @@ -0,0 +1,37 @@ +import os +import folder_paths +from ..modules import opus_translation + + + +# Prompt Translation +class PromptTranslation: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_name": (["Helsinki-NLP/opus-mt-zh-en", "Helsinki-NLP/opus-mt-en-zh"],), + "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, 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.") + + output_text = opus_translation.translation(model_path, prompt) + + return (output_text,) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4f9bdd7 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "ComfyUI-PromptWrapper" +description = "ComfyUI nodes for PromptWrapper, Mainly for prompt word translation, Chinese translation into English or English translation into Chinese." +version = "1.0.0" +license = {file = "LICENSE"} + +[project.urls] +Repository = "https://github.com/clouddreamfly/ComfyUI-PromptWrapper" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "clouddreamfly" +DisplayName = "ComfyUI-PromptWrapper" +Icon = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..39dab0f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +torch +transformers \ No newline at end of file