fix sdxl generate smoke tests.

Author: jfacevedo-google

src/maxdiffusion/checkpointing/base_stable_diffusion_checkpointer.py

@@ -88,11 +88,14 @@ def create_unet_state(self, pipeline, params, checkpoint_item_name, is_training)
         config=self.config,
         mesh=self.mesh,
         weights_init_fn=weights_init_fn,
-        model_params=None if self.config.train_new_unet else params.get("unet", None),
+        model_params=None,
         checkpoint_manager=self.checkpoint_manager,
         checkpoint_item=checkpoint_item_name,
         training=is_training,
     )
+    if not self.config.train_new_unet:
+      unet_state = unet_state.replace(params=params.get("unet", None))
+      unet_state = jax.device_put(unet_state, state_mesh_shardings)
     return unet_state, state_mesh_shardings, learning_rate_scheduler
 
   def create_vae_state(self, pipeline, params, checkpoint_item_name, is_training=False):
@@ -150,17 +153,20 @@ def create_text_encoder_2_state(self, pipeline, params, checkpoint_item_name, is
         input_shape=(self.total_train_batch_size, pipeline.tokenizer.model_max_length),
     )
 
-    return max_utils.setup_initial_state(
+    state, state_mesh_shardings = max_utils.setup_initial_state(
         model=pipeline.text_encoder_2,
         tx=tx,
         config=self.config,
         mesh=self.mesh,
         weights_init_fn=weights_init_fn,
-        model_params=params.get("text_encoder_2", None),
+        model_params=None,
         checkpoint_manager=self.checkpoint_manager,
         checkpoint_item=checkpoint_item_name,
         training=is_training,
     )
+    state = state.replace(params=params.get("text_encoder_2", None))
+    state = jax.device_put(state, state_mesh_shardings)
+    return state, state_mesh_shardings
 
   def restore_data_iterator_state(self, data_iterator):
     if (

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/generate_sdxl.py

@@ -249,9 +249,11 @@ def run(config):
         config=config,
         mesh=checkpoint_loader.mesh,
         weights_init_fn=weights_init_fn,
-        model_params=params.get("unet", None),
+        model_params=None,
         training=False,
     )
+    unet_state = unet_state.replace(params=params.get("unet", None))
+    unet_state = jax.device_put(unet_state, unet_state_shardings)
 
   vae_state, vae_state_shardings = checkpoint_loader.create_vae_state(
       pipeline, params, checkpoint_item_name="vae_state", is_training=False

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 jax.sharding import PositionalSharding
+from transformers import (
+  FlaxCLIPTextModel,
+  FlaxCLIPTextPreTrainedModel
+)
 from flax import struct
 from typing import (
     Callable,
@@ -315,7 +318,10 @@ def init_train_state(model, tx, weights_init_fn, params=None, training=True, eva
   Args: model_params, model, tx, training
   """
   if not params:
-    params = weights_init_fn(eval_only=eval_only)
+    if isinstance(model, FlaxCLIPTextModel) or isinstance(model, FlaxCLIPTextPreTrainedModel):
+      params = weights_init_fn()
+    else:
+      params = weights_init_fn(eval_only=eval_only)
   if training:
     state = train_state.TrainState.create(
         apply_fn=model.apply if hasattr(model, "apply") else model.__call__,

src/maxdiffusion/models/modeling_flax_pytorch_utils.py

@@ -163,6 +163,8 @@ def create_flax_params_from_pytorch_state(
       pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)
       flax_key_list = [*pt_tuple_key]
       flax_tensor = pt_tensor
+      if "lora" in flax_key_list:
+        flax_key_list[flax_key_list.index("lora")] = f"lora-{adapter_name}"
     else:
       flax_key_list = [*pt_tuple_key]
       if "text_encoder" in pt_tuple_key or "text_encoder_2" in pt_tuple_key: