refactor wan pipeline tests to the new mixin structure

tests/pipelines/wan/test_wan.py

@@ -13,46 +13,33 @@
 # limitations under the License.
 
 import gc
-import tempfile
-import unittest
 
-import numpy as np
+import pytest
 import torch
 from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
 
 from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanPipeline, WanTransformer3DModel
 
 from ...testing_utils import (
+    assert_tensors_close,
     backend_empty_cache,
-    enable_full_determinism,
     require_torch_accelerator,
     slow,
     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
+from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
 
 
-enable_full_determinism()
-
-
-class WanPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+class WanPipelineTesterConfig(BasePipelineTesterConfig):
     pipeline_class = WanPipeline
-    params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
-    batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
-    image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
-    image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
-    required_optional_params = frozenset(
-        [
-            "num_inference_steps",
-            "generator",
-            "latents",
-            "return_dict",
-            "callback_on_step_end",
-            "callback_on_step_end_tensor_inputs",
-        ]
+    required_input_params_in_call_signature = frozenset(
+        ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"]
+    )
+    batch_input_params = frozenset(["prompt"])
+    # Wan is a video pipeline: it exposes `num_videos_per_prompt`, not the base default `num_images_per_prompt`.
+    optional_input_params = frozenset(
+        ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"]
     )
-    test_xformers_attention = False
 
     def get_dummy_components(self):
         torch.manual_seed(0)
@@ -87,115 +74,94 @@ def get_dummy_components(self):
             rope_max_seq_len=32,
         )
 
-        components = {
+        return {
             "transformer": transformer,
             "vae": vae,
             "scheduler": scheduler,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "transformer_2": None,
         }
-        return components
-
-    def get_dummy_inputs(self, device, seed=0):
-        if str(device).startswith("mps"):
-            generator = torch.manual_seed(seed)
-        else:
-            generator = torch.Generator(device=device).manual_seed(seed)
-        inputs = {
+
+    def get_dummy_inputs(self):
+        return {
             "prompt": "dance monkey",
             "negative_prompt": "negative",  # TODO
-            "generator": generator,
+            "generator": self.get_generator(0),
             "num_inference_steps": 2,
             "guidance_scale": 6.0,
             "height": 16,
             "width": 16,
             "num_frames": 9,
             "max_sequence_length": 16,
+            # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
             "output_type": "pt",
         }
-        return inputs
 
-    def test_inference(self):
-        device = "cpu"
 
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe.to(device)
-        pipe.set_progress_bar_config(disable=None)
+class TestWanPipeline(WanPipelineTesterConfig, PipelineTesterMixin):
+    def test_inference(self):
+        # Run on CPU: the expected slice below is CPU-specific.
+        pipe = self.get_pipeline()
 
-        inputs = self.get_dummy_inputs(device)
+        inputs = self.get_dummy_inputs()
         video = pipe(**inputs).frames
         generated_video = video[0]
-        self.assertEqual(generated_video.shape, (9, 3, 16, 16))
+        assert generated_video.shape == (9, 3, 16, 16)
 
         # fmt: off
         expected_slice = torch.tensor([0.4525, 0.452, 0.4485, 0.4534, 0.4524, 0.4529, 0.454, 0.453, 0.5127, 0.5326, 0.5204, 0.5253, 0.5439, 0.5424, 0.5133, 0.5078])
         # fmt: on
 
         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))
+        assert torch.allclose(generated_slice, expected_slice, atol=1e-3)
 
-    @unittest.skip("Test not supported")
-    def test_attention_slicing_forward_pass(self):
-        pass
+    def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4):
+        # `_optional_components` lists both `transformer` and `transformer_2`, but only `transformer_2` is optional
+        # for this wan2.1 t2v pipeline. The base test nulls every optional component, which would drop the required
+        # `transformer` and leave no denoiser, so restrict this to `transformer_2`.
+        pipe = self.get_pipeline().to(torch_device)
+        pipe.transformer_2 = None
 
-    # _optional_components include transformer, transformer_2, but only transformer_2 is optional for this wan2.1 t2v pipeline
-    def test_save_load_optional_components(self, expected_max_difference=1e-4):
-        optional_component = "transformer_2"
-
-        components = self.get_dummy_components()
-        components[optional_component] = None
-        pipe = self.pipeline_class(**components)
-        for component in pipe.components.values():
-            if hasattr(component, "set_default_attn_processor"):
-                component.set_default_attn_processor()
-        pipe.to(torch_device)
-        pipe.set_progress_bar_config(disable=None)
-
-        generator_device = "cpu"
-        inputs = self.get_dummy_inputs(generator_device)
-        torch.manual_seed(0)
+        inputs = self.get_dummy_inputs()
         output = pipe(**inputs)[0]
 
-        with tempfile.TemporaryDirectory() as tmpdir:
-            pipe.save_pretrained(tmpdir, safe_serialization=False)
-            pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
-            for component in pipe_loaded.components.values():
-                if hasattr(component, "set_default_attn_processor"):
-                    component.set_default_attn_processor()
-            pipe_loaded.to(torch_device)
-            pipe_loaded.set_progress_bar_config(disable=None)
-
-        self.assertTrue(
-            getattr(pipe_loaded, optional_component) is None,
-            f"`{optional_component}` did not stay set to None after loading.",
-        )
+        pipe.save_pretrained(tmp_path, safe_serialization=False)
+        pipe_loaded = self.pipeline_class.from_pretrained(tmp_path)
+        pipe_loaded.to(torch_device)
+        pipe_loaded.set_progress_bar_config(disable=None)
 
-        inputs = self.get_dummy_inputs(generator_device)
-        torch.manual_seed(0)
+        assert pipe_loaded.transformer_2 is None, "`transformer_2` did not stay set to None after loading."
+
+        inputs = self.get_dummy_inputs()
         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)
+        assert_tensors_close(
+            output_loaded,
+            output,
+            atol=expected_max_difference,
+            msg="Output changed after dropping the optional component.",
+        )
+
+
+class TestWanPipelineMemory(WanPipelineTesterConfig, MemoryTesterMixin):
+    pass
 
 
 @slow
 @require_torch_accelerator
-class WanPipelineIntegrationTests(unittest.TestCase):
+class TestWanPipelineSlow:
     prompt = "A painting of a squirrel eating a burger."
 
-    def setUp(self):
-        super().setUp()
+    @pytest.fixture(autouse=True)
+    def cleanup(self):
         gc.collect()
         backend_empty_cache(torch_device)
-
-    def tearDown(self):
-        super().tearDown()
+        yield
         gc.collect()
         backend_empty_cache(torch_device)
 
-    @unittest.skip("TODO: test needs to be implemented")
-    def test_Wanx(self):
+    @pytest.mark.skip(reason="TODO: test needs to be implemented")
+    def test_wan(self):
         pass