chore: configure pre-commit hooks and format repository

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

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
 ```

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
-

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/""
-      ```
+      ```

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**
+🔗 **https://pantheon.corp.google.com/cluster-director/diagnostics**

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
+python3 end_to_end/tpu/eval_assert.py final_loss metrics.txt $LOSS_THRESHOLD

gpu_multi_process_run.sh

@@ -153,4 +153,4 @@ eval ${COMMAND} &
 PID=$!
 PIDS+=($PID)
 
-wait_all_success_or_exit "${PIDS[@]}"
+wait_all_success_or_exit "${PIDS[@]}"

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
+RUN rm -rf /root/.cache/pip

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
+WORKDIR /deps

maxdiffusion_runner.Dockerfile

@@ -7,4 +7,4 @@ WORKDIR /deps
 # Copy all files from local workspace into docker container
 COPY . .
 
-WORKDIR /deps
+WORKDIR /deps

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"
+line-ending = "auto"