Ruff format

Signed-off-by: Kunjan Patel <kunjanp@google.com>

Author: coolkp

src/maxdiffusion/models/attention_flax.py

@@ -171,7 +171,7 @@ def _pad_data_for_flash(tensor, heads, flash_block_size, num_shards: int = 1):
 
   return tensor, kv_size, seq_len
 
-def convert_to_tokamax_splash_config( block_sizes: BlockSizes, 
+def convert_to_tokamax_splash_config( block_sizes: BlockSizes,
                                       q_layout: tokamax_splash_attention_kernel.QKVLayout = tokamax_splash_attention_kernel.QKVLayout.HEAD_DIM_MINOR,
                                       k_layout: tokamax_splash_attention_kernel.QKVLayout = tokamax_splash_attention_kernel.QKVLayout.HEAD_DIM_MINOR,
                                       v_layout: tokamax_splash_attention_kernel.QKVLayout = tokamax_splash_attention_kernel.QKVLayout.HEAD_DIM_MINOR,
@@ -536,7 +536,7 @@ def _apply_attention(
     )
   elif attention_kernel == "ring":
     return _tpu_flash_attention(
-        query, key * scale, value, heads, mesh, axis_names_q, axis_names_kv, flash_block_sizes, dtype, attention_kernel, 
+        query, key * scale, value, heads, mesh, axis_names_q, axis_names_kv, flash_block_sizes, dtype, attention_kernel,
         mask_padding_tokens=mask_padding_tokens,
     )
   elif attention_kernel == "cudnn_flash_te":

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -567,7 +567,7 @@ def __call__(
         prompt = [prompt]
 
       batch_size = len(prompt)
-      
+
       with jax.named_scope("Encode-Prompt"):
         prompt_embeds, negative_prompt_embeds = self.encode_prompt(
             prompt=prompt,
@@ -578,7 +578,7 @@ def __call__(
         )
 
       num_channel_latents = self.transformer.config.in_channels
-      if latents is None: 
+      if latents is None:
         latents = self.prepare_latents(
             batch_size=batch_size,
             vae_scale_factor_temporal=self.vae_scale_factor_temporal,