Update pyconfig.py

Author: Serenagu525

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:
@@ -42,6 +41,21 @@ def string_to_bool(s: str) -> bool:
 config = None
 
 
+def create_parallelisms_list(raw_keys):
+  ici_parallelism = [
+      raw_keys["ici_data_parallelism"],
+      raw_keys["ici_fsdp_parallelism"],
+      raw_keys["ici_fsdp_transpose_parallelism"],
+      raw_keys["ici_sequence_parallelism"],
+      raw_keys["ici_tensor_parallelism"],
+      raw_keys["ici_tensor_transpose_parallelism"],
+      raw_keys["ici_expert_parallelism"],
+      raw_keys["ici_sequence_parallelism"],
+  ]
+  raw_keys["ici_parallelism"] = ici_parallelism
+  return raw_keys
+
+
 def print_system_information():
   max_logging.log(f"System Information: Jax Version: {jax.__version__}")
   max_logging.log(f"System Information: Jaxlib Version: {jax.lib.__version__}")
@@ -103,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]}")
@@ -112,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"""
@@ -176,6 +169,8 @@ def user_init(raw_keys):
     raw_keys["total_train_batch_size"] = max_utils.get_global_batch_size(raw_keys["per_device_batch_size"])
     raw_keys["num_slices"] = get_num_slices(raw_keys)
     raw_keys["quantization_local_shard_count"] = get_quantization_local_shard_count(raw_keys)
+    if "ici_fsdp_transpose_parallelism" in raw_keys:
+      raw_keys = create_parallelisms_list(raw_keys)
 
 
 def get_num_slices(raw_keys):