Skip to content

Commit d710e59

Browse files
committed
style
1 parent 1189a35 commit d710e59

File tree

3 files changed

+13
-5
lines changed

3 files changed

+13
-5
lines changed

tests/pipelines/wan/test_wan.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,10 @@
1313
# limitations under the License.
1414

1515
import gc
16+
import tempfile
1617
import unittest
1718

19+
import numpy as np
1820
import torch
1921
from transformers import AutoTokenizer, T5EncoderModel
2022

@@ -178,6 +180,7 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4):
178180
max_diff = np.abs(output.detach().cpu().numpy() - output_loaded.detach().cpu().numpy()).max()
179181
self.assertLess(max_diff, expected_max_difference)
180182

183+
181184
@slow
182185
@require_torch_accelerator
183186
class WanPipelineIntegrationTests(unittest.TestCase):

tests/pipelines/wan/test_wan_22_image_to_video.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import numpy as np
1919
import torch
20+
from PIL import Image
2021
from transformers import AutoTokenizer, T5EncoderModel
2122

2223
from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanImageToVideoPipeline, WanTransformer3DModel
@@ -27,7 +28,6 @@
2728

2829
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
2930
from ..test_pipelines_common import PipelineTesterMixin
30-
from PIL import Image
3131

3232

3333
enable_full_determinism()
@@ -157,7 +157,10 @@ def test_inference(self):
157157

158158
generated_slice = generated_video.flatten()
159159
generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
160-
self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3), f"generated_slice: {generated_slice}, expected_slice: {expected_slice}")
160+
self.assertTrue(
161+
torch.allclose(generated_slice, expected_slice, atol=1e-3),
162+
f"generated_slice: {generated_slice}, expected_slice: {expected_slice}",
163+
)
161164

162165
@unittest.skip("Test not supported")
163166
def test_attention_slicing_forward_pass(self):
@@ -371,7 +374,7 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4):
371374
self.assertTrue(
372375
getattr(pipe_loaded, component) is None,
373376
f"`{component}` did not stay set to None after loading.",
374-
)
377+
)
375378

376379
inputs = self.get_dummy_inputs(generator_device)
377380
torch.manual_seed(0)

tests/pipelines/wan/test_wan_image_to_video.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,10 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import tempfile
1516
import unittest
1617

18+
import numpy as np
1719
import torch
1820
from PIL import Image
1921
from transformers import (
@@ -25,7 +27,7 @@
2527
)
2628

2729
from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanImageToVideoPipeline, WanTransformer3DModel
28-
from diffusers.utils.testing_utils import enable_full_determinism
30+
from diffusers.utils.testing_utils import enable_full_determinism, torch_device
2931

3032
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
3133
from ..test_pipelines_common import PipelineTesterMixin
@@ -376,4 +378,4 @@ def test_save_load_optional_components(self, expected_max_difference=1e-4):
376378
output_loaded = pipe_loaded(**inputs)[0]
377379

378380
max_diff = np.abs(output.detach().cpu().numpy() - output_loaded.detach().cpu().numpy()).max()
379-
self.assertLess(max_diff, expected_max_difference)
381+
self.assertLess(max_diff, expected_max_difference)

0 commit comments

Comments
 (0)