split "Latent Caching" into "Image Caching" and "Text Caching"

modules/dataLoader/ChromaBaseDataLoader.py

@@ -51,7 +51,7 @@ def _preparation_modules(self, config: TrainConfig, model: ChromaModel):
         if not config.train_text_encoder_or_embedding():
             modules.append(encode_prompt)
 
-        if config.latent_caching and not config.train_text_encoder_or_embedding():
+        if config.text_caching and not config.train_text_encoder_or_embedding():
             modules.append(prune_masked_tokens)
 
         return modules
@@ -81,7 +81,7 @@ def _cache_modules(self, config: TrainConfig, model: ChromaModel, model_setup: B
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching = not config.train_text_encoder_or_embedding(),
+            text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
         )
 
     def _output_modules(self, config: TrainConfig, model: ChromaModel, model_setup: BaseChromaSetup):
@@ -111,7 +111,7 @@ def _output_modules(self, config: TrainConfig, model: ChromaModel, model_setup:
             train_dtype=model.train_dtype,
         )
 
-        if config.latent_caching and not config.train_text_encoder_or_embedding():
+        if config.text_caching and not config.train_text_encoder_or_embedding():
             output_module_list = [pad_masked_tokens] + output_module_list
 
         return output_module_list

modules/dataLoader/ErnieBaseDataLoader.py

@@ -43,7 +43,7 @@ def _preparation_modules(self, config: TrainConfig, model: ErnieModel):
             modules.append(downscale_mask)
         modules += [tokenize_prompt, encode_prompt]
 
-        if config.latent_caching:
+        if config.text_caching:
             modules.append(prune_masked_tokens)
 
         return modules
@@ -72,7 +72,7 @@ def _cache_modules(self, config: TrainConfig, model: ErnieModel, model_setup: Ba
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=True,
+            text_caching=config.text_caching,
         )
 
     def _output_modules(self, config: TrainConfig, model: ErnieModel, model_setup: BaseErnieSetup):
@@ -101,7 +101,7 @@ def _output_modules(self, config: TrainConfig, model: ErnieModel, model_setup: B
             train_dtype=model.train_dtype,
         )
 
-        if config.latent_caching:
+        if config.text_caching:
             output_module_list = [pad_masked_tokens] + output_module_list
 
         return output_module_list

modules/dataLoader/Flux2BaseDataLoader.py

@@ -86,7 +86,7 @@ def _cache_modules(self, config: TrainConfig, model: Flux2Model, model_setup: Ba
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=True,
+            text_caching=config.text_caching,
         )
 
     def _output_modules(self, config: TrainConfig, model: Flux2Model, model_setup: BaseFlux2Setup):

modules/dataLoader/FluxBaseDataLoader.py

@@ -105,7 +105,10 @@ def _cache_modules(self, config: TrainConfig, model: FluxModel, model_setup: Bas
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(),
+            text_caching=config.text_caching and (
+                not config.train_text_encoder_or_embedding()
+                or not config.train_text_encoder_2_or_embedding()
+            ),
         )
 
     def _output_modules(self, config: TrainConfig, model: FluxModel, model_setup: BaseFluxSetup):

modules/dataLoader/HiDreamBaseDataLoader.py

@@ -133,10 +133,12 @@ def _cache_modules(self, config: TrainConfig, model: HiDreamModel, model_setup:
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding() \
-                      or not config.train_text_encoder_2_or_embedding() \
-                      or not config.train_text_encoder_3_or_embedding() \
-                      or not config.train_text_encoder_4_or_embedding(),
+            text_caching=config.text_caching and (
+                not config.train_text_encoder_or_embedding()
+                or not config.train_text_encoder_2_or_embedding()
+                or not config.train_text_encoder_3_or_embedding()
+                or not config.train_text_encoder_4_or_embedding()
+            ),
         )
 
     def _output_modules(self, config: TrainConfig, model: HiDreamModel, model_setup: BaseHiDreamSetup):

modules/dataLoader/HunyuanVideoBaseDataLoader.py

