Merge ctk_copy: transformers 5.9, torch 2.12, compile thread-safety

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: dxqb

.github/pull_request_template.md

@@ -0,0 +1,32 @@
+<!--
+Thanks for contributing to OneTrainer.
+Please read docs/Contributing.md / docs/ProjectStructure.md for the project guide.
+-->
+
+## Summary
+
+<!-- What this PR changes and why. Link an issue or discussion if relevant. -->
+
+## Test plan
+
+<!--
+OneTrainer has no automated test suite. Manual verification is expected.
+Describe what you actually did, not what you "would do".
+-->
+
+- [ ] `pre-commit run --all-files` passes
+- [ ] Launched the affected UI or script and exercised the change
+- [ ] Tested with at least one real preset / config when relevant (note which: ____)
+
+## AI assistance
+
+<!--
+This project welcomes AI-assisted contributions but the reviewer should not have to be
+your co-author. Pick exactly one option below, delete the others.
+-->
+
+- No AI involvement
+- AI-assisted — I have read every line in this diff and can defend each change
+- Early AI prototype — opened for discussion, **not ready for review**
+
+<!-- By opening this PR (and not marking it as an early prototype, aka a draft) I confirm I have read every line in this diff and have tested it locally. -->

.gitignore

@@ -11,7 +11,8 @@ debug*.py
 
 # user data
 /workspace*
