transformer weight loading bug with scan layers = false

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -907,6 +907,8 @@ def init_block(rngs):
             norm_elementwise_affine=self.norm_elementwise_affine,
             norm_eps=self.norm_eps,
             rope_type=self.rope_type,
+            gated_attn=self.gated_attn,
+            cross_attn_mod=self.cross_attn_mod,
             dtype=self.dtype,
             weights_dtype=self.weights_dtype,
             mesh=self.mesh,

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -187,7 +187,11 @@ def create_model(rngs: nnx.Rngs, ltx2_config: dict):
   for path, val in flax.traverse_util.flatten_dict(params).items():
     if restored_checkpoint:
       path = path[:-1]
-    sharding = logical_state_sharding[path].value
+    try:
+      sharding = logical_state_sharding[path].value
+    except KeyError:
+      path_str = tuple(str(k) for k in path)
+      sharding = logical_state_sharding[path_str].value
     state[path].value = device_put_replicated(val, sharding)
   state = nnx.from_flat_state(state)