Skip to content

refactor wan vace pipeline tests to the new mixin structure #14231

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

Open
akshan-main wants to merge 1 commit into
base: main
Choose a base branch
from
+84 −142
Open

refactor wan vace pipeline tests to the new mixin structure #14231

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
226 changes: 84 additions & 142 deletions tests/pipelines/wan/test_wan_vace.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -88,147 +75,109 @@ def get_dummy_components(self):
vace_in_channels=96,
)

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)

def get_dummy_inputs(self):
num_frames = 17
height = 16
width = 16

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)

Expand All @@ -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
Loading