AlekPet
20377c8411
Split extension deeptranslator and googletranlate on folders Change initialize extensions when start comfyui Updates reamde files
126 lines
4.3 KiB
Python
126 lines
4.3 KiB
Python
import os
|
|
import requests
|
|
import json
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
from googletrans import Translator, LANGUAGES
|
|
|
|
### ===== GoogleTranslate Nodes [googletrans module] ===== ###
|
|
translator = Translator()
|
|
|
|
google_translation_key = os.environ.get("GOOGLE_TRANSLATION_API_KEY")
|
|
|
|
|
|
# Manual tranlsate prompts
|
|
@PromptServer.instance.routes.post("/alekpet/translate_manual")
|
|
async def translate_manual(request):
|
|
json_data = await request.json()
|
|
prompt = json_data.get("prompt", "")
|
|
|
|
if "prompt" in json_data and "srcTrans" in json_data and "toTrans" in json_data:
|
|
prompt = json_data.get("prompt")
|
|
srcTrans = json_data.get("srcTrans")
|
|
toTrans = json_data.get("toTrans")
|
|
|
|
translate_text_prompt = translate(prompt, srcTrans, toTrans)
|
|
|
|
return web.json_response({"translate_prompt": translate_text_prompt})
|
|
|
|
return web.json_response({"translate_prompt": prompt})
|
|
|
|
# Translate used Google API_KEY
|
|
class TranslationResult:
|
|
def __init__(self, text=''):
|
|
self.text = text
|
|
|
|
def translate_by_key(text, src, dest):
|
|
url = f"https://translation.googleapis.com/language/translate/v2?key={google_translation_key}"
|
|
|
|
data = {
|
|
'q': text,
|
|
'target': dest
|
|
}
|
|
|
|
resp = requests.post(url, data=data)
|
|
resp_data = json.loads(resp.text)
|
|
|
|
if 'translations' in resp_data.get('data', {}):
|
|
translations = resp_data['data']['translations']
|
|
if translations:
|
|
translated_text = translations[0]['translatedText']
|
|
return TranslationResult(translated_text)
|
|
|
|
return TranslationResult('')
|
|
|
|
|
|
def translate(prompt, srcTrans=None, toTrans=None):
|
|
if not srcTrans:
|
|
srcTrans = 'auto'
|
|
|
|
if not toTrans:
|
|
toTrans = 'en'
|
|
|
|
translate_text_prompt = ''
|
|
if prompt and prompt.strip() !="":
|
|
if not google_translation_key:
|
|
translate_text_prompt = translator.translate(prompt, src=srcTrans, dest=toTrans)
|
|
else:
|
|
translate_text_prompt = TranslationResult.translate_by_key(prompt, src=srcTrans, dest=toTrans)
|
|
|
|
return translate_text_prompt.text if hasattr(translate_text_prompt, 'text') else ''
|
|
|
|
class GoogleTranslateCLIPTextEncodeNode:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {
|
|
"required": {
|
|
"from_translate": (['auto']+list(LANGUAGES.keys()), {"default": "auto"}),
|
|
"to_translate": (list(LANGUAGES.keys()), {"default": "en"} ),
|
|
"manual_translate": ([True, False],),
|
|
"text": ("STRING", {"multiline": True,"placeholder":"Input prompt"}),
|
|
"clip": ("CLIP", )
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING","STRING",)
|
|
FUNCTION = "translate_text"
|
|
CATEGORY = "AlekPet Nodes/conditioning"
|
|
|
|
def translate_text(self, **kwargs):
|
|
from_translate = kwargs.get("from_translate")
|
|
to_translate = kwargs.get("to_translate")
|
|
manual_translate = kwargs.get("manual_translate", False)
|
|
text = kwargs.get("text")
|
|
clip = kwargs.get("clip")
|
|
|
|
text_tranlsated = translate(text, from_translate, to_translate) if not manual_translate else text
|
|
tokens = clip.tokenize(text_tranlsated)
|
|
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
|
|
return ([[cond, {"pooled_output": pooled}]], text_tranlsated)
|
|
|
|
|
|
class GoogleTranslateTextNode(GoogleTranslateCLIPTextEncodeNode):
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return_types = super().INPUT_TYPES()
|
|
del return_types["required"]["clip"]
|
|
return return_types
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "translate_text"
|
|
|
|
CATEGORY = "AlekPet Nodes/text"
|
|
|
|
def translate_text(self, **kwargs):
|
|
from_translate = kwargs.get("from_translate")
|
|
to_translate = kwargs.get("to_translate")
|
|
manual_translate = kwargs.get("manual_translate", False)
|
|
text = kwargs.get("text")
|
|
|
|
text_tranlsated = translate(text, from_translate, to_translate) if not manual_translate else text
|
|
return (text_tranlsated,)
|
|
|
|
### ===== GoogleTranslate Nodes [googletrans module] -> end ===== ### |