Move the MaxDiffusion backend to transformers 5

[olehtika]

Motivation

requirements-maxdiffusion.txt pins transformers==4.57.3, which carries CVE-2026-4372, CVE-2026-5241 and CVE-2026-9856. 5.10.0 is the first release fixing all three.

Because the pin is exact, it is not just that MaxDiffusion runs install a vulnerable transformers — a downstream image that upgrades transformers gets it resolved straight back down as soon as this file is installed. That is not hypothetical: it happened while fixing these CVEs in ROCm/MAD's primus_maxdiffusion image, which built an image that logged Successfully installed transformers-5.17.0 and then shipped 4.57.3, because this requirements file is installed afterwards.

The manifest is not the only place. .github/workflows/docker-release/Dockerfile.jax-v26.6 installs transformers==4.57.3 on top of the 5.x that jax_training ships, with a comment accepting all three CVEs as risk. So relaxing the manifest alone would have left the image this repo publishes still carrying them.

Technical Details

The comment above the pin explains why it was mandatory:

MUST stay on transformers 4.x: MaxDiffusion is a Flax codebase and imports FlaxCLIPTextModel/FlaxT5EncoderModel at module load

That is no longer true. MaxDiffusion now runs FLUX's CLIP-L and T5-XXL as PyTorch models under Torchax — the same mechanism it already uses for LTX2's Gemma3 and WAN's UMT5 encoders — and guards its remaining Flax imports, so import maxdiffusion no longer depends on classes that v5 deleted. Everything below follows from that one constraint being lifted.

The manifest and the config (first commit):

1. requirements-maxdiffusion.txt: transformers==4.57.3 becomes transformers>=5.10.0.
2. primus/configs/models/maxdiffusion/flux_dev.yaml: clip_model_name_or_path and t5xxl_model_name_or_path move off ariG23498/clip-vit-large-patch14-text-flax and ariG23498/t5-v1-1-xxl-flax onto the text_encoder and text_encoder_2 subfolders of black-forest-labs/FLUX.1-dev, with the matching tokenizer subfolders. Those mirrors hold Flax weights only, so the PyTorch encoders cannot read them at all. The replacement is the first-party repo the rest of this config already points at, and the same weights MaxDiffusion's own base_flux_dev.yml names.

The paths that kept 4.57.3 in reach (second commit):

3. Dockerfile.jax-v26.6: the transformers==4.57.3 install becomes a >=5.10.0 floor, and the accepted-risk comment goes with it. The sed immediately after it rewrote jnp.clip(..., a_max=) kwargs inside transformers/models/t5/modeling_flax_t5.py — a file v5 does not ship — so as an unconditional RUN it would have failed the build against a path that no longer exists; dropped. A closing assertion fails the build if any later layer resolves transformers back below the floor, rather than shipping an image that looks patched.
4. .github/dependabot.yml: the pip entry existed only to carry ignore: transformers >=5.0.0, and ignore.versions applies to security updates as well as version updates, so an advisory fixed only in 5.x could not raise a PR here at all. With the 4.x ceiling gone that entry has no purpose, so it is removed: pip version updates stay off (they need an explicit entry) and repo-wide security updates now reach this manifest unfiltered.
5. torchax==0.0.13 added to requirements-maxdiffusion.txt and to the image. It was in no manifest here — not the requirements, not setup_maxdiffusion_env.sh, not the image — because the two paths MaxDiffusion already uses it for, LTX2's Gemma3 and WAN's UMT5 encoders, are not exercised by the configs in this repo. FLUX's encoders now need it. Pinned because it dispatches against private torch overloads that torch >= 2.14 dropped.
6. The install hook comment stating that this file "pins transformers 4.x for MaxDiffusion's Flax code" is corrected, so the rationale for restoring the pin is not left lying around.

This needs a submodule bump to merge

Both parts depend on a third_party/maxdiffusion pin that carries the Torchax text encoders. main pins 68e069659f0af80694559e29939a17f879fe7f6a, which predates them, so merging these files alone would break FLUX rather than fix it — transformers 5 against that tree fails at import, and the repointed config hands PyTorch-only weights to Flax loaders. MAXDIFFUSION_COMMIT in the image Dockerfile has the same requirement and should move with the submodule, so an image build and a primus-cli launch agree on the tree.

