Modified wan_utils.py, dim errors corrected

Author: prishajain1

src/maxdiffusion/models/wan/wan_utils.py

@@ -255,10 +255,11 @@ def load_base_wan_transformer(
     del flattened_dict
     for pt_key, tensor in tensors.items():
       renamed_pt_key = rename_key(pt_key)
-      renamed_pt_key = renamed_pt_key.replace("image_embedder.ff.net.0", "image_embedder.ff.net_0")
-      renamed_pt_key = renamed_pt_key.replace("image_embedder.ff.net.2", "image_embedder.ff.net_2")
-      if "image_embedder.norm1.scale" in renamed_pt_key:
-        renamed_pt_key = renamed_pt_key.replace("norm1.scale", "norm1.kernel")
+      if "image_embedder" in renamed_pt_key:
+        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")
       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.py

@@ -532,8 +532,8 @@ def prepare_latents_i2v_base(
           encoded_output = self.vae.encode(video_condition, self.vae_cache)[0].mode()
 
       # Normalize latents
-      latents_mean = jnp.array(self.vae.latents_mean).reshape(1, self.vae.z_dim, 1, 1, 1)
-      latents_std = 1.0 / jnp.array(self.vae.latents_std).reshape(1, self.vae.z_dim, 1, 1, 1)
+      latents_mean = jnp.array(self.vae.latents_mean).reshape(1, 1, 1, 1, self.vae.z_dim)
+      latents_std = 1.0 / jnp.array(self.vae.latents_std).reshape(1, 1, 1, 1, self.vae.z_dim)
       latent_condition = (encoded_output - latents_mean) * latents_std
       latent_condition = latent_condition.astype(dtype)