FengWen
cb55519a84
* Unify service address to use BIZYAIR_SERVER_ADDRESS * Update path_manager.py * Update path_manager.py
95 lines
2.5 KiB
Python
95 lines
2.5 KiB
Python
import os
|
|
import uuid
|
|
|
|
import torch
|
|
|
|
from bizyair import BizyAirBaseNode, BizyAirNodeIO, create_node_data
|
|
from bizyair.common.env_var import BIZYAIR_SERVER_ADDRESS
|
|
from bizyair.data_types import CONDITIONING
|
|
from bizyair.image_utils import encode_data
|
|
|
|
from .utils import (
|
|
decode_and_deserialize,
|
|
get_api_key,
|
|
send_post_request,
|
|
serialize_and_encode,
|
|
)
|
|
|
|
CATEGORY_NAME = "☁️BizyAir/Kolors"
|
|
|
|
|
|
class BizyAirMZChatGLM3TextEncode:
|
|
API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/mzkolorschatglm3"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
|
|
FUNCTION = "encode"
|
|
CATEGORY = CATEGORY_NAME
|
|
|
|
def encode(self, text):
|
|
API_KEY = get_api_key()
|
|
assert len(text) <= 4096, f"the prompt is too long, length: {len(text)}"
|
|
|
|
payload = {
|
|
"text": text,
|
|
}
|
|
auth = f"Bearer {API_KEY}"
|
|
headers = {
|
|
"accept": "application/json",
|
|
"content-type": "application/json",
|
|
"authorization": auth,
|
|
}
|
|
|
|
response: str = send_post_request(
|
|
self.API_URL, payload=payload, headers=headers
|
|
)
|
|
tensors_np = decode_and_deserialize(response)
|
|
|
|
ret_conditioning = []
|
|
for item in tensors_np:
|
|
t, d = item
|
|
t_tensor = torch.from_numpy(t)
|
|
d_dict = {}
|
|
for k, v in d.items():
|
|
d_dict[k] = torch.from_numpy(v)
|
|
ret_conditioning.append([t_tensor, d_dict])
|
|
|
|
return (ret_conditioning,)
|
|
|
|
|
|
class BizyAir_MinusZoneChatGLM3TextEncode(BizyAirMZChatGLM3TextEncode, BizyAirBaseNode):
|
|
RETURN_TYPES = (CONDITIONING,)
|
|
|
|
FUNCTION = "mz_encode"
|
|
|
|
def mz_encode(self, text):
|
|
out = self.encode(text)[0]
|
|
node_data = create_node_data(
|
|
class_type="ComfyAirLoadData",
|
|
inputs={"conditioning": {"relay": out}},
|
|
outputs={"slot_index": 3},
|
|
)
|
|
node_data["is_changed"] = uuid.uuid4().hex
|
|
return (
|
|
BizyAirNodeIO(
|
|
self.assigned_id,
|
|
nodes={self.assigned_id: encode_data(node_data, old_version=True)},
|
|
),
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"BizyAir_MinusZoneChatGLM3TextEncode": BizyAir_MinusZoneChatGLM3TextEncode,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"BizyAir_MinusZoneChatGLM3TextEncode": "☁️BizyAir MinusZone ChatGLM3 Text Encode",
|
|
}
|