Skip to content

Commit 80cffd6

Browse files
dxqbdxqb
committed
Merge PR #1462 into preview
# Conflicts: # modules/ui/TrainUI.py # modules/util/config/TrainConfig.py
2 parents 8fe4080 + d5ac1e4 commit 80cffd6

60 files changed

Lines changed: 191 additions & 146 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎modules/dataLoader/ChromaBaseDataLoader.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ def _preparation_modules(self, config: TrainConfig, model: ChromaModel):
5050
if not config.train_text_encoder_or_embedding():
5151
modules.append(encode_prompt)
5252

53-
if config.latent_caching and not config.train_text_encoder_or_embedding():
53+
if config.text_caching and not config.train_text_encoder_or_embedding():
5454
modules.append(prune_masked_tokens)
5555

5656
return modules
@@ -80,7 +80,7 @@ def _cache_modules(self, config: TrainConfig, model: ChromaModel, model_setup: B
8080
text_split_names=text_split_names,
8181
sort_names=sort_names,
8282
config=config,
83-
text_caching = not config.train_text_encoder_or_embedding(),
83+
text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
8484
)
8585

8686
def _output_modules(self, config: TrainConfig, model: ChromaModel, model_setup: BaseChromaSetup):
@@ -110,7 +110,7 @@ def _output_modules(self, config: TrainConfig, model: ChromaModel, model_setup:
110110
train_dtype=model.train_dtype,
111111
)
112112

113-
if config.latent_caching and not config.train_text_encoder_or_embedding():
113+
if config.text_caching and not config.train_text_encoder_or_embedding():
114114
output_module_list = [pad_masked_tokens] + output_module_list
115115

116116
return output_module_list

‎modules/dataLoader/ErnieBaseDataLoader.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ def _cache_modules(self, config: TrainConfig, model: ErnieModel, model_setup: Ba
6464
text_split_names=text_split_names,
6565
sort_names=sort_names,
6666
config=config,
67-
text_caching=True,
67+
text_caching=config.text_caching,
6868
)
6969

7070
def _output_modules(self, config: TrainConfig, model: ErnieModel, model_setup: BaseErnieSetup):

‎modules/dataLoader/Flux2BaseDataLoader.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ def _cache_modules(self, config: TrainConfig, model: Flux2Model, model_setup: Ba
8585
text_split_names=text_split_names,
8686
sort_names=sort_names,
8787
config=config,
88-
text_caching=True,
88+
text_caching=config.text_caching,
8989
)
9090

9191
def _output_modules(self, config: TrainConfig, model: Flux2Model, model_setup: BaseFlux2Setup):

‎modules/dataLoader/FluxBaseDataLoader.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,10 @@ def _cache_modules(self, config: TrainConfig, model: FluxModel, model_setup: Bas
103103
text_split_names=text_split_names,
104104
sort_names=sort_names,
105105
config=config,
106-
text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(),
106+
text_caching=config.text_caching and (
107+
not config.train_text_encoder_or_embedding()
108+
or not config.train_text_encoder_2_or_embedding()
109+
),
107110
)
108111

109112
def _output_modules(self, config: TrainConfig, model: FluxModel, model_setup: BaseFluxSetup):

‎modules/dataLoader/HiDreamBaseDataLoader.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -132,10 +132,12 @@ def _cache_modules(self, config: TrainConfig, model: HiDreamModel, model_setup:
132132
text_split_names=text_split_names,
133133
sort_names=sort_names,
134134
config=config,
135-
text_caching=not config.train_text_encoder_or_embedding() \
136-
or not config.train_text_encoder_2_or_embedding() \
137-
or not config.train_text_encoder_3_or_embedding() \
138-
or not config.train_text_encoder_4_or_embedding(),
135+
text_caching=config.text_caching and (
136+
not config.train_text_encoder_or_embedding()
137+
or not config.train_text_encoder_2_or_embedding()
138+
or not config.train_text_encoder_3_or_embedding()
139+
or not config.train_text_encoder_4_or_embedding()
140+
),
139141
)
140142

