Fix transformer sharding, flash block sizing, and tests

Author: csgoogle

.github/workflows/UnitTests.yml

@@ -42,8 +42,6 @@ jobs:
         python-version: '3.12'
     - name: Install dependencies
       run: |
-        pip install -e .
-        pip uninstall jax jaxlib libtpu-nightly libtpu -y
         bash setup.sh MODE=stable
         export PATH=$PATH:$HOME/.local/bin
         pip install ruff
@@ -66,4 +64,4 @@ jobs:
 #      checks: read
 #      pull-requests: write
 #    needs: build
-#    uses: ./.github/workflows/AddLabel.yml
+#    uses: ./.github/workflows/AddLabel.yml

.github/workflows/UploadDockerImages.yml

@@ -28,15 +28,12 @@ jobs:
   build-image:
     runs-on: ["self-hosted", "e2", "cpu"]
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v5
     - name: Cleanup old docker images
       run: docker system prune --all --force
-    - name: build maxdiffusion jax ai image
+    - name: build maxdiffusion stable image
       run: |
-        bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
-    - name: build maxdiffusion w/ nightly jax ai image
-      run: |
-        bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack_nightly MODE=jax_ai_image PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/tpu/jax_nightly:latest
-    - name: build maxdiffusion jax nightly image
+        bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable MODE=stable PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_stable
+    - name: build maxdiffusion nightly image
       run: |
         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_nightly MODE=nightly PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_nightly

.github/workflows/pypi_release.yml

@@ -0,0 +1,48 @@
+# Copyright 2025 Google LLC
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This workflow will build, test and automatically release MaxDiffusion package to PyPI using Trusted Publishing (OIDC).
+
+name: Publish MaxDiffusion to PyPI
+
+# Triggers when a new "release" is published in the GitHub UI
+on:
+  release:
+    types: [published]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+  id-token: write
+
+jobs:
+  build_and_publish:
+    name: Build and Publish MaxDiffusion Package
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v5
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+    - name: Install build dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build hatchling hatch-requirements-txt
+    - name: Build package
+      run: python -m build
+    - name: Publish package
+      uses: pypa/gh-action-pypi-publish@release/v1
+      with:
+        packages-dir: dist/

Makefile

@@ -18,11 +18,11 @@ modified_only_fixup:
 # Update src/maxdiffusion/dependency_versions_table.py
 
 deps_table_update:
-	@python setup.py deps_table_update
+	@python utils/update_dependency_table.py
 
 deps_table_check_updated:
 	@md5sum src/maxdiffusion/dependency_versions_table.py > md5sum.saved
-	@python setup.py deps_table_update
+	@python utils/update_dependency_table.py
 	@md5sum -c --quiet md5sum.saved || (printf "\nError: the version dependency table is outdated.\nPlease run 'make fixup' or 'make style' and commit the changes.\n\n" && exit 1)
 	@rm md5sum.saved
 

README.md

@@ -17,6 +17,9 @@
 [![Unit Tests](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml/badge.svg)](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml)
 
 # What's new?
+- **`2026/03/31`**: Wan2.2 SenCache inference is now supported for T2V and I2V (up to 1.4x speedup)
+- **`2026/03/25`**: Wan2.1 and Wan2.2 Magcache inference is now supported
+- **`2026/03/25`**: LTX-2 Video Inference is now supported
 - **`2026/01/29`**: Wan LoRA for inference is now supported
 - **`2026/01/15`**: Wan2.1 and Wan2.2 Img2vid generation is now supported
 - **`2025/11/11`**: Wan2.2 txt2vid generation is now supported
@@ -49,6 +52,7 @@ MaxDiffusion supports
 * ControlNet inference (Stable Diffusion 1.4 & SDXL).
 * Dreambooth training support for Stable Diffusion 1.x,2.x.
 * LTX-Video text2vid, img2vid (inference).
+* LTX-2 Video text2vid (inference).
 * Wan2.1 text2vid (training and inference).
 * Wan2.2 text2vid (inference).
 
