diff --git a/.gitignore b/.gitignore
index 9b0e1abfa..7106e4f98 100644
--- a/.gitignore
+++ b/.gitignore
@@ -160,7 +160,8 @@ debug.env
 tags
 
 # pre-commit
-.pre-commit*
+.pre-commit-cache/
+!.pre-commit-config.yaml
 
 # .lock
 *.lock
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 000000000..8fec6736f
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,35 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: trailing-whitespace
+        exclude: ^\.github/
+      - id: end-of-file-fixer
+        exclude: ^\.github/
+      - id: check-yaml
+      - id: check-added-large-files
+        args: ['--maxkb=5000']
+      - id: check-merge-conflict
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.9.9
+    hooks:
+      - id: ruff
+        args: [--fix]
+
+  - repo: https://github.com/google/pyink
+    rev: 23.10.0
+    hooks:
+      - id: pyink
+        args: ["--pyink-indentation=2", "--line-length=125"]
+
+  - repo: local
+    hooks:
+      - id: pylint
+        name: pylint
+        entry: pylint --fail-under=7
+        language: system
+        types: [python]
+        files: ^(src/maxdiffusion|end_to_end)/
diff --git a/README.md b/README.md
index 0665ec174..6897441e2 100755
--- a/README.md
+++ b/README.md
@@ -181,7 +181,7 @@ After installation completes, run the training script.
   ```bash
   rm $TFRECORDS_DATASET_DIR/eval_timesteps/file_42-430.tfrec
   ```
-  
+
   ### Training on a Single VM
 
   Loading the data is supported both locally from the disk created above, or from `gcs`. In this guide, we'll be using a gcs bucket to train. First copy the data to the GCS bucket.
@@ -266,7 +266,7 @@ After installation completes, run the training script.
   - per_device_batch_size can be a fractional, but must be a whole number when multiplied by number of devices. In this example, 0.25 * 4 (devices) = effective global batch size = 1.
   - The step time in v5p-8 with global batch size = 1 is large due to using `FULL` remat. On larger number of chips we can run larger batch sizes greatly increasing MFU, as we will see in the next session of deploying with xpk.
   - To enable eval during training set `eval_every` to a value > 0.
-  - In Wan2.1, the ici_fsdp_parallelism axis is used for sequence parallelism, the ici_tensor_parallelism axis is used for head parallelism. 
+  - In Wan2.1, the ici_fsdp_parallelism axis is used for sequence parallelism, the ici_tensor_parallelism axis is used for head parallelism.
     - You can enable both, keeping in mind that Wan2.1 has 40 heads and 40 must be evenly divisible by ici_tensor_parallelism.
     - For Sequence parallelism, the code pads the sequence length to evenly divide the sequence. Try out different ici_fsdp_parallelism numbers, but we find 2 and 4 to be the best right now.
   - For use on GPU it is recommended to enable the cudnn_te_flash attention kernel for optimal performance.
@@ -309,7 +309,7 @@ After installation completes, run the training script.
   ### Deploying with XPK
 
   This assumes the user has already created an xpk cluster, installed all dependencies and the also created the dataset from the step above. For getting started with MaxDiffusion and xpk see [this guide](docs/getting_started/run_maxdiffusion_via_xpk.md).
-  
+
   Using v5p-256 Then the command to run on xpk is as follows:
 
   ```bash
@@ -523,8 +523,8 @@ To generate images, run the following command:
   ```bash
   python src/maxdiffusion/generate_ltx_video.py src/maxdiffusion/configs/ltx_video.yml output_dir="[SAME DIRECTORY]" config_path="src/maxdiffusion/models/ltx_video/ltxv-13B.json"
   ```
-  Img2video Generation: 
-  
+  Img2video Generation:
+
   Add conditioning image path as conditioning_media_paths in the form of ["IMAGE_PATH"] along with other generation parameters in the ltx_video.yml file. Then follow same instruction as above.
 
   ## LTX-2 Video
@@ -704,14 +704,14 @@ python src/maxdiffusion/generate_wan.py \
 ### Ring Attention
 We added ring attention support for Wan models. Below are the stats for one `720p` (81 frames) video generation (with CFG DP):
 | Accelerator |  Model | Attention Type | Inference Steps | Sharding | e2e Generation Time |
-| -- | -- | -- | -- | -- | -- | 
+| -- | -- | -- | -- | -- | -- |
 | v7x-8 | WAN 2.1 | Tokamax Flash | 50 | dp2-fsdp1-context4-tp1 | **249.3** |
 | v7x-8 | WAN 2.1 | Tokamax Ring | 50 | dp2-fsdp1-context4-tp1 | 252.4 |
 | v7x-8 | WAN 2.2 | Tokamax Flash | 40 | dp2-fsdp1-context4-tp1 | **194.4** |
 | v7x-8 | WAN 2.2 | Tokamax Ring | 40 | dp2-fsdp1-context4-tp1 | 201.7 |
 
 | Accelerator |  Model | Attention Type | Inference Steps | Sharding | e2e Generation Time |
-| -- | -- | -- | -- | -- | -- | 
+| -- | -- | -- | -- | -- | -- |
 | v7x-16 | WAN 2.1 | Tokamax Flash | 50 | dp2-fsdp1-context8-tp1 | **127.1** |
 | v7x-16 | WAN 2.1 | Tokamax Ring | 50 | dp2-fsdp1-context8-tp1 | 137.2 |
 | v7x-16 | WAN 2.2 | Tokamax Flash | 40 | dp2-fsdp1-context8-tp1 | **106.0** |
@@ -751,7 +751,7 @@ The optimal attention tile sizes (`block_q` / `block_kv`) depend on the sequence
 
   If you are using a TPU v6e (Trillium), you can use optimized flash block sizes for faster inference. Uncomment Flux-dev [config](src/maxdiffusion/configs/base_flux_dev.yml#60) and Flux-schnell [config](src/maxdiffusion/configs/base_flux_schnell.yml#68)
 
-  To keep text encoders, vae and transformer on HBM memory at all times, the following command shards the model across devices. 
+  To keep text encoders, vae and transformer on HBM memory at all times, the following command shards the model across devices.
 
   ```bash
   python src/maxdiffusion/generate_flux.py src/maxdiffusion/configs/base_flux_schnell.yml jax_cache_dir=/tmp/cache_dir run_name=flux_test output_dir=/tmp/ prompt="photograph of an electronics chip in the shape of a race car with trillium written on its side" per_device_batch_size=1 ici_data_parallelism=1 ici_fsdp_parallelism=-1 offload_encoders=False
