28 lines
889 B
Python
28 lines
889 B
Python
from transformers import CLIPTextModel, CLIPTextModelWithProjection
|
|
|
|
from .._mixin import OstrisTransformersMixin
|
|
|
|
|
|
class CLIPTextEncoder(CLIPTextModel, OstrisTransformersMixin):
|
|
"""CLIP-L text encoder (flux-style checkpoints: text_encoder/ +
|
|
tokenizer/)."""
|
|
|
|
aitk_subfolder = "text_encoder"
|
|
aitk_tokenizer_subfolder = "tokenizer"
|
|
|
|
@classmethod
|
|
def get_transformer_block_names(cls):
|
|
return ["text_model.encoder.layers"]
|
|
|
|
|
|
class CLIPTextEncoderWithProjection(CLIPTextModelWithProjection, OstrisTransformersMixin):
|
|
"""CLIP text encoder with the projection head (SDXL / SD3 / HiDream style
|
|
checkpoints; the second encoder lives at text_encoder_2/ + tokenizer_2/)."""
|
|
|
|
aitk_subfolder = "text_encoder"
|
|
aitk_tokenizer_subfolder = "tokenizer"
|
|
|
|
@classmethod
|
|
def get_transformer_block_names(cls):
|
|
return ["text_model.encoder.layers"]
|