That bump also has to clear AI-Hypercomputer/maxdiffusion#482. It fixes three independent FLUX regressions on MaxDiffusion main, introduced by 64450df4 (2026-07-16, the Flux.2-klein onboarding): timestep/guidance embeddings projected twice and failing at init on a broadcast error, a legacy BFL checkpoint loader still mapping to a renamed norm_out submodule, and norm_out silently inheriting the opposite shift/scale convention, which inflates the loss about 8x without failing. None of it affects the currently pinned 68e0696 (2026-03-03), which is an ancestor of that commit — but a pin moved forward to pick up the Torchax encoders without #482 in it would build a FLUX that trains wrong rather than not at all.

The MaxDiffusion side is upstream as AI-Hypercomputer/maxdiffusion#483. This PR stays draft until that lands and the submodule can be bumped to a commit containing it, at which point the bump belongs in this same change. Filing now so the CVE exposure and the ordering are visible; happy to carry the bump here or to hand these files to whoever owns the pin.

On three review suggestions, declined

• Bump the sibling HF packages (diffusers, tokenizers, accelerate, huggingface_hub, safetensors) alongside transformers. Nothing else is pinned in this file — transformers was the only pin in it, so there are no stale sibling pins to move, and adding some would newly constrain resolves that work today. In ROCm/MAD's build of this stack, pip resolved huggingface_hub 1.31.0 against transformers 5.17.0 on its own.
• Update setup_maxdiffusion_env.sh's jax/flax/orbax pins, since dropping Flax is the premise. The premise does not hold in either half. That script pins no jax, flax or orbax versions — they come from the base image — and MaxDiffusion remains a Flax codebase: only the two text encoders moved to PyTorch, while the transformer and VAE are still Flax. Its Flax T5 sed is already guarded by [ -f "$T5" ] && grep -q "a_max=", so it no-ops on v5 rather than failing, which is why it needed no change while the Dockerfile's unconditional one did.
• Fold the submodule bump in now. Agreed in principle, which is why this stays draft — the commit to bump to does not exist upstream yet.

Test Plan

FLUX.1-dev, WAN 2.1 1.3B and WAN 2.1 14B pretrain, MI355X configs, 8x MI355X (gfx950), 20 steps each at bf16, driven through ROCm/MAD's primus_maxdiffusion image on the rocm/jax-training:maxtext-v26.7 base, with Primus at dc3f4ba plus the manifest and config changes and a MaxDiffusion tree carrying the Torchax encoders.

Test Result

All three succeed, rc=0, with step times and losses matching runs on the pre-upgrade stack:

Config	Status	s/step	TFLOP/s/GPU	Loss @ step 18
flux_dev-pretrain	SUCCESS	1.548	583.3	572
wan2.1_1.3b-pretrain	SUCCESS	4.437	920.9	1.540
wan2.1_14b-pretrain	SUCCESS	26.868	789.3	1.551

Losses are identical to the pre-upgrade baseline, so the PyTorch encoders reproduce the Flax numbers rather than merely loading. FLUX was confirmed to read its encoders from the FLUX.1-dev text_encoder/text_encoder_2 subfolders, so the new path is the one exercised.

torchax==0.0.13 was checked separately against the torch this repo installs, since those runs went through ROCm/MAD's image, whose base carries torch 2.12.0 rather than the 2.8.0 named here and in setup_maxdiffusion_env.sh:

$ docker run --rm python:3.12-slim bash -c 'pip install -q --index-url https://download.pytorch.org/whl/cpu torch==2.8.0 && pip install -q torchax==0.0.13 "jax[cpu]" && python -c "import torch, torchax, torchax.interop; torchax.default_env(); print(torch.__version__, \"torchax OK\")"'
2.8.0+cpu torchax OK

Two things left for reviewers. Dockerfile.jax-v26.6 has not been built as changed — the pieces it changes are covered above and by that import check, but the build itself is unexercised, which is what the closing assertions are there to catch. And the image's MAXDIFFUSION_COMMIT still names the pre-Torchax commit, for the same reason the submodule does; both move together when the upstream change lands.

[pr-pundit-code-review]

Code Review: AMD-AGI/Primus

Summary

1 architecture concern, 2 files reviewed.
Rules checked: 2 (2 passed, 0 failed).

---

Architecture Concerns

These findings come from architecture principles distilled from the repo's PR lineage — patterns that caused PRs to be rejected or significantly reworked.

batch-related-dependency-bumps-atomically

When multiple dependencies belong to the same ecosystem/manifest and are managed by an automated dependency-update tool that groups them (e.g., a dependency group across a single directory), they should be bumped together in one coordinated PR rather than as a scattered series of one-off single-package PRs. Piecemeal single-package bumps each trigger a full CI/review cycle, leave the manifest in inconsistent intermediate states where interdependent packages are mismatched, and multiply merge-conflict and re-validation surface across the same set of files.