141143
def _output_modules(self, config: TrainConfig, model: HiDreamModel, model_setup: BaseHiDreamSetup):

‎modules/dataLoader/HunyuanVideoBaseDataLoader.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,10 @@ def _cache_modules(self, config: TrainConfig, model: HunyuanVideoModel, model_se
9797
text_split_names=text_split_names,
9898
sort_names=sort_names,
9999
config=config,
100-
text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(),
100+
text_caching=config.text_caching and (
101+
not config.train_text_encoder_or_embedding()
102+
or not config.train_text_encoder_2_or_embedding()
103+
),
101104
)
102105

103106
def _output_modules(self, config: TrainConfig, model: HunyuanVideoModel, model_setup: BaseHunyuanVideoSetup):

‎modules/dataLoader/PixArtAlphaBaseDataLoader.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ def _cache_modules(self, config: TrainConfig, model: PixArtAlphaModel, model_set
8585
text_split_names=text_split_names,
8686
sort_names=sort_names,
8787
config=config,
88-
text_caching=not config.train_text_encoder_or_embedding(),
88+
text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
8989
)
9090

9191
def _output_modules(self, config: TrainConfig, model: PixArtAlphaModel, model_setup: BasePixArtAlphaSetup):

‎modules/dataLoader/QwenBaseDataLoader.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ def _preparation_modules(self, config: TrainConfig, model: QwenModel):
5454
if not config.train_text_encoder_or_embedding():
5555
modules.append(encode_prompt)
5656

57-
if config.latent_caching and not config.train_text_encoder_or_embedding():
57+
if config.text_caching and not config.train_text_encoder_or_embedding():
5858
modules.append(prune_masked_tokens)
5959

6060
return modules
@@ -84,7 +84,7 @@ def _cache_modules(self, config: TrainConfig, model: QwenModel, model_setup: Bas
8484
text_split_names=text_split_names,
8585
sort_names=sort_names,
8686
config=config,
87-
text_caching=not config.train_text_encoder_or_embedding(),
87+
text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
8888
)
8989

9090
def _output_modules(self, config: TrainConfig, model: QwenModel, model_setup: BaseQwenSetup):
@@ -114,7 +114,7 @@ def _output_modules(self, config: TrainConfig, model: QwenModel, model_setup: Ba
114114
train_dtype=model.train_dtype,
115115
)
116116

117-
if config.latent_caching and not config.train_text_encoder_or_embedding():
117+
if config.text_caching and not config.train_text_encoder_or_embedding():
118118
output_module_list = [pad_masked_tokens] + output_module_list
119119

120120
return output_module_list

‎modules/dataLoader/SanaBaseDataLoader.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ def _cache_modules(self, config: TrainConfig, model: SanaModel, model_setup: Bas
7878
text_split_names=text_split_names,
7979
sort_names=sort_names,
8080
config=config,
81-
text_caching=not config.train_text_encoder_or_embedding(),
81+
text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
8282
)
8383

8484
def _output_modules(self, config: TrainConfig, model: SanaModel, model_setup: BaseSanaSetup):

‎modules/dataLoader/StableDiffusion3BaseDataLoader.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,11 @@ def _cache_modules(self, config: TrainConfig, model: StableDiffusion3Model, mode
117117
text_split_names=text_split_names,
118118
sort_names=sort_names,
119119
config=config,
120-
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(),
120+
text_caching=config.text_caching and (
121+
not config.train_text_encoder_or_embedding()
122+
or not config.train_text_encoder_2_or_embedding()
123+
or not config.train_text_encoder_3_or_embedding()
124+
),
121125
)
122126

123127
def _output_modules(self, config: TrainConfig, model: StableDiffusion3Model, model_setup: BaseStableDiffusion3Setup):

0 commit comments

Comments
 (0)