Fix transformers compatibility - handle FlaxCLIPTextModel import for v5.0+

Author: eltsai

src/maxdiffusion/generate_flux.py

@@ -29,7 +29,19 @@
 from chex import Array
 from einops import rearrange
 from flax.linen import partitioning as nn_partitioning
-from transformers import (CLIPTokenizer, FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel, AutoTokenizer)
+
+try:
+  from transformers import (CLIPTokenizer, FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel, AutoTokenizer)
+except ImportError:
+  # For transformers>=5.0, Flax models have different import paths
+  from transformers import CLIPTokenizer, T5EncoderModel, AutoTokenizer
+
+  try:
+    from transformers.models.clip.modeling_flax_clip import FlaxCLIPTextModel
+    from transformers.models.t5.modeling_flax_t5 import FlaxT5EncoderModel
+  except ImportError:
+    FlaxCLIPTextModel = None
+    FlaxT5EncoderModel = None
 
 from maxdiffusion import FlaxAutoencoderKL, pyconfig, max_logging
 from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel