From 147c99bb2c71edd6d4417946048bab2490d032d8 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Sat, 18 Jul 2026 05:15:12 -0700 Subject: [PATCH] refactor wan vace pipeline tests to the new mixin structure --- tests/pipelines/wan/test_wan_vace.py | 226 ++++++++++----------------- 1 file changed, 84 insertions(+), 142 deletions(-) diff --git a/tests/pipelines/wan/test_wan_vace.py b/tests/pipelines/wan/test_wan_vace.py index 7294cd368e0b..e4e5f24be844 100644 --- a/tests/pipelines/wan/test_wan_vace.py +++ b/tests/pipelines/wan/test_wan_vace.py @@ -12,10 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -import tempfile -import unittest -import numpy as np +import pytest import torch from PIL import Image from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -28,31 +26,20 @@ WanVACETransformer3DModel, ) -from ...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 +from ...testing_utils import assert_tensors_close, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -enable_full_determinism() - - -class WanVACEPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class WanVACEPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = WanVACEPipeline - 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"]) + # WanVACE 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) @@ -88,7 +75,7 @@ def get_dummy_components(self): vace_in_channels=96, ) - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -96,14 +83,8 @@ 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) + def get_dummy_inputs(self): num_frames = 17 height = 16 width = 16 @@ -111,124 +92,92 @@ def get_dummy_inputs(self, device, seed=0): video = [Image.new("RGB", (height, width))] * num_frames mask = [Image.new("L", (height, width), 0)] * num_frames - inputs = { + return { "video": video, "mask": mask, "prompt": "dance monkey", "negative_prompt": "negative", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "height": 16, "width": 16, "num_frames": num_frames, "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 TestWanVACEPipeline(WanVACEPipelineTesterConfig, PipelineTesterMixin): + @pytest.mark.skip(reason="Batching is not yet supported with this pipeline") + def test_inference_batch_consistent(self): + pass - inputs = self.get_dummy_inputs(device) - video = pipe(**inputs).frames[0] - self.assertEqual(video.shape, (17, 3, 16, 16)) + @pytest.mark.skip(reason="Batching is not yet supported with this pipeline") + def test_inference_batch_single_identical(self): + pass + + def test_inference(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() + + inputs = self.get_dummy_inputs() + video = pipe(**inputs).frames + generated_video = video[0] + assert generated_video.shape == (17, 3, 16, 16) # fmt: off - expected_slice = [0.4523, 0.45198, 0.44872, 0.45326, 0.45211, 0.45258, 0.45344, 0.453, 0.52431, 0.52572, 0.50701, 0.5118, 0.53717, 0.53093, 0.50557, 0.51402] + expected_slice = torch.tensor([0.4523, 0.45198, 0.44872, 0.45326, 0.45211, 0.45258, 0.45344, 0.453, 0.52431, 0.52572, 0.50701, 0.5118, 0.53717, 0.53093, 0.50557, 0.51402]) # fmt: on - video_slice = video.flatten() - video_slice = torch.cat([video_slice[:8], video_slice[-8:]]) - video_slice = [round(x, 5) for x in video_slice.tolist()] - self.assertTrue(np.allclose(video_slice, expected_slice, atol=1e-3)) + generated_slice = generated_video.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + assert torch.allclose(generated_slice, expected_slice, atol=1e-3) def test_inference_with_single_reference_image(self): - device = "cpu" + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["reference_images"] = Image.new("RGB", (16, 16)) - video = pipe(**inputs).frames[0] - self.assertEqual(video.shape, (17, 3, 16, 16)) + video = pipe(**inputs).frames + generated_video = video[0] + assert generated_video.shape == (17, 3, 16, 16) # fmt: off - expected_slice = [0.45247, 0.45214, 0.44874, 0.45314, 0.45171, 0.45299, 0.45428, 0.45317, 0.51378, 0.52658, 0.53361, 0.52303, 0.46204, 0.50435, 0.52555, 0.51342] + expected_slice = torch.tensor([0.45247, 0.45214, 0.44874, 0.45314, 0.45171, 0.45299, 0.45428, 0.45317, 0.51378, 0.52658, 0.53361, 0.52303, 0.46204, 0.50435, 0.52555, 0.51342]) # fmt: on - video_slice = video.flatten() - video_slice = torch.cat([video_slice[:8], video_slice[-8:]]) - video_slice = [round(x, 5) for x in video_slice.tolist()] - self.assertTrue(np.allclose(video_slice, expected_slice, atol=1e-3)) + generated_slice = generated_video.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + assert torch.allclose(generated_slice, expected_slice, atol=1e-3) def test_inference_with_multiple_reference_image(self): - device = "cpu" + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["reference_images"] = [[Image.new("RGB", (16, 16))] * 2] - video = pipe(**inputs).frames[0] - self.assertEqual(video.shape, (17, 3, 16, 16)) + video = pipe(**inputs).frames + generated_video = video[0] + assert generated_video.shape == (17, 3, 16, 16) # fmt: off - expected_slice = [0.45321, 0.45221, 0.44818, 0.45375, 0.45268, 0.4519, 0.45271, 0.45253, 0.51244, 0.52223, 0.51253, 0.51321, 0.50743, 0.51177, 0.51626, 0.50983] + expected_slice = torch.tensor([0.45321, 0.45221, 0.44818, 0.45375, 0.45268, 0.4519, 0.45271, 0.45253, 0.51244, 0.52223, 0.51253, 0.51321, 0.50743, 0.51177, 0.51626, 0.50983]) # fmt: on - video_slice = video.flatten() - video_slice = torch.cat([video_slice[:8], video_slice[-8:]]) - video_slice = [round(x, 5) for x in video_slice.tolist()] - self.assertTrue(np.allclose(video_slice, expected_slice, atol=1e-3)) - - @unittest.skip("Test not supported") - def test_attention_slicing_forward_pass(self): - pass - - @unittest.skip("Errors out because passing multiple prompts at once is not yet supported by this pipeline.") - def test_encode_prompt_works_in_isolation(self): - pass - - @unittest.skip("Batching is not yet supported with this pipeline") - def test_inference_batch_consistent(self): - pass - - @unittest.skip("Batching is not yet supported with this pipeline") - def test_inference_batch_single_identical(self): - return super().test_inference_batch_single_identical() - - @unittest.skip( - "AutoencoderKLWan encoded latents are always in FP32. This test is not designed to handle mixed dtype inputs" - ) - def test_float16_inference(self): - pass - - @unittest.skip( - "AutoencoderKLWan encoded latents are always in FP32. This test is not designed to handle mixed dtype inputs" - ) - def test_save_load_float16(self): - pass + generated_slice = generated_video.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + assert torch.allclose(generated_slice, expected_slice, atol=1e-3) def test_inference_with_only_transformer(self): components = self.get_dummy_components() components["transformer_2"] = None components["boundary_ratio"] = 0.0 - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline(**components).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() video = pipe(**inputs).frames[0] assert video.shape == (17, 3, 16, 16) @@ -244,56 +193,49 @@ def test_inference_with_only_transformer_2(self): ) components["boundary_ratio"] = 1.0 - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline(**components).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() video = pipe(**inputs).frames[0] assert video.shape == (17, 3, 16, 16) - def test_save_load_optional_components(self, expected_max_difference=1e-4): - optional_component = ["transformer"] - + def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4): + # `_optional_components` lists both `transformer` and `transformer_2`. Here we drop the (optional) + # `transformer` and denoise with `transformer_2` only, which needs `boundary_ratio=1.0` and a scheduler that + # can run the low-noise stage on its own (FlowMatchEuler can't, since its starting timestep equals the + # boundary). components = self.get_dummy_components() components["transformer_2"] = components["transformer"] - # FlowMatchEulerDiscreteScheduler doesn't support running low noise only scheduler - # because starting timestep t == 1000 == boundary_timestep + components["transformer"] = None components["scheduler"] = UniPCMultistepScheduler( prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0 ) - for component in optional_component: - components[component] = None - components["boundary_ratio"] = 1.0 - 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) + pipe = self.get_pipeline(**components).to(torch_device) - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() torch.manual_seed(0) 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) + 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) - for component in optional_component: - assert getattr(pipe_loaded, component) is None, f"`{component}` did not stay set to None after loading." + assert pipe_loaded.transformer is None, "`transformer` did not stay set to None after loading." - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() torch.manual_seed(0) output_loaded = pipe_loaded(**inputs)[0] - max_diff = np.abs(output.detach().cpu().numpy() - output_loaded.detach().cpu().numpy()).max() - assert max_diff < expected_max_difference, "Outputs exceed expecpted maximum difference" + assert_tensors_close( + output_loaded, + output, + atol=expected_max_difference, + msg="Output changed after dropping the optional component.", + ) + + +class TestWanVACEPipelineMemory(WanVACEPipelineTesterConfig, MemoryTesterMixin): + pass