@@ -73,6 +77,7 @@ MaxDiffusion supports
   - [Inference](#inference)
     - [Wan](#wan-models)
     - [LTX-Video](#ltx-video)
+    - [LTX-2 Video](#ltx-2-video)
     - [Flux](#flux)
       - [Fused Attention for GPU](#fused-attention-for-gpu)
     - [SDXL](#stable-diffusion-xl)
@@ -497,6 +502,33 @@ To generate images, run the following command:
 
   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
+
+  Although not required, attaching an external disk is recommended as weights take up a lot of disk space. [Follow these instructions if you would like to attach an external disk](https://cloud.google.com/tpu/docs/attach-durable-block-storage).
+
+  The following command will run LTX-2 T2V:
+
+   ```bash
+  HF_HUB_CACHE=/mnt/disks/external_disk/maxdiffusion_hf_cache/ \
+  LIBTPU_INIT_ARGS="--xla_tpu_enable_async_collective_fusion=true \
+  --xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true \
+  --xla_tpu_enable_async_collective_fusion_multiple_steps=true \
+  --xla_tpu_overlap_compute_collective_tc=true \
+  --xla_enable_async_all_reduce=true" \
+  HF_HUB_ENABLE_HF_TRANSFER=1 \
+  python src/maxdiffusion/generate_ltx2.py \
+  src/maxdiffusion/configs/ltx2_video.yml \
+  attention="flash" \
+  num_inference_steps=40 \
+  num_frames=121 \
+  width=768 \
+  height=512 \
+  per_device_batch_size=.125 \
+  ici_data_parallelism=2 \
+  ici_context_parallelism=4 \
+  run_name=ltx2-inference
+  ```
+
   ## Wan Models
 
   Although not required, attaching an external disk is recommended as weights take up a lot of disk space. [Follow these instructions if you would like to attach an external disk](https://cloud.google.com/tpu/docs/attach-durable-block-storage).
@@ -540,6 +572,31 @@ To generate images, run the following command:
   * For Wan2.2 T2V, use `base_wan_27b.yml`.
   * For Wan2.2 I2V, use `base_wan_i2v_27b.yml`.
 
+  ### Caching Mechanisms
+
+  Wan 2.x pipelines support several caching strategies to accelerate inference by skipping redundant transformer forward passes. These are **mutually exclusive** — enable only one at a time.
+
+  | Cache Type | Config Flag | Supported Pipelines | Speedup | Description |
+  | --- | --- | --- | --- | --- |
+  | **CFG Cache** | `use_cfg_cache: True` | Wan 2.1 T2V, Wan 2.2 T2V/I2V | ~1.2x | FasterCache-style: caches the unconditional branch and applies FFT frequency-domain compensation on skipped steps. |
+  | **SenCache** | `use_sen_cache: True` | Wan 2.2 T2V/I2V | ~1.4x | Sensitivity-Aware Caching ([arXiv:2602.24208](https://arxiv.org/abs/2602.24208)): predicts output change via first-order sensitivity S = α_x·‖Δx‖ + α_t·\|Δt\|. Skips the full CFG forward pass when predicted change is below tolerance ε. |
+
+  To enable a caching mechanism, set the corresponding flag in your config YAML or pass it as a command-line override:
+
+  ```bash
+  # Example: enable SenCache for Wan 2.2 T2V
+  python src/maxdiffusion/generate_wan.py \
+    src/maxdiffusion/configs/base_wan_27b.yml \
+    use_sen_cache=True \
+    ...
+
+  # Example: enable CFG Cache for Wan 2.2 I2V
+  python src/maxdiffusion/generate_wan.py \
+    src/maxdiffusion/configs/base_wan_i2v_27b.yml \
+    use_cfg_cache=True \
+    ...
+  ```
+
   ## Flux
 
   First make sure you have permissions to access the Flux repos in Huggingface.

dependencies/requirements/base_requirements/requirements.txt

@@ -0,0 +1,43 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+absl-py
+aqtp
+chex
+datasets
+einops
+flax
+ftfy
+google-cloud-storage
+grain
+hf_transfer
+huggingface_hub
+imageio-ffmpeg
+imageio
+jax
+jaxlib
+Jinja2
+opencv-python-headless
+optax
+orbax-checkpoint
+parameterized
+Pillow
+pyink
+pylint
+pytest
+ruff
+scikit-image
+sentencepiece
+tensorboard-plugin-profile
+tensorboard
+tensorboardx
+tensorflow-datasets
+tensorflow
+tokamax
+tokenizers
+transformers<5.0.0
+
+# pinning torch and torchvision to specific versions to avoid
+# installing GPU versions from PyPI when running seed-env
+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
+