adds causvid

Author: Juan Acevedo

src/maxdiffusion/configs/base_wan_14b.yml

@@ -28,8 +28,9 @@ save_config_to_gcs: False
 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: 'lightx2v/Wan2.1-T2V-14B-CausVid'
+transformer_pretrained_model_name_or_path: ''
 
 unet_checkpoint: ''
 revision: ''

src/maxdiffusion/models/wan/wan_utils.py

@@ -8,6 +8,8 @@
 from flax.traverse_util import unflatten_dict, flatten_dict
 from ..modeling_flax_pytorch_utils import (rename_key, rename_key_and_reshape_tensor, torch2jax, validate_flax_state_dict)
 
+CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH = "lightx2v/Wan2.1-T2V-14B-CausVid"
+
 
 def _tuple_str_to_int(in_tuple):
   out_list = []
@@ -25,13 +27,14 @@ def rename_for_nnx(key):
     new_key = key[:-1] + ("scale",)
   return new_key
 
+
 def load_causvid_transformer(pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True):
   device = jax.devices(device)[0]
   with jax.default_device(device):
     if hf_download:
       ckpt_shard_path = hf_hub_download(pretrained_model_name_or_path, filename="causal_model.pt")
       loaded_state_dict = torch.load(ckpt_shard_path)
-      
+
       tensors = {}
       flax_state_dict = {}
       cpu = jax.local_devices(backend="cpu")[0]
@@ -77,13 +80,15 @@ def load_causvid_transformer(pretrained_model_name_or_path: str, eval_shapes: di
       jax.clear_caches()
       return flax_state_dict
 
+
 def load_wan_transformer(pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True):
-  
-  if "CausVid" in pretrained_model_name_or_path:
+
+  if pretrained_model_name_or_path == CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH:
     return load_causvid_transformer(pretrained_model_name_or_path, eval_shapes, device, hf_download)
   else:
     return load_base_wan_transformer(pretrained_model_name_or_path, eval_shapes, device, hf_download)
 
+
 def load_base_wan_transformer(pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True):
   device = jax.devices(device)[0]
   with jax.default_device(device):

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -71,9 +71,7 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
     return wan_transformer
 
   # 1. Load config.
-  wan_config = WanModel.load_config(
-    config.pretrained_model_name_or_path,
-    subfolder="transformer")
+  wan_config = WanModel.load_config(config.pretrained_model_name_or_path, subfolder="transformer")
   wan_config["mesh"] = mesh
   wan_config["dtype"] = config.activations_dtype
   wan_config["weights_dtype"] = config.weights_dtype
@@ -97,9 +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 or config.pretrained_model_name_or_path,
-      params, "cpu")
+  params = load_wan_transformer(config.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

@@ -25,6 +25,7 @@
 import yaml
 from . import max_logging
 from . import max_utils
+from .models.wan.wan_utils import CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH
 
 
 def string_to_bool(s: str) -> bool:
@@ -102,6 +103,7 @@ def __init__(self, argv: list[str], **kwargs):
       jax.config.update("jax_compilation_cache_dir", raw_keys["jax_cache_dir"])
 
     _HyperParameters.user_init(raw_keys)
+    _HyperParameters.wan_init(raw_keys)
     self.keys = raw_keys
     for k in sorted(raw_keys.keys()):
       max_logging.log(f"Config param {k}: {raw_keys[k]}")
@@ -110,6 +112,22 @@ def _load_kwargs(self, argv: list[str]):
     args_dict = dict(a.split("=", 1) for a in argv[2:])
     return args_dict
 
+  @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")
+
   @staticmethod
   def user_init(raw_keys):
     """Transformations between the config data and configs used at runtime"""