| """ |
| Ernie-Image TextEncoder for DiffSynth-Studio. |
| |
| Wraps transformers Ministral3Model to output text embeddings. |
| Pattern: lazy import + manual config dict + torch.nn.Module wrapper. |
| Only loads the text (language) model, ignoring vision components. |
| """ |
|
|
| import torch |
|
|
|
|
| class ErnieImageTextEncoder(torch.nn.Module): |
| """ |
| Text encoder using Ministral3Model (transformers). |
| Only the text_config portion of the full Mistral3Model checkpoint. |
| Uses the base model (no lm_head) since the checkpoint only has embeddings. |
| """ |
|
|
| def __init__(self): |
| super().__init__() |
| from transformers import Ministral3Config, Ministral3Model |
|
|
| text_config = { |
| "attention_dropout": 0.0, |
| "bos_token_id": 1, |
| "dtype": "bfloat16", |
| "eos_token_id": 2, |
| "head_dim": 128, |
| "hidden_act": "silu", |
| "hidden_size": 3072, |
| "initializer_range": 0.02, |
| "intermediate_size": 9216, |
| "max_position_embeddings": 262144, |
| "model_type": "ministral3", |
| "num_attention_heads": 32, |
| "num_hidden_layers": 26, |
| "num_key_value_heads": 8, |
| "pad_token_id": 11, |
| "rms_norm_eps": 1e-05, |
| "rope_parameters": { |
| "beta_fast": 32.0, |
| "beta_slow": 1.0, |
| "factor": 16.0, |
| "llama_4_scaling_beta": 0.1, |
| "mscale": 1.0, |
| "mscale_all_dim": 1.0, |
| "original_max_position_embeddings": 16384, |
| "rope_theta": 1000000.0, |
| "rope_type": "yarn", |
| "type": "yarn", |
| }, |
| "sliding_window": None, |
| "tie_word_embeddings": True, |
| "use_cache": True, |
| "vocab_size": 131072, |
| } |
| config = Ministral3Config(**text_config) |
| self.model = Ministral3Model(config) |
| self.config = config |
|
|
| def forward( |
| self, |
| input_ids=None, |
| attention_mask=None, |
| position_ids=None, |
| **kwargs, |
| ): |
| outputs = self.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| output_hidden_states=True, |
| return_dict=True, |
| **kwargs, |
| ) |
| return (outputs.hidden_states,) |
|
|
