47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
import re
|
|
from googletrans import Translator, LANGUAGES
|
|
|
|
translator = Translator()
|
|
empty_str = re.compile('^\s*$', re.I | re.M)
|
|
|
|
class TranslateCLIPTextEncodeNode:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def translate(self, prompt, srcTrans=None, toTrans=None):
|
|
if not srcTrans:
|
|
srcTrans = 'auto'
|
|
|
|
if not toTrans:
|
|
toTrans = 'en'
|
|
|
|
translate_text_prompt = ''
|
|
if prompt and not empty_str.match(prompt):
|
|
translate_text_prompt = vars(translator.translate(prompt, src=srcTrans, dest=toTrans))
|
|
|
|
return translate_text_prompt.get('text', '') if 'text' in translate_text_prompt else ''
|
|
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
|
|
return {
|
|
"required": {
|
|
"from_translate": (['auto']+list(LANGUAGES.keys()), {"default": "auto"}),
|
|
"to_translate": (list(LANGUAGES.keys()), {"default": "en"} ),
|
|
"text": ("STRING", {"multiline": True}),
|
|
"clip": ("CLIP", )
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "translate_text"
|
|
|
|
CATEGORY = "conditioning"
|
|
|
|
def translate_text(self, from_translate, to_translate, text, clip):
|
|
text = self.translate(text, from_translate, to_translate)
|
|
return ([[clip.encode(text), {}]],)
|
|
|