chore: configure pre-commit hooks and format repository

- Add .pre-commit-config.yaml with trailing-whitespace, end-of-file-fixer, check-yaml, ruff, pyink, and pylint
- Fix duplicate keys in base_xl.yml, base21.yml, and ltx2_video.yml
- Update .gitignore to track .pre-commit-config.yaml
- Format files and fix end of file newlines across the repository
- Add pre-commit setup and venv instructions to README.md

Author: mbohlool

.gitignore

@@ -160,7 +160,8 @@ debug.env
 tags
 
 # pre-commit
-.pre-commit*
+.pre-commit-cache/
+!.pre-commit-config.yaml
 
 # .lock
 *.lock

.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)/

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).
+To learn how to enable ML Diagnostics metrics tracking for your runs, please see our [Metrics Guide](docs/metrics.md).

_typos.toml

@@ -10,4 +10,4 @@ parms="parms" # parms is used in scripts/convert_original_stable_diffusion_to_di
 
 
 [files]
-extend-exclude = ["_typos.toml"]
+extend-exclude = ["_typos.toml"]

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
-

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
+zstandard>=0.25.0

docker_build_dependency_image.sh

@@ -76,4 +76,4 @@ else
     --build-arg BASEIMAGE=${BASEIMAGE} \
     -t ${LOCAL_IMAGE_NAME} \
     -f maxdiffusion_dependencies.Dockerfile .
-fi
+fi

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}"
+echo "All done, check out your artifacts at: gcr.io/$PROJECT/${CLOUD_IMAGE_NAME}"

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
 

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,18 +13,18 @@
 
 ## 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)
 
 > "tokamax_flash" uses the splash attention implementation in [tokamax-repo](https://github.com/openxla/tokamax/blob/main/tokamax/_src/ops/experimental/tpu/splash_attention/splash_attention_kernel.py) This kernel only supports fused backward pass where gradients for q,k,v are computed in a single kernel so "block_q_dq" and "block_kv_dq" are not used
 
 ## 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.
+> 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.