Fix issue with split_mode and train_blocks

bmaltais · Aug 21, 2024 · 33937c2 · 33937c2
1 parent 0ee4053
commit 33937c2
Show file tree

Hide file tree

Showing 6 changed files with 283 additions and 139 deletions.
diff --git a/kohya_gui/class_flux1.py b/kohya_gui/class_flux1.py
@@ -112,19 +112,10 @@ def noise_offset_type_change(
                         value=self.config.get("flux1.timestep_sampling", "sigma"),
                         interactive=True,
                     )
-
-                    self.flux1_cache_text_encoder_outputs = gr.Checkbox(
-                        label="Cache Text Encoder Outputs",
-                        value=self.config.get("flux1.cache_text_encoder_outputs", False),
-                        info="Cache text encoder outputs to speed up inference",
-                        interactive=True,
-                    )
-                    self.flux1_cache_text_encoder_outputs_to_disk = gr.Checkbox(
-                        label="Cache Text Encoder Outputs to Disk",
-                        value=self.config.get(
-                            "flux1.cache_text_encoder_outputs_to_disk", False
-                        ),
-                        info="Cache text encoder outputs to disk to speed up inference",
+                    self.apply_t5_attn_mask = gr.Checkbox(
+                        label="Apply T5 Attention Mask",
+                        value=self.config.get("flux1.apply_t5_attn_mask", False),
+                        info="Apply attention mask to T5-XXL encode and FLUX double blocks ",
                         interactive=True,
                     )
                 with gr.Row():
@@ -158,12 +149,29 @@ def noise_offset_type_change(
                         step=1,
                         interactive=True,
                     )
+
+                with gr.Row():
+                    self.flux1_cache_text_encoder_outputs = gr.Checkbox(
+                        label="Cache Text Encoder Outputs",
+                        value=self.config.get("flux1.cache_text_encoder_outputs", False),
+                        info="Cache text encoder outputs to speed up inference",
+                        interactive=True,
+                    )
+                    self.flux1_cache_text_encoder_outputs_to_disk = gr.Checkbox(
+                        label="Cache Text Encoder Outputs to Disk",
+                        value=self.config.get(
+                            "flux1.cache_text_encoder_outputs_to_disk", False
+                        ),
+                        info="Cache text encoder outputs to disk to speed up inference",
+                        interactive=True,
+                    )
                     self.mem_eff_save = gr.Checkbox(
                         label="Memory Efficient Save",
                         value=self.config.get("flux1.mem_eff_save", False),
                         info="[Experimentsl] Enable memory efficient save. We do not recommend using it unless you are familiar with the code.",
                         interactive=True,
                     )
+
                 with gr.Row(visible=True if finetuning else False):
                     self.blockwise_fused_optimizers = gr.Checkbox(
                         label="Blockwise Fused Optimizer",