Skip to content

refactor wan 2.2 pipeline tests to the new mixin structure #14240

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sayakpaul merged 2 commits into from
Jul 20, 2026
+95 −162
Merged

refactor wan 2.2 pipeline tests to the new mixin structure #14240

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8e6827a
refactor wan 2.2 pipeline tests to the new mixin structure
akshan-main Jul 19, 2026
3ee2ee6
Merge branch 'main' into refactor-wan22-pipeline-tests
sayakpaul Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
257 changes: 95 additions & 162 deletions tests/pipelines/wan/test_wan_22.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,40 +12,26 @@
# 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 transformers import AutoConfig, AutoTokenizer, T5EncoderModel

from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanPipeline, WanTransformer3DModel

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 Wan22PipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class Wan22PipelineTesterConfig(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)
Expand Down Expand Up @@ -95,7 +81,7 @@ def get_dummy_components(self):
rope_max_seq_len=32,
)

components = {
return {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
Expand All @@ -104,112 +90,83 @@ def get_dummy_components(self):
"transformer_2": transformer_2,
"boundary_ratio": 0.875,
}
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",
"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 TestWan22Pipeline(Wan22PipelineTesterConfig, 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))

@unittest.skip("Test not supported")
def test_attention_slicing_forward_pass(self):
pass
Comment on lines -151 to -153

@sayakpaul sayakpaul Jul 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How it this being treated?

@akshan-main akshan-main Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new PipelineTesterMixin has no attention-slicing test, so there's nothing to skip; the old override was just disabling the old base mixin's version, same as the flux reference.

@akshan-main akshan-main Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wan doesn't support attention slicing anyways


def test_save_load_optional_components(self, expected_max_difference=1e-4):
optional_component = "transformer"
assert torch.allclose(generated_slice, expected_slice, atol=1e-3)

def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4):
# For wan 2.2 14B, `transformer` is not used when `boundary_ratio` is 1.0, so only then is it optional.
components = self.get_dummy_components()
components[optional_component] = None
components["boundary_ratio"] = 1.0 # for wan 2.2 14B, transformer is not used when boundary_ratio is 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)

generator_device = "cpu"
inputs = self.get_dummy_inputs(generator_device)
torch.manual_seed(0)
components["transformer"] = None
components["boundary_ratio"] = 1.0
pipe = self.get_pipeline(**components).to(torch_device)

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, "transformer") is None,
"`transformer` 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 is None, "`transformer` 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 TestWan22PipelineMemory(Wan22PipelineTesterConfig, MemoryTesterMixin):
pass

class Wan225BPipelineFastTests(PipelineTesterMixin, unittest.TestCase):

class Wan225BPipelineTesterConfig(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)
Expand Down Expand Up @@ -251,7 +208,7 @@ def get_dummy_components(self):
rope_max_seq_len=32,
)

components = {
return {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
Expand All @@ -261,41 +218,32 @@ def get_dummy_components(self):
"boundary_ratio": None,
"expand_timesteps": True,
}
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": 32,
"width": 32,
"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 TestWan225BPipeline(Wan225BPipelineTesterConfig, 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, 32, 32))
assert generated_video.shape == (9, 3, 32, 32)

# fmt: off
expected_slice = torch.tensor([[[0.4814, 0.4298, 0.5094, 0.4289, 0.5061, 0.4301, 0.5043, 0.4284, 0.5375,
Expand All @@ -304,61 +252,46 @@ 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}",
assert 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):
pass

def test_components_function(self):
init_components = self.get_dummy_components()
init_components.pop("boundary_ratio")
init_components.pop("expand_timesteps")
pipe = self.pipeline_class(**init_components)
pipe = self.get_pipeline(**init_components)

self.assertTrue(hasattr(pipe, "components"))
self.assertTrue(set(pipe.components.keys()) == set(init_components.keys()))
assert hasattr(pipe, "components")
assert set(pipe.components.keys()) == set(init_components.keys())

def test_save_load_optional_components(self, expected_max_difference=1e-4):
optional_component = "transformer_2"
def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4):
pipe = self.get_pipeline().to(torch_device)
pipe.transformer_2 = None

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.",
)

def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=2e-3):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)


def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(expected_max_diff=2e-3)
class TestWan225BPipelineMemory(Wan225BPipelineTesterConfig, MemoryTesterMixin):
pass
Loading