Skip to content

Commit f7b4145

Browse files
eltsaieltsai
committed
Merge origin/main: resolve VAE autoencoder and imports conflicts
2 parents 0fa8678 + 6de9d57 commit f7b4145

34 files changed

Lines changed: 1436 additions & 208 deletions

‎.github/workflows/UnitTests.yml‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,6 @@ jobs:
4242
python-version: '3.12'
4343
- name: Install dependencies
4444
run: |
45-
pip install -e .
46-
pip uninstall jax jaxlib libtpu-nightly libtpu -y
4745
bash setup.sh MODE=stable
4846
export PATH=$PATH:$HOME/.local/bin
4947
pip install ruff
@@ -59,8 +57,8 @@ jobs:
5957
- name: PyTest
6058
run: | #--deselect=src/maxdiffusion/tests/input_pipeline_interface_test.py
6159
export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
62-
HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ TOKENIZERS_PARALLELISM=false python3 -m pytest --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py --ignore=src/maxdiffusion/kernels/splash_attention -x
63-
# add_pull_ready:
60+
HF_HUB_CACHE=/mnt/disks/github-runner-disk/ HF_HOME=/mnt/disks/github-runner-disk/ TOKENIZERS_PARALLELISM=false python3 -m pytest --deselect=src/maxdiffusion/tests/ltx_transformer_step_test.py -x
61+
# add_pull_ready:q
6462
# if: github.ref != 'refs/heads/main'
6563
# permissions:
6664
# checks: read

‎.github/workflows/UploadDockerImages.yml‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,15 +28,12 @@ jobs:
2828
build-image:
2929
runs-on: ["self-hosted", "e2", "cpu"]
3030
steps:
31-
- uses: actions/checkout@v3
31+
- uses: actions/checkout@v5
3232
- name: Cleanup old docker images
3333
run: docker system prune --all --force
34-
- name: build maxdiffusion jax ai image
34+
- name: build maxdiffusion stable image
3535
run: |
36-
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
37-
- name: build maxdiffusion w/ nightly jax ai image
38-
run: |
39-
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
40-
- name: build maxdiffusion jax nightly image
36+
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
37+
- name: build maxdiffusion nightly image
4138
run: |
4239
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‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
# Copyright 2025 Google LLC
2+
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# This workflow will build, test and automatically release MaxDiffusion package to PyPI using Trusted Publishing (OIDC).
16+
17+
name: Publish MaxDiffusion to PyPI
18+
19+
# Triggers when a new "release" is published in the GitHub UI
20+
on:
21+
release:
22+
types: [published]
23+
workflow_dispatch:
24+
25+
permissions:
26+
contents: read
27+
id-token: write
28+
29+
jobs:
30+
build_and_publish:
31+
name: Build and Publish MaxDiffusion Package
32+
runs-on: ubuntu-latest
33+
steps:
34+
- uses: actions/checkout@v5
35+
- name: Set up Python
36+
uses: actions/setup-python@v5
37+
with:
38+
python-version: '3.12'
39+
- name: Install build dependencies
40+
run: |
41+
python -m pip install --upgrade pip
42+
pip install build hatchling hatch-requirements-txt
43+
- name: Build package
44+
run: python -m build
45+
- name: Publish package
46+
uses: pypa/gh-action-pypi-publish@release/v1
47+
with:
48+
packages-dir: dist/

‎Makefile‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,11 @@ modified_only_fixup:
1818
# Update src/maxdiffusion/dependency_versions_table.py
1919

2020
deps_table_update:
21-
@python setup.py deps_table_update
21+
@python utils/update_dependency_table.py
2222

2323
deps_table_check_updated:
2424
@md5sum src/maxdiffusion/dependency_versions_table.py > md5sum.saved
25-
@python setup.py deps_table_update
25+
@python utils/update_dependency_table.py
2626
@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)
2727
@rm md5sum.saved
2828