@@ -98,7 +98,10 @@ def _cache_modules(self, config: TrainConfig, model: HunyuanVideoModel, model_se
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(),
+            text_caching=config.text_caching and (
+                not config.train_text_encoder_or_embedding()
+                or not config.train_text_encoder_2_or_embedding()
+            ),
         )
 
     def _output_modules(self, config: TrainConfig, model: HunyuanVideoModel, model_setup: BaseHunyuanVideoSetup):

modules/dataLoader/IdeogramBaseDataLoader.py

@@ -55,7 +55,7 @@ def _preparation_modules(self, config: TrainConfig, model: IdeogramModel):
             modules.append(downscale_mask)
         modules += [tokenize_prompt, encode_prompt]
 
-        if config.latent_caching:
+        if config.text_caching:
             modules.append(prune_masked_tokens)
 
         return modules
@@ -84,7 +84,7 @@ def _cache_modules(self, config: TrainConfig, model: IdeogramModel, model_setup:
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=True,
+            text_caching=config.text_caching,
         )
 
     def _output_modules(self, config: TrainConfig, model: IdeogramModel, model_setup: BaseIdeogramSetup):
@@ -113,7 +113,7 @@ def _output_modules(self, config: TrainConfig, model: IdeogramModel, model_setup
             train_dtype=model.train_dtype,
         )
 
-        if config.latent_caching:
+        if config.text_caching:
             output_module_list = [pad_masked_tokens] + output_module_list
 
         return output_module_list

modules/dataLoader/Krea2BaseDataLoader.py

@@ -62,7 +62,7 @@ def _preparation_modules(self, config: TrainConfig, model: Krea2Model):
         if not config.train_text_encoder_or_embedding():
             modules.append(encode_prompt)
 
-        if config.latent_caching and not config.train_text_encoder_or_embedding():
+        if config.text_caching and not config.train_text_encoder_or_embedding():
             modules.append(prune_masked_tokens)
 
         return modules
@@ -92,7 +92,7 @@ def _cache_modules(self, config: TrainConfig, model: Krea2Model, model_setup: Ba
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding(),
+            text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
         )
 
     def _output_modules(self, config: TrainConfig, model: Krea2Model, model_setup: BaseKrea2Setup):
@@ -122,7 +122,7 @@ def _output_modules(self, config: TrainConfig, model: Krea2Model, model_setup: B
             train_dtype=model.train_dtype,
         )
 
-        if config.latent_caching and not config.train_text_encoder_or_embedding():
+        if config.text_caching and not config.train_text_encoder_or_embedding():
             output_module_list = [pad_masked_tokens] + output_module_list
 
         return output_module_list

modules/dataLoader/PixArtAlphaBaseDataLoader.py

@@ -87,7 +87,7 @@ def _cache_modules(self, config: TrainConfig, model: PixArtAlphaModel, model_set
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding(),
+            text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
         )
 
     def _output_modules(self, config: TrainConfig, model: PixArtAlphaModel, model_setup: BasePixArtAlphaSetup):

modules/dataLoader/QwenBaseDataLoader.py

@@ -55,7 +55,7 @@ def _preparation_modules(self, config: TrainConfig, model: QwenModel):
         if not config.train_text_encoder_or_embedding():
             modules.append(encode_prompt)
 
-        if config.latent_caching and not config.train_text_encoder_or_embedding():
+        if config.text_caching and not config.train_text_encoder_or_embedding():
             modules.append(prune_masked_tokens)
 
         return modules
@@ -85,7 +85,7 @@ def _cache_modules(self, config: TrainConfig, model: QwenModel, model_setup: Bas
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding(),
+            text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
         )
 
     def _output_modules(self, config: TrainConfig, model: QwenModel, model_setup: BaseQwenSetup):
@@ -115,7 +115,7 @@ def _output_modules(self, config: TrainConfig, model: QwenModel, model_setup: Ba
             train_dtype=model.train_dtype,
         )
 
-        if config.latent_caching and not config.train_text_encoder_or_embedding():
+        if config.text_caching and not config.train_text_encoder_or_embedding():
             output_module_list = [pad_masked_tokens] + output_module_list
 
         return output_module_list

modules/dataLoader/SanaBaseDataLoader.py

@@ -79,7 +79,7 @@ def _cache_modules(self, config: TrainConfig, model: SanaModel, model_setup: Bas
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding(),
+            text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
         )
 
     def _output_modules(self, config: TrainConfig, model: SanaModel, model_setup: BaseSanaSetup):

