Changed state dict

concept_attention/flux/src/flux/util.py

@@ -7,7 +7,7 @@ from huggingface_hub import hf_hub_download
 from imwatermark import WatermarkEncoder
 from safetensors.torch import load_file as load_sft
 
-from transformers import T5EncoderModel, AutoConfig, AutoModel
+from transformers import T5EncoderModel, AutoConfig, AutoModel, T5Tokenizer
 
 from concept_attention.flux.src.flux.model import Flux, FluxParams
 from concept_attention.flux.src.flux.modules.autoencoder import AutoEncoder, AutoEncoderParams
@@ -139,7 +139,12 @@ def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmb
     state_dict.update(load_sft(safe_tensor_1, device=str(device)))
     state_dict.update(load_sft(safe_tensor_2, device=str(device)))
     # Load the state dict
-    t5_encoder = AutoModel.from_config(config=model_config, state_dict=state_dict)
+    t5_encoder = T5EncoderModel(config=model_config)
+    t5_encoder.load_state_dict(state_dict, strict=False)
+
+    # Load the tokenizer
+    tokenizer = T5Tokenizer.from_pretrained("google/t5-v1_1-xxl")
+    t5_encoder.tokenizer = tokenizer
 
     # max length 64, 128, 256 and 512 should work (if your sequence is short enough)
     # Load the safe tensors model