提交版本库

.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 

README.md

@@ -0,0 +1,2 @@
+# ComfyUI-PromptWrapper
+

__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'] 

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)

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("<pad> "):]
+
+
+
+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)

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,)

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 = ""

requirements.txt

@@ -0,0 +1,2 @@
+torch
+transformers