-/models*
+/models/*
+!/models/.gitkeep
 /training_concepts
 /training_samples
 /training_user_settings

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
     - id: check-yaml
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.15.14
+  rev: v0.15.15
   hooks:
     # Run the Ruff linter, but not the formatter.
     - id: ruff

models/.gitkeep

@@ -7,7 +7,6 @@
 from modules.util import factory
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.enum.ModelType import ModelType
-from modules.util.thread_safety import apply_thread_safe_forward
 from modules.util.TrainProgress import TrainProgress
 
 from mgds.pipelineModules.DecodeTokens import DecodeTokens
@@ -32,8 +31,6 @@ def _preparation_modules(self, config: TrainConfig, model: ErnieModel):
         image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean')
         downscale_mask = ScaleImage(in_name='mask', out_name='latent_mask', factor=0.125)
         tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH)
-        if config.dataloader_threads > 1:
-            apply_thread_safe_forward(model.text_encoder)  # workaround for transformers#42673, unclear if Mistral is affected
         encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask',
                                           text_encoder=model.text_encoder, hidden_state_output_index=HIDDEN_STATES_LAYER, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype())
 

modules/dataLoader/ErnieBaseDataLoader.py

@@ -14,7 +14,6 @@
 from modules.util import factory
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.enum.ModelType import ModelType
-from modules.util.thread_safety import apply_thread_safe_forward
 from modules.util.TrainProgress import TrainProgress
 
 from mgds.pipelineModules.DecodeTokens import DecodeTokens
@@ -43,8 +42,6 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model):
             tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length,
                                         apply_chat_template = lambda caption: mistral_format_input([caption], MISTRAL_SYSTEM_MESSAGE), apply_chat_template_kwargs = {'add_generation_prompt': False},
                                       )
-            if config.dataloader_threads > 1:
-                apply_thread_safe_forward(model.text_encoder)  # workaround for transformers#42673
             encode_prompt = EncodeMistralText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask',
                                               text_encoder=model.text_encoder, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(),
                                               hidden_state_output_index=MISTRAL_HIDDEN_STATES_LAYERS,
@@ -53,8 +50,6 @@ def _preparation_modules(self, config: TrainConfig, model: Flux2Model):
             tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=config.text_encoder_sequence_length,
                                         apply_chat_template = lambda caption: qwen3_format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': False}
                                       )
-            if config.dataloader_threads > 1:
-                apply_thread_safe_forward(model.text_encoder)  # workaround for transformers#42673
             encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask',
                                            text_encoder=model.text_encoder, hidden_state_output_index=QWEN3_HIDDEN_STATES_LAYERS, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype())
 

modules/dataLoader/Flux2BaseDataLoader.py

@@ -9,7 +9,6 @@
 from modules.util import factory
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.enum.ModelType import ModelType
-from modules.util.thread_safety import apply_thread_safe_forward
 from modules.util.TrainProgress import TrainProgress
 
 from mgds.pipelineModules.DecodeTokens import DecodeTokens
@@ -38,8 +37,6 @@ def _preparation_modules(self, config: TrainConfig, model: ZImageModel):
         tokenize_prompt = Tokenize(in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH,
                                     apply_chat_template = lambda caption: format_input(caption), apply_chat_template_kwargs = {'add_generation_prompt': True, 'enable_thinking': True}
                                   )
-        if config.dataloader_threads > 1:
-            apply_thread_safe_forward(model.text_encoder)  # workaround for transformers#42673
         encode_prompt = EncodeQwenText(tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask',
                                        text_encoder=model.text_encoder, hidden_state_output_index=-2, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype())
         prune_masked_tokens = PruneMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state')

modules/dataLoader/ZImageBaseDataLoader.py

@@ -40,6 +40,7 @@ def __init__(
 class ChromaModel(BaseModel):
     # base model data
     tokenizer: T5Tokenizer | None
+    orig_tokenizer: T5Tokenizer | None
     noise_scheduler: FlowMatchEulerDiscreteScheduler | None
     text_encoder: T5EncoderModel | None
     vae: AutoencoderKL | None
@@ -72,6 +73,7 @@ def __init__(
         )
 
         self.tokenizer = None
+        self.orig_tokenizer = None
         self.noise_scheduler = None
         self.text_encoder = None
         self.vae = None
@@ -141,13 +143,13 @@ def eval(self):
             self.text_encoder.eval()
         self.transformer.eval()
 
-    def create_pipeline(self) -> DiffusionPipeline:
+    def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline:
         return ChromaPipeline(
             transformer=self.transformer,
             scheduler=self.noise_scheduler,
             vae=self.vae,
             text_encoder=self.text_encoder,
-            tokenizer=self.tokenizer,
+            tokenizer=self.orig_tokenizer if use_original_tokenizers else self.tokenizer,
         )
 
     def add_text_encoder_embeddings_to_prompt(self, prompt: str) -> str:

modules/model/ChromaModel.py

@@ -50,7 +50,9 @@ def __init__(
 class FluxModel(BaseModel):
     # base model data
     tokenizer_1: CLIPTokenizer | None
+    orig_tokenizer_1: CLIPTokenizer | None
     tokenizer_2: T5Tokenizer | None
+    orig_tokenizer_2: T5Tokenizer | None
     noise_scheduler: FlowMatchEulerDiscreteScheduler | None
     text_encoder_1: CLIPTextModel | None
     text_encoder_2: T5EncoderModel | None
@@ -86,7 +88,9 @@ def __init__(
         )
 
         self.tokenizer_1 = None
+        self.orig_tokenizer_1 = None
         self.tokenizer_2 = None
+        self.orig_tokenizer_2 = None
         self.noise_scheduler = None
         self.text_encoder_1 = None
         self.text_encoder_2 = None
@@ -177,15 +181,15 @@ def eval(self):
             self.text_encoder_2.eval()
         self.transformer.eval()
 
-    def create_pipeline(self) -> DiffusionPipeline:
+    def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline:
         return FluxPipeline(
             transformer=self.transformer,
             scheduler=self.noise_scheduler,
             vae=self.vae,
             text_encoder=self.text_encoder_1,
-            tokenizer=self.tokenizer_1,
+            tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1,
             text_encoder_2=self.text_encoder_2,
-            tokenizer_2=self.tokenizer_2,
+            tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2,
         )
 
     def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str:

modules/model/FluxModel.py

@@ -267,19 +267,19 @@ def eval(self):
             self.text_encoder_4.eval()
         self.transformer.eval()
 
-    def create_pipeline(self, use_original_modules: bool) -> DiffusionPipeline:
+    def create_pipeline(self, use_original_tokenizers: bool = False) -> DiffusionPipeline:
         return HiDreamImagePipeline(
             transformer=self.transformer,
             scheduler=self.noise_scheduler,
             vae=self.vae,
             text_encoder=self.text_encoder_1,
-            tokenizer=self.orig_tokenizer_1 if use_original_modules else self.tokenizer_1,
+            tokenizer=self.orig_tokenizer_1 if use_original_tokenizers else self.tokenizer_1,
             text_encoder_2=self.text_encoder_2,
-            tokenizer_2=self.orig_tokenizer_2 if use_original_modules else self.tokenizer_2,
+            tokenizer_2=self.orig_tokenizer_2 if use_original_tokenizers else self.tokenizer_2,
             text_encoder_3=self.text_encoder_3,
-            tokenizer_3=self.orig_tokenizer_3 if use_original_modules else self.tokenizer_3,
+            tokenizer_3=self.orig_tokenizer_3 if use_original_tokenizers else self.tokenizer_3,
             text_encoder_4=self.text_encoder_4,
-            tokenizer_4=self.orig_tokenizer_4 if use_original_modules else self.tokenizer_4,
+            tokenizer_4=self.orig_tokenizer_4 if use_original_tokenizers else self.tokenizer_4,
         )
 
     def add_text_encoder_1_embeddings_to_prompt(self, prompt: str) -> str: