fix and ltx2 backward compatibility

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -379,6 +379,8 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
     logical_state_spec = nnx.get_partition_spec(state)
     logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, config.logical_axis_rules)
     logical_state_sharding = dict(nnx.to_flat_state(logical_state_sharding))
+    params = state.to_pure_dict()
+    state = dict(nnx.to_flat_state(state))
     filename = "ltx-2.3-22b-dev.safetensors" if getattr(config, "model_name", "") == "ltx2.3" else None
     params = load_connectors_weights(
         config.pretrained_model_name_or_path,