Findings:

• This PR relaxes transformers to >=5.10.0 while its own comment states 'Bump the submodule and this line together' — but third_party/maxdiffusion is not touched in this changeset. Add the third_party/maxdiffusion submodule pointer bump (to the commit that carries the Torchax/PyTorch text-encoder change and guarded Flax imports) to this same PR so the requirement and the code that satisfies it land atomically.
• Bump the rest of the HuggingFace stack that moves in lockstep with transformers 5.x (diffusers, tokenizers, accelerate, huggingface_hub, safetensors) in this same requirements-maxdiffusion.txt change rather than letting them resolve to whatever 4.x-era pins currently apply; a lone transformers major-version jump leaves sibling HF packages stale and can produce an unresolvable or silently broken intermediate state.
• The header of requirements-maxdiffusion.txt points at examples/maxdiffusion/setup_maxdiffusion_env.sh as the place where the jax/flax/--find-links deps are installed. Since dropping the Flax code path is the premise of this bump, update that script's jax/flax/orbax pins in the same PR so the two manifests are not left describing different transformers/Flax expectations.
• Consolidate the flux_dev.yaml text-encoder repo/subfolder retarget, the transformers pin change, and the maxdiffusion submodule bump into one coordinated group PR — each of these three is a no-op or a breakage on its own (PyTorch encoders need the new submodule; the new submodule needs transformers 5.x; transformers 5.x needs the non-Flax encoder paths), so splitting them forces multiple CI+review cycles through known-broken intermediate commits.

---

Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.

[Copilot]

🟡 Changes recommended

Critical compatibility and installation-path issues remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Updates MaxDiffusion for Transformers 5 and PyTorch-based FLUX text encoders to address vulnerable dependency versions.

Changes:

• Raises the Transformers requirement to >=5.10.0.
• Points FLUX encoders and tokenizers to FLUX.1-dev subfolders.
• Requires coordinated updates to the MaxDiffusion submodule and image installation paths.

File summaries

File	Summary	Findings
requirements-maxdiffusion.txt	Updates the Transformers floor and compatibility rationale.	Critical (2 votes): The baked image still installs Transformers 4.57.3 and applies an unconditional Flax-T5 patch. Moderate (1 vote): Dependabot ignores all Transformers 5 updates. Nit (1 vote): The installation hook comment still describes the old Transformers 4 requirement.
primus/configs/models/maxdiffusion/flux_dev.yaml	Configures PyTorch FLUX encoder and tokenizer subfolders.	Critical (2 votes): This requires the pending Torchax loader and matching MaxDiffusion submodule revision; the currently pinned revision will ignore these settings and attempt incompatible Flax loading.

Review details

Suppressed comments (2)

requirements-maxdiffusion.txt:20

• The security automation still ignores every Transformers version >=5.0.0 in .github/dependabot.yml:37-45. After this floor is raised to 5.10.0, that rule suppresses updates in the supported major, including future security fixes; remove or narrow the ignore and update its Flax-based rationale.

transformers>=5.10.0

requirements-maxdiffusion.txt:20

• The normal primus-cli install hook still says this file pins Transformers 4.x for MaxDiffusion's Flax code (runner/helpers/hooks/train/pretrain/maxdiffusion/00_install_requirements.sh:19-21). That is now false and conflicts with this migration, so update the hook comment and related installation documentation with the requirement change to avoid restoring the vulnerable pin later.

# maxdiffusion` outright. MaxDiffusion now runs FLUX's text encoders as PyTorch
# models under Torchax and guards its remaining Flax imports, which lifts that
# constraint -- but only from a third_party/maxdiffusion pin that carries those
# changes. Bump the submodule and this line together.
transformers>=5.10.0

• Files reviewed: 2/2 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[pr-pundit-code-review]

Code Review: AMD-AGI/Primus

Summary

1 architecture concern, 5 files reviewed.
Rules checked: 2 (2 passed, 0 failed).

---

Architecture Concerns

These findings come from architecture principles distilled from the repo's PR lineage — patterns that caused PRs to be rejected or significantly reworked.

batch-related-dependency-bumps-atomically

