style

Author: yiyixuxu

tests/pipelines/wan/test_wan.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 
 import gc
+import tempfile
 import unittest
 
+import numpy as np
 import torch
 from transformers import AutoTokenizer, T5EncoderModel
 
@@ -178,6 +180,7 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4):
         max_diff = np.abs(output.detach().cpu().numpy() - output_loaded.detach().cpu().numpy()).max()
         self.assertLess(max_diff, expected_max_difference)
 
+
 @slow
 @require_torch_accelerator
 class WanPipelineIntegrationTests(unittest.TestCase):

tests/pipelines/wan/test_wan_22_image_to_video.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 import torch
+from PIL import Image
 from transformers import AutoTokenizer, T5EncoderModel
 
 from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanImageToVideoPipeline, WanTransformer3DModel
@@ -27,7 +28,6 @@
 
 from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
 from ..test_pipelines_common import PipelineTesterMixin
-from PIL import Image
 
 
 enable_full_determinism()
@@ -157,7 +157,10 @@ def test_inference(self):
 
         generated_slice = generated_video.flatten()
         generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
-        self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3), f"generated_slice: {generated_slice}, expected_slice: {expected_slice}")
+        self.assertTrue(
+            torch.allclose(generated_slice, expected_slice, atol=1e-3),
+            f"generated_slice: {generated_slice}, expected_slice: {expected_slice}",
+        )
 
     @unittest.skip("Test not supported")
     def test_attention_slicing_forward_pass(self):
@@ -371,7 +374,7 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4):
             self.assertTrue(
                 getattr(pipe_loaded, component) is None,
                 f"`{component}` did not stay set to None after loading.",
-                )
+            )
 
         inputs = self.get_dummy_inputs(generator_device)
         torch.manual_seed(0)

tests/pipelines/wan/test_wan_image_to_video.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import tempfile
 import unittest
 
+import numpy as np
 import torch
 from PIL import Image
 from transformers import (
@@ -25,7 +27,7 @@
 )
 
 from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanImageToVideoPipeline, WanTransformer3DModel
-from diffusers.utils.testing_utils import enable_full_determinism
+from diffusers.utils.testing_utils import enable_full_determinism, torch_device
 
 from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
 from ..test_pipelines_common import PipelineTesterMixin
@@ -376,4 +378,4 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4):
         output_loaded = pipe_loaded(**inputs)[0]
 
         max_diff = np.abs(output.detach().cpu().numpy() - output_loaded.detach().cpu().numpy()).max()
-        self.assertLess(max_diff, expected_max_difference)
+        self.assertLess(max_diff, expected_max_difference)