Rename cvt and get_w to descriptive helper function names in util.py

Author: amepas

src/maxdiffusion/models/flux/util.py

@@ -428,43 +428,47 @@ def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_lay
   first_leaf = jax.tree_util.tree_leaves(params)[0]
   target_dtype = first_leaf.dtype
 
-  def cvt(tensor, transpose=False):
+  def convert_and_transpose_tensor(tensor, transpose=False):
     if transpose and len(tensor.shape) == 2:
       tensor = tensor.T
     return jnp.array(tensor, dtype=target_dtype)
 
   # Global layers
-  params["context_embedder"]["kernel"] = cvt(pt_state_dict.pop("context_embedder.weight"), transpose=True)
-  params["x_embedder"]["kernel"] = cvt(pt_state_dict.pop("x_embedder.weight"), transpose=True)
-  params["double_stream_modulation_img"]["kernel"] = cvt(
+  params["context_embedder"]["kernel"] = convert_and_transpose_tensor(
+      pt_state_dict.pop("context_embedder.weight"), transpose=True
+  )
+  params["x_embedder"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("x_embedder.weight"), transpose=True)
+  params["double_stream_modulation_img"]["kernel"] = convert_and_transpose_tensor(
       pt_state_dict.pop("double_stream_modulation_img.linear.weight"), transpose=True
   )
-  params["double_stream_modulation_txt"]["kernel"] = cvt(
+  params["double_stream_modulation_txt"]["kernel"] = convert_and_transpose_tensor(
       pt_state_dict.pop("double_stream_modulation_txt.linear.weight"), transpose=True
   )
-  params["single_stream_modulation"]["kernel"] = cvt(
+  params["single_stream_modulation"]["kernel"] = convert_and_transpose_tensor(
       pt_state_dict.pop("single_stream_modulation.linear.weight"), transpose=True
   )
-  params["proj_out"]["kernel"] = cvt(pt_state_dict.pop("proj_out.weight"), transpose=True)
+  params["proj_out"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("proj_out.weight"), transpose=True)
 
   # norm_out
-  params["norm_out"]["linear"]["kernel"] = cvt(pt_state_dict.pop("norm_out.linear.weight"), transpose=True)
+  params["norm_out"]["linear"]["kernel"] = convert_and_transpose_tensor(
+      pt_state_dict.pop("norm_out.linear.weight"), transpose=True
+  )
 
   # time_text_embed (Timestep Embedding)
   if "time_guidance_embed.timestep_embedder.linear_1.weight" in pt_state_dict:
-    params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_1"]["kernel"] = cvt(
+    params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_1"]["kernel"] = convert_and_transpose_tensor(
         pt_state_dict.pop("time_guidance_embed.timestep_embedder.linear_1.weight"), transpose=True
     )
   if "time_guidance_embed.timestep_embedder.linear_1.bias" in pt_state_dict:
-    params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_1"]["bias"] = cvt(
+    params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_1"]["bias"] = convert_and_transpose_tensor(
         pt_state_dict.pop("time_guidance_embed.timestep_embedder.linear_1.bias")
     )
   if "time_guidance_embed.timestep_embedder.linear_2.weight" in pt_state_dict:
-    params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_2"]["kernel"] = cvt(
+    params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_2"]["kernel"] = convert_and_transpose_tensor(
         pt_state_dict.pop("time_guidance_embed.timestep_embedder.linear_2.weight"), transpose=True
     )
   if "time_guidance_embed.timestep_embedder.linear_2.bias" in pt_state_dict:
-    params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_2"]["bias"] = cvt(
+    params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_2"]["bias"] = convert_and_transpose_tensor(
         pt_state_dict.pop("time_guidance_embed.timestep_embedder.linear_2.bias")
     )
 
@@ -486,22 +490,34 @@ def cvt(tensor, transpose=False):
     jax_db["attn"]["e_qkv"]["kernel"] = jnp.array(np.concatenate([add_q, add_k, add_v], axis=1), dtype=target_dtype)
 
     # Projections out
-    jax_db["attn"]["i_proj"]["kernel"] = cvt(pt_state_dict.pop(prefix + "attn.to_out.0.weight"), transpose=True)
-    jax_db["attn"]["e_proj"]["kernel"] = cvt(pt_state_dict.pop(prefix + "attn.to_add_out.weight"), transpose=True)
+    jax_db["attn"]["i_proj"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "attn.to_out.0.weight"), transpose=True
+    )
+    jax_db["attn"]["e_proj"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "attn.to_add_out.weight"), transpose=True
+    )
 
     # Norm scales
