cleaned up implementation, some potential merge conflicts to be resolved in next commit

Author: amepas

src/maxdiffusion/models/flux/util.py

@@ -302,16 +302,20 @@ def prepare_text_ids(batch_size, seq_len):
 # -----------------------------------------------------------------------------
 
 
-def cast_dict_to_bfloat16_inplace(d):
+def cast_dict_to_bfloat16_inplace(d, device=None):
   """Casts a nested dictionary of JAX/numpy arrays to bfloat16 in-place, freeing memory immediately."""
   import gc
   import jax.numpy as jnp
+  import jax
 
   for k, v in list(d.items()):
     if isinstance(v, dict):
-      cast_dict_to_bfloat16_inplace(v)
+      cast_dict_to_bfloat16_inplace(v, device=device)
     elif hasattr(v, "astype"):
-      d[k] = jnp.array(v, dtype=jnp.bfloat16)
+      if device is not None:
+        d[k] = jax.device_put(jnp.array(v, dtype=jnp.bfloat16), device=device)
+      else:
+        d[k] = jnp.array(v, dtype=jnp.bfloat16)
       if hasattr(d[k], "block_until_ready"):
         d[k].block_until_ready()
       del v
@@ -358,8 +362,8 @@ def cvt(tensor, transpose=False):
     return jnp.array(tensor.to(torch.float32).cpu().numpy(), dtype=target_dtype)
 
   # Global layers
-  params["txt_in"]["kernel"] = cvt(pt_state_dict.pop("context_embedder.weight"), transpose=True)
-  params["img_in"]["kernel"] = cvt(pt_state_dict.pop("x_embedder.weight"), transpose=True)
+  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(
       pt_state_dict.pop("double_stream_modulation_img.linear.weight"), transpose=True
   )
@@ -372,7 +376,7 @@ def cvt(tensor, transpose=False):
   params["proj_out"]["kernel"] = cvt(pt_state_dict.pop("proj_out.weight"), transpose=True)
 
   # norm_out
-  params["norm_out"]["Dense_0"]["kernel"] = cvt(pt_state_dict.pop("norm_out.linear.weight"), transpose=True)
+  params["norm_out"]["linear"]["kernel"] = cvt(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:
@@ -420,10 +424,10 @@ def cvt(tensor, transpose=False):
     jax_db["attn"]["encoder_key_norm"]["scale"] = cvt(pt_state_dict.pop(prefix + "attn.norm_added_k.weight"))
 
     # SwiGLU MLPs
-    jax_db["img_mlp"]["linear_in"]["kernel"] = cvt(pt_state_dict.pop(prefix + "ff.linear_in.weight"), transpose=True)
-    jax_db["img_mlp"]["linear_out"]["kernel"] = cvt(pt_state_dict.pop(prefix + "ff.linear_out.weight"), transpose=True)
-    jax_db["txt_mlp"]["linear_in"]["kernel"] = cvt(pt_state_dict.pop(prefix + "ff_context.linear_in.weight"), transpose=True)
-    jax_db["txt_mlp"]["linear_out"]["kernel"] = cvt(
+    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(pt_state_dict.pop(prefix + "ff_context.linear_in.weight"), transpose=True)
+    jax_db["ff_context"]["linear_out"]["kernel"] = cvt(
         pt_state_dict.pop(prefix + "ff_context.linear_out.weight"), transpose=True
     )