‎README.md‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
[![Unit Tests](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml/badge.svg)](https://github.com/AI-Hypercomputer/maxdiffusion/actions/workflows/UnitTests.yml)
1818

1919
# What's new?
20+
- **`2026/03/31`**: Wan2.2 SenCache inference is now supported for T2V and I2V (up to 1.4x speedup)
2021
- **`2026/03/25`**: Wan2.1 and Wan2.2 Magcache inference is now supported
2122
- **`2026/03/25`**: LTX-2 Video Inference is now supported
2223
- **`2026/01/29`**: Wan LoRA for inference is now supported
@@ -571,6 +572,34 @@ To generate images, run the following command:
571572
* For Wan2.2 T2V, use `base_wan_27b.yml`.
572573
* For Wan2.2 I2V, use `base_wan_i2v_27b.yml`.
573574

575+
<<<<<<< HEAD
576+
=======
577+
### Caching Mechanisms
578+
579+
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.
580+
581+
| Cache Type | Config Flag | Supported Pipelines | Speedup | Description |
582+
| --- | --- | --- | --- | --- |
583+
| **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. |
584+
| **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 ε. |
585+
586+
To enable a caching mechanism, set the corresponding flag in your config YAML or pass it as a command-line override:
587+
588+
```bash
589+
# Example: enable SenCache for Wan 2.2 T2V
590+
python src/maxdiffusion/generate_wan.py \
591+
src/maxdiffusion/configs/base_wan_27b.yml \
592+
use_sen_cache=True \
593+
...
594+
595+
# Example: enable CFG Cache for Wan 2.2 I2V
596+
python src/maxdiffusion/generate_wan.py \
597+
src/maxdiffusion/configs/base_wan_i2v_27b.yml \
598+
use_cfg_cache=True \
599+
...
600+
```
601+
602+
>>>>>>> origin/main
574603
## Flux
575604

576605
First make sure you have permissions to access the Flux repos in Huggingface.

‎dependencies/requirements/base_requirements/requirements.txt‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
--extra-index-url https://download.pytorch.org/whl/cpu
2+
absl-py
3+
aqtp
4+
chex
5+
datasets
6+
einops
7+
flax
8+
ftfy
9+
google-cloud-storage
10+
grain
11+
hf_transfer
12+
huggingface_hub
13+
imageio-ffmpeg
14+
imageio
15+
jax
16+
jaxlib
17+
Jinja2
18+
opencv-python-headless
19+
optax
20+
orbax-checkpoint
21+
parameterized
22+
Pillow
23+
pyink
24+
pylint
25+
pytest
26+
ruff
27+
scikit-image
28+
sentencepiece
29+
tensorboard-plugin-profile
30+
tensorboard
31+
tensorboardx
32+
tensorflow-datasets
33+
tensorflow
34+
tokamax
35+
tokenizers
36+
transformers<5.0.0
37+
38+
# pinning torch and torchvision to specific versions to avoid
39+
# installing GPU versions from PyPI when running seed-env
40+
torch @ https://download.pytorch.org/whl/cpu/torch-2.10.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl
41+
torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.25.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl
42+
qwix @ https://github.com/google/qwix/archive/408a0f48f988b6c5b180e07f0cb1d05997bf0dcc.zip
43+

‎dependencies/requirements/generated_requirements/requirements.txt‎

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,200 @@
1+
# Generated by seed-env. Do not edit manually.
2+
# If you need to modify dependencies, please do so in the host requirements file and run seed-env again.
3+
4+
absl-py>=2.3.1
5+
aiofiles>=25.1.0
6+
aiohappyeyeballs>=2.6.1
7+
aiohttp>=3.13.3
8+
aiosignal>=1.4.0
9+
annotated-types>=0.7.0
10+
anyio>=4.13.0
11+
aqtp>=0.9.0
12+
array-record>=0.8.3 ; sys_platform != 'win32'
13+
astroid>=4.0.4
14+
astunparse>=1.6.3
15+
attrs>=25.4.0
16+
auditwheel>=6.6.0
17+
black>=25.12.0
18+
build>=1.4.0
19+
certifi>=2026.1.4
20+
cffi>=2.0.0 ; platform_python_implementation != 'PyPy'
21+
charset-normalizer>=3.4.4
22+
cheroot>=11.1.2
23+
chex>=0.1.91
24+
click>=8.3.1
25+
cloudpickle>=3.1.2
26+
colorama>=0.4.6
27+
contourpy>=1.3.3
28+
cryptography>=46.0.6
29+
cycler>=0.12.1
30+
dataclasses-json>=0.6.7
31+
datasets>=4.8.4
32+
decorator>=5.2.1
33+
dill>=0.4.1
34+
dm-tree>=0.1.9
35+
docstring-parser>=0.17.0
36+
einops>=0.8.2
37+
einshape>=1.0
38+
etils>=1.13.0
39+
execnet>=2.1.2
40+
filelock>=3.20.3
41+
flatbuffers>=25.12.19
42+
flax>=0.12.6
43+
fonttools>=4.61.1
44+
frozenlist>=1.8.0
45+
fsspec>=2026.1.0
46+
ftfy>=6.3.1
47+
gast>=0.7.0
48+
gcsfs>=2026.1.0
49+
google-api-core>=2.30.0
50+
google-auth-oauthlib>=1.3.0
51+
google-auth>=2.49.1
52+
google-cloud-core>=2.5.0
53+
google-cloud-storage-control>=1.11.0
54+
google-cloud-storage>=3.10.1
55+
google-crc32c>=1.8.0
56+
google-pasta>=0.2.0
57+
google-resumable-media>=2.8.0
58+
googleapis-common-protos>=1.73.1
59+
grain>=0.2.16
60+
grpc-google-iam-v1>=0.14.3
61+
grpcio-status>=1.76.0
62+
grpcio>=1.76.0
63+
gviz-api>=1.10.0
64+
h11>=0.16.0
65+
h5py>=3.15.1
66+
hf-transfer>=0.1.9
67+
hf-xet>=1.4.2 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
68+
httpcore>=1.0.9
69+
httpx>=0.28.1
70+
huggingface-hub>=0.36.2
71+
humanize>=4.15.0
72+
hypothesis>=6.142.1
73+
idna>=3.11
74+
imageio-ffmpeg>=0.6.0
75+
imageio>=2.37.3
76+
immutabledict>=4.3.1
77+
importlib-resources>=6.5.2
78+
iniconfig>=2.3.0
79+
isort>=8.0.1
80+
jaraco-functools>=4.4.0
81+
jax>=0.9.0
82+
jaxlib>=0.9.0
83+
jaxtyping>=0.3.9
84+
jinja2>=3.1.6
85+
keras>=3.13.1
86+
kiwisolver>=1.4.9
87+
lazy-loader>=0.5
88+
libclang>=18.1.1
89+
libtpu>=0.0.34 ; platform_machine == 'x86_64' and sys_platform == 'linux'
90+
markdown-it-py>=4.0.0
91+
markdown>=3.10.1
92+
markupsafe>=3.0.3
93+
marshmallow>=3.26.2
94+
matplotlib>=3.10.8
95+
mccabe>=0.7.0
96+
mdurl>=0.1.2
97+
ml-dtypes>=0.5.4
98+
more-itertools>=10.8.0
99+
mpmath>=1.3.0
100+
msgpack>=1.1.2
101+
multidict>=6.7.1
102+
multiprocess>=0.70.19
103+
mypy-extensions>=1.1.0
104+
namex>=0.1.0
105+
networkx>=3.6.1
106+
numpy-typing-compat>=20251206.2.0
107+
numpy>=2.0.2
108+
nvidia-cuda-cccl>=13.1.115
109+
oauthlib>=3.3.1
110+
opencv-python-headless>=4.13.0.92
111+
opt-einsum>=3.4.0
112+
optax>=0.2.8
113+
optree>=0.18.0
114+
optype>=0.15.0
115+
orbax-checkpoint>=0.11.33
116+
orbax-export>=0.0.8
117+
packaging>=26.0
118+
pandas>=3.0.1
119+
parameterized>=0.9.0
120+
pathspec>=1.0.4
121+
pillow>=12.1.0
122+
platformdirs>=4.9.4
123+
pluggy>=1.6.0
124+
portpicker>=1.6.0
125+
promise>=2.3
126+
propcache>=0.4.1
127+
proto-plus>=1.27.2
128+
protobuf>=6.33.6
129+
psutil>=7.2.1
130+
pyarrow>=23.0.1
131+
pyasn1-modules>=0.4.2
132+
pyasn1>=0.6.3
133+
pycparser>=3.0 ; implementation_name != 'PyPy' and platform_python_implementation != 'PyPy'
134+
pydantic-core>=2.41.5
135+
pydantic>=2.12.5
136+
pyelftools>=0.32
137+
pygments>=2.19.2
138+
pyink>=25.12.0
139+
pylint>=4.0.5
140+
pyparsing>=3.3.2
141+
pyproject-hooks>=1.2.0
142+
pytest-xdist>=3.8.0
143+
pytest>=8.4.2
144+
python-dateutil>=2.9.0.post0
145+
pytokens>=0.4.1
146+
pyyaml>=6.0.3
147+
qwix @ https://github.com/google/qwix/archive/408a0f48f988b6c5b180e07f0cb1d05997bf0dcc.zip
148+
regex>=2026.2.28
149+
requests-oauthlib>=2.0.0
150+
requests>=2.32.5
151+
rich>=14.2.0
152+
ruff>=0.15.8
153+
safetensors>=0.7.0
154+
scikit-image>=0.26.0
155+
scipy-stubs>=1.17.0.1
156+
scipy>=1.17.0
157+
sentencepiece>=0.2.1
158+
setuptools>=80.10.1
159+
simple-parsing>=0.1.8
160+
simplejson>=3.20.2
161+
six>=1.17.0
162+
sortedcontainers>=2.4.0
163+
sympy>=1.14.0
164+
tensorboard-data-server>=0.7.2
165+
tensorboard-plugin-profile>=2.22.0
166+
tensorboard>=2.20.0
167+
tensorboardx>=2.6.4
168+
tensorflow-datasets>=4.9.9
169+
tensorflow-metadata>=1.17.3
170+
tensorflow>=2.20.0
171+
tensorstore>=0.1.80
172+
termcolor>=3.3.0
173+
tifffile>=2026.3.3
174+
tokamax>=0.0.10
175+
tokenizers>=0.22.2
176+
toml>=0.10.2
177+
tomlkit>=0.14.0
178+
toolz>=1.1.0
179+
torch @ https://download.pytorch.org/whl/cpu/torch-2.10.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl
180+
torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.25.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl
181+
tqdm>=4.67.3
182+
transformers>=4.57.6
183+
treescope>=0.1.10
184+
typeguard>=2.13.3
185+
typing-extensions>=4.15.0
186+
typing-inspect>=0.9.0
187+
typing-inspection>=0.4.2
188+
tzdata>=2025.3 ; sys_platform == 'emscripten' or sys_platform == 'win32'
189+
urllib3>=2.6.3
190+
uvloop>=0.22.1
191+
wadler-lindig>=0.1.7
192+
wcwidth>=0.6.0
193+
werkzeug>=3.1.5
194+
wheel>=0.46.2
195+
wrapt>=2.1.2
196+
xprof>=2.22.0
197+
xxhash>=3.6.0
198+
yarl>=1.23.0
199+
zipp>=3.23.0
200+
zstandard>=0.25.0

0 commit comments

Comments
 (0)