transformer file changed

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -115,6 +115,10 @@ def get_key_and_value(pt_tuple_key, tensor, flax_state_dict, random_flax_state_d
   # Also check 'weight' because rename_key might not have converted it to kernel if it wasn't a known Linear
   flax_key_str = [str(k) for k in flax_key]
 
+  # DEBUG: Check specific keys
+  if "norm_k" in flax_key_str or "audio_caption_projection" in flax_key_str:
+       print(f"DEBUG: get_key_and_value mapping: {pt_tuple_key} -> {flax_key_str}")
+  
   if flax_key_str[-1] in ["kernel", "weight"]:
        # Try replacing with scale and check if it exists in random_flax_state_dict
        temp_key_str = flax_key_str[:-1] + ["scale"]

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -807,7 +807,7 @@ def init_block(rngs):
     # 6. Output layers
     self.gradient_checkpoint = GradientCheckpointType.from_str(remat_policy)
     self.norm_out = nnx.LayerNorm(
-        inner_dim, epsilon=1e-6, use_scale=False, rngs=rngs, dtype=jnp.float32, param_dtype=jnp.float32
+        inner_dim, epsilon=1e-6, use_scale=False, use_bias=False, rngs=rngs, dtype=jnp.float32, param_dtype=jnp.float32
     )
     self.proj_out = nnx.Linear(
         inner_dim,
@@ -820,7 +820,7 @@ def init_block(rngs):
     )
 
     self.audio_norm_out = nnx.LayerNorm(
-        audio_inner_dim, epsilon=1e-6, use_scale=False, rngs=rngs, dtype=jnp.float32, param_dtype=jnp.float32
+        audio_inner_dim, epsilon=1e-6, use_scale=False, use_bias=False, rngs=rngs, dtype=jnp.float32, param_dtype=jnp.float32
     )
     self.audio_proj_out = nnx.Linear(
         audio_inner_dim,