modules/dataLoader/StableDiffusion3BaseDataLoader.py

@@ -119,7 +119,11 @@ def _cache_modules(self, config: TrainConfig, model: StableDiffusion3Model, mode
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding() or not config.train_text_encoder_3_or_embedding(),
+            text_caching=config.text_caching and (
+                not config.train_text_encoder_or_embedding()
+                or not config.train_text_encoder_2_or_embedding()
+                or not config.train_text_encoder_3_or_embedding()
+            ),
         )
 
     def _output_modules(self, config: TrainConfig, model: StableDiffusion3Model, model_setup: BaseStableDiffusion3Setup):

modules/dataLoader/StableDiffusionBaseDataLoader.py

@@ -94,7 +94,7 @@ def _cache_modules(self, config: TrainConfig, model: StableDiffusionModel, model
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding(),
+            text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
         )
 
     def _output_modules(self, config: TrainConfig, model: StableDiffusionModel, model_setup: BaseStableDiffusionSetup):

modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py

@@ -205,7 +205,7 @@ def before_cache_fun():
 
         modules = []
 
-        if config.latent_caching:
+        if config.image_caching:
             modules.append(disk_cache)
             modules.append(variation_sorting)
 
@@ -229,7 +229,7 @@ def __output_modules(self, config: TrainConfig):
 
         image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean')
 
-        if config.latent_caching:
+        if config.image_caching:
             batch_sorting = AspectBatchSorting(resolution_in_name='crop_resolution', names=sort_names, batch_size=config.batch_size)
         else:
             batch_sorting = InlineAspectBatchSorting(resolution_in_name='crop_resolution', names=sort_names, batch_size=config.batch_size)

modules/dataLoader/StableDiffusionXLBaseDataLoader.py

@@ -101,7 +101,10 @@ def _cache_modules(self, config: TrainConfig, model: StableDiffusionXLModel, mod
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(),
+            text_caching=config.text_caching and (
+                not config.train_text_encoder_or_embedding()
+                or not config.train_text_encoder_2_or_embedding()
+            ),
         )
 
     def _output_modules(self, config: TrainConfig, model: StableDiffusionXLModel, model_setup: BaseStableDiffusionXLSetup):

modules/dataLoader/WuerstchenBaseDataLoader.py

@@ -86,7 +86,7 @@ def before_cache_image_fun():
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=not config.train_text_encoder_or_embedding(),
+            text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
             before_cache_image_fun=before_cache_image_fun
         )
 

modules/dataLoader/ZImageBaseDataLoader.py

@@ -49,7 +49,7 @@ def _preparation_modules(self, config: TrainConfig, model: ZImageModel):
 
         modules += [tokenize_prompt, encode_prompt]
 
-        if config.latent_caching:
+        if config.text_caching:
             modules.append(prune_masked_tokens)
 
         return modules
@@ -78,7 +78,7 @@ def _cache_modules(self, config: TrainConfig, model: ZImageModel, model_setup: B
             text_split_names=text_split_names,
             sort_names=sort_names,
             config=config,
-            text_caching=True,
+            text_caching=config.text_caching,
         )
 
     def _output_modules(self, config: TrainConfig, model: ZImageModel, model_setup: BaseZImageSetup):
@@ -107,7 +107,7 @@ def _output_modules(self, config: TrainConfig, model: ZImageModel, model_setup:
             train_dtype=model.train_dtype,
         )
 
-        if config.latent_caching:
+        if config.text_caching:
             output_module_list = [pad_masked_tokens] + output_module_list
 
         return output_module_list

modules/dataLoader/mixin/DataLoaderText2ImageMixin.py

@@ -301,7 +301,7 @@ def prepare_vae():
         )
 
         world_size = multi.world_size() if config.multi_gpu else 1  #world_size can be 1 for validation dataloader, even if multi.world_size() returns > 1
-        if config.latent_caching:
+        if config.image_caching:
             batch_sorting = AspectBatchSorting(resolution_in_name='crop_resolution', names=sort_names, batch_size=config.batch_size * world_size)
             distributed_sampler = DistributedSampler(names=sort_names, world_size=world_size, rank=multi.rank())
         else:
