Hugging Face's logo

thisisAce
/
runpod_custom_nodes

TensorBoard
ONNX
Model card Files Metrics Community
runpod_custom_nodes
File size: 4,460 Bytes 
1e3b872
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
 
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 ===== ###