clean up code and lint

Author: jfacevedo-google

requirements.txt

@@ -9,7 +9,7 @@ optax>=0.2.3
 torch>=2.3.1
 torchvision>=0.18.1
 ftfy
-tensorboard==2.17.0
+tensorboard>=2.17.0
 tensorboardx==2.6.2.2
 tensorboard-plugin-profile==2.15.2
 Jinja2
@@ -25,7 +25,7 @@ ruff>=0.1.5,<=0.2
 git+https://github.com/mlperf/logging.git
 opencv-python-headless==4.10.0.84
 orbax-checkpoint==0.10.2
-tokenizers==0.20.0
+tokenizers==0.21.0
 huggingface_hub==0.24.7
 transformers==4.48.1
 einops==0.8.0

src/maxdiffusion/checkpointing/base_stable_diffusion_checkpointer.py

@@ -330,10 +330,7 @@ def load_checkpoint(self, step=None, scheduler_class=None):
       if self.checkpoint_type == STABLE_DIFFUSION_XL_CHECKPOINT:
         te_pretrained_2_config = CLIPTextConfig(**model_configs[0]["text_encoder_2_config"])
         text_encoder_2 = FlaxCLIPTextModelWithProjection(
-            te_pretrained_2_config,
-            seed=self.config.seed,
-            dtype=self.config.activations_dtype,
-            _do_init=False
+            te_pretrained_2_config, seed=self.config.seed, dtype=self.config.activations_dtype, _do_init=False
         )
         pipeline_kwargs["text_encoder_2"] = text_encoder_2
         # both tokenizers in sdxl are the same.

src/maxdiffusion/configs/base_flux_schnell.yml

@@ -54,17 +54,27 @@ precision: "DEFAULT"
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash
-flash_block_sizes: {}
-# Use the following flash_block_sizes on v6e (Trillium) due to larger vmem.
+flash_block_sizes: {
+  "block_q" : 256,
+  "block_kv_compute" : 256,
+  "block_kv" : 256,
+  "block_q_dkv" : 256,
+  "block_kv_dkv" : 256,
+  "block_kv_dkv_compute" : 256,
+  "block_q_dq" : 256,
+  "block_kv_dq" : 256
+}
+
+# Use the following flash_block_sizes on v6e (Trillium).
 # flash_block_sizes: {
-#   "block_q" : 1536,
-#   "block_kv_compute" : 1536,
-#   "block_kv" : 1536,
-#   "block_q_dkv" : 1536,
-#   "block_kv_dkv" : 1536,
-#   "block_kv_dkv_compute" : 1536,
-#   "block_q_dq" : 1536,
-#   "block_kv_dq" : 1536
+#   "block_q" : 2176,
+#   "block_kv_compute" : 2176,
+#   "block_kv" : 2176,
+#   "block_q_dkv" : 2176,
+#   "block_kv_dkv" : 2176,
+#   "block_kv_dkv_compute" : 2176,
+#   "block_q_dq" : 2176,
+#   "block_kv_dq" : 2176
 # }
 # GroupNorm groups
 norm_num_groups: 32