@@ -358,15 +358,15 @@ def before_cache_text_fun():
 
         modules = []
 
-        if config.latent_caching:
+        if config.image_caching:
             modules.append(image_disk_cache)
 
             sort_names = [x for x in sort_names if x not in image_aggregate_names]
             sort_names = [x for x in sort_names if x not in image_split_names]
 
-            if text_caching:
-                modules.append(text_disk_cache)
-                sort_names = [x for x in sort_names if x not in text_split_names]
+        if text_caching:
+            modules.append(text_disk_cache)
+            sort_names = [x for x in sort_names if x not in text_split_names]
 
         if len(sort_names) > 0:
             variation_sorting = VariationSorting(names=sort_names, balancing_in_name='concept.balancing', balancing_strategy_in_name='concept.balancing_strategy',

modules/modelSetup/AnimaFineTuneSetup.py

@@ -69,9 +69,9 @@ def setup_train_device(
             model: AnimaModel,
             config: TrainConfig,
     ):
-        vae_on_train_device = not config.latent_caching
+        vae_on_train_device = not config.image_caching
 
-        model.text_encoder_to(self.temp_device if config.latent_caching else self.train_device)
+        model.text_encoder_to(self.train_device if not config.text_caching else self.temp_device)
         model.vae_to(self.train_device if vae_on_train_device else self.temp_device)
         model.transformer_to(self.train_device)
 

modules/modelSetup/AnimaLoRASetup.py

@@ -81,9 +81,9 @@ def setup_train_device(
             model: AnimaModel,
             config: TrainConfig,
     ):
-        vae_on_train_device = not config.latent_caching
+        vae_on_train_device = not config.image_caching
 
-        model.text_encoder_to(self.temp_device if config.latent_caching else self.train_device)
+        model.text_encoder_to(self.train_device if not config.text_caching else self.temp_device)
         model.vae_to(self.train_device if vae_on_train_device else self.temp_device)
         model.transformer_to(self.train_device)
 

modules/modelSetup/ChromaEmbeddingSetup.py

@@ -72,9 +72,10 @@ def setup_train_device(
             model: ChromaModel,
             config: TrainConfig,
     ):
-        vae_on_train_device = not config.latent_caching
+        vae_on_train_device = not config.image_caching
+        text_encoder_on_train_device = config.train_text_encoder_or_embedding() or not config.text_caching
 
-        model.text_encoder_to(self.train_device if config.text_encoder.train_embedding else self.temp_device)
+        model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device)
         model.vae_to(self.train_device if vae_on_train_device else self.temp_device)
         model.transformer_to(self.train_device)
 

modules/modelSetup/ChromaFineTuneSetup.py

@@ -82,10 +82,10 @@ def setup_train_device(
             model: ChromaModel,
             config: TrainConfig,
     ):
-        vae_on_train_device = not config.latent_caching
+        vae_on_train_device = not config.image_caching
         text_encoder_on_train_device = \
             config.train_text_encoder_or_embedding() \
-            or not config.latent_caching
+            or not config.text_caching
 
         model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device)
         model.vae_to(self.train_device if vae_on_train_device else self.temp_device)

modules/modelSetup/ChromaLoRASetup.py

@@ -109,10 +109,10 @@ def setup_train_device(
             model: ChromaModel,
             config: TrainConfig,
     ):
-        vae_on_train_device = not config.latent_caching
+        vae_on_train_device = not config.image_caching
         text_encoder_on_train_device = \
             config.train_text_encoder_or_embedding() \
-            or not config.latent_caching
+            or not config.text_caching
 
         model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device)
         model.vae_to(self.train_device if vae_on_train_device else self.temp_device)

modules/modelSetup/ErnieFineTuneSetup.py

@@ -64,8 +64,8 @@ def setup_train_device(
             model: ErnieModel,
             config: TrainConfig,
     ):
-        vae_on_train_device = not config.latent_caching
-        text_encoder_on_train_device = not config.latent_caching
+        vae_on_train_device = not config.image_caching
+        text_encoder_on_train_device = not config.text_caching
 
         model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device)
         model.vae_to(self.train_device if vae_on_train_device else self.temp_device)