From 1cbc9280f76c4412004c6a0dd9fae34e02e1eb06 Mon Sep 17 00:00:00 2001 From: Marksusu Date: Fri, 19 Jul 2024 19:27:44 +0700 Subject: [PATCH] Initial commit --- __init__.py | 2 + __pycache__/__init__.cpython-312.pyc | Bin 0 -> 286 bytes __pycache__/nodes.cpython-312.pyc | Bin 0 -> 4413 bytes nodes.py | 119 +++++++++++++++++++++++++++ requirements.txt | 1 + 5 files changed, 122 insertions(+) create mode 100644 __init__.py create mode 100644 __pycache__/__init__.cpython-312.pyc create mode 100644 __pycache__/nodes.cpython-312.pyc create mode 100644 nodes.py create mode 100644 requirements.txt diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..c554b37 --- /dev/null +++ b/__init__.py @@ -0,0 +1,2 @@ +from .nodes import NODE_CLASS_MAPPINGS +__all__ = ['NODE_CLASS_MAPPINGS'] \ No newline at end of file diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cde9e5511549ebf7d1df1b8ee66939bd3c4d9241 GIT binary patch literal 286 zcmX@j%ge<81huET)4YN7V-N=hn4pZ$VnD`ph7^Vr#vF!R#wbQch7_iB#weyrW=)ot zj6g|E##_RE{w}WZ&OVO8!STM10Rf(V?!iUOKsi56=3A_J`6;QzVD2sU`1r(}ocQ>a z44*;f|H{?R$j?pH&rK}K*4N9)PfpCyFV09TO4ZLU$}CCLPt1%@$xKTtEzZo((=Sdf zDoZWWch1jEs|@wjPcAJk$kodsN$jEq?LG1yz!UZnHM)o3hpdbKpoKPzO literal 0 HcmV?d00001 diff --git a/__pycache__/nodes.cpython-312.pyc b/__pycache__/nodes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da09b3b21c4c78a2de7493a2f3e084aa796fad4e GIT binary patch literal 4413 zcma)9TWk~A8J-#2$?2xLoyDz+jrBbDR>VL*|ZVO$Gl5@^~ zJ^%lmGw1yNyTxK8p#5s-4~d^U3E~?})SI?8d7XsG0>Kh2nIuwVm`qV&io`UP)Pyza zs10k?Q5V*!BORt;)FkyOCd`n;4T9DFi3}Sk;w-`H?hq{fn5udY8(Eq&@?_Y=(ku<5 z`8MU(zr;d6rC3H{h>xX5LL(WLODeW18OqFXX+F=Ppr2GUDfWa)}sF3OyFOJQO?| z%6DWDm*8)8NF5@IrG;cnc=>%+_HwWE+9vC| zhPE5m3EE(M4YowQzWJEIu&P^i0wxQDxUr2Q=BjPF{z=psFtVgX)(`XJnhk6m70fqB z#&ZX`ef#!~=Z+*2iEGmMJJ4GgpTy`-W->{ILYzViTnxphU`w?EmrNiBcolGD z9UaI8gW})y#){}yZdovfKHxxNV#M~zyXbOgTGJTUZNK_uQ3P#aVx1F}uXPcvBta~Y# zM0`Zm_9L2Rijd6H4R!6!k@fHd$V`%+HL|9^n(I0Infx_but=t9_oi-rCRtbuKF6SM z-8engA3?|l=v(XF3s#9n7KuTo8V=SCG;{?+x)G$zrUWCDHiI@JZR%kjcq{;th-U4o zw&Q^&4c|b)#x_WL@CbCl4s&dU=MM0JbpdrutGW%Mr=AziiS&nHvDMcW?FSp*4(R#7>vp|Oe_Qpe+g1jh2h=-7=XhguXWtSXXm5+| zzz?F^^tUy-_HEH&FI%EZ+a=Bay(J9WtPR_(C2|@ct{qJ>)Uz)*L1WhrF_4lI&ZE12 z9>?JZD7Xr4$;CEGZl1hB-5@1Hf`S*p8_e4#BI82uwH~QDbSgx?sVg!eUF-IbOFJho z@niB5-zYEf)tYx;W2t71$Id*}+mW}NQVzPbs<^D3i-7IzdEy7PWvF;rC0iB@omYQ2fA436&9xFmwuYgG+M=Y9NJ%`V0yX zMVCdm;9=Cvp_Eh_pezJIRJE_-;crBEzmW}hxZqGIJ-N&9oK}T}=+Ezf78i@KwwOTmf{vu~kbDCIu%61C)1ft2yfIHn0~Q z0QM~LADXZ?{O0BtH|55zqOZtSYzOAFP^38nbNV%V!vb@cS@b;Wtk^%8)BW3IUv)Px zMD9i&Ies^>;`W!_{))S2?#ybd?_uO#r0BRmvC_J~+`7NgdSEX2W?fjY(YtWa$+wiQw2sUL*W8|+S`e1YyW$L#oq>||sq^^idmlVv zid6B#Vwc=@;J2n1t-Zy3x%J4rajm6irRC#t%g0OSD=lZ|4Qt-rE8at8@1fFk#XB&s zd(qmnc;(^5y@_(`%vb!9RsLB-=J=<5(?HJ~SA9JzzW%bWzx0!(*x!!ISHrTezv7F` z+aUP&cmHM2ANR=qqf2P{c*Qq1Z+it{HQ!6D?(TVXu6VNeOS$hj$o4FIW#9R`7ru7y z+A5KKCu+Qda_6DaCLv9{@=DJju$o^w;7lyBQxveUxEri)} z^NDA!0XPM_+8=h`>%QMRcL6drSm(`)k;M_&+WnO2dEsy`wBBu9{H)^WnWI-dyH-4X zWlvx6=5k2(^i@3Pzc9~f=Mpa(8W-D(dt~o^+0ifC50oOZZ9rxQ)Z!X$Wb@#~;Hja| z(8VD9Kts5|LZ_e$`i%+`Rj-6-R53=Qa2Mo~9L&wp=%=|@5?>ZTnc~TqAfz~PD#IdN z991tp>_RbLJ{`JpIavM8q%dc$1PAdt3NttmI(_!ykJ5*3ZvXJz_t8 zi(*zUqeKR!I3%dou6iR^uWt1>EGLM9Lc`BdHj`4c(?TY#=%-@BRFa=W@H2xzxK32V zQk=nyr%p!)hX=;Sq9X&NqX5~M(yTI14ULTs4}2004vd`MknKqDd*-KA7>}V4_{0yX zFa+JYh9t><5+6M$49^MMbHehRF#Ut@e5cjgs9Dp0tt?5ly@KB}WZQaBL+G6I!|NK# UL&=15-H6E`kZn|%XjU`-A9})o9smFU literal 0 HcmV?d00001 diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..d71b063 --- /dev/null +++ b/nodes.py @@ -0,0 +1,119 @@ +from transformers import MarianMTModel, MarianTokenizer +import re + +marian_list = [ + "opus-mt-zh-en", + "opus-mt-ru-en", + "opus-mt-th-en", +] +# https://huggingface.co/Helsinki-NLP 在这个地址里可以找到更多语种的模型，添加到如上列表后，下次使用时会自动下载模型 + +class MTCLIPEncode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "clip": ("CLIP", {}), + "checkpoint": (marian_list, {"multiline": False,"default": "opus-mt-zh-en"}), + "text": ("STRING", {"multiline": True,"default": "prefix | translate_part | suffix"}), + } + } + + def is_valid_translate_part(self, text): + # 简单检查文本是否包含至少一个汉字或一个长度大于2的非英文单词 + return bool(re.search(r'[\u4e00-\u9fff]|\b\w{3,}\b', text)) + + def mtencode(self, clip, checkpoint, text): + + # 使用正则表达式修剪首尾可能存在的一个或多个 `,`、[空格] + text = re.sub(r'^[\s,]+|[\s,]+$', '', text) + + # 使用正则表达式分离出需要翻译和不需要翻译的部分 + pattern = r"([^|]*)\|([^|]*)\|([^|]*)" + match = re.match(pattern, text) + + if match: + prefix = match.group(1).strip() + translate_part = match.group(2).strip() + suffix = match.group(3).strip() + + # 修剪三个部分的文本，首尾可能存在的一个或多个`|`、`,`、[空格] + prefix = re.sub(r'^[\s|,]+|[\s|,]+$', '', prefix) + translate_part = re.sub(r'^[\s|,]+|[\s|,]+$', '', translate_part) + suffix = re.sub(r'^[\s|,]+|[\s|,]+$', '', suffix) + + # 确保 translate_part 有效 + if not self.is_valid_translate_part(translate_part): + prompt_text = f"{prefix}, {suffix}" + else: + model_name = 'Helsinki-NLP/' + checkpoint + tokenizer = MarianTokenizer.from_pretrained(model_name) + model = MarianMTModel.from_pretrained(model_name) + + translated = model.generate(**tokenizer(translate_part, return_tensors="pt", padding=True)) + translated_text = tokenizer.decode(translated[0], skip_special_tokens=True) + + # 修剪 translated_text 首尾可能存在的空格、逗号和句号 + translated_text = re.sub(r'^[\s,.]+|[\s,.]+$', '', translated_text) + + # 合并不需要翻译的部分和翻译后的部分 + if prefix and suffix: + prompt_text = f"{prefix}, {translated_text}, {suffix}" + elif prefix: + prompt_text = f"{prefix}, {translated_text}" + elif suffix: + prompt_text = f"{translated_text}, {suffix}" + else: + prompt_text = translated_text + + # ANSI 转义序列用于颜色输出 + color_prefix = f"\033[94m{prefix}\033[0m" if prefix else "" + color_translate_part = f"\033[92m{translate_part}\033[0m" + color_translated_text = f"\033[92m{translated_text}\033[0m" + color_suffix = f"\033[94m{suffix}\033[0m" if suffix else "" + + # 打印格式化输出 + print(f"　　　　　　🫐　🫐　🫐　🫐　🫐　🫐") + if prefix and suffix: + print(f"　　　　{color_prefix}, {color_translate_part}, {color_suffix}") + elif prefix: + print(f"　　　　{color_prefix}, {color_translate_part}") + elif suffix: + print(f"　　　　{color_translate_part}, {color_suffix}") + else: + print(f"　　　　{color_translate_part}") + + if prefix and suffix: + print(f"　　　　{color_prefix}, {color_translated_text}, {color_suffix}") + elif prefix: + print(f"　　　　{color_prefix}, {color_translated_text}") + elif suffix: + print(f"　　　　{color_translated_text}, {color_suffix}") + else: + print(f"　　　　{color_translated_text}") + + else: + prompt_text = text.strip() # 如果没有匹配项，使用原文本 + + # 打印格式化输出 + print(f"　　　　{prompt_text}") + + tokens = clip.tokenize(prompt_text) + cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True) + + return ([[cond, {"pooled_output": pooled}]], prompt_text) + + RETURN_TYPES = ( + "CONDITIONING", + "STRING", + ) + FUNCTION = "mtencode" + CATEGORY = "MTCLIPEncode" + +NODE_CLASS_MAPPINGS = { + "MTCLIPEncode": MTCLIPEncode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MTCLIPEncode": "MTCLIPEncode", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..747b7aa --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +transformers \ No newline at end of file