Run FLUX text encoders under Torchax instead of Flax

[olehtika]

Motivation

FLUX's text encoders are the last thing in this repo that needs transformers 4.x. They import FlaxCLIPTextModel and FlaxT5EncoderModel, and transformers 5.0 removed every Flax implementation, so import maxdiffusion dies at module load on any 5.x.

That matters beyond keeping current: transformers 4.57.3 and earlier carry CVE-2026-4372, CVE-2026-5241 and CVE-2026-9856, and 5.10.0 is the first release that fixes all three. There is no 4.x backport, so as long as FLUX needs the Flax classes, every downstream image built on this repo is pinned to a version with three open CVEs and no way out. Downstreams are already hitting this — ROCm/MAD's primus_maxdiffusion image and AMD-AGI/Primus both pin 4.57.3 for exactly this reason.

Technical Details

Rather than look for Flax replacements that no longer exist, this runs FLUX's PyTorch CLIP-L and T5-XXL under JAX through Torchax, which is already how this repo runs LTX2's Gemma3 and WAN's UMT5 encoders. FLUX's text encoders are frozen and only run during embedding precompute — the trainer deletes them before the training loop — so tracing them costs nothing on the training path.

• models/flux/text_encoders/torchax_text_encoders.py: TorchaxCLIPTextEncoder and TorchaxT5TextEncoder, wrapping the PyTorch encoders as JittableModules. Parameters are converted once at load and placed across devices with make_array_from_callback, so a multi-device mesh does not hold a full copy per host.
• The callers follow: flux_checkpointer, flux_pipeline, flux_trainer, generate_flux and generate_flux_multi_res.
• The FLUX configs move clip_model_name_or_path and t5xxl_model_name_or_path off the third-party Flax mirrors (ariG23498/clip-vit-large-patch14-text-flax, ariG23498/t5-v1-1-xxl-flax) onto the text_encoder and text_encoder_2 subfolders of black-forest-labs/FLUX.1-dev. 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 these configs already name.
• The new *_subfolder keys are read through a small helper that tolerates their absence, since pyconfig raises ValueError rather than AttributeError for a key it does not hold. A config written before these keys existed falls back to the repository root instead of crashing.

The result runs on both 4.57.3 and 5.x: dtype= replaces the torch_dtype= kwarg that 4.56 deprecated and v5 removed, and attn_implementation="eager" is accepted by both. Nobody still on the old pin is stranded.

Relationship to #482

On current main, FLUX training is broken independently of this change, by three regressions from 64450df4 (the Flux.2-klein onboarding) that #482 fixes. They touch different files from this PR — embeddings_flax.py, flux/util.py, transformer_flux_flax.py — so the two do not overlap, but FLUX cannot be exercised end to end on main without both. The GPU validation below therefore carries this change and #482 together.

Test Plan

1. A new CPU-only parity test, tests/flux_text_encoder_parity_test.py, covering encoder output and parameter placement/offload with tiny randomly-initialized models (vocab 99, hidden 32, 2 layers), so it needs no hub access and runs in seconds. tests/text_encoders_test.py is updated for the new classes.
2. FLUX-dev on 8x MI355X (gfx950), synthetic data, 512 resolution, batch 14/device, 20 steps, train_new_flux=False, comparing main plus this change plus Fix three FLUX regressions from the Flux.2-Klein onboarding #482 against 68e0696 — six months behind main, and the last commit where FLUX training was known-good here.
3. The same change rebased onto 68e0696 and run through ROCm/MAD's primus_maxdiffusion image on 8x MI355X, to check the FLUX and WAN configs downstreams actually run, against transformers 5.17.0.

Test Result

(2) reproduces the known-good regime:

	step time	loss, steps 0-5
68e0696 baseline	1.555-1.571 s	18432, 10816, 11712, 17536, 12160, 11072
main + this PR + #482	1.563-1.855 s	20864, 12288, 12544, 16064, 9856, 8512

The losses are not expected to match exactly, with 200+ commits between the two trees, but they sit in the same regime.

(3) succeeds on all three configs the downstream image ships, rc=0, with step times and losses matching its pre-upgrade baseline:

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

Losses there are identical to the pre-upgrade baseline, so the PyTorch encoders reproduce the Flax numbers rather than merely loading. That image carries transformers 5.17.0, huggingface_hub 1.31.0 and torchax 0.0.13, and FLUX was confirmed to read its encoders from the FLUX.1-dev text_encoder/text_encoder_2 subfolders, so the Torchax path is the one exercised rather than a silent fallback.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request replaces the Flax-based CLIP and T5 text encoders for FLUX with PyTorch implementations run under JAX using Torchax, adapting to the removal of Flax implementations in transformers v5. It introduces Torchax-based text encoder wrappers, updates configurations, pipelines, and trainers to support loading, sharding, and offloading these encoders, and adds parity tests. Feedback on the changes suggests updating the place_params method in _TorchaxTextEncoder to also place self.buffers on the target device/sharding to prevent potential device mismatch errors or silent host-to-device transfers.

[gemini-code-assist]

The place_params method only places self.params on the target device/sharding, but leaves self.buffers untouched. For models like CLIPTextModel that contain persistent buffers (e.g., position_ids or attention masks), these buffers will remain on the host CPU. This can lead to device mismatch errors or silent host-to-device transfers during functional calls on accelerators (TPU/GPU).\n\nPlease update place_params to also place self.buffers on the target device/sharding.

  def place_params(self, device_or_sharding) -> None:\n    \"\"\"Move the frozen weights onto device_or_sharding.\n\n    Torchax holds the weights as torch views over JAX arrays, so the placement\n    goes through jax_view/torch_view. A Sharding target is built with\n    make_array_from_callback (as max_utils.device_put_replicated does),\n    because each host only holds its own full copy of these weights and\n    device_put cannot span non-addressable devices.\n    \"\"\"\n\n    def place(leaf):\n      if isinstance(device_or_sharding, jax.sharding.Sharding):\n        return jax.make_array_from_callback(leaf.shape, device_or_sharding, lambda index: leaf[index])\n      return jax.device_put(leaf, device_or_sharding)\n\n    self.params = interop.torch_view(jax.tree_util.tree_map(place, interop.jax_view(self.params)))\n    self.buffers = interop.torch_view(jax.tree_util.tree_map(place, interop.jax_view(self.buffers)))