diff --git a/.github/workflows/publish_action.yml b/.github/workflows/publish_action.yml new file mode 100644 index 0000000..ef9d8d9 --- /dev/null +++ b/.github/workflows/publish_action.yml @@ -0,0 +1,22 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - master + - main + paths: + - "pyproject.toml" + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/README-CN.md b/README-CN.md new file mode 100644 index 0000000..07692cb --- /dev/null +++ b/README-CN.md @@ -0,0 +1,27 @@ +[中文](README-CN.md)|[English](README.md) + +# 小米 GemmaX 的 ComfyUI 节点: 基于 Gemma 多语言翻译, 支持 28 种语言. + +![](https://github.com/billwuhao/ComfyUI_gemmax/blob/main/images/2025-03-23_07-12-01.png) + +## 📣 更新 + +[2025-03-23]⚒️: 发布版本 v1.0.0. + +## 安装 + +``` +cd ComfyUI/custom_nodes +git clone https://github.com/billwuhao/ComfyUI_gemmax.git +``` + +## 模型下载 + +将下列模型人选其一, 下载放到 `ComfyUI/models/TTS` 目录下. + +[GemmaX2-28-2B-v0.1](https://huggingface.co/ModelSpace/GemmaX2-28-2B-v0.1) +[GemmaX2-28-9B-v0.1](https://huggingface.co/ModelSpace/GemmaX2-28-9B-v0.1) + +## 鸣谢 + +[gemmax](https://github.com/xiaomi-research/gemmax) diff --git a/README.md b/README.md index 55c3f70..fd7b94c 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,27 @@ -# ComfyUI_gemmax -XiaoMi GemmaX: Support 28 languages, Multilingual Translator based on Gemma. A node for ComfyUI. +[中文](README-CN.md)|[English](README.md) + +# ComfyUI Nodes for Xiaomi GemmaX: Multi-language translation based on Gemma, supporting 28 languages. + +![](https://github.com/billwuhao/ComfyUI_gemmax/blob/main/images/2025-03-23_07-12-01.png) + +## 📣 Updates + +[2025-03-23]⚒️: Released version v1.0.0. + +## Installation + +``` +cd ComfyUI/custom_nodes +git clone https://github.com/billwuhao/ComfyUI_gemmax.git +``` + +## Model Download + +Choose one of the following models and download it to the `ComfyUI/models/TTS` directory. + +[GemmaX2-28-2B-v0.1](https://huggingface.co/ModelSpace/GemmaX2-28-2B-v0.1) +[GemmaX2-28-9B-v0.1](https://huggingface.co/ModelSpace/GemmaX2-28-9B-v0.1) + +## Acknowledgements + +[gemmax](https://github.com/xiaomi-research/gemmax) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..90f92f8 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .gemmaxnode import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/gemmaxnode.py b/gemmaxnode.py new file mode 100644 index 0000000..9cd9ec0 --- /dev/null +++ b/gemmaxnode.py @@ -0,0 +1,69 @@ +from transformers import AutoModelForCausalLM, AutoTokenizer +import folder_paths +import os +import torch + +models_dir = folder_paths.models_dir +model_path = os.path.join(models_dir, "TTS") + +LANGUAGES = ["Arabic", "Bengali", "Czech", "German", "English", "Spanish", "Persian", "French", "Hebrew", "Hindi", + "Indonesian", "Italian", "Japanese", "Khmer", "Korean", "Lao", "Malay", "Burmese", "Dutch", "Polish", + "Portuguese", "Russian", "Thai", "Tagalog", "Turkish", "Urdu", "Vietnamese", "中文"] + +class GemmaxRun: + tokenizer = None + model_cache = None + device = "cuda" if torch.cuda.is_available() else "cpu" + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model":(["GemmaX2-28-2B-v0.1", "GemmaX2-28-9B-v0.1"],{"default": "GemmaX2-28-2B-v0.1"}), + "source_language": (LANGUAGES, {"default": "English"}), + "target_language": (LANGUAGES, {"default": "中文"}), + "text": ("STRING",), + "max_new_tokens": ("INT", {"default": 200, "min": 1,}), + "unload_model": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("translations",) + FUNCTION = "translate" + CATEGORY = "MW/MW-gemmax" + def translate(self, model, source_language, target_language, text, max_new_tokens, unload_model): + model_id = model_path + "/" + model + if self.model_cache is None: + model = AutoModelForCausalLM.from_pretrained(model_id).eval().to(self.device) + tokenizer = AutoTokenizer.from_pretrained(model_id) + self.tokenizer = tokenizer + self.model_cache = model + else: + model = self.model_cache + tokenizer = self.tokenizer + + text = "将文本从{}翻译成{}:\n\n{}:{}\n\n{}:".format(source_language, target_language, source_language, text, target_language) + + inputs = tokenizer(text, return_tensors="pt").to(self.device) + outputs = model.generate(**inputs, max_new_tokens=max_new_tokens) + translations = tokenizer.decode(outputs[0], skip_special_tokens=True) + + translations = translations.split(f"\n\n{target_language}:")[-1] + + if unload_model: + import gc + self.tokenizer = None + self.model_cache = None + gc.collect() + torch.cuda.empty_cache() + + return (translations,) + + +NODE_CLASS_MAPPINGS = { + "GemmaxRun": GemmaxRun, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "GemmaxRun": "Gemmax Run", +} \ No newline at end of file diff --git a/images/2025-03-23_07-12-01.png b/images/2025-03-23_07-12-01.png new file mode 100644 index 0000000..d70adb8 Binary files /dev/null and b/images/2025-03-23_07-12-01.png differ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..bf458a6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "gemmax-mw" +description = "XiaoMi GemmaX: Support 28 languages, Multilingual Translator based on Gemma. A node for ComfyUI." +version = "1.0.0" +license = {file = "LICENSE"} +dependencies = [] + +[project.urls] +Repository = "https://github.com/billwuhao/ComfyUI_gemmax" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "mw" +DisplayName = "MW-ComfyUI_gemmax" +Icon = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e69de29