47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
import torch
|
|
|
|
from .kolors.kolors_core import chatglm3_text_encode
|
|
|
|
class KolorsMultiTextEncode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"chatglm3_model": ("CHATGLM3MODEL", ),
|
|
"text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
"hid_proj": ("TorchLinear", ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
|
|
FUNCTION = "encode"
|
|
CATEGORY = "emojiiii"
|
|
|
|
def encode(self, chatglm3_model, text, hid_proj):
|
|
text_list = text.split('\n')
|
|
|
|
if len(text_list)==0:
|
|
raise Exception('No text found in input')
|
|
|
|
pooled_out = []
|
|
cond_out = []
|
|
for i in range(len(text_list)):
|
|
if text_list[i]=='' or text_list[i]=='\n':
|
|
continue
|
|
prompt_embeds, pooled_output = chatglm3_text_encode(self, chatglm3_model, prompt=text_list[i])\
|
|
|
|
if hid_proj.weight.dtype != prompt_embeds.dtype:
|
|
with torch.cuda.amp.autocast(dtype=hid_proj.weight.dtype):
|
|
prompt_embeds = hid_proj(prompt_embeds)
|
|
else:
|
|
prompt_embeds = hid_proj(prompt_embeds)
|
|
|
|
cond_out.append(prompt_embeds)
|
|
pooled_out.append(pooled_output)
|
|
|
|
final_pooled_output = torch.cat(pooled_out, dim=0)
|
|
final_conditioning = torch.cat(cond_out, dim=0)
|
|
|
|
return ([[final_conditioning, {"pooled_output": final_pooled_output}]],)
|