Modified wan_utils.py, dim errors corrected

Author: prishajain1

src/maxdiffusion/models/wan/wan_utils.py

@@ -259,7 +259,8 @@ def load_base_wan_transformer(
         renamed_pt_key = renamed_pt_key.replace("net.0", "net_0")
         renamed_pt_key = renamed_pt_key.replace("net.2", "net_2")
         renamed_pt_key = renamed_pt_key.replace("norm1", "norm1.layer_norm")
-        renamed_pt_key = renamed_pt_key.replace("norm2", "norm2.layer_norm")
+        if "norm1" in renamed_pt_key or "norm2" in renamed_pt_key:
+          renamed_pt_key = renamed_pt_key.replace("kernel", "scale")
       renamed_pt_key = renamed_pt_key.replace("blocks_", "blocks.")
       renamed_pt_key = renamed_pt_key.replace(".scale_shift_table", ".adaln_scale_shift_table")
       renamed_pt_key = renamed_pt_key.replace("to_out_0", "proj_attn")

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -112,15 +112,17 @@ def prepare_latents(
         if last_image is None:
             mask_lat_size = mask_lat_size.at[:, :, 1:, :, :].set(0)
         else:
-            mask_lat_size = mask_lat_size.at[:, :, 1:-1, :, :].set(0)
+            mask_lat_size = mask_lat_size.at[:, :, 1:-1, :, :].set(0)     
         first_frame_mask = mask_lat_size[:, :, 0:1]
         first_frame_mask = jnp.repeat(first_frame_mask, self.vae_scale_factor_temporal, axis=2)
         mask_lat_size = jnp.concatenate([first_frame_mask, mask_lat_size[:, :, 1:]], axis=2)
         mask_lat_size = mask_lat_size.reshape(
             batch_size, -1, self.vae_scale_factor_temporal, latent_height, latent_width
         )
         mask_lat_size = jnp.swapaxes(mask_lat_size, 1, 2)
-        condition = jnp.concatenate([mask_lat_size, latent_condition], axis=1)
+        mask_lat_size = jnp.transpose(mask_lat_size, (0, 2, 3, 4, 1))
+        condition = jnp.concatenate([mask_lat_size, latent_condition], axis=-1)
+
         return latents, condition, None