added test

Author: Serenagu525

src/maxdiffusion/configs/ltx_video.yml

@@ -22,19 +22,6 @@ weights_dtype: 'bfloat16'
 activations_dtype: 'bfloat16'
 
 
-run_name: ''
-output_dir: 'ltx-video-output'
-save_config_to_gcs: False
-
-#hardware
-hardware: 'tpu'
-skip_jax_distributed_system: False
-
-jax_cache_dir: ''
-weights_dtype: 'bfloat16'
-activations_dtype: 'bfloat16'
-
-
 run_name: ''
 output_dir: 'ltx-video-output'
 save_config_to_gcs: False
@@ -78,4 +65,4 @@ per_device_batch_size: 1
 compile_topology_num_slices: -1 
 quantization_local_shard_count: -1
 jit_initializers: True 
-enable_single_replica_ckpt_restoring: False
+enable_single_replica_ckpt_restoring: False

src/maxdiffusion/models/ltx_video/transformers/attention.py

@@ -622,6 +622,12 @@ def partial_flash_attention(q, k, v, q_segment_ids, kv_segment_ids):
         raise ValueError(f"Expected mask with 2 dims, got {q_segment_ids.ndim}.")
       # Based on: ("activation_kv_batch", "activation_kv_heads", "activation_length", "activation_kv_head_dim")
       # Computation of the spec based on the logical constraints can be found in logical_axes_to_spec.py.
+      # qkvo_sharding_spec = jax.sharding.PartitionSpec(
+      #     ("data", "fsdp", "fsdp_transpose", "expert"),
+      #     ("tensor", "tensor_transpose", "sequence", "tensor_sequence"),
+      #     None,
+      #     None,
+      # )
       qkvo_sharding_spec = jax.sharding.PartitionSpec(
           "data",
           "fsdp",

src/maxdiffusion/pyconfig.py

@@ -25,7 +25,6 @@
 import yaml
 from . import max_logging
 from . import max_utils
-from .models.wan.wan_utils import CAUSVID_TRANSFORMER_MODEL_NAME_OR_PATH, WAN_21_FUSION_X_MODEL_NAME_OR_PATH
 
 
 def string_to_bool(s: str) -> bool:
@@ -118,7 +117,6 @@ 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]}")
@@ -127,26 +125,6 @@ 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):
-    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
-          or transformer_pretrained_model_name_or_path == WAN_21_FUSION_X_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"""