fix lint issues

Author: csgoogle

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

@@ -735,38 +735,32 @@ def __call__(
     # frame-wise reshape so each query batch element still corresponds to a full frame.
     query = jax.lax.with_sharding_constraint(
         query,
-        nn.logical_to_mesh_axes(
-            (
-                common_types.BATCH,
-                None,
-                common_types.CROSS_ATTN_HEAD,
-                common_types.D_KV,
-            )
-        ),
+        nn.logical_to_mesh_axes((
+            common_types.BATCH,
+            None,
+            common_types.CROSS_ATTN_HEAD,
+            common_types.D_KV,
+        )),
     )
     key = jax.lax.with_sharding_constraint(
         key,
-        nn.logical_to_mesh_axes(
-            (
-                common_types.BATCH,
-                None,
-                None,
-                common_types.CROSS_ATTN_HEAD,
-                common_types.D_KV,
-            )
-        ),
+        nn.logical_to_mesh_axes((
+            common_types.BATCH,
+            None,
+            None,
+            common_types.CROSS_ATTN_HEAD,
+            common_types.D_KV,
+        )),
     )
     value = jax.lax.with_sharding_constraint(
         value,
-        nn.logical_to_mesh_axes(
-            (
-                common_types.BATCH,
-                None,
-                None,
-                common_types.CROSS_ATTN_HEAD,
-                common_types.D_KV,
-            )
-        ),
+        nn.logical_to_mesh_axes((
+            common_types.BATCH,
+            None,
+            None,
+            common_types.CROSS_ATTN_HEAD,
+            common_types.D_KV,
+        )),
     )
 
     query_S = query.shape[1]
@@ -803,13 +797,11 @@ def __call__(
     hidden_states = self.to_out(attn_output)
     hidden_states = jax.lax.with_sharding_constraint(
         hidden_states,
-        nn.logical_to_mesh_axes(
-            (
-                common_types.BATCH,
-                common_types.LENGTH,
-                common_types.EMBED,
-            )
-        ),
+        nn.logical_to_mesh_axes((
+            common_types.BATCH,
+            common_types.LENGTH,
+            common_types.EMBED,
+        )),
     )
 
     if attention_mask is not None:

src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py

@@ -610,9 +610,7 @@ def prepare_prev_segment_cond_latents(
       remaining_frames = segment_frame_length - prev_segment_cond_frames
       remaining = jnp.zeros((batch_size, 3, remaining_frames, height, width), dtype=vae_dtype)
 
-    full_cond_video = jnp.concatenate(
-        [prev_segment_cond_video.astype(vae_dtype), remaining], axis=2
-    )  # (B, C, T_seg, H, W)
+    full_cond_video = jnp.concatenate([prev_segment_cond_video.astype(vae_dtype), remaining], axis=2)  # (B, C, T_seg, H, W)
 
     cond_latents = self._encode_video_to_latents(full_cond_video, dtype)
     # (B, T_lat, H_lat, W_lat, z_dim)
@@ -785,9 +783,7 @@ def __call__(
           f"`segment_frame_length - 1` must be divisible by {self.vae_scale_factor_temporal}. "
           f"Rounding {segment_frame_length}."
       )
-      segment_frame_length = (
-          segment_frame_length // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
-      )
+      segment_frame_length = segment_frame_length // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
     segment_frame_length = max(segment_frame_length, 1)
 
     do_classifier_free_guidance = guidance_scale > 1.0
@@ -994,9 +990,7 @@ def __call__(
             noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
 
           noise_pred = noise_pred.astype(seg_latents.dtype)
-          seg_latents, scheduler_state = self.scheduler.step(
-              scheduler_state, noise_pred, t, seg_latents, return_dict=False
-          )
+          seg_latents, scheduler_state = self.scheduler.step(scheduler_state, noise_pred, t, seg_latents, return_dict=False)
 
         # Decode this segment (skip reference frame at index 0).
         out_frames_cf = self._decode_segment_to_pixels(seg_latents[:, 1:, :, :, :])

src/maxdiffusion/tests/wan_animate_diffusers_parity_test.py

@@ -503,9 +503,7 @@ def fake_encode(video, dtype):
       latent_t = (video.shape[2] - 1) // self.max_pipeline.vae_scale_factor_temporal + 1
       latent_h = video.shape[3] // self.max_pipeline.vae_scale_factor_spatial
       latent_w = video.shape[4] // self.max_pipeline.vae_scale_factor_spatial
-      return jnp.zeros(
-          (video.shape[0], latent_t, latent_h, latent_w, self.max_pipeline.vae.z_dim), dtype=jnp.float32
-      )
+      return jnp.zeros((video.shape[0], latent_t, latent_h, latent_w, self.max_pipeline.vae.z_dim), dtype=jnp.float32)
 
     self.max_pipeline._encode_video_to_latents = fake_encode
 
@@ -722,9 +720,7 @@ def __call__(
     np.testing.assert_allclose(to_numpy(capture["pose_hidden_states"]), to_numpy(expected_pose), atol=0.0, rtol=0.0)
     np.testing.assert_allclose(to_numpy(capture["face_pixel_values"]), to_numpy(face_video), atol=0.0, rtol=0.0)
     np.testing.assert_allclose(to_numpy(capture["encoder_hidden_states"]), to_numpy(prompt_embeds), atol=0.0, rtol=0.0)
-    np.testing.assert_allclose(
-        to_numpy(capture["encoder_hidden_states_image"]), to_numpy(image_embeds), atol=0.0, rtol=0.0
-    )
+    np.testing.assert_allclose(to_numpy(capture["encoder_hidden_states_image"]), to_numpy(image_embeds), atol=0.0, rtol=0.0)
     self.assertEqual(capture["motion_encode_batch_size"], 7)
     self.assertFalse(capture["return_dict"])
     np.testing.assert_allclose(to_numpy(noise_pred), to_numpy(latents), atol=0.0, rtol=0.0)
@@ -872,9 +868,7 @@ def test_flax_unipc_flow_sigmas_match_diffusers(self):
       max_model_output = jnp.array(to_numpy(hf_model_output))
 
       hf_sample = hf_scheduler.step(hf_model_output, int(timestep), hf_sample, return_dict=False)[0]
-      max_sample, max_state = max_scheduler.step(
-          max_state, max_model_output, int(timestep), max_sample, return_dict=False
-      )
+      max_sample, max_state = max_scheduler.step(max_state, max_model_output, int(timestep), max_sample, return_dict=False)
 
       np.testing.assert_allclose(to_numpy(max_sample), to_numpy(hf_sample), atol=1e-4, rtol=1e-5)