Merge branch 'ctk_copy' into ctk_abstraction

# Conflicts:
#	modules/ui/BaseConceptWindowView.py
#	modules/ui/BaseTimestepDistributionWindowView.py
#	modules/ui/BaseTrainingTabView.py
#	modules/ui/CtkConceptWindowView.py
#	modules/ui/CtkTimestepDistributionWindowView.py
#	modules/ui/CtkTrainingTabView.py

Author: dxqb

.pre-commit-config.yaml

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

modules/dataLoader/ChromaBaseDataLoader.py

@@ -26,6 +26,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.CHROMA_1)
 class ChromaBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -154,5 +155,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-factory.register(BaseDataLoader, ChromaBaseDataLoader, ModelType.CHROMA_1)

modules/dataLoader/ErnieBaseDataLoader.py

@@ -7,13 +7,14 @@
 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
 from mgds.pipelineModules.DecodeVAE import DecodeVAE
 from mgds.pipelineModules.EncodeMistralText import EncodeMistralText
 from mgds.pipelineModules.EncodeVAE import EncodeVAE
+from mgds.pipelineModules.PadMaskedTokens import PadMaskedTokens
+from mgds.pipelineModules.PruneMaskedTokens import PruneMaskedTokens
 from mgds.pipelineModules.RescaleImageChannels import RescaleImageChannels
 from mgds.pipelineModules.SampleVAEDistribution import SampleVAEDistribution
 from mgds.pipelineModules.SaveImage import SaveImage
@@ -22,6 +23,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.ERNIE)
 class ErnieBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -32,15 +34,18 @@ 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())
+        prune_masked_tokens = PruneMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state')
 
         modules = [rescale_image, encode_image, image_sample]
         if config.masked_training or config.model_type.has_mask_input():
             modules.append(downscale_mask)
         modules += [tokenize_prompt, encode_prompt]
+
+        if config.latent_caching:
+            modules.append(prune_masked_tokens)
+
         return modules
 
     def _cache_modules(self, config: TrainConfig, model: ErnieModel, model_setup: BaseErnieSetup):
@@ -71,6 +76,8 @@ def _cache_modules(self, config: TrainConfig, model: ErnieModel, model_setup: Ba
         )
 
     def _output_modules(self, config: TrainConfig, model: ErnieModel, model_setup: BaseErnieSetup):
+        pad_masked_tokens = PadMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state', max_length=PROMPT_MAX_LENGTH)
+
         output_names = [
             'image_path', 'latent_image',
             'prompt',
@@ -84,7 +91,7 @@ def _output_modules(self, config: TrainConfig, model: ErnieModel, model_setup: B
 
         output_names.append('text_encoder_hidden_state')
 
-        return self._output_modules_from_out_names(
+        output_module_list = self._output_modules_from_out_names(
             model, model_setup,
             output_names=output_names,
             config=config,
@@ -94,6 +101,11 @@ def _output_modules(self, config: TrainConfig, model: ErnieModel, model_setup: B
             train_dtype=model.train_dtype,
         )
 
+        if config.latent_caching:
+            output_module_list = [pad_masked_tokens] + output_module_list
+
+        return output_module_list
+
     def _debug_modules(self, config: TrainConfig, model: ErnieModel):
         debug_dir = os.path.join(config.debug_dir, "dataloader")
 
@@ -133,6 +145,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-
-factory.register(BaseDataLoader, ErnieBaseDataLoader, ModelType.ERNIE)

modules/dataLoader/Flux2BaseDataLoader.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
@@ -30,6 +29,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.FLUX_2)
 class Flux2BaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -43,8 +43,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 +51,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())
 
@@ -161,6 +157,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-
-factory.register(BaseDataLoader, Flux2BaseDataLoader, ModelType.FLUX_2)

modules/dataLoader/FluxBaseDataLoader.py

@@ -26,6 +26,8 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.FLUX_DEV_1)
+@factory.register(BaseDataLoader, ModelType.FLUX_FILL_DEV_1)
 class FluxBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -181,6 +183,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-factory.register(BaseDataLoader, FluxBaseDataLoader, ModelType.FLUX_DEV_1)
-factory.register(BaseDataLoader, FluxBaseDataLoader, ModelType.FLUX_FILL_DEV_1)

modules/dataLoader/HiDreamBaseDataLoader.py

@@ -27,6 +27,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.HI_DREAM_FULL)
 class HiDreamBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -219,5 +220,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=64,
         )
-
-factory.register(BaseDataLoader, HiDreamBaseDataLoader, ModelType.HI_DREAM_FULL)

modules/dataLoader/HunyuanVideoBaseDataLoader.py

@@ -29,6 +29,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.HUNYUAN_VIDEO)
 class HunyuanVideoBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -183,5 +184,3 @@ def _create_dataset(
             allow_video_files=True,
             vae_frame_dim=True,
         )
-
-factory.register(BaseDataLoader, HunyuanVideoBaseDataLoader, ModelType.HUNYUAN_VIDEO)

modules/dataLoader/PixArtAlphaBaseDataLoader.py

@@ -24,6 +24,8 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.PIXART_ALPHA)
+@factory.register(BaseDataLoader, ModelType.PIXART_SIGMA)
 class PixArtAlphaBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -159,6 +161,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=16,
         )
-
-factory.register(BaseDataLoader, PixArtAlphaBaseDataLoader, ModelType.PIXART_ALPHA)
-factory.register(BaseDataLoader, PixArtAlphaBaseDataLoader, ModelType.PIXART_SIGMA)

modules/dataLoader/QwenBaseDataLoader.py

@@ -30,6 +30,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.QWEN)
 class QwenBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -166,5 +167,3 @@ def _create_dataset(
             allow_video_files=False, #don't allow video files, but...
             vae_frame_dim=True,  #...Qwen has a video-capable VAE. convert images to video dimensions
         )
-
-factory.register(BaseDataLoader, QwenBaseDataLoader, ModelType.QWEN)

modules/dataLoader/SanaBaseDataLoader.py

@@ -23,6 +23,7 @@
 from mgds.pipelineModules.Tokenize import Tokenize
 
 
+@factory.register(BaseDataLoader, ModelType.SANA)
 class SanaBaseDataLoader(
     BaseDataLoader,
     DataLoaderText2ImageMixin,
@@ -152,5 +153,3 @@ def _create_dataset(
             config, model, model_setup, train_progress, is_validation,
             aspect_bucketing_quantization=32,
         )
-
-factory.register(BaseDataLoader, SanaBaseDataLoader, ModelType.SANA)