From 8159d66c345154716e6a7b3332c24d1ec9fbf2e8 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Fri, 17 Jul 2026 22:34:48 -0700 Subject: [PATCH 1/3] refactor wan pipeline tests to the new mixin structure --- tests/pipelines/wan/test_wan.py | 142 ++++++++++++-------------------- 1 file changed, 54 insertions(+), 88 deletions(-) diff --git a/tests/pipelines/wan/test_wan.py b/tests/pipelines/wan/test_wan.py index bd2d059ff52c..1a504831de5f 100644 --- a/tests/pipelines/wan/test_wan.py +++ b/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 - - -enable_full_determinism() +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, +) -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"] ) - test_xformers_attention = False + batch_input_params = frozenset(["prompt"]) def get_dummy_components(self): torch.manual_seed(0) @@ -87,7 +74,7 @@ def get_dummy_components(self): rope_max_seq_len=32, ) - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -95,39 +82,32 @@ def get_dummy_components(self): "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]) @@ -135,67 +115,53 @@ 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)) + 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 From 46059e066670e98ceed3d80263affb1567aca599 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Fri, 17 Jul 2026 22:39:56 -0700 Subject: [PATCH 2/3] use num_videos_per_prompt in wan pipeline test config --- tests/pipelines/wan/test_wan.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/pipelines/wan/test_wan.py b/tests/pipelines/wan/test_wan.py index 1a504831de5f..382cadb07715 100644 --- a/tests/pipelines/wan/test_wan.py +++ b/tests/pipelines/wan/test_wan.py @@ -40,6 +40,10 @@ class WanPipelineTesterConfig(BasePipelineTesterConfig): ["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"] + ) def get_dummy_components(self): torch.manual_seed(0) From dcaf802b7462beccc35cdae3271f91bf8d7a52d3 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Sat, 18 Jul 2026 03:32:45 -0700 Subject: [PATCH 3/3] testing_utils import single line --- tests/pipelines/wan/test_wan.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/pipelines/wan/test_wan.py b/tests/pipelines/wan/test_wan.py index 382cadb07715..7a91f914f702 100644 --- a/tests/pipelines/wan/test_wan.py +++ b/tests/pipelines/wan/test_wan.py @@ -27,11 +27,7 @@ slow, torch_device, ) -from ..testing_utils import ( - BasePipelineTesterConfig, - MemoryTesterMixin, - PipelineTesterMixin, -) +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin class WanPipelineTesterConfig(BasePipelineTesterConfig):