fix linting

Author: csgoogle

src/maxdiffusion/tests/wan_animate/test_transformer_wan_animate.py

@@ -1,5 +1,5 @@
 """
-Copyright 2025 Google LLC
+Copyright 2026 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -16,14 +16,9 @@
 
 import os
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
 import numpy as np
 import jax.numpy as jnp
 from flax import nnx
-import math
-import sys
-import pytest
 
 # Import from the codebase
 # Make sure to set Python path or import correctly
@@ -42,9 +37,6 @@
 from jax.sharding import Mesh
 
 
-
-
-
 def transfer_conv_weights(pt_conv, jax_conv):
   if hasattr(jax_conv, "weight"):
     jax_conv.weight[...] = jnp.array(pt_conv.weight.detach().numpy())
@@ -227,8 +219,6 @@ def transfer_transformer_weights(pt_model, jax_model):
   jax_model.scale_shift_table[...] = jnp.array(pt_model.scale_shift_table.detach().numpy())
 
 
-
-
 class TestWanAnimateTransformer:
 
   def test_motion_conv_equivalence(self):
@@ -244,6 +234,7 @@ def test_motion_conv_equivalence(self):
     x_jax = jnp.array(x_np)
 
     from diffusers.models.transformers.transformer_wan_animate import MotionConv2d
+
     pt_model = MotionConv2d(
         in_channels=C_in,
         out_channels=C_out,
@@ -697,5 +688,3 @@ def test_equivalence_wan_animate_transformer(self):
 
       assert np_pt.shape == np_jax.shape
       np.testing.assert_allclose(np_pt, np_jax, rtol=1e-4, atol=1e-4)
-
-