88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
class PixArtResolutionCond:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"cond": ("CONDITIONING", ),
|
|
"width": ("INT", {"default": 1024.0, "min": 0, "max": 8192}),
|
|
"height": ("INT", {"default": 1024.0, "min": 0, "max": 8192}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
RETURN_NAMES = ("cond",)
|
|
FUNCTION = "add_cond"
|
|
CATEGORY = "ExtraModels/PixArt"
|
|
TITLE = "PixArt Resolution Conditioning"
|
|
|
|
def add_cond(self, cond, width, height):
|
|
for c in range(len(cond)):
|
|
cond[c][1].update({
|
|
"img_hw": [[height, width]],
|
|
"aspect_ratio": [[height/width]],
|
|
})
|
|
return (cond,)
|
|
|
|
class PixArtT5FromSD3CLIP:
|
|
"""
|
|
Split the T5 text encoder away from SD3
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"sd3_clip": ("CLIP",),
|
|
"padding": ("INT", {"default": 1, "min": 1, "max": 300}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CLIP",)
|
|
RETURN_NAMES = ("t5",)
|
|
FUNCTION = "split"
|
|
CATEGORY = "ExtraModels/PixArt"
|
|
TITLE = "PixArt T5 from SD3 CLIP"
|
|
|
|
def split(self, sd3_clip, padding):
|
|
try:
|
|
from comfy.text_encoders.sd3_clip import SD3Tokenizer, SD3ClipModel
|
|
except ImportError:
|
|
# fallback for older ComfyUI versions
|
|
from comfy.sd3_clip import SD3Tokenizer, SD3ClipModel # type: ignore
|
|
import copy
|
|
|
|
clip = sd3_clip.clone()
|
|
assert clip.cond_stage_model.t5xxl is not None, "CLIP must have T5 loaded!"
|
|
|
|
# remove transformer
|
|
transformer = clip.cond_stage_model.t5xxl.transformer
|
|
clip.cond_stage_model.t5xxl.transformer = None
|
|
|
|
# clone object
|
|
tmp = SD3ClipModel(clip_l=False, clip_g=False, t5=False)
|
|
tmp.t5xxl = copy.deepcopy(clip.cond_stage_model.t5xxl)
|
|
# put transformer back
|
|
clip.cond_stage_model.t5xxl.transformer = transformer
|
|
tmp.t5xxl.transformer = transformer
|
|
|
|
# override special tokens
|
|
tmp.t5xxl.special_tokens = copy.deepcopy(clip.cond_stage_model.t5xxl.special_tokens)
|
|
tmp.t5xxl.special_tokens.pop("end") # make sure empty tokens match
|
|
|
|
# add attn mask opt if present in original
|
|
if hasattr(sd3_clip.cond_stage_model, "t5_attention_mask"):
|
|
tmp.t5_attention_mask = False
|
|
|
|
# tokenizer
|
|
tok = SD3Tokenizer()
|
|
tok.t5xxl.min_length = padding
|
|
|
|
clip.cond_stage_model = tmp
|
|
clip.tokenizer = tok
|
|
|
|
return (clip, )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PixArtResolutionCond" : PixArtResolutionCond,
|
|
"PixArtT5FromSD3CLIP": PixArtT5FromSD3CLIP,
|
|
}
|