146 lines
4.5 KiB
Python
146 lines
4.5 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": ("BOOLEAN", {"default": False}),
|
|
"text": ("STRING", {"multiline": True, "placeholder": "Input prompt"}),
|
|
"clip": ("CLIP",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"CONDITIONING",
|
|
"STRING",
|
|
)
|
|
FUNCTION = "translate_text"
|
|
DESCRIPTION = "This is a node that translates the prompt into another language using Google Translate."
|
|
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 ===== ###
|