v1.0.0

.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 }}

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)

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)

__init__.py

@@ -0,0 +1,3 @@
+from .gemmaxnode import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

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",
+}

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