ruff fixed

Author: Serenagu525

src/maxdiffusion/generate_ltx_video.py

@@ -64,10 +64,6 @@ def run_inference(
       segment_ids=segment_ids,
       encoder_attention_segment_ids=encoder_attention_segment_ids,
   )
-  prof = profiler.Profiler(config)
-  prof.activate(optional_postfix="transformer step")
-  prof.deactivate()
-  
 
   with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
     noise_pred, transformer_state, _ = jax.lax.fori_loop(0, 1, loop_body_p, (latents, transformer_state, timestep))
@@ -176,8 +172,8 @@ def run(config):
       in_shardings=(state_shardings,),
       out_shardings=None,
   )
-  with jax.profiler.trace("/tmp/jax-trace", create_perfetto_link=True):
-    noise_pred = p_run_inference(states).block_until_ready()
+
+  noise_pred = p_run_inference(states).block_until_ready()
   print(noise_pred)  # (4, 256, 128)
 
 

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:
@@ -417,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(
@@ -628,4 +609,4 @@ def maybe_initialize_jax_distributed_system(raw_keys):
     initialize_jax_for_gpu()
     max_logging.log("Jax distributed system initialized on GPU!")
   else:
-    jax.distributed.initialize()
+    jax.distributed.initialize()

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):

src/maxdiffusion/tests/ltx_transformer_step_test.py

@@ -191,7 +191,7 @@ def test_one_step_transformer(self):
         in_shardings=(state_shardings,),
         out_shardings=None,
     )
-   
+
     noise_pred = p_run_inference(states).block_until_ready()
     noise_pred = torch.from_numpy(np.array(noise_pred))