feat(flux2klein): Faster Loading + Compliation

src/maxdiffusion/configs/base_flux2klein.yml

@@ -203,7 +203,7 @@ num_train_epochs: 1
 seed: 0
 output_dir: 'output/'
 output_name: "flux2klein_generated_image.png"
-per_device_batch_size: 1
+per_device_batch_size: 1.0
 
 warmup_steps_fraction: 0.1
 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.
@@ -231,6 +231,7 @@ do_classifier_free_guidance: True
 guidance_scale: 4.0
 guidance_rescale: 0.0
 num_inference_steps: 4
+num_reps: 1
 save_final_checkpoint: False
 
 # SDXL Lightning parameters

src/maxdiffusion/configs/base_flux2klein_9B.yml

@@ -203,7 +203,7 @@ num_train_epochs: 1
 seed: 0
 output_dir: 'output/'
 output_name: "flux2klein_generated_image.png"
-per_device_batch_size: 1
+per_device_batch_size: 1.0
 
 warmup_steps_fraction: 0.1
 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.
@@ -231,6 +231,7 @@ do_classifier_free_guidance: True
 guidance_scale: 4.0
 guidance_rescale: 0.0
 num_inference_steps: 4
+num_reps: 1
 save_final_checkpoint: False
 
 # SDXL Lightning parameters

src/maxdiffusion/models/flux/util.py

