From bfce7b6521133365e9ca5d237adf086477932033 Mon Sep 17 00:00:00 2001 From: olehtika Date: Tue, 15 Sep 2026 11:55:01 +0000 Subject: [PATCH 1/2] Run FLUX text encoders under Torchax instead of Flax transformers v5 removed every Flax implementation, so FlaxCLIPTextModel and FlaxT5EncoderModel are gone and FLUX cannot import its text encoders. Staying below v5 is not an option for downstreams that run FLUX: 4.57.3 and earlier carry CVE-2026-4372, CVE-2026-5241 and CVE-2026-9856. Wrap the PyTorch CLIP-L and T5-XXL encoders with Torchax instead, which maxdiffusion already uses for LTX2's Gemma3 encoder. The wrappers subclass interop.JittableModule and trace a static forward through functional_call, returning plain tensors so callers no longer index into model output dicts or thread a separate params tree. FLUX's encoders are frozen and only run during embedding precompute, and the trainer deletes them before the training loop, so tracing them costs nothing on the training path. place_params uses make_array_from_callback rather than device_put so that parameters can be replicated across non-addressable devices on multi-host runs, matching max_utils.device_put_replicated. The configs now point at the text_encoder and text_encoder_2 subfolders of black-forest-labs/FLUX.1-dev rather than third-party Flax mirrors, so the weights come from the official repository. from_pretrained calls pass dtype= rather than torch_dtype=, which 4.56 deprecated and v5 still warns about, and attn_implementation="eager" to keep the models traceable. Both work on 4.57.3 and 5.x, so this runs on either. Add a CPU-only parity test covering both wrappers against eager PyTorch, plus param placement and offload, using tiny randomly-initialized models so it needs no hub access. --- .../checkpointing/flux_checkpointer.py | 25 +-- src/maxdiffusion/configs/base_flux_dev.yml | 10 +- .../configs/base_flux_dev_multi_res.yml | 10 +- .../configs/base_flux_schnell.yml | 10 +- src/maxdiffusion/generate_flux.py | 51 +++--- src/maxdiffusion/generate_flux_multi_res.py | 53 +++--- .../models/flux/text_encoders/__init__.py | 15 ++ .../text_encoders/torchax_text_encoders.py | 168 ++++++++++++++++++ .../pipelines/flux/flux_pipeline.py | 33 ++-- .../tests/flux_text_encoder_parity_test.py | 117 ++++++++++++ src/maxdiffusion/tests/text_encoders_test.py | 23 ++- src/maxdiffusion/trainers/flux_trainer.py | 7 +- 12 files changed, 405 insertions(+), 117 deletions(-) create mode 100644 src/maxdiffusion/models/flux/text_encoders/__init__.py create mode 100644 src/maxdiffusion/models/flux/text_encoders/torchax_text_encoders.py create mode 100644 src/maxdiffusion/tests/flux_text_encoder_parity_test.py diff --git a/src/maxdiffusion/checkpointing/flux_checkpointer.py b/src/maxdiffusion/checkpointing/flux_checkpointer.py index 70b54d08d..d09101d27 100644 --- a/src/maxdiffusion/checkpointing/flux_checkpointer.py +++ b/src/maxdiffusion/checkpointing/flux_checkpointer.py @@ -30,7 +30,10 @@ from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel from ..pipelines.flux.flux_pipeline import FluxPipeline -from transformers import (CLIPTokenizer, FlaxCLIPTextModel, FlaxT5EncoderModel, AutoTokenizer) +from maxdiffusion.models.flux.text_encoders.torchax_text_encoders import ( + load_clip_encoder_and_tokenizer, + load_t5_encoder_and_tokenizer, +) from maxdiffusion.checkpointing.checkpointing_utils import (create_orbax_checkpoint_manager) from maxdiffusion.models.flux.util import load_flow_model @@ -187,12 +190,8 @@ def load_diffusers_checkpoint(self): context = nullcontext() with context: - clip_encoder = FlaxCLIPTextModel.from_pretrained(self.config.clip_model_name_or_path, dtype=self.config.weights_dtype) - clip_tokenizer = CLIPTokenizer.from_pretrained(self.config.clip_model_name_or_path, max_length=77, use_fast=True) - t5_encoder = FlaxT5EncoderModel.from_pretrained(self.config.t5xxl_model_name_or_path, dtype=self.config.weights_dtype) - t5_tokenizer = AutoTokenizer.from_pretrained( - self.config.t5xxl_model_name_or_path, model_max_length=self.config.max_sequence_length, use_fast=True - ) + clip_encoder, clip_tokenizer = load_clip_encoder_and_tokenizer(self.config) + t5_encoder, t5_tokenizer = load_t5_encoder_and_tokenizer(self.config) vae, vae_params = FlaxAutoencoderKL.from_pretrained( self.config.pretrained_model_name_or_path, @@ -256,16 +255,8 @@ def load_checkpoint(self, step=None, scheduler_class=None): context = nullcontext() with context: - clip_encoder = FlaxCLIPTextModel.from_pretrained( - self.config.clip_model_name_or_path, dtype=self.config.weights_dtype - ) - clip_tokenizer = CLIPTokenizer.from_pretrained(self.config.clip_model_name_or_path, max_length=77, use_fast=True) - t5_encoder = FlaxT5EncoderModel.from_pretrained( - self.config.t5xxl_model_name_or_path, dtype=self.config.weights_dtype - ) - t5_tokenizer = AutoTokenizer.from_pretrained( - self.config.t5xxl_model_name_or_path, model_max_length=self.config.max_sequence_length, use_fast=True - ) + clip_encoder, clip_tokenizer = load_clip_encoder_and_tokenizer(self.config) + t5_encoder, t5_tokenizer = load_t5_encoder_and_tokenizer(self.config) vae = FlaxAutoencoderKL.from_config( model_configs[0]["vae_config"], diff --git a/src/maxdiffusion/configs/base_flux_dev.yml b/src/maxdiffusion/configs/base_flux_dev.yml index 5761eda0a..5e9076032 100644 --- a/src/maxdiffusion/configs/base_flux_dev.yml +++ b/src/maxdiffusion/configs/base_flux_dev.yml @@ -28,8 +28,14 @@ save_config_to_gcs: False log_period: 100 pretrained_model_name_or_path: 'black-forest-labs/FLUX.1-dev' -clip_model_name_or_path: 'ariG23498/clip-vit-large-patch14-text-flax' -t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax' +# The text encoders are the PyTorch checkpoints, run under JAX through Torchax: +# transformers 5 ships no Flax implementations. +clip_model_name_or_path: 'black-forest-labs/FLUX.1-dev' +clip_model_subfolder: 'text_encoder' +clip_tokenizer_subfolder: 'tokenizer' +t5xxl_model_name_or_path: 'black-forest-labs/FLUX.1-dev' +t5xxl_model_subfolder: 'text_encoder_2' +t5xxl_tokenizer_subfolder: 'tokenizer_2' # Flux params flux_name: "flux-dev" diff --git a/src/maxdiffusion/configs/base_flux_dev_multi_res.yml b/src/maxdiffusion/configs/base_flux_dev_multi_res.yml index 5bc034ec5..9a395bc3b 100644 --- a/src/maxdiffusion/configs/base_flux_dev_multi_res.yml +++ b/src/maxdiffusion/configs/base_flux_dev_multi_res.yml @@ -28,8 +28,14 @@ save_config_to_gcs: False log_period: 100 pretrained_model_name_or_path: 'black-forest-labs/FLUX.1-dev' -clip_model_name_or_path: 'ariG23498/clip-vit-large-patch14-text-flax' -t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax' +# The text encoders are the PyTorch checkpoints, run under JAX through Torchax: +# transformers 5 ships no Flax implementations. +clip_model_name_or_path: 'black-forest-labs/FLUX.1-dev' +clip_model_subfolder: 'text_encoder' +clip_tokenizer_subfolder: 'tokenizer' +t5xxl_model_name_or_path: 'black-forest-labs/FLUX.1-dev' +t5xxl_model_subfolder: 'text_encoder_2' +t5xxl_tokenizer_subfolder: 'tokenizer_2' # Flux params flux_name: "flux-dev" diff --git a/src/maxdiffusion/configs/base_flux_schnell.yml b/src/maxdiffusion/configs/base_flux_schnell.yml index 1206a2516..555736fca 100644 --- a/src/maxdiffusion/configs/base_flux_schnell.yml +++ b/src/maxdiffusion/configs/base_flux_schnell.yml @@ -28,8 +28,14 @@ save_config_to_gcs: False log_period: 100 pretrained_model_name_or_path: 'black-forest-labs/FLUX.1-schnell' -clip_model_name_or_path: 'ariG23498/clip-vit-large-patch14-text-flax' -t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax' +# The text encoders are the PyTorch checkpoints, run under JAX through Torchax: +# transformers 5 ships no Flax implementations. +clip_model_name_or_path: 'black-forest-labs/FLUX.1-schnell' +clip_model_subfolder: 'text_encoder' +clip_tokenizer_subfolder: 'tokenizer' +t5xxl_model_name_or_path: 'black-forest-labs/FLUX.1-schnell' +t5xxl_model_subfolder: 'text_encoder_2' +t5xxl_tokenizer_subfolder: 'tokenizer_2' # Flux params flux_name: "flux-schnell" diff --git a/src/maxdiffusion/generate_flux.py b/src/maxdiffusion/generate_flux.py index 8d2153c65..73a579833 100644 --- a/src/maxdiffusion/generate_flux.py +++ b/src/maxdiffusion/generate_flux.py @@ -29,13 +29,18 @@ from chex import Array from einops import rearrange from flax.linen import partitioning as nn_partitioning -from transformers import (CLIPTokenizer, FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel, AutoTokenizer) +from transformers import (CLIPTokenizer, AutoTokenizer) from maxdiffusion import FlaxAutoencoderKL, pyconfig, max_logging, max_utils +from maxdiffusion.models.flux.text_encoders.torchax_text_encoders import ( + TorchaxCLIPTextEncoder, + TorchaxT5TextEncoder, + load_clip_encoder_and_tokenizer, + load_t5_encoder_and_tokenizer, +) from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel from maxdiffusion.train_utils import transformer_engine_context from maxdiffusion.max_utils import ( - device_put_replicated, get_memory_allocations, create_device_mesh, get_flash_block_sizes, @@ -193,7 +198,7 @@ def prepare_latents( def get_clip_prompt_embeds( - prompt: Union[str, List[str]], num_images_per_prompt: int, tokenizer: CLIPTokenizer, text_encoder: FlaxCLIPTextModel + prompt: Union[str, List[str]], num_images_per_prompt: int, tokenizer: CLIPTokenizer, text_encoder: TorchaxCLIPTextEncoder ): prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) @@ -207,10 +212,9 @@ def get_clip_prompt_embeds( return_tensors="np", ) - text_input_ids = text_inputs.input_ids + text_input_ids = jnp.asarray(text_inputs.input_ids, dtype=jnp.int32) - prompt_embeds = text_encoder(text_input_ids, params=text_encoder.params, train=False) - prompt_embeds = prompt_embeds.pooler_output + prompt_embeds = text_encoder(text_input_ids) prompt_embeds = jnp.tile(prompt_embeds, (batch_size * num_images_per_prompt, 1)) return prompt_embeds @@ -219,7 +223,7 @@ def get_t5_prompt_embeds( prompt: Union[str, List[str]], num_images_per_prompt: int, tokenizer: AutoTokenizer, - text_encoder: T5EncoderModel, + text_encoder: TorchaxT5TextEncoder, max_sequence_length: int = 512, ): prompt = [prompt] if isinstance(prompt, str) else prompt @@ -233,10 +237,8 @@ def get_t5_prompt_embeds( padding="max_length", return_tensors="np", ) - text_input_ids = text_inputs.input_ids - prompt_embeds = text_encoder(text_input_ids, attention_mask=None, output_hidden_states=False)["last_hidden_state"] - dtype = text_encoder.dtype - prompt_embeds = prompt_embeds.astype(dtype) + text_input_ids = jnp.asarray(text_inputs.input_ids, dtype=jnp.int32) + prompt_embeds = text_encoder(text_input_ids) _, seq_len, _ = prompt_embeds.shape # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method prompt_embeds = jnp.tile(prompt_embeds, (1, num_images_per_prompt, 1)) @@ -248,9 +250,9 @@ def encode_prompt( prompt: Union[str, List[str]], prompt_2: Union[str, List[str]], clip_tokenizer: CLIPTokenizer, - clip_text_encoder: FlaxCLIPTextModel, + clip_text_encoder: TorchaxCLIPTextEncoder, t5_tokenizer: AutoTokenizer, - t5_text_encoder: T5EncoderModel, + t5_text_encoder: TorchaxT5TextEncoder, num_images_per_prompt: int = 1, max_sequence_length: int = 512, ): @@ -331,24 +333,12 @@ def run(config): ) # LOAD TEXT ENCODERS - clip_text_encoder = FlaxCLIPTextModel.from_pretrained( - config.pretrained_model_name_or_path, subfolder="text_encoder", from_pt=True, dtype=config.weights_dtype - ) - clip_tokenizer = CLIPTokenizer.from_pretrained( - config.pretrained_model_name_or_path, subfolder="tokenizer", dtype=config.weights_dtype - ) - - t5_encoder = FlaxT5EncoderModel.from_pretrained(config.t5xxl_model_name_or_path, dtype=config.weights_dtype) - t5_tokenizer = AutoTokenizer.from_pretrained( - config.t5xxl_model_name_or_path, max_length=config.max_sequence_length, use_fast=True - ) + clip_text_encoder, clip_tokenizer = load_clip_encoder_and_tokenizer(config) + t5_encoder, t5_tokenizer = load_t5_encoder_and_tokenizer(config) encoders_sharding = NamedSharding(mesh, P()) - partial_device_put_replicated = functools.partial(device_put_replicated, sharding=encoders_sharding) - clip_text_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), clip_text_encoder.params) - clip_text_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, clip_text_encoder.params) - t5_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), t5_encoder.params) - t5_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, t5_encoder.params) + clip_text_encoder.place_params(encoders_sharding) + t5_encoder.place_params(encoders_sharding) prompt_embeds, pooled_prompt_embeds, text_ids = encode_prompt( prompt=config.prompt, @@ -382,8 +372,7 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep pooled_prompt_embeds = jax.device_put(pooled_prompt_embeds, data_sharding) if config.offload_encoders: - cpus = jax.devices("cpu") - t5_encoder.params = jax.device_put(t5_encoder.params, device=cpus[0]) + t5_encoder.offload_params() get_memory_allocations() # evaluate shapes diff --git a/src/maxdiffusion/generate_flux_multi_res.py b/src/maxdiffusion/generate_flux_multi_res.py index c74533d42..d0ff79ad0 100644 --- a/src/maxdiffusion/generate_flux_multi_res.py +++ b/src/maxdiffusion/generate_flux_multi_res.py @@ -28,12 +28,17 @@ import flax.linen as nn from chex import Array from flax.linen import partitioning as nn_partitioning -from transformers import (CLIPTokenizer, FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel, AutoTokenizer) +from transformers import (CLIPTokenizer, AutoTokenizer) from maxdiffusion import FlaxAutoencoderKL, pyconfig, max_logging, max_utils +from maxdiffusion.models.flux.text_encoders.torchax_text_encoders import ( + TorchaxCLIPTextEncoder, + TorchaxT5TextEncoder, + load_clip_encoder_and_tokenizer, + load_t5_encoder_and_tokenizer, +) from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel from maxdiffusion.max_utils import ( - device_put_replicated, get_memory_allocations, create_device_mesh, get_flash_block_sizes, @@ -225,7 +230,7 @@ def tokenize_clip(prompt: Union[str, List[str]], tokenizer: CLIPTokenizer): def get_clip_prompt_embeds( - prompt: Union[str, List[str]], num_images_per_prompt: int, tokenizer: CLIPTokenizer, text_encoder: FlaxCLIPTextModel + prompt: Union[str, List[str]], num_images_per_prompt: int, tokenizer: CLIPTokenizer, text_encoder: TorchaxCLIPTextEncoder ): prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) @@ -239,10 +244,9 @@ def get_clip_prompt_embeds( return_tensors="np", ) - text_input_ids = text_inputs.input_ids + text_input_ids = jnp.asarray(text_inputs.input_ids, dtype=jnp.int32) - prompt_embeds = text_encoder(text_input_ids, params=text_encoder.params, train=False) - prompt_embeds = prompt_embeds.pooler_output + prompt_embeds = text_encoder(text_input_ids) prompt_embeds = jnp.tile(prompt_embeds, (batch_size * num_images_per_prompt, 1)) return prompt_embeds @@ -265,7 +269,7 @@ def get_t5_prompt_embeds( prompt: Union[str, List[str]], num_images_per_prompt: int, tokenizer: AutoTokenizer, - text_encoder: T5EncoderModel, + text_encoder: TorchaxT5TextEncoder, max_sequence_length: int = 512, ): prompt = [prompt] if isinstance(prompt, str) else prompt @@ -279,10 +283,8 @@ def get_t5_prompt_embeds( padding="max_length", return_tensors="np", ) - text_input_ids = text_inputs.input_ids - prompt_embeds = text_encoder(text_input_ids, attention_mask=None, output_hidden_states=False)["last_hidden_state"] - dtype = text_encoder.dtype - prompt_embeds = prompt_embeds.astype(dtype) + text_input_ids = jnp.asarray(text_inputs.input_ids, dtype=jnp.int32) + prompt_embeds = text_encoder(text_input_ids) _, seq_len, _ = prompt_embeds.shape # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method prompt_embeds = jnp.tile(prompt_embeds, (1, num_images_per_prompt, 1)) @@ -294,9 +296,9 @@ def encode_prompt( prompt: Union[str, List[str]], prompt_2: Union[str, List[str]], clip_tokenizer: CLIPTokenizer, - clip_text_encoder: FlaxCLIPTextModel, + clip_text_encoder: TorchaxCLIPTextEncoder, t5_tokenizer: AutoTokenizer, - t5_text_encoder: T5EncoderModel, + t5_text_encoder: TorchaxT5TextEncoder, num_images_per_prompt: int = 1, max_sequence_length: int = 512, ): @@ -365,24 +367,12 @@ def run(config): num_channels_latents = transformer.in_channels // 4 # LOAD TEXT ENCODERS - clip_text_encoder = FlaxCLIPTextModel.from_pretrained( - config.pretrained_model_name_or_path, subfolder="text_encoder", from_pt=True, dtype=config.weights_dtype - ) - clip_tokenizer = CLIPTokenizer.from_pretrained( - config.pretrained_model_name_or_path, subfolder="tokenizer", dtype=config.weights_dtype - ) - - t5_encoder = FlaxT5EncoderModel.from_pretrained(config.t5xxl_model_name_or_path, dtype=config.weights_dtype) - t5_tokenizer = AutoTokenizer.from_pretrained( - config.t5xxl_model_name_or_path, max_length=config.max_sequence_length, use_fast=True - ) + clip_text_encoder, clip_tokenizer = load_clip_encoder_and_tokenizer(config) + t5_encoder, t5_tokenizer = load_t5_encoder_and_tokenizer(config) encoders_sharding = NamedSharding(mesh, P()) - partial_device_put_replicated = functools.partial(device_put_replicated, sharding=encoders_sharding) - clip_text_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), clip_text_encoder.params) - clip_text_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, clip_text_encoder.params) - t5_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), t5_encoder.params) - t5_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, t5_encoder.params) + clip_text_encoder.place_params(encoders_sharding) + t5_encoder.place_params(encoders_sharding) def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timesteps, guidance, pooled_prompt_embeds): print("latents.shape: ", latents.shape, latents.dtype) @@ -460,7 +450,7 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep for _ in range(2): s0 = time.perf_counter() if config.offload_encoders: - t5_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, t5_encoder.params) + t5_encoder.place_params(encoders_sharding) max_logging.log(f"Moving encoder to TPU time: {(time.perf_counter() - s0)}") prompt_embeds, pooled_prompt_embeds, text_ids = encode_prompt( prompt=config.prompt, @@ -474,8 +464,7 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep ) if config.offload_encoders: s1 = time.perf_counter() - cpus = jax.devices("cpu") - t5_encoder.params = jax.device_put(t5_encoder.params, device=cpus[0]) + t5_encoder.offload_params() max_logging.log(f"Text encoding offload time: {(time.perf_counter() - s1)}") text_encoding_time_final = time.perf_counter() - s0 max_logging.log(f"text encoding time: {text_encoding_time_final}") diff --git a/src/maxdiffusion/models/flux/text_encoders/__init__.py b/src/maxdiffusion/models/flux/text_encoders/__init__.py new file mode 100644 index 000000000..11f31009e --- /dev/null +++ b/src/maxdiffusion/models/flux/text_encoders/__init__.py @@ -0,0 +1,15 @@ +""" +Copyright 2026 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. +""" diff --git a/src/maxdiffusion/models/flux/text_encoders/torchax_text_encoders.py b/src/maxdiffusion/models/flux/text_encoders/torchax_text_encoders.py new file mode 100644 index 000000000..78be4fecb --- /dev/null +++ b/src/maxdiffusion/models/flux/text_encoders/torchax_text_encoders.py @@ -0,0 +1,168 @@ +# Copyright 2026 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. + +"""FLUX text encoders (CLIP-L, T5-XXL) running under JAX through Torchax. + +transformers dropped its Flax implementations in v5, so FLUX loads the +PyTorch encoders and traces them into JAX here, the same way LTX2 wraps +Gemma 3 and WAN wraps UMT5. Both encoders are frozen: training only calls +them to precompute prompt embeddings before the transformer sees a batch. +""" + +from typing import Optional, Tuple + +import jax +import jax.numpy as jnp +import torch +from torchax import default_env, interop +from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer, T5EncoderModel + +CLIP_TOKENIZER_MAX_LENGTH = 77 + +TORCH_DTYPE_MAP = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, +} + + +def _torch_dtype(jax_dtype) -> torch.dtype: + name = jax_dtype.name if hasattr(jax_dtype, "name") else str(jax_dtype) + if name not in TORCH_DTYPE_MAP: + raise ValueError(f"Unsupported text encoder dtype: {name}. Supported values are: {list(TORCH_DTYPE_MAP.keys())}") + return TORCH_DTYPE_MAP[name] + + +class _TorchaxTextEncoder(interop.JittableModule): + """Common Torchax plumbing for the frozen FLUX text encoders.""" + + def __init__(self, model: torch.nn.Module, jax_dtype: jnp.dtype): + super().__init__(model) + self.jax_dtype = jax_dtype + + @classmethod + def from_torch(cls, model: torch.nn.Module, jax_dtype: jnp.dtype) -> "_TorchaxTextEncoder": + model.eval() + with default_env(): + return cls(model.to("jax"), jax_dtype) + + def place_params(self, device_or_sharding) -> None: + """Move the frozen weights onto `device_or_sharding`. + + Torchax holds the weights as torch views over JAX arrays, so the placement + goes through `jax_view`/`torch_view`. A `Sharding` target is built with + `make_array_from_callback` (as `max_utils.device_put_replicated` does), + because each host only holds its own full copy of these weights and + `device_put` cannot span non-addressable devices. + """ + + def place(leaf): + if isinstance(device_or_sharding, jax.sharding.Sharding): + return jax.make_array_from_callback(leaf.shape, device_or_sharding, lambda index: leaf[index]) + return jax.device_put(leaf, device_or_sharding) + + self.params = interop.torch_view(jax.tree_util.tree_map(place, interop.jax_view(self.params))) + + def offload_params(self) -> None: + """Park the weights in host memory until the next encode call.""" + self.place_params(jax.devices("cpu")[0]) + + +class TorchaxCLIPTextEncoder(_TorchaxTextEncoder): + """Wraps `transformers.CLIPTextModel`, returning FLUX's pooled embedding.""" + + def __call__(self, input_ids: jax.Array) -> jax.Array: + with default_env(): + pooler_output = self.functional_call( + self._pooler_output, + params=self.params, + buffers=self.buffers, + input_ids=interop.torch_view(input_ids), + ) + return interop.jax_view(pooler_output).astype(self.jax_dtype) + + @staticmethod + def _pooler_output(model, input_ids): + # Returning the tensor rather than the output dataclass keeps the result a + # plain pytree that `interop.jax_view` can convert. + return model(input_ids=input_ids).pooler_output + + +class TorchaxT5TextEncoder(_TorchaxTextEncoder): + """Wraps `transformers.T5EncoderModel`, returning the last hidden state.""" + + def __call__(self, input_ids: jax.Array, attention_mask: Optional[jax.Array] = None) -> jax.Array: + with default_env(): + last_hidden_state = self.functional_call( + self._last_hidden_state, + params=self.params, + buffers=self.buffers, + input_ids=interop.torch_view(input_ids), + attention_mask=None if attention_mask is None else interop.torch_view(attention_mask), + ) + return interop.jax_view(last_hidden_state).astype(self.jax_dtype) + + @staticmethod + def _last_hidden_state(model, input_ids, attention_mask): + return model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state + + +def _optional_subfolder(config, key) -> str: + """Read a subfolder key that a config predating the PyTorch weights may not carry. + + pyconfig raises ValueError for a key it does not hold, which getattr's default + cannot absorb, so an older config would fail here rather than fall back to the + repository root. + """ + try: + return getattr(config, key) or "" + except (AttributeError, KeyError, ValueError): + return "" + + +def load_clip_encoder_and_tokenizer(config) -> Tuple[TorchaxCLIPTextEncoder, CLIPTokenizer]: + """Load FLUX's CLIP-L text encoder and its tokenizer from PyTorch weights.""" + subfolder = _optional_subfolder(config, "clip_model_subfolder") + # `dtype` rather than the deprecated `torch_dtype`: transformers 5 documents + # only the former. Eager attention keeps the graph traceable under Torchax. + encoder = CLIPTextModel.from_pretrained( + config.clip_model_name_or_path, + subfolder=subfolder, + dtype=_torch_dtype(config.weights_dtype), + attn_implementation="eager", + ) + tokenizer = CLIPTokenizer.from_pretrained( + config.clip_model_name_or_path, + subfolder=_optional_subfolder(config, "clip_tokenizer_subfolder"), + model_max_length=CLIP_TOKENIZER_MAX_LENGTH, + ) + return TorchaxCLIPTextEncoder.from_torch(encoder, config.weights_dtype), tokenizer + + +def load_t5_encoder_and_tokenizer(config) -> Tuple[TorchaxT5TextEncoder, AutoTokenizer]: + """Load FLUX's T5-XXL text encoder and its tokenizer from PyTorch weights.""" + subfolder = _optional_subfolder(config, "t5xxl_model_subfolder") + encoder = T5EncoderModel.from_pretrained( + config.t5xxl_model_name_or_path, + subfolder=subfolder, + dtype=_torch_dtype(config.weights_dtype), + attn_implementation="eager", + ) + tokenizer = AutoTokenizer.from_pretrained( + config.t5xxl_model_name_or_path, + subfolder=_optional_subfolder(config, "t5xxl_tokenizer_subfolder"), + model_max_length=config.max_sequence_length, + use_fast=True, + ) + return TorchaxT5TextEncoder.from_torch(encoder, config.weights_dtype), tokenizer diff --git a/src/maxdiffusion/pipelines/flux/flux_pipeline.py b/src/maxdiffusion/pipelines/flux/flux_pipeline.py index 15b2c4f5e..d8887de38 100644 --- a/src/maxdiffusion/pipelines/flux/flux_pipeline.py +++ b/src/maxdiffusion/pipelines/flux/flux_pipeline.py @@ -18,7 +18,7 @@ import jax import jax.numpy as jnp import math -from transformers import (CLIPTokenizer, FlaxCLIPTextModel, FlaxT5EncoderModel, AutoTokenizer) +from transformers import (CLIPTokenizer, AutoTokenizer) from einops import rearrange from jax.typing import DTypeLike from chex import Array @@ -28,6 +28,7 @@ from maxdiffusion.utils import logging from ...models import FlaxAutoencoderKL +from ...models.flux.text_encoders.torchax_text_encoders import (TorchaxCLIPTextEncoder, TorchaxT5TextEncoder) from ...schedulers import (FlaxEulerDiscreteScheduler) from ..pipeline_flax_utils import FlaxDiffusionPipeline from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel @@ -43,10 +44,10 @@ class FluxPipeline(FlaxDiffusionPipeline): def __init__( self, - t5_encoder: FlaxCLIPTextModel, - clip_encoder: FlaxCLIPTextModel, + t5_encoder: TorchaxT5TextEncoder, + clip_encoder: TorchaxCLIPTextEncoder, vae: FlaxAutoencoderKL, - t5_tokenizer: FlaxT5EncoderModel, + t5_tokenizer: AutoTokenizer, clip_tokenizer: CLIPTokenizer, flux: FluxTransformer2DModel, scheduler: FlaxEulerDiscreteScheduler, @@ -163,7 +164,7 @@ def get_clip_prompt_embeds( prompt: Union[str, List[str]], num_images_per_prompt: int, tokenizer: CLIPTokenizer, - text_encoder: FlaxCLIPTextModel, + text_encoder: TorchaxCLIPTextEncoder, ): prompt = [prompt] if isinstance(prompt, str) else prompt text_inputs = tokenizer( @@ -176,10 +177,9 @@ def get_clip_prompt_embeds( return_tensors="np", ) - text_input_ids = text_inputs.input_ids + text_input_ids = jnp.asarray(text_inputs.input_ids, dtype=jnp.int32) - prompt_embeds = text_encoder(text_input_ids, params=text_encoder.params, train=False) - prompt_embeds = prompt_embeds.pooler_output + prompt_embeds = text_encoder(text_input_ids) prompt_embeds = jnp.tile(prompt_embeds, (num_images_per_prompt, 1)) return prompt_embeds @@ -188,7 +188,7 @@ def get_t5_prompt_embeds( prompt: Union[str, List[str]], num_images_per_prompt: int, tokenizer: AutoTokenizer, - text_encoder: FlaxT5EncoderModel, + text_encoder: TorchaxT5TextEncoder, max_sequence_length: int = 512, encode_in_batches=False, encode_batch_size=None, @@ -205,27 +205,22 @@ def get_t5_prompt_embeds( padding="max_length", return_tensors="np", ) - text_input_ids = text_inputs.input_ids + text_input_ids = jnp.asarray(text_inputs.input_ids, dtype=jnp.int32) if encode_in_batches: prompt_embeds = None for i in range(0, text_input_ids.shape[0], encode_batch_size): - batch_prompt_embeds = text_encoder( - text_input_ids[i : i + encode_batch_size], attention_mask=None, output_hidden_states=False - )["last_hidden_state"] + batch_prompt_embeds = text_encoder(text_input_ids[i : i + encode_batch_size]) if prompt_embeds is None: prompt_embeds = batch_prompt_embeds else: prompt_embeds = jnp.concatenate([prompt_embeds, batch_prompt_embeds]) else: - prompt_embeds = text_encoder(text_input_ids, attention_mask=None, output_hidden_states=False)["last_hidden_state"] + prompt_embeds = text_encoder(text_input_ids) _, seq_len, _ = prompt_embeds.shape # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method prompt_embeds = jnp.tile(prompt_embeds, (1, num_images_per_prompt, 1)) prompt_embeds = jnp.reshape(prompt_embeds, (batch_size * num_images_per_prompt, seq_len, -1)) - dtype = text_encoder.dtype - prompt_embeds = prompt_embeds.astype(dtype) - return prompt_embeds def encode_prompt( @@ -233,9 +228,9 @@ def encode_prompt( prompt: Union[str, List[str]], prompt_2: Union[str, List[str]], clip_tokenizer: CLIPTokenizer, - clip_text_encoder: FlaxCLIPTextModel, + clip_text_encoder: TorchaxCLIPTextEncoder, t5_tokenizer: AutoTokenizer, - t5_text_encoder: FlaxT5EncoderModel, + t5_text_encoder: TorchaxT5TextEncoder, num_images_per_prompt: int = 1, max_sequence_length: int = 512, encode_in_batches: bool = False, diff --git a/src/maxdiffusion/tests/flux_text_encoder_parity_test.py b/src/maxdiffusion/tests/flux_text_encoder_parity_test.py new file mode 100644 index 000000000..c0822b6c6 --- /dev/null +++ b/src/maxdiffusion/tests/flux_text_encoder_parity_test.py @@ -0,0 +1,117 @@ +""" +Copyright 2026 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. +""" + +# Parity between eager PyTorch and the Torchax-traced FLUX text encoders. The +# models here are small randomly-initialized stand-ins, so the test runs on CPU +# without touching the hub: what is under test is the Torchax wrapper, not the +# CLIP or T5 weights. + +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +import torch +from absl.testing import absltest +from transformers import CLIPTextConfig, CLIPTextModel, T5Config, T5EncoderModel + +from ..models.flux.text_encoders.torchax_text_encoders import TorchaxCLIPTextEncoder, TorchaxT5TextEncoder + +SEQ_LEN = 16 +BATCH_SIZE = 2 +VOCAB_SIZE = 99 + + +def _input_ids(): + ids = np.random.default_rng(0).integers(3, VOCAB_SIZE, size=(BATCH_SIZE, SEQ_LEN), dtype=np.int32) + # CLIP pools at the EOS position, so every sequence needs one. + ids[:, -1] = 2 + return ids + + +def _tiny_clip_text_model(): + torch.manual_seed(0) + config = CLIPTextConfig( + vocab_size=VOCAB_SIZE, + hidden_size=32, + intermediate_size=37, + num_hidden_layers=2, + num_attention_heads=4, + max_position_embeddings=SEQ_LEN, + bos_token_id=0, + pad_token_id=1, + eos_token_id=2, + attn_implementation="eager", + ) + return CLIPTextModel(config).eval() + + +def _tiny_t5_encoder_model(): + torch.manual_seed(0) + config = T5Config( + vocab_size=VOCAB_SIZE, + d_model=32, + d_ff=37, + d_kv=8, + num_layers=2, + num_heads=4, + is_encoder_decoder=False, + attn_implementation="eager", + ) + return T5EncoderModel(config).eval() + + +class FluxTextEncoderParityTest(unittest.TestCase): + """Compares Torchax outputs against the eager PyTorch reference.""" + + def test_clip_pooled_output_matches_eager(self): + input_ids = _input_ids() + model = _tiny_clip_text_model() + with torch.no_grad(): + expected = model(input_ids=torch.from_numpy(input_ids).long()).pooler_output.numpy() + + encoder = TorchaxCLIPTextEncoder.from_torch(model, jnp.float32) + actual = encoder(jnp.asarray(input_ids, dtype=jnp.int32)) + + self.assertEqual(actual.shape, expected.shape) + np.testing.assert_allclose(np.asarray(actual), expected, atol=1e-4, rtol=1e-4) + + def test_t5_last_hidden_state_matches_eager(self): + input_ids = _input_ids() + model = _tiny_t5_encoder_model() + with torch.no_grad(): + expected = model(input_ids=torch.from_numpy(input_ids).long()).last_hidden_state.numpy() + + encoder = TorchaxT5TextEncoder.from_torch(model, jnp.float32) + actual = encoder(jnp.asarray(input_ids, dtype=jnp.int32)) + + self.assertEqual(actual.shape, expected.shape) + np.testing.assert_allclose(np.asarray(actual), expected, atol=1e-4, rtol=1e-4) + + def test_place_and_offload_params_preserve_outputs(self): + input_ids = jnp.asarray(_input_ids(), dtype=jnp.int32) + encoder = TorchaxT5TextEncoder.from_torch(_tiny_t5_encoder_model(), jnp.float32) + before = np.asarray(encoder(input_ids)) + + encoder.place_params(jax.devices()[0]) + np.testing.assert_allclose(np.asarray(encoder(input_ids)), before, atol=1e-6, rtol=1e-6) + + encoder.offload_params() + np.testing.assert_allclose(np.asarray(encoder(input_ids)), before, atol=1e-6, rtol=1e-6) + + +if __name__ == "__main__": + absltest.main() diff --git a/src/maxdiffusion/tests/text_encoders_test.py b/src/maxdiffusion/tests/text_encoders_test.py index c91bca9a1..aaa06ad71 100644 --- a/src/maxdiffusion/tests/text_encoders_test.py +++ b/src/maxdiffusion/tests/text_encoders_test.py @@ -19,14 +19,19 @@ import pytest from absl.testing import absltest -from transformers import CLIPTokenizer, FlaxCLIPTextModel -from transformers import T5TokenizerFast, FlaxT5EncoderModel +import jax.numpy as jnp +import torch +from transformers import CLIPTokenizer, CLIPTextModel +from transformers import T5TokenizerFast, T5EncoderModel from ..generate_flux import get_clip_prompt_embeds, get_t5_prompt_embeds +from ..models.flux.text_encoders.torchax_text_encoders import TorchaxCLIPTextEncoder, TorchaxT5TextEncoder IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +FLUX_DEV = "black-forest-labs/FLUX.1-dev" + class TextEncoderTest(unittest.TestCase): """Test text encoders""" @@ -36,9 +41,12 @@ def setUp(self): @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions") def test_flux_t5_text_encoder(self): - text_encoder = FlaxT5EncoderModel.from_pretrained("ariG23498/t5-v1-1-xxl-flax") + text_encoder = TorchaxT5TextEncoder.from_torch( + T5EncoderModel.from_pretrained(FLUX_DEV, subfolder="text_encoder_2", dtype=torch.bfloat16), + jnp.bfloat16, + ) - tokenizer_2 = T5TokenizerFast.from_pretrained("ariG23498/t5-v1-1-xxl-flax") + tokenizer_2 = T5TokenizerFast.from_pretrained(FLUX_DEV, subfolder="tokenizer_2") embeds = get_t5_prompt_embeds("A dog on a skateboard", 2, tokenizer_2, text_encoder) @@ -46,10 +54,11 @@ def test_flux_t5_text_encoder(self): @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions") def test_flux_clip_text_encoder(self): - text_encoder = FlaxCLIPTextModel.from_pretrained( - "black-forest-labs/FLUX.1-dev", subfolder="text_encoder", from_pt=True, dtype="bfloat16" + text_encoder = TorchaxCLIPTextEncoder.from_torch( + CLIPTextModel.from_pretrained(FLUX_DEV, subfolder="text_encoder", dtype=torch.bfloat16), + jnp.bfloat16, ) - tokenizer = CLIPTokenizer.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="tokenizer", dtype="bfloat16") + tokenizer = CLIPTokenizer.from_pretrained(FLUX_DEV, subfolder="tokenizer") embeds = get_clip_prompt_embeds("A cat riding a skateboard", 2, tokenizer, text_encoder) assert embeds.shape == (2, 768) diff --git a/src/maxdiffusion/trainers/flux_trainer.py b/src/maxdiffusion/trainers/flux_trainer.py index 747b62930..ea1eeb0cf 100644 --- a/src/maxdiffusion/trainers/flux_trainer.py +++ b/src/maxdiffusion/trainers/flux_trainer.py @@ -92,11 +92,8 @@ def start_training(self): # move params to accelerator encoders_sharding = NamedSharding(self.mesh, P(None)) - partial_device_put_replicated = partial(max_utils.device_put_replicated, sharding=encoders_sharding) - pipeline.clip_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), pipeline.clip_encoder.params) - pipeline.clip_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, pipeline.clip_encoder.params) - pipeline.t5_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), pipeline.t5_encoder.params) - pipeline.t5_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, pipeline.t5_encoder.params) + pipeline.clip_encoder.place_params(encoders_sharding) + pipeline.t5_encoder.place_params(encoders_sharding) vae_state, vae_state_mesh_shardings = self.create_vae_state( pipeline=pipeline, params=params, checkpoint_item_name=VAE_STATE_KEY, is_training=False From a16fc1f85a63bda7d14418068683b4b60058e0ad Mon Sep 17 00:00:00 2001 From: olehtika Date: Thu, 22 Oct 2026 06:32:15 +0000 Subject: [PATCH 2/2] Place Torchax text encoder buffers alongside params place_params moved only self.params, so buffers such as CLIP's position_ids stayed on the default device while the weights were replicated across a mesh or offloaded to the host. Move self.buffers with the same placement, and add a CLIP test, since T5 carries no buffers and the existing placement test cannot see them. --- .../text_encoders/torchax_text_encoders.py | 11 ++++++++-- .../tests/flux_text_encoder_parity_test.py | 22 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/src/maxdiffusion/models/flux/text_encoders/torchax_text_encoders.py b/src/maxdiffusion/models/flux/text_encoders/torchax_text_encoders.py index 78be4fecb..238a6a9a5 100644 --- a/src/maxdiffusion/models/flux/text_encoders/torchax_text_encoders.py +++ b/src/maxdiffusion/models/flux/text_encoders/torchax_text_encoders.py @@ -58,13 +58,16 @@ def from_torch(cls, model: torch.nn.Module, jax_dtype: jnp.dtype) -> "_TorchaxTe return cls(model.to("jax"), jax_dtype) def place_params(self, device_or_sharding) -> None: - """Move the frozen weights onto `device_or_sharding`. + """Move the frozen weights and buffers onto `device_or_sharding`. Torchax holds the weights as torch views over JAX arrays, so the placement goes through `jax_view`/`torch_view`. A `Sharding` target is built with `make_array_from_callback` (as `max_utils.device_put_replicated` does), because each host only holds its own full copy of these weights and `device_put` cannot span non-addressable devices. + + Buffers move too: CLIP's `position_ids` is one, and Torchax files any + parameter with `requires_grad=False` under buffers rather than params. """ def place(leaf): @@ -72,7 +75,11 @@ def place(leaf): return jax.make_array_from_callback(leaf.shape, device_or_sharding, lambda index: leaf[index]) return jax.device_put(leaf, device_or_sharding) - self.params = interop.torch_view(jax.tree_util.tree_map(place, interop.jax_view(self.params))) + def place_tree(tree): + return interop.torch_view(jax.tree_util.tree_map(place, interop.jax_view(tree))) + + self.params = place_tree(self.params) + self.buffers = place_tree(self.buffers) def offload_params(self) -> None: """Park the weights in host memory until the next encode call.""" diff --git a/src/maxdiffusion/tests/flux_text_encoder_parity_test.py b/src/maxdiffusion/tests/flux_text_encoder_parity_test.py index c0822b6c6..174bf9eb8 100644 --- a/src/maxdiffusion/tests/flux_text_encoder_parity_test.py +++ b/src/maxdiffusion/tests/flux_text_encoder_parity_test.py @@ -26,6 +26,7 @@ import numpy as np import torch from absl.testing import absltest +from torchax import interop from transformers import CLIPTextConfig, CLIPTextModel, T5Config, T5EncoderModel from ..models.flux.text_encoders.torchax_text_encoders import TorchaxCLIPTextEncoder, TorchaxT5TextEncoder @@ -112,6 +113,27 @@ def test_place_and_offload_params_preserve_outputs(self): encoder.offload_params() np.testing.assert_allclose(np.asarray(encoder(input_ids)), before, atol=1e-6, rtol=1e-6) + def test_place_and_offload_params_move_buffers(self): + input_ids = jnp.asarray(_input_ids(), dtype=jnp.int32) + # CLIP, unlike T5, carries a buffer (`position_ids`), so only it can catch + # buffers left behind on the default device. + encoder = TorchaxCLIPTextEncoder.from_torch(_tiny_clip_text_model(), jnp.float32) + self.assertTrue(encoder.buffers) + before = np.asarray(encoder(input_ids)) + + mesh = jax.sharding.Mesh(np.array(jax.devices()), ("data",)) + replicated = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) + encoder.place_params(replicated) + for leaf in jax.tree_util.tree_leaves(interop.jax_view(encoder.buffers)): + self.assertEqual(leaf.sharding, replicated) + np.testing.assert_allclose(np.asarray(encoder(input_ids)), before, atol=1e-6, rtol=1e-6) + + encoder.offload_params() + cpu = jax.devices("cpu")[0] + for leaf in jax.tree_util.tree_leaves(interop.jax_view(encoder.buffers)): + self.assertEqual(leaf.devices(), {cpu}) + np.testing.assert_allclose(np.asarray(encoder(input_ids)), before, atol=1e-6, rtol=1e-6) + if __name__ == "__main__": absltest.main()