@@ -900,13 +900,40 @@ MaxDiffusion started as a fork of [Diffusers](https://github.com/huggingface/dif
 Whether you are forking MaxDiffusion for your own needs or intending to contribute back to the community, a full suite of tests can be found in `tests` and `src/maxdiffusion/tests`.
 
 To run unit tests simply run:
-```
+```bash
 python -m pytest
 ```
 
-This project uses `pylint` and `pyink` to enforce code style. Before submitting a pull request, please ensure your code passes these checks by running:
+### Pre-commit Hooks
+
+We use [pre-commit](https://pre-commit.com/) to automatically check and format code before each commit (using `pyink`, `ruff`, `pylint`, and general git hygiene checks).
+
+> **Important:** Make sure you are in your active virtual environment (e.g. `maxdiffusion_venv` or your active venv) before running `pre-commit install`, so that hooks run using the environment's installed dependencies.
+
+```bash
+# 1. Activate your virtual environment first
+source <path-to-venv>/bin/activate
+
+# 2. Install pre-commit (if not already installed)
+pip install pre-commit
 
+# 3. Install git pre-commit hooks
+pre-commit install
 ```
+
+Once installed, pre-commit will automatically run on staged files whenever you run `git commit`.
+
+You can also run all pre-commit checks manually across the entire repository at any time:
+
+```bash
+pre-commit run --all-files
+```
+
+### Code Style
+
+This project uses `pylint` and `pyink` to enforce code style. Before submitting a pull request, please ensure your code passes these checks by running:
+
+```bash
 bash code_style.sh
 ```
 
@@ -919,4 +946,4 @@ The full suite of -end-to end tests is in `tests` and `src/maxdiffusion/tests`.
 To learn how to enable ML Diagnostics and XProf profiling for your runs, please see our [ML Diagnostics Guide](docs/profiling.md).
 
 ## Metrics
-To learn how to enable ML Diagnostics metrics tracking for your runs, please see our [Metrics Guide](docs/metrics.md).
\ No newline at end of file
+To learn how to enable ML Diagnostics metrics tracking for your runs, please see our [Metrics Guide](docs/metrics.md).
diff --git a/_typos.toml b/_typos.toml
index 94be96203..551099f98 100644
--- a/_typos.toml
+++ b/_typos.toml
@@ -10,4 +10,4 @@ parms="parms" # parms is used in scripts/convert_original_stable_diffusion_to_di
 
 
 [files]
-extend-exclude = ["_typos.toml"]
\ No newline at end of file
+extend-exclude = ["_typos.toml"]
diff --git a/dependencies/requirements/base_requirements/requirements.txt b/dependencies/requirements/base_requirements/requirements.txt
index d9b53811f..b69bb1f0d 100644
--- a/dependencies/requirements/base_requirements/requirements.txt
+++ b/dependencies/requirements/base_requirements/requirements.txt
@@ -44,4 +44,3 @@ transformers<5.0.0
 torch @ https://download.pytorch.org/whl/cpu/torch-2.10.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl
 torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.25.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl
 qwix @ https://github.com/google/qwix/archive/408a0f48f988b6c5b180e07f0cb1d05997bf0dcc.zip
-
diff --git a/dependencies/requirements/generated_requirements/requirements.txt b/dependencies/requirements/generated_requirements/requirements.txt
index 03e9e7217..742f10cc0 100644
--- a/dependencies/requirements/generated_requirements/requirements.txt
+++ b/dependencies/requirements/generated_requirements/requirements.txt
@@ -201,4 +201,4 @@ xprof>=2.22.0
 xxhash>=3.6.0
 yarl>=1.23.0
 zipp>=3.23.0
-zstandard>=0.25.0
\ No newline at end of file
+zstandard>=0.25.0
diff --git a/docker_build_dependency_image.sh b/docker_build_dependency_image.sh
index 5c2c0f8d2..48685b56d 100644
--- a/docker_build_dependency_image.sh
+++ b/docker_build_dependency_image.sh
@@ -76,4 +76,4 @@ else
     --build-arg BASEIMAGE=${BASEIMAGE} \
     -t ${LOCAL_IMAGE_NAME} \
     -f maxdiffusion_dependencies.Dockerfile .
-fi
\ No newline at end of file
+fi
diff --git a/docker_upload_runner.sh b/docker_upload_runner.sh
index dc365c665..969233135 100644
--- a/docker_upload_runner.sh
+++ b/docker_upload_runner.sh
@@ -49,4 +49,4 @@ docker build --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} -f ./maxdiffusion_runner.
 docker tag ${LOCAL_IMAGE_NAME_RUNNER} gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}:latest
 docker push gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}:latest
 
-echo "All done, check out your artifacts at: gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}"
\ No newline at end of file
+echo "All done, check out your artifacts at: gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}"
diff --git a/docs/README.md b/docs/README.md
index aaad1f4e5..543b95943 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -14,7 +14,7 @@ This folder contains documentation for getting started with and using MaxDiffusi
 
 ## Training
 
-* **[Common Training Guide](train_README.md)** -  Provides a comprehensive guide to training MaxDiffusion models, including script usage, configuration options, and sharding strategies. 
+* **[Common Training Guide](train_README.md)** -  Provides a comprehensive guide to training MaxDiffusion models, including script usage, configuration options, and sharding strategies.
 
 ## Data Input
 
diff --git a/docs/attention_blocks_flowchart.md b/docs/attention_blocks_flowchart.md
index 69816ac79..d60fb75bd 100644
--- a/docs/attention_blocks_flowchart.md
+++ b/docs/attention_blocks_flowchart.md
@@ -4,7 +4,7 @@
 - "block_q": Block sizes (HBM TO VMEM and VREG) to tile along Q sequence in forward pass
 - "block_kv_compute" : Sub Block size (VMEM to VREG) of "block_kv" where compute is performed in forward pass. It must be factor or same as "block_kv"
 - "block_kv" : Block sizes (HBM TO VMEM) to tile along KV sequence in forward pass
-- "block_q_dkv" : Block sizes along Q sequence in backward pass with fused kernel to compute gradient of q, k , v. It must be factor or same as block_q 
+- "block_q_dkv" : Block sizes along Q sequence in backward pass with fused kernel to compute gradient of q, k , v. It must be factor or same as block_q
 - "block_kv_dkv" : Block sizes along KV sequence in backward pass. It must be factor or same as block_kv
 - "block_kv_dkv_compute" : Sub Block Sizes of block_kv_dkv, must be factor or same as "block_kv_dkv"
 - "block_q_dq" : Block sizes along Q sequence in backward pass with unfused kernel to compute gradient of just q. it must be factor or same as "block_q"
@@ -13,7 +13,7 @@
 
 ## Flowchart
 
-Maxdiffusion automatically adheres to this flowchart to ensure working, and there is a log that will inform you on the modifications that maxdiffusion makes to the specified block sizes. 
+Maxdiffusion automatically adheres to this flowchart to ensure working, and there is a log that will inform you on the modifications that maxdiffusion makes to the specified block sizes.
 
 ![alt text](attention_blocks_flowchart.png)
 
@@ -21,10 +21,10 @@ Maxdiffusion automatically adheres to this flowchart to ensure working, and ther
 
 ## How block sizes matter for perfomance and accuracy
 
-Block sizes key to saturating HBM bandwidth and ensuring maximum possible overlap of computation on cores with HBM use and VMEM to VREG. It is highly recommended to tune them. 
+Block sizes key to saturating HBM bandwidth and ensuring maximum possible overlap of computation on cores with HBM use and VMEM to VREG. It is highly recommended to tune them.
 
 Block sizes also have an effect on the sequence length. Sequence length is multiple of resolution and number of frames (video), along with VAE scale down factors and patchifying ratios. This sequence length or shard of this sequence length needs to be multiple of the block sizes specified. Therefore maxdiffusion pads the sequence lengths to the nearest multiple of the block sizes. It is advisable to choose block sizes which are factor of sequence length, atleast for the Q block sizes.
 
 > In cross attention Image or Video tokens are attending to text tokens sequence length of text tokens is really small and potentially smaller than specified block size so KV block sizes are overwritten to safe values.
 
-> KV block sizes must be multiple of 128 since the size of register is 8x128 and in attention KV sequence dim lies on 128 for the multiplications as K is transposed.
\ No newline at end of file
+> KV block sizes must be multiple of 128 since the size of register is 8x128 and in attention KV sequence dim lies on 128 for the multiplications as K is transposed.
diff --git a/docs/data_README.md b/docs/data_README.md
index 5459f2906..02c11ef09 100644
--- a/docs/data_README.md
+++ b/docs/data_README.md
@@ -4,7 +4,7 @@
 Currently MaxDiffusion supports 3 data input pipelines, controlled by the flag `dataset_type`
 | Pipeline | Dataset Location | Dataset formats | Features or limitations |
 | -------- | ---------------- | --------------- | ----------------------- |
-| HuggingFace (hf) | datasets in HuggingFace Hub or local/Cloud Storage | Formats supported in HF Hub: parquet, arrow, json, csv, txt | data are not loaded in memory but streamed from the saved location, good for big dataset | 
+| HuggingFace (hf) | datasets in HuggingFace Hub or local/Cloud Storage | Formats supported in HF Hub: parquet, arrow, json, csv, txt | data are not loaded in memory but streamed from the saved location, good for big dataset |
 | tf | dataset will be downloaded form HuggingFace Hub to disk | Formats supported in HF Hub: parquet, arrow, json, csv, txt | Will read the whole dataset into memory, works for small dataset |
 | tfrecord | local/Cloud Storage | TFRecord | data are not loaded in memory but streamed from the saved location, good for big dataset |
 | Grain | local/Cloud Storage | ArrayRecord (or any random access format) | data are not loaded in memory but streamed from the saved location, good for big dataset, supports global shuffle and data iterator checkpoint for determinism (see details in [doc](https://github.com/AI-Hypercomputer/maxtext/blob/main/getting_started/Data_Input_Pipeline.md#grain-pipeline---for-determinism)) |
@@ -14,16 +14,16 @@ Currently MaxDiffusion supports 3 data input pipelines, controlled by the flag `
 ### HuggingFace Streaming (dataset_type=hf)
 #### Example config for streaming from HuggingFace Hub (no download needed):
 ```
-dataset_type: hf 
+dataset_type: hf
 dataset_name: BleachNick/UltraEdit_500k  # for using https://huggingface.co/datasets/BleachNick/UltraEdit_500k
-image_column: source_image 
-caption_column: source_caption 
+image_column: source_image
+caption_column: source_caption
 train_split: FreeForm
 hf_access_token: ''  # provide token if using gated dataset or tokenizer
 ```
 #### Example config for streaming from downloaded data in a GCS bucket:
 ```
-dataset_type: hf 
+dataset_type: hf
 dataset_name: parquet  # or json, arrow, etc.
 hf_train_files: gs://<bucket>/<folder>/*-train-*.parquet  # match the train files
 ```
diff --git a/docs/getting_started/first_run.md b/docs/getting_started/first_run.md
index 7ba0d106d..54531a36e 100644
--- a/docs/getting_started/first_run.md
+++ b/docs/getting_started/first_run.md
@@ -11,7 +11,7 @@ multiple hosts.
 1. [Create and SSH to a single-host TPU (v6-8). ](https://cloud.google.com/tpu/docs/users-guide-tpu-vm#creating_a_cloud_tpu_vm_with_gcloud)
 * You can find here [here](https://cloud.google.com/tpu/docs/regions-zones) the list of zones that support the v6(Trillium) TPUs
 * We recommend using the base VM image "v2-alpha-tpuv6e", which meets the version requirements: Ubuntu Version 22.04, Python 3.12 and Tensorflow >= 2.12.0
-   
+
 1. Clone MaxDiffusion in your TPU VM.
 ```
 git clone https://github.com/AI-Hypercomputer/maxdiffusion.git
@@ -34,4 +34,3 @@ source ~/$venv_name/bin/activate
 ## Getting Starting: Multihost development
 
 [GKE, recommended] [Running MaxDiffusion with xpk](run_maxdiffusion_via_xpk.md) - Quick Experimentation and Production support
-
diff --git a/docs/getting_started/run_maxdiffusion_via_xpk.md b/docs/getting_started/run_maxdiffusion_via_xpk.md
index e243135a4..78204a481 100644
--- a/docs/getting_started/run_maxdiffusion_via_xpk.md
+++ b/docs/getting_started/run_maxdiffusion_via_xpk.md
@@ -66,7 +66,7 @@ after which log out and log back in to the machine.
 
     See details on docker images in xpk here: https://github.com/google/xpk/blob/main/README.md#how-to-add-docker-images-to-a-xpk-workload
 
-    **Note:** When using the XPK command, ensure you include `pip install .` to install the package from the current directory. This is necessary because the container is created from a copy of your local directory, and `pip install .` ensures any local changes you've made are applied within the container. 
+    **Note:** When using the XPK command, ensure you include `pip install .` to install the package from the current directory. This is necessary because the container is created from a copy of your local directory, and `pip install .` ensures any local changes you've made are applied within the container.
 
     __Using xpk to upload image to your gcp project and run MaxDiffusion__
 
@@ -104,4 +104,4 @@ after which log out and log back in to the machine.
       --tpu-type=v4-8 \
       --num-slices=1  \
       --command "pip install . && python src/maxdiffusion/train.py src/maxdiffusion/configs/base_2_base.yml run_name="my_run" output_dir="gs://your-bucket/""
-      ```
\ No newline at end of file
+      ```
diff --git a/docs/profiling.md b/docs/profiling.md
index b43e8ee7d..3a8c010bc 100644
--- a/docs/profiling.md
+++ b/docs/profiling.md
@@ -31,4 +31,4 @@ If permissions are not configured correctly, your job will fail with an error si
 ## 4. Viewing Your Runs
 Once your job is running with diagnostics enabled, you can monitor the profiles, execution times, and metrics in the Cluster Director console here:
 
-🔗 **https://pantheon.corp.google.com/cluster-director/diagnostics**
\ No newline at end of file
+🔗 **https://pantheon.corp.google.com/cluster-director/diagnostics**
diff --git a/end_to_end/tpu/test_sdxl_training_loss.sh b/end_to_end/tpu/test_sdxl_training_loss.sh
index 78913f6c0..6418c1d60 100755
--- a/end_to_end/tpu/test_sdxl_training_loss.sh
+++ b/end_to_end/tpu/test_sdxl_training_loss.sh
@@ -22,4 +22,4 @@ export LIBTPU_INIT_ARGS=""
 $TRAIN_CMD
 
 # Assert training loss is smaller than input LOSS_THRESHOLD
-python3 end_to_end/tpu/eval_assert.py final_loss metrics.txt $LOSS_THRESHOLD
\ No newline at end of file
+python3 end_to_end/tpu/eval_assert.py final_loss metrics.txt $LOSS_THRESHOLD
diff --git a/gpu_multi_process_run.sh b/gpu_multi_process_run.sh
index b689347f9..efe6b9c1f 100644
--- a/gpu_multi_process_run.sh
+++ b/gpu_multi_process_run.sh
@@ -153,4 +153,4 @@ eval ${COMMAND} &
 PID=$!
 PIDS+=($PID)
 
-wait_all_success_or_exit "${PIDS[@]}"
\ No newline at end of file
+wait_all_success_or_exit "${PIDS[@]}"
diff --git a/maxdiffusion_dependencies.Dockerfile b/maxdiffusion_dependencies.Dockerfile
index 9a9598271..a87430169 100644
--- a/maxdiffusion_dependencies.Dockerfile
+++ b/maxdiffusion_dependencies.Dockerfile
@@ -51,4 +51,4 @@ RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VER
 RUN --mount=type=cache,target=/root/.cache/pip bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION}
 
 # Cleanup
-RUN rm -rf /root/.cache/pip
\ No newline at end of file
+RUN rm -rf /root/.cache/pip
diff --git a/maxdiffusion_gpu_dependencies.Dockerfile b/maxdiffusion_gpu_dependencies.Dockerfile
index 45f033549..eebb09212 100644
--- a/maxdiffusion_gpu_dependencies.Dockerfile
+++ b/maxdiffusion_gpu_dependencies.Dockerfile
@@ -1,6 +1,6 @@
 # syntax=docker/dockerfile:experimental
 # Note: This pulls in the lastest of jax:base
-ARG BASEIMAGE=ghcr.io/nvidia/jax:base 
+ARG BASEIMAGE=ghcr.io/nvidia/jax:base
 FROM $BASEIMAGE
 
 # Stopgaps measure to circumvent gpg key setup issue.
@@ -45,4 +45,4 @@ RUN ls .
 RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}"
 RUN --mount=type=cache,target=/root/.cache/pip bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE}
 
-WORKDIR /deps
\ No newline at end of file
+WORKDIR /deps
diff --git a/maxdiffusion_runner.Dockerfile b/maxdiffusion_runner.Dockerfile
index 25bb66ae3..676f13eee 100644
--- a/maxdiffusion_runner.Dockerfile
+++ b/maxdiffusion_runner.Dockerfile
@@ -7,4 +7,4 @@ WORKDIR /deps
 # Copy all files from local workspace into docker container
 COPY . .
 
-WORKDIR /deps
\ No newline at end of file
+WORKDIR /deps
diff --git a/pyproject.toml b/pyproject.toml
index 6065c8a3e..25d1cc31e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -67,9 +67,9 @@ test = [
 torch = ["torch", "accelerate"]
 flax = ["jax", "jaxlib", "flax"]
 dev = [
-    "urllib3", "black", "isort", "ruff", "hf-doc-builder", "compel", "datasets", "Jinja2", 
-    "invisible-watermark", "k-diffusion", "librosa", "omegaconf", "parameterized", "pytest", 
-    "pytest-timeout", "pytest-xdist", "requests-mock", "safetensors", "sentencepiece", "scipy", 
+    "urllib3", "black", "isort", "ruff", "hf-doc-builder", "compel", "datasets", "Jinja2",
+    "invisible-watermark", "k-diffusion", "librosa", "omegaconf", "parameterized", "pytest",
+    "pytest-timeout", "pytest-xdist", "requests-mock", "safetensors", "sentencepiece", "scipy",
     "torchvision", "accelerate", "protobuf", "tensorboard", "torch", "jax", "jaxlib", "flax"
 ]
 
@@ -119,4 +119,4 @@ indent-style = "space"
 skip-magic-trailing-comma = false
 
 # Like Black, automatically detect the appropriate line ending.
-line-ending = "auto"
\ No newline at end of file
+line-ending = "auto"
diff --git a/setup.sh b/setup.sh
index 3f1141888..d6d5e089d 100644
--- a/setup.sh
+++ b/setup.sh
@@ -172,4 +172,4 @@ else
 fi
 
 # Install maxdiffusion
-python3 -m uv pip install --no-deps .
\ No newline at end of file
+python3 -m uv pip install --no-deps .
diff --git a/setup_gcsfuse.sh b/setup_gcsfuse.sh
index 2d1d57b8d..62eb1957b 100755
--- a/setup_gcsfuse.sh
+++ b/setup_gcsfuse.sh
@@ -15,7 +15,7 @@
 # limitations under the License.
 
 # Description:
-# DATASET_GCS_BUCKET and MOUNT_PATH are required，FILE_PATH is optional, if provided, 
+# DATASET_GCS_BUCKET and MOUNT_PATH are required，FILE_PATH is optional, if provided,
 # will use ls -R to improve first-time read performance by prefetching metadata
 # https://cloud.google.com/storage/docs/cloud-storage-fuse/performance#improve-first-time-reads
 # bash setup_gcsfuse.sh DATASET_GCS_BUCKET=maxdiffusion-github-runner-test-assets MOUNT_PATH=/tmp/gcsfuse FILE_PATH=/tmp/gcsfuse/my_dataset
@@ -52,7 +52,7 @@ gcsfuse -o ro --implicit-dirs --log-severity=debug \
         --type-cache-max-size-mb=-1 --stat-cache-max-size-mb=-1 --kernel-list-cache-ttl-secs=-1 --metadata-cache-ttl-secs=-1 \
         --log-file=$HOME/gcsfuse_$TIMESTAMP.json "$DATASET_GCS_BUCKET" "$MOUNT_PATH"
 
-if [[ ! -z ${FILE_PATH} ]] ; then 
+if [[ ! -z ${FILE_PATH} ]] ; then
   FILE_COUNT=$(ls -R $FILE_PATH | wc -l)
   echo $FILE_COUNT files found in $FILE_PATH
 fi
diff --git a/src/maxdiffusion/configs/README.md b/src/maxdiffusion/configs/README.md
index d57d8b0c7..376f69594 100644
--- a/src/maxdiffusion/configs/README.md
+++ b/src/maxdiffusion/configs/README.md
@@ -27,4 +27,4 @@ base_xl_lightning.yml - used to run inference using [SDXL-Lightning](https://hug
 
 base_flux_dev.yml - used for training and inference using [Flux Dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)
 
-base_flux_schnell.yml - used for training and inference using [Flux Schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell)
\ No newline at end of file
+base_flux_schnell.yml - used for training and inference using [Flux Schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell)
diff --git a/src/maxdiffusion/configs/base14.yml b/src/maxdiffusion/configs/base14.yml
index 55a19c8b2..a688fccca 100644
--- a/src/maxdiffusion/configs/base14.yml
+++ b/src/maxdiffusion/configs/base14.yml
@@ -21,7 +21,7 @@ write_metrics: True
 gcs_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
@@ -53,12 +53,12 @@ attention: 'dot_product' # Supported attention: dot_product, flash
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 flash_block_sizes: {}
 # GroupNorm groups
 norm_num_groups: 32
@@ -245,7 +245,7 @@ cache_dreambooth_dataset: False
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
+use_qwix_quantization: False
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
 
 # ML Diagnostics settings
diff --git a/src/maxdiffusion/configs/base21.yml b/src/maxdiffusion/configs/base21.yml
index e5b575503..48fec99e2 100644
--- a/src/maxdiffusion/configs/base21.yml
+++ b/src/maxdiffusion/configs/base21.yml
@@ -21,7 +21,7 @@ write_metrics: True
 gcs_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
@@ -52,12 +52,12 @@ attention: 'dot_product' # Supported attention: dot_product, flash
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 
 flash_block_sizes: {}
 # GroupNorm groups
@@ -185,11 +185,6 @@ checkpoint_every: -1
 # enables one replica to read the ckpt then broadcast to the rest
 enable_single_replica_ckpt_restoring: False
 
-# Prepare image latents and text encoder outputs
-# during dataset creation to reduce memory consumption.
-cache_latents_text_encoder_outputs: True
-
-
 # Training loop
 learning_rate: 1.e-7
 scale_lr: False
@@ -247,7 +242,7 @@ quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
-use_qwix_quantization: False 
+use_qwix_quantization: False
 
 # ML Diagnostics settings
 enable_ml_diagnostics: False
diff --git a/src/maxdiffusion/configs/base_2_base.yml b/src/maxdiffusion/configs/base_2_base.yml
index 3ae2fb14b..808bbd5ca 100644
--- a/src/maxdiffusion/configs/base_2_base.yml
+++ b/src/maxdiffusion/configs/base_2_base.yml
@@ -21,7 +21,7 @@ write_metrics: True
 gcs_metrics: False
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
@@ -53,12 +53,12 @@ attention: 'flash' # Supported attention: dot_product, flash
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 
 flash_block_sizes: {}
 # to override default block sizes for flash attention
@@ -261,10 +261,10 @@ cache_dreambooth_dataset: False
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
+use_qwix_quantization: False
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
 
 # ML Diagnostics settings
 enable_ml_diagnostics: False
 profiler_gcs_path: ""
-enable_ondemand_xprof: False
\ No newline at end of file
+enable_ondemand_xprof: False
diff --git a/src/maxdiffusion/configs/base_flux2klein.yml b/src/maxdiffusion/configs/base_flux2klein.yml
index f1a7e0fa4..36f098c7d 100644
--- a/src/maxdiffusion/configs/base_flux2klein.yml
+++ b/src/maxdiffusion/configs/base_flux2klein.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 write_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
@@ -67,12 +67,12 @@ attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 
 flash_block_sizes: {
   "block_q": 4608,
@@ -270,7 +270,7 @@ controlnet_conditioning_scale: 0.5
 controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png'
 quantization: ''
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
+use_qwix_quantization: False
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
 
 # ML Diagnostics settings
@@ -286,4 +286,3 @@ interactive: False
 
 # Note: Architecture dimensions (depth, num_double_layers, num_attention_heads) are
 # automatically inferred from pretrained_model_name_or_path (transformer/config.json).
-
diff --git a/src/maxdiffusion/configs/base_flux2klein_9B.yml b/src/maxdiffusion/configs/base_flux2klein_9B.yml
index a3a0afeac..3d9c1848e 100644
--- a/src/maxdiffusion/configs/base_flux2klein_9B.yml
+++ b/src/maxdiffusion/configs/base_flux2klein_9B.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 write_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
@@ -67,12 +67,12 @@ attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 
 flash_block_sizes: {
   "block_q": 4608,
@@ -270,7 +270,7 @@ controlnet_conditioning_scale: 0.5
 controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png'
 quantization: ''
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
+use_qwix_quantization: False
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
 
 # ML Diagnostics settings
@@ -286,4 +286,3 @@ interactive: False
 
 # Note: Architecture dimensions (depth, num_double_layers, num_attention_heads) are
 # automatically inferred from pretrained_model_name_or_path (transformer/config.json).
-
diff --git a/src/maxdiffusion/configs/base_flux_dev.yml b/src/maxdiffusion/configs/base_flux_dev.yml
index 20f1fcc8f..5761eda0a 100644
--- a/src/maxdiffusion/configs/base_flux_dev.yml
+++ b/src/maxdiffusion/configs/base_flux_dev.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 write_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
@@ -68,12 +68,12 @@ use_experimental_scheduler: False
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 
 #flash_block_sizes: {}
 # Use the following flash_block_sizes on v6e (Trillium) due to larger vmem.
@@ -322,7 +322,7 @@ controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Goo
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
+use_qwix_quantization: False
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
 
 # ML Diagnostics settings
diff --git a/src/maxdiffusion/configs/base_flux_dev_multi_res.yml b/src/maxdiffusion/configs/base_flux_dev_multi_res.yml
index 548b6164f..5bc034ec5 100644
--- a/src/maxdiffusion/configs/base_flux_dev_multi_res.yml
+++ b/src/maxdiffusion/configs/base_flux_dev_multi_res.yml
@@ -21,7 +21,7 @@ write_metrics: True
 gcs_metrics: False
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
@@ -66,12 +66,12 @@ attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 
 #flash_block_sizes: {}
 # Use the following flash_block_sizes on v6e (Trillium) due to larger vmem.
@@ -288,7 +288,7 @@ controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Goo
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
+use_qwix_quantization: False
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
 
 # ML Diagnostics settings
diff --git a/src/maxdiffusion/configs/base_flux_schnell.yml b/src/maxdiffusion/configs/base_flux_schnell.yml
index 5d3f7f2fa..1206a2516 100644
--- a/src/maxdiffusion/configs/base_flux_schnell.yml
+++ b/src/maxdiffusion/configs/base_flux_schnell.yml
@@ -21,7 +21,7 @@ write_metrics: True
 gcs_metrics: False
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
@@ -65,12 +65,12 @@ attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 flash_block_sizes: {
   "block_q" : 256,
   "block_kv_compute" : 256,
@@ -297,7 +297,7 @@ controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Goo
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
+use_qwix_quantization: False
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
 
 save_final_checkpoint: False
diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml
index d690b4893..fd529c15b 100644
--- a/src/maxdiffusion/configs/base_wan_14b.yml
+++ b/src/maxdiffusion/configs/base_wan_14b.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 write_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml
index 185b01277..35b17f9af 100644
--- a/src/maxdiffusion/configs/base_wan_27b.yml
+++ b/src/maxdiffusion/configs/base_wan_27b.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 write_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
@@ -464,4 +464,4 @@ enable_ssim: False
 # ML Diagnostics settings
 enable_ml_diagnostics: False
 profiler_gcs_path: ""
-enable_ondemand_xprof: False
\ No newline at end of file
+enable_ondemand_xprof: False
diff --git a/src/maxdiffusion/configs/base_wan_animate.yml b/src/maxdiffusion/configs/base_wan_animate.yml
index 2b547dd3d..e753df31e 100644
--- a/src/maxdiffusion/configs/base_wan_animate.yml
+++ b/src/maxdiffusion/configs/base_wan_animate.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 write_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
@@ -96,12 +96,12 @@ flash_min_seq_length: 4096
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 dropout: 0.0
 
 # Tuned for 720p (720x1280), 81 frames, CP=8 on Trillium (32MB VMEM):
diff --git a/src/maxdiffusion/configs/base_wan_i2v_14b.yml b/src/maxdiffusion/configs/base_wan_i2v_14b.yml
index 32f95620c..d3ca82d85 100644
--- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml
+++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 write_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
diff --git a/src/maxdiffusion/configs/base_wan_i2v_27b.yml b/src/maxdiffusion/configs/base_wan_i2v_27b.yml
index c97be88ec..43ae00e62 100644
--- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml
+++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 write_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
diff --git a/src/maxdiffusion/configs/base_xl.yml b/src/maxdiffusion/configs/base_xl.yml
index 031a018ea..e76849237 100644
--- a/src/maxdiffusion/configs/base_xl.yml
+++ b/src/maxdiffusion/configs/base_xl.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 write_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
@@ -53,12 +53,12 @@ attention: 'dot_product' # Supported attention: dot_product, flash
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 flash_block_sizes: {}
 # GroupNorm groups
 norm_num_groups: 32
@@ -257,12 +257,11 @@ controlnet_model_name_or_path: 'diffusers/controlnet-canny-sdxl-1.0'
 controlnet_from_pt: True
 controlnet_conditioning_scale: 0.5
 controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png'
-enable_mllog: False
 
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
+use_qwix_quantization: False
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
 
 # ML Diagnostics settings
diff --git a/src/maxdiffusion/configs/base_xl_lightning.yml b/src/maxdiffusion/configs/base_xl_lightning.yml
index d8f7801b4..9d06436d5 100644
--- a/src/maxdiffusion/configs/base_xl_lightning.yml
+++ b/src/maxdiffusion/configs/base_xl_lightning.yml
@@ -20,7 +20,7 @@ metrics_file: "" # for testing, local file that stores scalar metrics. If empty,
 gcs_metrics: True
 
 timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written.
-write_timing_metrics: True 
+write_timing_metrics: True
 
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
@@ -51,12 +51,12 @@ attention: 'flash' # Supported attention: dot_product, flash
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: True
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 flash_block_sizes: {}
 # GroupNorm groups
 norm_num_groups: 32
@@ -212,7 +212,7 @@ enable_mllog: False
 quantization: ''
 # Shard the range finding operation for quantization. By default this is set to number of slices.
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
+use_qwix_quantization: False
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
 
 # ML Diagnostics settings
diff --git a/src/maxdiffusion/configs/ltx2_3_video.yml b/src/maxdiffusion/configs/ltx2_3_video.yml
index b47ee927f..0aeaca647 100644
--- a/src/maxdiffusion/configs/ltx2_3_video.yml
+++ b/src/maxdiffusion/configs/ltx2_3_video.yml
@@ -15,7 +15,7 @@ ulysses_shards: -1
 ulysses_attention_chunks: 1
 a2v_attention_kernel: 'flash'
 v2a_attention_kernel: 'dot_product'
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 precision: 'bf16'
 scan_layers: True
 names_which_can_be_saved: []
@@ -142,14 +142,14 @@ train_split: 'train'
 dataset_type: 'tfrecord'
 cache_latents_text_encoder_outputs: True
 per_device_batch_size: 1.0
-compile_topology_num_slices: -1 
+compile_topology_num_slices: -1
 quantization_local_shard_count: -1
 use_qwix_quantization: False
 weight_quantization_calibration_method: "absmax"
 act_quantization_calibration_method: "absmax"
 bwd_quantization_calibration_method: "absmax"
 qwix_module_path: ".*"
-jit_initializers: True 
+jit_initializers: True
 enable_single_replica_ckpt_restoring: False
 seed: 10
 audio_format: "s16"
diff --git a/src/maxdiffusion/configs/ltx2_video.yml b/src/maxdiffusion/configs/ltx2_video.yml
index 271b07deb..f953eb8a0 100644
--- a/src/maxdiffusion/configs/ltx2_video.yml
+++ b/src/maxdiffusion/configs/ltx2_video.yml
@@ -15,7 +15,7 @@ ulysses_shards: -1
 ulysses_attention_chunks: 1
 a2v_attention_kernel: 'dot_product'
 v2a_attention_kernel: 'dot_product'
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 precision: 'bf16'
 
 # For scanning transformer layers
@@ -147,16 +147,15 @@ train_split: 'train'
 dataset_type: 'tfrecord'
 cache_latents_text_encoder_outputs: True
 per_device_batch_size: 1.0
-compile_topology_num_slices: -1 
+compile_topology_num_slices: -1
 quantization_local_shard_count: -1
 use_qwix_quantization: False
 weight_quantization_calibration_method: "absmax"
 act_quantization_calibration_method: "absmax"
 bwd_quantization_calibration_method: "absmax"
 qwix_module_path: ".*"
-jit_initializers: True 
+jit_initializers: True
 enable_single_replica_ckpt_restoring: False
-seed: 0
 audio_format: "s16"
 
 # LoRA parameters
diff --git a/src/maxdiffusion/configs/ltx_video.yml b/src/maxdiffusion/configs/ltx_video.yml
index d70154e0e..4b32c65a3 100644
--- a/src/maxdiffusion/configs/ltx_video.yml
+++ b/src/maxdiffusion/configs/ltx_video.yml
@@ -2,7 +2,7 @@
 hardware: 'tpu'
 skip_jax_distributed_system: False
 attention: 'flash'
-attention_sharding_uniform: True 
+attention_sharding_uniform: True
 
 jax_cache_dir: ''
 weights_dtype: 'bfloat16'
@@ -99,10 +99,10 @@ train_split: 'train'
 dataset_type: 'tf'
 cache_latents_text_encoder_outputs: True
 per_device_batch_size: 1
-compile_topology_num_slices: -1 
+compile_topology_num_slices: -1
 quantization_local_shard_count: -1
-use_qwix_quantization: False 
-jit_initializers: True 
+use_qwix_quantization: False
+jit_initializers: True
 enable_single_replica_ckpt_restoring: False
 
 # ML Diagnostics settings
diff --git a/src/maxdiffusion/models/ltx_video/ltxv-13B.json b/src/maxdiffusion/models/ltx_video/ltxv-13B.json
index 75b16b011..b23177b2a 100644
--- a/src/maxdiffusion/models/ltx_video/ltxv-13B.json
+++ b/src/maxdiffusion/models/ltx_video/ltxv-13B.json
@@ -22,4 +22,4 @@
     "positional_embedding_max_pos": [20, 2048, 2048],
     "timestep_scale_multiplier": 1000,
     "in_channels": 128
-}
\ No newline at end of file
+}
diff --git a/src/maxdiffusion/utils/model_card_template.md b/src/maxdiffusion/utils/model_card_template.md
index f19c85b0f..8dc0fcfcb 100644
--- a/src/maxdiffusion/utils/model_card_template.md
+++ b/src/maxdiffusion/utils/model_card_template.md
@@ -9,7 +9,7 @@ should probably proofread and complete it, then remove this comment. -->
 
 ## Model description
 
-This diffusion model is trained with the [🤗 Diffusers](https://github.com/huggingface/diffusers) library 
+This diffusion model is trained with the [🤗 Diffusers](https://github.com/huggingface/diffusers) library
 on the `{{ dataset_name }}` dataset.
 
 ## Intended uses & limitations
@@ -46,5 +46,3 @@ The following hyperparameters were used during training:
 ### Training results
 
 📈 [TensorBoard logs](https://huggingface.co/{{ repo_name }}/tensorboard?#scalars)
-
-
diff --git a/unit_test_and_lint.sh b/unit_test_and_lint.sh
index a908e06a0..0fd1d88f7 100755
--- a/unit_test_and_lint.sh
+++ b/unit_test_and_lint.sh
@@ -17,4 +17,4 @@
 # python3 -m pylint $(git ls-files '*.py')
 
 cd src/maxdiffusion
-python3 -m pytest
\ No newline at end of file
+python3 -m pytest
diff --git a/utils/update_dependency_table.py b/utils/update_dependency_table.py
index af0c840df..0fd7a04fe 100644
--- a/utils/update_dependency_table.py
+++ b/utils/update_dependency_table.py
@@ -20,42 +20,44 @@
 import os
 import re
 
+
 def main():
-    requirements_file = "dependencies/requirements/generated_requirements/requirements.txt"
-    target_file = "src/maxdiffusion/dependency_versions_table.py"
-
-    if not os.path.exists(requirements_file):
-        print(f"Error: {requirements_file} not found.")
-        return
-
-    deps = {}
-    with open(requirements_file, "r", encoding="utf-8") as f:
-        for line in f:
-            line = line.strip()
-            if not line or line.startswith(("#", "--")):
-                continue
-
-            # Simple regex to extract package name and version
-            # This covers: pkg, pkg>=1.0, pkg==1.0, pkg @ https://...
-            match = re.match(r"^([^!=<>~@\s]+)(?:.*)$", line)
-            if match:
-                pkg_name = match.group(1)
-                deps[pkg_name] = line
-
-    entries = "\n".join([f'    "{k}": "{v}",' for k, v in sorted(deps.items())])
-    content = [
-        "# THIS FILE HAS BEEN AUTOGENERATED. To update:",
-        "# 1. modify the requirements in dependencies/requirements/generated_requirements/requirements.txt",
-        "# 2. run `make deps_table_update` or `python utils/update_dependency_table.py`",
-        "deps = {",
-        entries,
-        "}",
-        "",
-    ]
-
-    print(f"Updating {target_file}")
-    with open(target_file, "w", encoding="utf-8", newline="\n") as f:
-        f.write("\n".join(content))
+  requirements_file = "dependencies/requirements/generated_requirements/requirements.txt"
+  target_file = "src/maxdiffusion/dependency_versions_table.py"
+
+  if not os.path.exists(requirements_file):
+    print(f"Error: {requirements_file} not found.")
+    return
+
+  deps = {}
+  with open(requirements_file, "r", encoding="utf-8") as f:
+    for line in f:
+      line = line.strip()
+      if not line or line.startswith(("#", "--")):
+        continue
+
+      # Simple regex to extract package name and version
+      # This covers: pkg, pkg>=1.0, pkg==1.0, pkg @ https://...
+      match = re.match(r"^([^!=<>~@\s]+)(?:.*)$", line)
+      if match:
+        pkg_name = match.group(1)
+        deps[pkg_name] = line
+
+  entries = "\n".join([f'    "{k}": "{v}",' for k, v in sorted(deps.items())])
+  content = [
+      "# THIS FILE HAS BEEN AUTOGENERATED. To update:",
+      "# 1. modify the requirements in dependencies/requirements/generated_requirements/requirements.txt",
+      "# 2. run `make deps_table_update` or `python utils/update_dependency_table.py`",
+      "deps = {",
+      entries,
+      "}",
+      "",
+  ]
+
+  print(f"Updating {target_file}")
+  with open(target_file, "w", encoding="utf-8", newline="\n") as f:
+    f.write("\n".join(content))
+
 
 if __name__ == "__main__":
-    main()
+  main()