feat(flux2klein): eliminate CPU parameter recasting during single-pass safetensors conversion

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -360,20 +360,12 @@ def unbox_fn(x):
       max_logging.log(f" -> [SUB-TIMING 1/3] PyTree unboxing template setup: {time.time() - t_sub0:.2f}s")
       t_sub1 = time.time()
 
-      params = load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, depth)
-      vae_params, vae_bn_mean, vae_bn_std = load_and_convert_vae_weights(vae_safetensors_path, vae_params)
-      qwen3_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params, qwen3_config)
-      max_logging.log(f" -> [SUB-TIMING 2/3] Safetensors loading & key mapping: {time.time() - t_sub1:.2f}s")
-
-      t_sub2 = time.time()
-      if config.weights_dtype == "bfloat16":
-        max_logging.log("Casting JAX parameters to bfloat16 in-place...")
-        cast_dict_to_bfloat16_inplace(params, exclude_keywords=("norm",))
-        cast_dict_to_bfloat16_inplace(vae_params, exclude_keywords=("norm",))
-        cast_dict_to_bfloat16_inplace(qwen3_params, exclude_keywords=("norm",))
-        vae_bn_mean = vae_bn_mean.astype(jnp.bfloat16)
-        vae_bn_std = vae_bn_std.astype(jnp.bfloat16)
-        max_logging.log(f" -> [SUB-TIMING 2b/3] bfloat16 in-place casting: {time.time() - t_sub2:.2f}s")
+      weight_dtype = jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32
+
+      params = load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, depth, dtype=weight_dtype)
+      vae_params, vae_bn_mean, vae_bn_std = load_and_convert_vae_weights(vae_safetensors_path, vae_params, dtype=weight_dtype)
+      qwen3_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params, qwen3_config, dtype=weight_dtype)
+      max_logging.log(f" -> [SUB-TIMING 2/3] Safetensors loading & key mapping (in target dtype): {time.time() - t_sub1:.2f}s")
 
       params = flax.core.freeze(params)
       vae_params = flax.core.freeze(vae_params)

src/maxdiffusion/models/flux/util.py

@@ -411,7 +411,7 @@ def cast_dict_to_bfloat16_inplace(d, device=None, exclude_keywords=None, parent_
 # -----------------------------------------------------------------------------
 
 
-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):
   """
   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.
@@ -439,12 +439,13 @@ def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_lay
   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(
@@ -563,7 +564,7 @@ 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):
   """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
@@ -576,11 +577,13 @@ def load_and_convert_vae_weights(safetensors_path, jax_params):
   jax_params = flax.core.unfreeze(jax_params)
 
   first_leaf = jax.tree_util.tree_leaves(jax_params)[0]
-  target_dtype = first_leaf.dtype
+  target_dtype = dtype if dtype is not None else first_leaf.dtype
 
-  def get_pytorch_weight_tensor(key, dtype=target_dtype):
+  def get_pytorch_weight_tensor(key, dtype_val=target_dtype):
     tensor = pt_state_dict[key]
-    return jnp.array(tensor, dtype=dtype)
+    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/qwen3_flax.py

@@ -626,7 +626,7 @@ def __call__(
 # -----------------------------------------------------------------------------
 
 
-def load_and_convert_qwen3_weights(safetensors_path: str, jax_params: dict, config: FlaxQwen3Config) -> dict:
+def load_and_convert_qwen3_weights(safetensors_path: str, jax_params: dict, config: FlaxQwen3Config, dtype=None) -> dict:
   """
   Loads weights from safetensors via zero-copy safetensors.numpy and converts them to JAX parameter dictionary.
   """
@@ -649,7 +649,7 @@ def load_and_convert_qwen3_weights(safetensors_path: str, jax_params: dict, conf
   max_logging.log("Safetensors weights loaded successfully. Starting JAX parameter mapping...")
 
   first_leaf = jax.tree_util.tree_leaves(jax_params)[0]
-  target_dtype = first_leaf.dtype
+  target_dtype = dtype if dtype is not None else first_leaf.dtype
 
   # Helper to transpose and cast weight directly to target_dtype
   def get_w(name: str, transpose: bool = True):
@@ -659,7 +659,9 @@ def get_w(name: str, transpose: bool = True):
     t = torch_weights[name]
     if len(t.shape) == 2 and transpose:
       t = t.T
-    return jnp.array(t, dtype=target_dtype)
+    is_norm = any(kw in name.lower() for kw in ("norm", "layernorm", "rmsnorm", "groupnorm"))
+    leaf_dtype = jnp.float32 if is_norm else target_dtype
+    return jnp.array(t, dtype=leaf_dtype)
 
   # Create mutable copy of JAX params to populate
   import flax