When multiple dependencies belong to the same ecosystem/manifest and are managed by an automated dependency-update tool that groups them (e.g., a dependency group across a single directory), they should be bumped together in one coordinated PR rather than as a scattered series of one-off single-package PRs. Piecemeal single-package bumps each trigger a full CI/review cycle, leave the manifest in inconsistent intermediate states where interdependent packages are mismatched, and multiply merge-conflict and re-validation surface across the same set of files.

Findings:

• This PR raises transformers to >=5.10.0 in both requirements-maxdiffusion.txt and Dockerfile.jax-v26.6 while leaving ARG MAXDIFFUSION_COMMIT=68e069659f0af80694559e29939a17f879fe7f6a untouched and third_party/maxdiffusion out of the changed-file set — even though the comments state the floor only holds 'from a third_party/maxdiffusion pin that carries those changes' and 'Bump the submodule and this line together'. Include the third_party/maxdiffusion submodule bump and the matching MAXDIFFUSION_COMMIT update in this same PR so the transformers floor, the torchax add, and the MaxDiffusion tree land atomically instead of as two CI/review cycles with an intermediate broken state.
• MAXDIFFUSION_COMMIT in .github/workflows/docker-release/Dockerfile.jax-v26.6 and the third_party/maxdiffusion submodule SHA are cross-referenced versions of the same dependency (the Dockerfile comment explicitly asks that 'an image build and a primus-cli launch agree on the tree'). Update both pins in one commit rather than letting the image pin drift from the vendored checkout.
• torchax==0.0.13 is added with a stated hard coupling to torch ('imports cleanly on torch 2.12.0 and fails on torch >= 2.14'), but the sibling torch/torchvision pins that constrain it are not touched in this PR. Add the corresponding torch/torchvision constraint (e.g. torch<2.14) alongside the torchax pin in requirements-maxdiffusion.txt and the Dockerfile so the torchax/torch pair is bumped and bounded as one group instead of relying on a build-time import assertion to catch a future drift.
• The transformers floor is now asserted in three places (requirements-maxdiffusion.txt, the Dockerfile pip install, and the Dockerfile import assertion) with the version literal 5.10.0 duplicated in each. Hoist the floor to a single source (a shared constraints file or a Docker ARG consumed by both the install and the assertion) so a future security bump is one coordinated edit rather than three that can silently diverge.
• requirements-maxdiffusion.txt and requirements-jax.txt are described as deliberately separate but both feed transformers into the same image lineage; verify requirements-jax.txt (and any other root manifest) carries no conflicting transformers pin and, if it does, bump it in this same PR so a MaxText-then-MaxDiffusion install order cannot re-resolve transformers back below 5.10.0.

---

Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.

[Copilot]

🟡 Changes recommended

The stale MaxDiffusion commit/submodule pin is a critical blocker, and Torchax compatibility with torch 2.14+ is not enforced.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (2)

.github/workflows/docker-release/Dockerfile.jax-v26.6:592

• The build guard checks only the Transformers version and torchax; it never imports MaxDiffusion. Thus the current stale checkout can pass this guard and ship an image whose MaxDiffusion import fails as soon as training starts. Add import maxdiffusion to this final assertion (or exercise the actual encoder imports) so the image build catches an incompatible checkout.

RUN python3 -c "import transformers; from packaging.version import Version; \
v = Version(transformers.__version__); \
assert v >= Version('5.10.0'), f'transformers {v} still carries CVE-2026-4372, CVE-2026-5241 and CVE-2026-9856'; \
print('transformers', v, 'OK')" \
    && python3 -c "import torchax; print('torchax OK')"

requirements-maxdiffusion.txt:28

• torchax==0.0.13 is documented as incompatible with torch >=2.14, but this manifest imposes no torch constraint and setup_maxdiffusion_env.sh skips its ROCm install whenever any importable torch is already present (lines 66-73). A host/image with torch 2.14+ therefore passes dependency installation and only fails when torchax is imported; add a version check before accepting the existing torch or otherwise enforce a compatible torch/torchax pair.

torchax==0.0.13

• Files reviewed: 5/5 changed files
• Comments generated: 1
• Review effort level: Lite

[pr-pundit-code-review]

Code Review: AMD-AGI/Primus

Summary

1 architecture concern, 5 files reviewed.
Rules checked: 2 (2 passed, 0 failed).

---

Architecture Concerns

These findings come from architecture principles distilled from the repo's PR lineage — patterns that caused PRs to be rejected or significantly reworked.

batch-related-dependency-bumps-atomically

