Skip to content

Commit 00f380a

Browse files
hameerabbasihameerabbasi
committed
Merge remote-tracking branch 'origin/preview' into pixi-preview
2 parents e6aa51b + 263ea76 commit 00f380a

108 files changed

Lines changed: 1169 additions & 355 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
@@ -51,7 +51,7 @@ def _preparation_modules(self, config: TrainConfig, model: ChromaModel):
5151
if not config.train_text_encoder_or_embedding():
5252
modules.append(encode_prompt)
5353

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

5757
return modules
@@ -81,7 +81,7 @@ def _cache_modules(self, config: TrainConfig, model: ChromaModel, model_setup: B
8181
text_split_names=text_split_names,
8282
sort_names=sort_names,
8383
config=config,
84-
text_caching = not config.train_text_encoder_or_embedding(),
84+
text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
8585
)
8686

8787
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:
111111
train_dtype=model.train_dtype,
112112
)
113113

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

117117
return output_module_list

‎modules/dataLoader/ErnieBaseDataLoader.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ def _preparation_modules(self, config: TrainConfig, model: ErnieModel):
4343
modules.append(downscale_mask)
4444
modules += [tokenize_prompt, encode_prompt]
4545

46-
if config.latent_caching:
46+
if config.text_caching:
4747
modules.append(prune_masked_tokens)
4848

4949
return modules
@@ -72,7 +72,7 @@ def _cache_modules(self, config: TrainConfig, model: ErnieModel, model_setup: Ba
7272
text_split_names=text_split_names,
7373
sort_names=sort_names,
7474
config=config,
75-
text_caching=True,
75+
text_caching=config.text_caching,
7676
)
7777

7878
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
101101
train_dtype=model.train_dtype,
102102
)
103103

104-
if config.latent_caching:
104+
if config.text_caching:
105105
output_module_list = [pad_masked_tokens] + output_module_list
106106

107107
return output_module_list

‎modules/dataLoader/Flux2BaseDataLoader.py‎

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

9292
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
@@ -105,7 +105,10 @@ def _cache_modules(self, config: TrainConfig, model: FluxModel, model_setup: Bas
105105
text_split_names=text_split_names,
106106
sort_names=sort_names,
107107
config=config,
108-
text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(),
108+
text_caching=config.text_caching and (
109+
not config.train_text_encoder_or_embedding()
110+
or not config.train_text_encoder_2_or_embedding()
111+
),
109112
)
110113

111114
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
@@ -133,10 +133,12 @@ def _cache_modules(self, config: TrainConfig, model: HiDreamModel, model_setup:
133133
text_split_names=text_split_names,
134134
sort_names=sort_names,
135135
config=config,
136-
text_caching=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(),
136+
text_caching=config.text_caching and (
137+
not config.train_text_encoder_or_embedding()
138+
or not config.train_text_encoder_2_or_embedding()
139+
or not config.train_text_encoder_3_or_embedding()
140+
or not config.train_text_encoder_4_or_embedding()
141+
),
140142
)
141143

142144
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
@@ -98,7 +98,10 @@ def _cache_modules(self, config: TrainConfig, model: HunyuanVideoModel, model_se
9898
text_split_names=text_split_names,
9999
sort_names=sort_names,
100100
config=config,
101-
text_caching=not config.train_text_encoder_or_embedding() or not config.train_text_encoder_2_or_embedding(),
101+
text_caching=config.text_caching and (
102+
not config.train_text_encoder_or_embedding()
103+
or not config.train_text_encoder_2_or_embedding()
104+
),
102105
)
103106

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

‎modules/dataLoader/IdeogramBaseDataLoader.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ def _preparation_modules(self, config: TrainConfig, model: IdeogramModel):
5555
modules.append(downscale_mask)
5656
modules += [tokenize_prompt, encode_prompt]
5757

58-
if config.latent_caching:
58+
if config.text_caching:
5959
modules.append(prune_masked_tokens)
6060

6161
return modules
@@ -84,7 +84,7 @@ def _cache_modules(self, config: TrainConfig, model: IdeogramModel, model_setup:
8484
text_split_names=text_split_names,
8585
sort_names=sort_names,
8686
config=config,
87-
text_caching=True,
87+
text_caching=config.text_caching,
8888
)
8989

9090
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
113113
train_dtype=model.train_dtype,
114114
)
115115

116-
if config.latent_caching:
116+
if config.text_caching:
117117
output_module_list = [pad_masked_tokens] + output_module_list
118118

119119
return output_module_list

‎modules/dataLoader/Krea2BaseDataLoader.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ def _preparation_modules(self, config: TrainConfig, model: Krea2Model):
6262
if not config.train_text_encoder_or_embedding():
6363
modules.append(encode_prompt)
6464

65-
if config.latent_caching and not config.train_text_encoder_or_embedding():
65+
if config.text_caching and not config.train_text_encoder_or_embedding():
6666
modules.append(prune_masked_tokens)
6767

6868
return modules
@@ -92,7 +92,7 @@ def _cache_modules(self, config: TrainConfig, model: Krea2Model, model_setup: Ba
9292
text_split_names=text_split_names,
9393
sort_names=sort_names,
9494
config=config,
95-
text_caching=not config.train_text_encoder_or_embedding(),
95+
text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
9696
)
9797

9898
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
122122
train_dtype=model.train_dtype,
123123
)
124124

125-
if config.latent_caching and not config.train_text_encoder_or_embedding():
125+
if config.text_caching and not config.train_text_encoder_or_embedding():
126126
output_module_list = [pad_masked_tokens] + output_module_list
127127

128128
return output_module_list

‎modules/dataLoader/PixArtAlphaBaseDataLoader.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ def _cache_modules(self, config: TrainConfig, model: PixArtAlphaModel, model_set
8787
text_split_names=text_split_names,
8888
sort_names=sort_names,
8989
config=config,
90-
text_caching=not config.train_text_encoder_or_embedding(),
90+
text_caching=config.text_caching and not config.train_text_encoder_or_embedding(),
9191
)
9292

9393
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
@@ -55,7 +55,7 @@ def _preparation_modules(self, config: TrainConfig, model: QwenModel):
5555
if not config.train_text_encoder_or_embedding():
5656
modules.append(encode_prompt)
5757

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

6161
return modules
@@ -85,7 +85,7 @@ def _cache_modules(self, config: TrainConfig, model: QwenModel, model_setup: Bas
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: QwenModel, model_setup: BaseQwenSetup):
@@ -115,7 +115,7 @@ def _output_modules(self, config: TrainConfig, model: QwenModel, model_setup: Ba
115115
train_dtype=model.train_dtype,
116116
)
117117

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

121121
return output_module_list

0 commit comments

Comments
 (0)