73 lines
2.6 KiB
Python
73 lines
2.6 KiB
Python
import logging
|
|
from enum import Enum
|
|
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import comfy.text_encoders
|
|
|
|
from .pixart.tenc import pixart_te, PixArtTokenizer
|
|
|
|
class TencType(Enum):
|
|
# offset in case we ever integrate w/ original
|
|
PixArt = 1001
|
|
MiaoBi = 1002
|
|
# HunYuan = 1003 # deprecated
|
|
Sana = 1004
|
|
|
|
tenc_names = {
|
|
# for node readout
|
|
"PixArt": TencType.PixArt,
|
|
"MiaoBi": TencType.MiaoBi,
|
|
# "HunYuan": TencType.HunYuan,
|
|
"Sana": TencType.Sana,
|
|
}
|
|
|
|
|
|
def load_text_encoder(ckpt_paths, embedding_directory=None, clip_type=TencType.PixArt, model_options={}):
|
|
# Partial duplicate of ComfyUI/comfy/sd:load_clip
|
|
clip_data = []
|
|
for p in ckpt_paths:
|
|
if p.lower().endswith(".gguf"):
|
|
# TODO: cross-node call w/o code duplication
|
|
raise NotImplementedError("Planned!")
|
|
else:
|
|
clip_data.append(comfy.utils.load_torch_file(p, safe_load=True))
|
|
return load_text_encoder_state_dicts(clip_data, embedding_directory=embedding_directory, clip_type=clip_type, model_options=model_options)
|
|
|
|
def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip_type=TencType.PixArt, model_options={}):
|
|
# Partial duplicate of ComfyUI/comfy/sd:load_text_encoder_state_dicts
|
|
clip_data = state_dicts
|
|
|
|
class EmptyClass:
|
|
pass
|
|
|
|
for i in range(len(clip_data)):
|
|
if "transformer.resblocks.0.ln_1.weight" in clip_data[i]:
|
|
clip_data[i] = comfy.utils.clip_text_transformers_convert(clip_data[i], "", "")
|
|
else:
|
|
if "text_projection" in clip_data[i]:
|
|
clip_data[i]["text_projection.weight"] = clip_data[i]["text_projection"].transpose(0, 1) #old models saved with the CLIPSave node
|
|
|
|
clip_target = EmptyClass()
|
|
clip_target.params = {}
|
|
|
|
if clip_type == TencType.PixArt:
|
|
clip_target.clip = pixart_te(**comfy.sd.t5xxl_detect(clip_data))
|
|
clip_target.tokenizer = PixArtTokenizer
|
|
|
|
parameters = 0
|
|
tokenizer_data = {}
|
|
for c in clip_data:
|
|
parameters += comfy.utils.calculate_parameters(c)
|
|
tokenizer_data, model_options = comfy.text_encoders.long_clipl.model_options_long_clip(c, tokenizer_data, model_options)
|
|
|
|
clip = comfy.sd.CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, tokenizer_data=tokenizer_data, model_options=model_options)
|
|
for c in clip_data:
|
|
m, u = clip.load_sd(c)
|
|
if len(m) > 0:
|
|
logging.warning("clip missing: {}".format(m))
|
|
|
|
if len(u) > 0:
|
|
logging.debug("clip unexpected: {}".format(u))
|
|
return clip
|