fix unit tests.

Author: Juan Acevedo

src/maxdiffusion/configs/base_wan_14b.yml

@@ -30,7 +30,7 @@ log_period: 100
 pretrained_model_name_or_path: 'Wan-AI/Wan2.1-T2V-14B-Diffusers'
 
 # Overrides the transformer from pretrained_model_name_or_path
-transformer_pretrained_model_name_or_path: ''
+wan_transformer_pretrained_model_name_or_path: ''
 
 unet_checkpoint: ''
 revision: ''

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -95,7 +95,7 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
   # 4. Load pretrained weights and move them to device using the state shardings from (3) above.
   # This helps with loading sharded weights directly into the accelerators without fist copying them
   # all to one device and then distributing them, thus using low HBM memory.
-  params = load_wan_transformer(config.transformer_pretrained_model_name_or_path, params, "cpu")
+  params = load_wan_transformer(config.wan_transformer_pretrained_model_name_or_path, params, "cpu")
   params = jax.tree_util.tree_map(lambda x: x.astype(config.weights_dtype), params)
   for path, val in flax.traverse_util.flatten_dict(params).items():
     sharding = logical_state_sharding[path].value

src/maxdiffusion/pyconfig.py

@@ -114,19 +114,20 @@ def _load_kwargs(self, argv: list[str]):
 
   @staticmethod
   def wan_init(raw_keys):
-    transformer_pretrained_model_name_or_path = raw_keys["transformer_pretrained_model_name_or_path"]
-    if transformer_pretrained_model_name_or_path == "":
-      raw_keys["transformer_pretrained_model_name_or_path"] = raw_keys["pretrained_model_name_or_path"]
-    elif transformer_pretrained_model_name_or_path == CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH:
-      # Set correct parameters for CausVid in case of user error.
-      raw_keys["guidance_scale"] = 1.0
-      num_inference_steps = raw_keys["num_inference_steps"]
-      if num_inference_steps > 10:
-        max_logging.log(
-            f"Warning: Try setting num_inference_steps to less than 8 steps when using CausVid, currently you are setting {num_inference_steps} steps."
-        )
-    else:
-      raise ValueError(f"{transformer_pretrained_model_name_or_path} transformer model is not supported for Wan 2.1")
+    if "wan_transformer_pretrained_model_name_or_path" in raw_keys:
+      transformer_pretrained_model_name_or_path = raw_keys["wan_transformer_pretrained_model_name_or_path"]
+      if transformer_pretrained_model_name_or_path == "":
+        raw_keys["wan_transformer_pretrained_model_name_or_path"] = raw_keys["pretrained_model_name_or_path"]
+      elif transformer_pretrained_model_name_or_path == CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH:
+        # Set correct parameters for CausVid in case of user error.
+        raw_keys["guidance_scale"] = 1.0
+        num_inference_steps = raw_keys["num_inference_steps"]
+        if num_inference_steps > 10:
+          max_logging.log(
+              f"Warning: Try setting num_inference_steps to less than 8 steps when using CausVid, currently you are setting {num_inference_steps} steps."
+          )
+      else:
+        raise ValueError(f"{transformer_pretrained_model_name_or_path} transformer model is not supported for Wan 2.1")
 
   @staticmethod
   def user_init(raw_keys):