fix sdxl generate smoke tests.

Author: jfacevedo-google

src/maxdiffusion/generate_flux.py

@@ -31,7 +31,7 @@
 
 from maxdiffusion import FlaxAutoencoderKL, pyconfig, max_logging
 from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel
-from max_utils import (
+from maxdiffusion.max_utils import (
     device_put_replicated,
     get_memory_allocations,
     create_device_mesh,
@@ -52,9 +52,6 @@ def unpack(x: Array, height: int, width: int) -> Array:
   )
 
 
-from einops import rearrange
-
-
 def vae_decode(latents, vae, state, config):
   img = unpack(x=latents, height=config.resolution, width=config.resolution)
   img = img / vae.config.scaling_factor + vae.config.shift_factor

src/maxdiffusion/max_utils.py

@@ -46,7 +46,10 @@
 from flax.linen import partitioning as nn_partitioning
 from flax.training import train_state
 from jax.experimental import mesh_utils
-from transformers import (FlaxCLIPTextModel, FlaxCLIPTextPreTrainedModel)
+from transformers import (
+  FlaxCLIPTextModel,
+  FlaxCLIPTextPreTrainedModel
+)
 from flax import struct
 from typing import (
     Callable,