Modified wan_utils.py, chnaged added_kv_proj_dim to 5120, errors corrected

Author: prishajain1

src/maxdiffusion/models/wan/wan_utils.py

@@ -255,6 +255,10 @@ 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")
       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

@@ -108,7 +108,7 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
   else:
     wan_config = WanModel.load_config(config.pretrained_model_name_or_path, subfolder=subfolder)
   if config.model_type == "I2V":
-    wan_config["added_kv_proj_dim"] = 1024
+    wan_config["added_kv_proj_dim"] = 5120
   wan_config["mesh"] = mesh
   wan_config["dtype"] = config.activations_dtype
   wan_config["weights_dtype"] = config.weights_dtype
@@ -529,7 +529,7 @@ def prepare_latents_i2v_base(
       video_condition = video_condition.astype(vae_dtype)
 
       with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-          encoded_output = self.vae.encode(video_condition, self.vae_cache)[0]
+          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)