changed back pyconfig

Author: Serenagu525

src/maxdiffusion/max_utils.py

@@ -257,21 +257,6 @@ def create_device_mesh(config, devices=None, logging=True):
   if devices is None:
     devices = jax.devices()
   num_devices = len(devices)
-  ##special case for ltx-video
-  if "fsdp_transpose" in config.mesh_axes:
-    num_slices = 1
-    # if config.inference_benchmark_test else config.num_slices
-    num_devices_per_slice = num_devices // num_slices
-    # Find possible unspecified parallelisms
-    ici_parallelism = fill_unspecified_mesh_axes(config.ici_parallelism.copy(), num_devices_per_slice, "ICI")
-    mesh = mesh_utils.create_device_mesh(
-        ici_parallelism,
-        devices,
-    )
-    max_logging.log(f"Num_devices: {num_devices}, shape {mesh.shape}")
-
-    return mesh
-    
   try:
     num_slices = 1 + max([d.slice_index for d in devices])
   except:
@@ -303,66 +288,9 @@ def create_device_mesh(config, devices=None, logging=True):
   if logging:
     max_logging.log(f"Decided on mesh: {mesh}")
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
   return mesh
 
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
 def unbox_logicallypartioned_trainstate(boxed_train_state: train_state.TrainState):
   """Unboxes the flax.LogicallyPartitioned pieces in a train state.
 
@@ -474,11 +402,7 @@ def setup_initial_state(
           config.enable_single_replica_ckpt_restoring,
       )
       if state:
-        ###!Edited
-        if checkpoint_item == " ":
-          state = state
-        else:
-          state = state[checkpoint_item]
+        state = state[checkpoint_item]
     if not state:
       max_logging.log(f"Could not find the item in orbax, creating state...")
       init_train_state_partial = functools.partial(

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, WAN_21_FUSION_X_MODEL_NAME_OR_PATH
 
 
 def string_to_bool(s: str) -> bool:
@@ -41,21 +42,6 @@ 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__}")
@@ -117,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]}")
@@ -125,6 +112,26 @@ 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"""
@@ -169,8 +176,6 @@ 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):
@@ -221,4 +226,4 @@ def initialize(argv, **kwargs):
 if __name__ == "__main__":
   initialize(sys.argv)
   print(config.steps)
-  r = range(config.steps)
+  r = range(config.steps)