transformer weight loading bug with scan layers = false

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -144,10 +144,7 @@ def get_key_and_value(pt_tuple_key, tensor, flax_state_dict, random_flax_state_d
 
   flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, scan_layers)
 
-  # Transpose back caption projections for LTX-2.3 as they are already in JAX format or shouldn't be transposed
-  if ("caption_projection" in flax_key or "audio_caption_projection" in flax_key) and "timestep_embedder" not in flax_key:
-    if "kernel" in flax_key and flax_tensor.ndim == 2:
-      flax_tensor = flax_tensor.T
+
 
   flax_key_str = [str(k) for k in flax_key]
 

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -124,8 +124,8 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict):
         "audio_attention_head_dim": 64,
         "audio_cross_attention_dim": 4096,
         "num_layers": 48,
-        "caption_channels": 8192,
-        "audio_caption_channels": 4096,
+        "caption_channels": 4096,
+        "audio_caption_channels": 2048,
     }
   else:
     ltx2_config = LTX2VideoTransformer3DModel.load_config(config.pretrained_model_name_or_path, subfolder=subfolder)