-    jax_db["attn"]["query_norm"]["scale"] = cvt(pt_state_dict.pop(prefix + "attn.norm_q.weight"))
-    jax_db["attn"]["key_norm"]["scale"] = cvt(pt_state_dict.pop(prefix + "attn.norm_k.weight"))
-    jax_db["attn"]["encoder_query_norm"]["scale"] = cvt(pt_state_dict.pop(prefix + "attn.norm_added_q.weight"))
-    jax_db["attn"]["encoder_key_norm"]["scale"] = cvt(pt_state_dict.pop(prefix + "attn.norm_added_k.weight"))
+    jax_db["attn"]["query_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_q.weight"))
+    jax_db["attn"]["key_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_k.weight"))
+    jax_db["attn"]["encoder_query_norm"]["scale"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "attn.norm_added_q.weight")
+    )
+    jax_db["attn"]["encoder_key_norm"]["scale"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "attn.norm_added_k.weight")
+    )
 
     # SwiGLU MLPs
-    jax_db["ff"]["linear_in"]["kernel"] = cvt(pt_state_dict.pop(prefix + "ff.linear_in.weight"), transpose=True)
-    jax_db["ff"]["linear_out"]["kernel"] = cvt(pt_state_dict.pop(prefix + "ff.linear_out.weight"), transpose=True)
-    jax_db["ff_context"]["linear_in"]["kernel"] = cvt(
+    jax_db["ff"]["linear_in"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "ff.linear_in.weight"), transpose=True
+    )
+    jax_db["ff"]["linear_out"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "ff.linear_out.weight"), transpose=True
+    )
+    jax_db["ff_context"]["linear_in"]["kernel"] = convert_and_transpose_tensor(
         pt_state_dict.pop(prefix + "ff_context.linear_in.weight"), transpose=True
     )