When multiple dependencies belong to the same ecosystem/manifest and are managed by an automated dependency-update tool that groups them (e.g., a dependency group across a single directory), they should be bumped together in one coordinated PR rather than as a scattered series of one-off single-package PRs. Piecemeal single-package bumps each trigger a full CI/review cycle, leave the manifest in inconsistent intermediate states where interdependent packages are mismatched, and multiply merge-conflict and re-validation surface across the same set of files.

Findings:

• This PR is not a scattered single-package bump (transformers>=5.10.0 and the new torchax==0.0.13 pin are moved together in requirements-maxdiffusion.txt and Dockerfile.jax-v26.6), but it leaves the third member of the same coordinated group out: the vendored third_party/maxdiffusion checkout. Both the requirements comment ("only from a third_party/maxdiffusion pin that carries those changes. Bump the submodule and this line together.") and the Dockerfile comment ("it should track third_party/maxdiffusion: move the two together") state the rule, yet no submodule bump appears in the changed-file list. Add the third_party/maxdiffusion submodule pointer update to this same PR so the transformers floor, torchax pin, and MaxDiffusion tree land atomically instead of in two CI/review cycles with an intermediate state where import maxdiffusion resolves Flax classes that transformers 5.10 no longer ships.
• MAXDIFFUSION_COMMIT is still 68e069659f0af80694559e29939a17f879fe7f6a in Dockerfile.jax-v26.6 while the surrounding comment now asserts "This commit has to carry the Torchax text encoders". Bump MAXDIFFUSION_COMMIT to the Torchax-carrying commit in this PR, matching the third_party/maxdiffusion submodule bump, so the image build and a primus-cli launch agree on the same tree rather than requiring a follow-up PR.
• requirements-maxdiffusion.txt's new torchax comment ties the 0.0.13 pin to "torch 2.8.0, the version setup_maxdiffusion_env.sh installs", but examples/maxdiffusion/setup_maxdiffusion_env.sh is not in the changed files. Include that script (or an explicit torch<2.14 constraint alongside the torchax pin in the same manifest) in this PR so the torch/torchax pairing is pinned as one group instead of relying on a comment that a later isolated torch bump can silently invalidate.

---

Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.

[Copilot]

🔵 Needs a closer look

Update the MaxDiffusion commit, strengthen Docker compatibility checks, and validate the Torchax/Torch pairing.

Review details

Suppressed comments (4)

.github/workflows/docker-release/Dockerfile.jax-v26.6:538

• The image still clones the old 68e0696 tree at the ARG below, so the new transformers floor makes this Dockerfile fail at MaxDiffusion import/install even if third_party/maxdiffusion is bumped separately. Update MAXDIFFUSION_COMMIT together with the vendored submodule to a commit containing the Torchax encoder changes before merging.

# Clone and install MaxDiffusion. This commit has to carry the Torchax text
# encoders, since the transformers floor above rules out the Flax classes the
# older trees import, and it should track third_party/maxdiffusion: move the two
# together so an image build and a primus-cli launch agree on the tree.

.github/workflows/docker-release/Dockerfile.jax-v26.6:593

• import torchax alone does not exercise the interop registration or the default_env() path that the FLUX encoders use, so this assertion can pass while the pinned torchax/torch pair is unusable. Match the validation used in the test plan by importing torchax.interop and calling torchax.default_env() here.

    && python3 -c "import torchax; print('torchax OK')"

.github/workflows/docker-release/Dockerfile.jax-v26.6:593

• This build guard never imports MaxDiffusion; it only imports the two newly installed dependencies. An incompatible or stale checkout can therefore pass the image build and fail only when the launcher imports maxdiffusion (which is exactly what the current commit does with Transformers 5). Add import maxdiffusion to this assertion so the build checks the actual compatibility boundary.

RUN python3 -c "import transformers; from packaging.version import Version; \
v = Version(transformers.__version__); \
assert v >= Version('5.10.0'), f'transformers {v} still carries CVE-2026-4372, CVE-2026-5241 and CVE-2026-9856'; \
print('transformers', v, 'OK')" \
    && python3 -c "import torchax; print('torchax OK')"

requirements-maxdiffusion.txt:29

• The new pin documents that torchax==0.0.13 fails with torch >=2.14, but the bare-checkout installer accepts any already-importable torch and skips its 2.8 installation (examples/maxdiffusion/setup_maxdiffusion_env.sh:66-73). On such an environment, this requirements install succeeds and the FLUX encoder fails later at runtime. Add a compatible-torch version check or installation path before treating an existing torch as valid.

torchax==0.0.13

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Lite