Wan training: Add checks for 'shape' and 'dtype' attributes during checkpoint loading.

Co-authored-by: martinarroyo <martinarroyo@google.com>

Author: ninatu

src/maxdiffusion/checkpointing/wan_checkpointer_2_1.py

@@ -45,9 +45,12 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
     state = metadatas.wan_state
 
     def add_sharding_to_struct(leaf_struct, sharding):
-      return jax.ShapeDtypeStruct(
-          shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding
-      )
+      struct = ocp.utils.to_shape_dtype_struct(leaf_struct)
+      if hasattr(struct, "shape") and hasattr(struct, "dtype"):
+        return jax.ShapeDtypeStruct(
+            shape=struct.shape, dtype=struct.dtype, sharding=sharding
+        )
+      return struct
 
     target_shardings = jax.tree_util.tree_map(
         lambda x: replicated_sharding, state

src/maxdiffusion/checkpointing/wan_checkpointer_i2v_2p1.py

@@ -45,9 +45,12 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
     state = metadatas.wan_state
 
     def add_sharding_to_struct(leaf_struct, sharding):
-      return jax.ShapeDtypeStruct(
-          shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding
-      )
+      struct = ocp.utils.to_shape_dtype_struct(leaf_struct)
+      if hasattr(struct, "shape") and hasattr(struct, "dtype"):
+        return jax.ShapeDtypeStruct(
+            shape=struct.shape, dtype=struct.dtype, sharding=sharding
+        )
+      return struct
 
     target_shardings = jax.tree_util.tree_map(
         lambda x: replicated_sharding, state

src/maxdiffusion/checkpointing/wan_vace_checkpointer_2_1.py

@@ -43,9 +43,12 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
     state = metadatas.wan_state
 
     def add_sharding_to_struct(leaf_struct, sharding):
-      return jax.ShapeDtypeStruct(
-          shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding
-      )
+      struct = ocp.utils.to_shape_dtype_struct(leaf_struct)
+      if hasattr(struct, "shape") and hasattr(struct, "dtype"):
+        return jax.ShapeDtypeStruct(
+            shape=struct.shape, dtype=struct.dtype, sharding=sharding
+        )
+      return struct
 
     target_shardings = jax.tree_util.tree_map(
         lambda x: replicated_sharding, state