@@ -300,17 +300,17 @@ def unpack_latents(latents, batch_size, num_channels_latents, height, width):
   Unpacks packed sequence of shape (batch_size, (height//16)*(width//16), channels*4)
   back to the unpacked spatial grid shape (batch_size, channels, height//8, width//8).
   """
-  import numpy as np
+  import jax.numpy as jnp
 
   h_latent = height // 8
   w_latent = width // 8
 
   # 1. Reshape to split spatial grid and packed channel blocks
-  latents = np.reshape(latents, (batch_size, h_latent // 2, w_latent // 2, num_channels_latents, 2, 2))
+  latents = jnp.reshape(latents, (batch_size, h_latent // 2, w_latent // 2, num_channels_latents, 2, 2))
   # 2. Permute dimensions back to unpacked order
-  latents = np.transpose(latents, (0, 3, 1, 4, 2, 5))
+  latents = jnp.transpose(latents, (0, 3, 1, 4, 2, 5))
   # 3. Flatten back to 4D unpacked latent shape
-  latents = np.reshape(latents, (batch_size, num_channels_latents, h_latent, w_latent))
+  latents = jnp.reshape(latents, (batch_size, num_channels_latents, h_latent, w_latent))
   return latents
 
 
@@ -398,19 +398,22 @@ def cast_dict_to_bfloat16_inplace(d, device=None, exclude_keywords=None, parent_
       is_excluded = exclude_keywords and any(kw.lower() in current_key.lower() for kw in exclude_keywords)
       target_dtype = jnp.float32 if is_excluded else jnp.bfloat16
 
-      d[k] = v.astype(target_dtype)
-      if hasattr(d[k], "block_until_ready"):
-        d[k].block_until_ready()
-      del v
-      gc.collect()
+      if v.dtype != target_dtype:
+        d[k] = v.astype(target_dtype)
+        if hasattr(d[k], "block_until_ready"):
+          d[k].block_until_ready()
+        del v
+        gc.collect()
 
 
 # -----------------------------------------------------------------------------
 # Safetensors Weight Loader & Key Converter Functions
 # -----------------------------------------------------------------------------
 
 
-def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, num_single_layers):
+def load_and_convert_flux_klein_weights(
+    safetensors_path, params, num_double_layers, num_single_layers, dtype=None, pt_state_dict=None
+):
   """
   Loads weights from safetensors via zero-copy safetensors.numpy and converts them to JAX parameter dictionary.
   Supports dynamic layer counts (double and single stream blocks) and sharded safetensors directories.
@@ -422,28 +425,30 @@ def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_lay
   import os
   import gc
 
-  pt_state_dict = {}
-  if os.path.isdir(safetensors_path):
-    shards = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
-    max_logging.log(f"Loading sharded weights from directory: {safetensors_path} (Found {len(shards)} shards)...")
-    for shard in sorted(shards):
-      max_logging.log(f"Loading shard: {shard}...")
-      pt_state_dict.update(load_file(shard))
-  else:
-    max_logging.log(f"Loading weights from: {safetensors_path}")
-    pt_state_dict = load_file(safetensors_path)
+  if pt_state_dict is None:
+    pt_state_dict = {}
+    if os.path.isdir(safetensors_path):
+      shards = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
+      max_logging.log(f"Loading sharded weights from directory: {safetensors_path} (Found {len(shards)} shards)...")
+      for shard in sorted(shards):
+        max_logging.log(f"Loading shard: {shard}...")
+        pt_state_dict.update(load_file(shard))
+    else:
+      max_logging.log(f"Loading weights from: {safetensors_path}")
+      pt_state_dict = load_file(safetensors_path)
 
   max_logging.log("Mapping weights to JAX parameters...")
 
   expected_pytree = jax.tree_util.tree_map(lambda leaf: leaf, params)
 
   first_leaf = jax.tree_util.tree_leaves(params)[0]
-  target_dtype = first_leaf.dtype
+  target_dtype = dtype if dtype is not None else first_leaf.dtype
 
-  def convert_and_transpose_tensor(tensor, transpose=False):
+  def convert_and_transpose_tensor(tensor, transpose=False, is_norm=False):
     if transpose and len(tensor.shape) == 2:
       tensor = tensor.T
-    return jnp.array(tensor, dtype=target_dtype)
+    leaf_dtype = jnp.float32 if is_norm else target_dtype
+    return jnp.array(tensor, dtype=leaf_dtype)
 
   # Global layers
   params["context_embedder"]["kernel"] = convert_and_transpose_tensor(
@@ -562,21 +567,28 @@ def convert_and_transpose_tensor(tensor, transpose=False):
   return params
 
 
-def load_and_convert_vae_weights(safetensors_path, jax_params):
+def load_and_convert_vae_weights(safetensors_path, jax_params, dtype=None, pt_state_dict=None):
   """Loads VAE weights from safetensors via zero-copy safetensors.numpy, maps them to JAX, and extracts BN stats."""
   from safetensors.numpy import load_file
   import flax
   import jax.numpy as jnp
 
-  max_logging.log(f"Loading VAE weights from: {safetensors_path}")
-  pt_state_dict = load_file(safetensors_path)
-
-  def get_pytorch_weight_tensor(key):
-    return pt_state_dict[key]
+  if pt_state_dict is None:
+    max_logging.log(f"Loading VAE weights from: {safetensors_path}")
+    pt_state_dict = load_file(safetensors_path)
 
   # Unfreeze JAX params so we can load the weights
   jax_params = flax.core.unfreeze(jax_params)
 
+  first_leaf = jax.tree_util.tree_leaves(jax_params)[0]
+  target_dtype = dtype if dtype is not None else first_leaf.dtype
+
+  def get_pytorch_weight_tensor(key, dtype_val=target_dtype):
+    tensor = pt_state_dict[key]
+    is_norm = any(kw in key.lower() for kw in ("norm", "layernorm", "rmsnorm", "groupnorm"))
+    leaf_dtype = jnp.float32 if is_norm else dtype_val
+    return jnp.array(tensor, dtype=leaf_dtype)
+
   # Map weights
   max_logging.log("Mapping VAE decoder weights to JAX parameters...")
 

src/maxdiffusion/models/resnet_flax.py

@@ -57,9 +57,8 @@ def setup(self):
   @nn.compact
   def __call__(self, hidden_states):
     batch, height, width, channels = hidden_states.shape
-    hidden_states = jax.image.resize(
-        hidden_states, shape=(batch, height * 2, width * 2, channels), method="nearest", precision=self.precision
-    )
+    hidden_states = jnp.broadcast_to(hidden_states[:, :, None, :, None, :], (batch, height, 2, width, 2, channels))
+    hidden_states = jnp.reshape(hidden_states, (batch, height * 2, width * 2, channels))
 
     hidden_states = nn.with_logical_constraint(hidden_states, ("conv_batch", "height", "keep_2", "out_channels"))
 

src/maxdiffusion/models/vae_flax.py

@@ -87,11 +87,8 @@ def setup(self):
 
   def __call__(self, hidden_states):
     batch, height, width, channels = hidden_states.shape
-    hidden_states = jax.image.resize(
-        hidden_states,
-        shape=(batch, height * 2, width * 2, channels),
-        method="nearest",
-    )
+    hidden_states = jnp.broadcast_to(hidden_states[:, :, None, :, None, :], (batch, height, 2, width, 2, channels))
+    hidden_states = jnp.reshape(hidden_states, (batch, height * 2, width * 2, channels))
     hidden_states = self.conv(hidden_states)
     return hidden_states