Performance fixes and bug fixes

extensions_built_in/diffusion_models/flux2/flux2_model.py

@@ -98,12 +98,21 @@ class Flux2Model(BaseModel):
         self.print_and_status_update("Loading Mistral")
 
         # load + quantize + offload + placement, all driven by model_config
+        # tie_word_embeddings=False: the checkpoint carries both embed_tokens
+        # and lm_head with different values; the config's tie claim is wrong
         text_encoder = Mistral3TextEncoder.load(
-            MISTRAL_PATH, subfolder="", **self.component_load_kwargs("te")
+            MISTRAL_PATH,
+            subfolder="",
+            tie_word_embeddings=False,
+            **self.component_load_kwargs("te"),
         )
         flush()
 
-        tokenizer = AutoProcessor.from_pretrained(MISTRAL_PATH)
+        # fix_mistral_regex=False: keep the exact tokenization flux2 has always
+        # used (True would change the pre-tokenizer and shift conditioning)
+        tokenizer = AutoProcessor.from_pretrained(
+            MISTRAL_PATH, fix_mistral_regex=False
+        )
         return text_encoder, tokenizer
 
     def load_model(self):

extensions_built_in/diffusion_models/flux2/src/pipeline.py

@@ -109,15 +109,19 @@ class Flux2Pipeline(DiffusionPipeline):
         # Process all messages at once
         # with image processing a too short max length can throw an error in here.
         try:
+            # tokenization kwargs ride in processor_kwargs (same values end up
+            # in the same place; loose **kwargs just warn on new transformers)
             inputs = self.tokenizer.apply_chat_template(
                 messages_batch,
                 add_generation_prompt=False,
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt",
-                padding="max_length",
-                truncation=True,
-                max_length=max_sequence_length,
+                processor_kwargs={
+                    "padding": "max_length",
+                    "truncation": True,
+                    "max_length": max_sequence_length,
+                },
             )
         except ValueError as e:
             print(