26 lines
793 B
Python
26 lines
793 B
Python
from transformers import Mistral3ForConditionalGeneration, Mistral3Model
|
|
|
|
from .._mixin import OstrisTransformersMixin
|
|
|
|
|
|
class Mistral3TextEncoder(Mistral3ForConditionalGeneration, OstrisTransformersMixin):
|
|
"""Mistral3 conditioning stack (flux2)."""
|
|
|
|
aitk_subfolder = "text_encoder"
|
|
aitk_tokenizer_subfolder = "tokenizer"
|
|
|
|
@classmethod
|
|
def get_transformer_block_names(cls):
|
|
return ["model.language_model.layers", "language_model.model.layers"]
|
|
|
|
|
|
class Mistral3ModelEncoder(Mistral3Model, OstrisTransformersMixin):
|
|
"""The inner Mistral3 base model (ernie_image's text encoder)."""
|
|
|
|
aitk_subfolder = "text_encoder"
|
|
aitk_tokenizer_subfolder = "tokenizer"
|
|
|
|
@classmethod
|
|
def get_transformer_block_names(cls):
|
|
return ["language_model.layers"]
|