linting.

Author: entrpn

src/maxdiffusion/checkpointing/wan_checkpointer.py

@@ -15,7 +15,6 @@
 """
 
 from abc import ABC
-from flax import nnx
 from maxdiffusion.checkpointing.checkpointing_utils import (create_orbax_checkpoint_manager)
 from ..pipelines.wan.wan_pipeline import WanPipeline
 from .. import max_logging, max_utils

src/maxdiffusion/models/attention_flax.py

@@ -202,7 +202,7 @@ def _tpu_flash_attention(
   def wrap_flash_attention(query, key, value):
     mask = splash_attention_mask.FullMask(_shape=(query.shape[2], key.shape[2]))
     multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1])
-    # make_splash_mha is wrapped around shardmap and seq and head is already 
+    # make_splash_mha is wrapped around shardmap and seq and head is already
     # sharded based on in_specs, therefore setting head_shards=1 and q_seq_shards=1.
     splash_kernel = splash_attention_kernel.make_splash_mha(
         mask=multi_head_mask,

src/maxdiffusion/models/gradient_checkpoint.py

@@ -6,7 +6,8 @@
 
 SKIP_GRADIENT_CHECKPOINT_KEY = "skip"
 
-# This class only works with NNX modules. 
+
+# This class only works with NNX modules.
 class GradientCheckpointType(Enum):
   """
   Defines the type of the gradient checkpoint we will have

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -364,7 +364,7 @@ def __init__(
       weights_dtype: jnp.dtype = jnp.float32,
       precision: jax.lax.Precision = None,
       attention: str = "dot_product",
-      remat_policy: str = "None"
+      remat_policy: str = "None",
   ):
     inner_dim = num_attention_heads * attention_head_dim
     out_channels = out_channels or in_channels
@@ -383,13 +383,7 @@ def __init__(
         precision=precision,
         kernel_init=nnx.with_partitioning(
             nnx.initializers.xavier_uniform(),
-            (
-                None,
-                None,
-                None,
-                None,
-                "conv_out"
-            ),
+            (None, None, None, None, "conv_out"),
         ),
     )
 

src/maxdiffusion/multihost_dataloading.py

@@ -114,4 +114,4 @@ def __iter__(self):
     return self
 
   def __next__(self):
-    return get_next_batch_sharded(self.local_iterator, self.global_mesh)
+    return get_next_batch_sharded(self.local_iterator, self.global_mesh)

src/maxdiffusion/trainers/wan_trainer.py

@@ -37,15 +37,18 @@
 from skimage.metrics import structural_similarity as ssim
 from flax.training import train_state
 
+
 class TrainState(train_state.TrainState):
   graphdef: nnx.GraphDef
   rest_of_state: nnx.State
 
+
 def _to_array(x):
   if not isinstance(x, jax.Array):
     x = jnp.asarray(x)
   return x
 
+
 def generate_sample(config, pipeline, filename_prefix):
   """
   Generates a video to validate training did not corrupt the model
@@ -79,7 +82,6 @@ def __init__(self, config):
     if config.train_text_encoder:
       raise ValueError("this script currently doesn't support training text_encoders")
 
-    #self.global_batch_size = self.config.per_device_batch_size * jax.device_count()
     self.global_batch_size = config.per_device_batch_size * jax.device_count()
 
   def post_training_steps(self, pipeline, params, train_states, msg=""):
@@ -95,13 +97,10 @@ def create_scheduler(self):
   def calculate_tflops(self, pipeline):
     max_logging.log("WARNING : Calculting tflops is not implemented in Wan 2.1. Returning 0...")
     return 0
-  
+
   def get_data_shardings(self, mesh):
     data_sharding = jax.sharding.NamedSharding(mesh, P(*self.config.data_sharding))
-    data_sharding = {
-      "latents" : data_sharding,
-      "encoder_hidden_states" : data_sharding
-    }
+    data_sharding = {"latents": data_sharding, "encoder_hidden_states": data_sharding}
     return data_sharding
 
   def load_dataset(self, mesh):
@@ -167,11 +166,7 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, data_itera
 
     with mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
       state = TrainState.create(
-        apply_fn=graphdef.apply,
-        params=params,
-        tx=optimizer,
-        graphdef=graphdef,
-        rest_of_state=rest_of_state
+          apply_fn=graphdef.apply, params=params, tx=optimizer, graphdef=graphdef, rest_of_state=rest_of_state
       )
       state = jax.tree.map(_to_array, state)
       state_spec = nnx.get_partition_spec(state)
@@ -196,8 +191,8 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, data_itera
 
     p_train_step = jax.jit(
         functools.partial(train_step, scheduler=pipeline.scheduler, config=self.config),
-        in_shardings = (state_shardings, data_shardings, None, None),
-        out_shardings = (state_shardings, None, None, None),
+        in_shardings=(state_shardings, data_shardings, None, None),
+        out_shardings=(state_shardings, None, None, None),
         donate_argnums=(0,),
     )
     rng = jax.random.key(self.config.seed)
@@ -284,6 +279,7 @@ def loss_fn(params):
     loss = jnp.mean(loss)
 
     return loss
+
   grad_fn = nnx.value_and_grad(loss_fn)
   loss, grads = grad_fn(state.params)
   new_state = state.apply_gradients(grads=grads)