Run FLUX text encoders under Torchax instead of Flax

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"],

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"

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"

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"

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

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}")

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.
+"""