-    jax_db["ff_context"]["linear_out"]["kernel"] = cvt(
+    jax_db["ff_context"]["linear_out"]["kernel"] = convert_and_transpose_tensor(
         pt_state_dict.pop(prefix + "ff_context.linear_out.weight"), transpose=True
     )
 
@@ -512,12 +528,16 @@ def cvt(tensor, transpose=False):
     s_prefix = f"single_transformer_blocks.{block_idx}."
 
     # Joint projections
-    jax_sb["linear1"]["kernel"] = cvt(pt_state_dict.pop(s_prefix + "attn.to_qkv_mlp_proj.weight"), transpose=True)
-    jax_sb["linear2"]["kernel"] = cvt(pt_state_dict.pop(s_prefix + "attn.to_out.weight"), transpose=True)
+    jax_sb["linear1"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(s_prefix + "attn.to_qkv_mlp_proj.weight"), transpose=True
+    )
+    jax_sb["linear2"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(s_prefix + "attn.to_out.weight"), transpose=True
+    )
 
     # Norm scales
-    jax_sb["attn"]["query_norm"]["scale"] = cvt(pt_state_dict.pop(s_prefix + "attn.norm_q.weight"))
-    jax_sb["attn"]["key_norm"]["scale"] = cvt(pt_state_dict.pop(s_prefix + "attn.norm_k.weight"))
+    jax_sb["attn"]["query_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(s_prefix + "attn.norm_q.weight"))
+    jax_sb["attn"]["key_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(s_prefix + "attn.norm_k.weight"))
 
   params = jax.tree_util.tree_map(
       lambda leaf: jnp.zeros(leaf.shape, dtype=leaf.dtype) if isinstance(leaf, jax.ShapeDtypeStruct) else leaf, params
@@ -537,7 +557,7 @@ def load_and_convert_vae_weights(safetensors_path, jax_params):
   max_logging.log(f"Loading VAE weights from: {safetensors_path}")
   pt_state_dict = load_file(safetensors_path)
 
-  def get_w(key):
+  def get_pytorch_weight_tensor(key):
     return pt_state_dict[key]
 
   # Unfreeze JAX params so we can load the weights
@@ -547,45 +567,49 @@ def get_w(key):
   max_logging.log("Mapping VAE decoder weights to JAX parameters...")
 
   # post_quant_conv
-  jax_params["post_quant_conv"]["kernel"] = jnp.array(get_w("post_quant_conv.weight").transpose(2, 3, 1, 0))
-  jax_params["post_quant_conv"]["bias"] = jnp.array(get_w("post_quant_conv.bias"))
+  jax_params["post_quant_conv"]["kernel"] = jnp.array(
+      get_pytorch_weight_tensor("post_quant_conv.weight").transpose(2, 3, 1, 0)
+  )
+  jax_params["post_quant_conv"]["bias"] = jnp.array(get_pytorch_weight_tensor("post_quant_conv.bias"))
 
   # decoder.conv_in
-  jax_params["decoder"]["conv_in"]["kernel"] = jnp.array(get_w("decoder.conv_in.weight").transpose(2, 3, 1, 0))
-  jax_params["decoder"]["conv_in"]["bias"] = jnp.array(get_w("decoder.conv_in.bias"))
+  jax_params["decoder"]["conv_in"]["kernel"] = jnp.array(
+      get_pytorch_weight_tensor("decoder.conv_in.weight").transpose(2, 3, 1, 0)
+  )
+  jax_params["decoder"]["conv_in"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.bias"))
 
   # decoder.mid_block
   # resnets
   for idx in [0, 1]:
     res_jax = jax_params["decoder"]["mid_block"][f"resnets_{idx}"]
     res_pt_prefix = f"decoder.mid_block.resnets.{idx}"
 
-    res_jax["norm1"]["scale"] = jnp.array(get_w(f"{res_pt_prefix}.norm1.weight"))
-    res_jax["norm1"]["bias"] = jnp.array(get_w(f"{res_pt_prefix}.norm1.bias"))
-    res_jax["conv1"]["kernel"] = jnp.array(get_w(f"{res_pt_prefix}.conv1.weight").transpose(2, 3, 1, 0))
-    res_jax["conv1"]["bias"] = jnp.array(get_w(f"{res_pt_prefix}.conv1.bias"))
+    res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.weight"))
+    res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm1.bias"))
+    res_jax["conv1"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.weight").transpose(2, 3, 1, 0))
+    res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv1.bias"))
 
-    res_jax["norm2"]["scale"] = jnp.array(get_w(f"{res_pt_prefix}.norm2.weight"))
-    res_jax["norm2"]["bias"] = jnp.array(get_w(f"{res_pt_prefix}.norm2.bias"))
-    res_jax["conv2"]["kernel"] = jnp.array(get_w(f"{res_pt_prefix}.conv2.weight").transpose(2, 3, 1, 0))
-    res_jax["conv2"]["bias"] = jnp.array(get_w(f"{res_pt_prefix}.conv2.bias"))
+    res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.weight"))
+    res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.norm2.bias"))
+    res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.weight").transpose(2, 3, 1, 0))
+    res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt_prefix}.conv2.bias"))
 
   # attentions
   attn_pt_prefix = "decoder.mid_block.attentions.0"
   attn_jax = jax_params["decoder"]["mid_block"]["attentions_0"]
 
-  attn_jax["group_norm"]["scale"] = jnp.array(get_w(f"{attn_pt_prefix}.group_norm.weight"))
-  attn_jax["group_norm"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.group_norm.bias"))
+  attn_jax["group_norm"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.weight"))
+  attn_jax["group_norm"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.group_norm.bias"))
 
-  attn_jax["query"]["kernel"] = jnp.array(get_w(f"{attn_pt_prefix}.to_q.weight").T)
-  attn_jax["query"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.to_q.bias"))
-  attn_jax["key"]["kernel"] = jnp.array(get_w(f"{attn_pt_prefix}.to_k.weight").T)
-  attn_jax["key"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.to_k.bias"))
-  attn_jax["value"]["kernel"] = jnp.array(get_w(f"{attn_pt_prefix}.to_v.weight").T)
-  attn_jax["value"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.to_v.bias"))
+  attn_jax["query"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.weight").T)
+  attn_jax["query"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_q.bias"))
+  attn_jax["key"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.weight").T)
+  attn_jax["key"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_k.bias"))
+  attn_jax["value"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.weight").T)
+  attn_jax["value"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_v.bias"))
 
-  attn_jax["proj_attn"]["kernel"] = jnp.array(get_w(f"{attn_pt_prefix}.to_out.0.weight").T)
-  attn_jax["proj_attn"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.to_out.0.bias"))
+  attn_jax["proj_attn"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.weight").T)
+  attn_jax["proj_attn"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{attn_pt_prefix}.to_out.0.bias"))
 
   # decoder.up_blocks
   for b_idx in range(4):
@@ -596,33 +620,37 @@ def get_w(key):
       res_jax = up_block_jax[f"resnets_{r_idx}"]
       res_pt = f"{up_block_pt}.resnets.{r_idx}"
 
-      res_jax["norm1"]["scale"] = jnp.array(get_w(f"{res_pt}.norm1.weight"))
-      res_jax["norm1"]["bias"] = jnp.array(get_w(f"{res_pt}.norm1.bias"))
-      res_jax["conv1"]["kernel"] = jnp.array(get_w(f"{res_pt}.conv1.weight").transpose(2, 3, 1, 0))
-      res_jax["conv1"]["bias"] = jnp.array(get_w(f"{res_pt}.conv1.bias"))
+      res_jax["norm1"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.weight"))
+      res_jax["norm1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm1.bias"))
+      res_jax["conv1"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.weight").transpose(2, 3, 1, 0))
+      res_jax["conv1"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv1.bias"))
 
-      res_jax["norm2"]["scale"] = jnp.array(get_w(f"{res_pt}.norm2.weight"))
-      res_jax["norm2"]["bias"] = jnp.array(get_w(f"{res_pt}.norm2.bias"))
-      res_jax["conv2"]["kernel"] = jnp.array(get_w(f"{res_pt}.conv2.weight").transpose(2, 3, 1, 0))
-      res_jax["conv2"]["bias"] = jnp.array(get_w(f"{res_pt}.conv2.bias"))
+      res_jax["norm2"]["scale"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.weight"))
+      res_jax["norm2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.norm2.bias"))
+      res_jax["conv2"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.weight").transpose(2, 3, 1, 0))
+      res_jax["conv2"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv2.bias"))
 
       shortcut_key = f"{res_pt}.conv_shortcut.weight"
       if shortcut_key in pt_state_dict:
-        res_jax["conv_shortcut"]["kernel"] = jnp.array(get_w(shortcut_key).transpose(2, 3, 1, 0))
-        res_jax["conv_shortcut"]["bias"] = jnp.array(get_w(f"{res_pt}.conv_shortcut.bias"))
+        res_jax["conv_shortcut"]["kernel"] = jnp.array(get_pytorch_weight_tensor(shortcut_key).transpose(2, 3, 1, 0))
+        res_jax["conv_shortcut"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{res_pt}.conv_shortcut.bias"))
 
     if b_idx < 3:
       upsampler_jax = up_block_jax["upsamplers_0"]
       upsampler_pt = f"{up_block_pt}.upsamplers.0"
 
-      upsampler_jax["conv"]["kernel"] = jnp.array(get_w(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0))
-      upsampler_jax["conv"]["bias"] = jnp.array(get_w(f"{upsampler_pt}.conv.bias"))
+      upsampler_jax["conv"]["kernel"] = jnp.array(
+          get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0)
+      )
+      upsampler_jax["conv"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{upsampler_pt}.conv.bias"))
 
   # decoder.conv_norm_out & conv_out
-  jax_params["decoder"]["conv_norm_out"]["scale"] = jnp.array(get_w("decoder.conv_norm_out.weight"))
-  jax_params["decoder"]["conv_norm_out"]["bias"] = jnp.array(get_w("decoder.conv_norm_out.bias"))
-  jax_params["decoder"]["conv_out"]["kernel"] = jnp.array(get_w("decoder.conv_out.weight").transpose(2, 3, 1, 0))
-  jax_params["decoder"]["conv_out"]["bias"] = jnp.array(get_w("decoder.conv_out.bias"))
+  jax_params["decoder"]["conv_norm_out"]["scale"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.weight"))
+  jax_params["decoder"]["conv_norm_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.bias"))
+  jax_params["decoder"]["conv_out"]["kernel"] = jnp.array(
+      get_pytorch_weight_tensor("decoder.conv_out.weight").transpose(2, 3, 1, 0)
+  )
+  jax_params["decoder"]["conv_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.bias"))
 
   jax_params = jax.tree_util.tree_map(
       lambda leaf: jnp.zeros(leaf.shape, dtype=leaf.dtype) if isinstance(leaf, jax.ShapeDtypeStruct) else leaf, jax_params
@@ -632,8 +660,8 @@ def get_w(key):
 
   # Extract Batch Normalization running stats
   max_logging.log("Extracting VAE Batch Normalization running stats...")
-  bn_mean = jnp.array(get_w("bn.running_mean")).reshape(1, -1, 1, 1)
-  bn_var = jnp.array(get_w("bn.running_var")).reshape(1, -1, 1, 1)
+  bn_mean = jnp.array(get_pytorch_weight_tensor("bn.running_mean")).reshape(1, -1, 1, 1)
+  bn_var = jnp.array(get_pytorch_weight_tensor("bn.running_var")).reshape(1, -1, 1, 1)
   batch_norm_eps = 0.0001
   bn_std = jnp.sqrt(bn_var + batch_norm_eps)