diff --git a/.ai/testing.md b/.ai/testing.md index d13f0c744530..24d39da3bd68 100644 --- a/.ai/testing.md +++ b/.ai/testing.md @@ -22,6 +22,7 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers - **These are pytest-style, not `unittest`** — no `unittest.TestCase` subclassing, no `setUp`/`tearDown` (a `cleanup` fixture handles VRAM), and skips use `pytest.skip` / `@pytest.mark.skip`, never `@unittest.skip`. Fixtures like `tmp_path` and the cached `base_pipe_output` are injected into test methods as arguments. - **Define one config class**, `PipelineTesterConfig`, subclassing `BasePipelineTesterConfig` (from `..testing_utils`). It holds the whole testing contract and performs no assertions: - Set `pipeline_class`, `required_input_params_in_call_signature` (params that must appear in `__call__`'s signature), and `batch_input_params` (params that get batched). Use the canonical sets in `..pipeline_params` where one fits, or an inline `frozenset([...])`. + - Set `output_shape` — the per-sample output shape for `get_dummy_inputs()`, i.e. `(channels, height, width)` for an image pipeline and `(num_frames, channels, height, width)` for a video one. Assert against `self.output_shape` in pipeline-specific tests instead of repeating the literal. - Implement `get_dummy_components(...)` — build every sub-module from the **real classes** at tiny config, each preceded by `torch.manual_seed(0)`. - Implement `get_dummy_inputs()` — **no `device` / `seed` arguments** (unlike the old style). Use `self.get_generator(0)` for the generator, keep sizes tiny, and set `output_type="pt"` so tests compare torch tensors directly with `assert_tensors_close` (no numpy round-trip). Remember `"pt"` images are `(batch, channels, height, width)`. - **Compose the config with one mixin per concern**, one test class each, named `Test...`. Add only the mixins that apply: diff --git a/tests/pipelines/cogvideo/test_cogvideox.py b/tests/pipelines/cogvideo/test_cogvideox.py index efcb03482886..5f5755a08e35 100644 --- a/tests/pipelines/cogvideo/test_cogvideox.py +++ b/tests/pipelines/cogvideo/test_cogvideox.py @@ -46,6 +46,7 @@ class CogVideoXPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (8, 3, 16, 16) # CogVideoX 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"] @@ -134,7 +135,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (8, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4370, 0.3687, 0.3268, 0.3554, 0.3712, 0.3622, 0.3604, 0.3981, 0.5380, 0.5220, 0.5235, 0.5247, 0.5405, 0.5487, 0.5489, 0.5326]) diff --git a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py index 0cae00f6f7a0..9aa8d1e34c7e 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py +++ b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py @@ -35,6 +35,7 @@ class CogVideoXFunControlPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt", "control_video"]) + output_shape = (8, 3, 16, 16) # CogVideoX 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"] @@ -128,7 +129,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (8, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.5921, 0.6076, 0.6015, 0.6024, 0.6140, 0.5966, 0.5870, 0.6079, 0.5049, 0.5086, 0.4891, 0.4783, 0.4770, 0.4738, 0.4537, 0.4508]) diff --git a/tests/pipelines/cogvideo/test_cogvideox_image2video.py b/tests/pipelines/cogvideo/test_cogvideox_image2video.py index 2ef82d90cba6..4e5cae9c193b 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_image2video.py +++ b/tests/pipelines/cogvideo/test_cogvideox_image2video.py @@ -54,6 +54,7 @@ class CogVideoXImageToVideoPipelineTesterConfig(BasePipelineTesterConfig): ] ) batch_input_params = frozenset(["prompt", "image"]) + output_shape = (8, 3, 16, 16) # CogVideoX 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"] @@ -154,7 +155,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (8, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4367, 0.4802, 0.5403, 0.5509, 0.5595, 0.5698, 0.5206, 0.5207, 0.5930, 0.5178, 0.4597, 0.4430, 0.4488, 0.4766, 0.5003, 0.4865]) diff --git a/tests/pipelines/cogvideo/test_cogvideox_video2video.py b/tests/pipelines/cogvideo/test_cogvideox_video2video.py index c418d8babc02..9b6cb76e4e6d 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_video2video.py +++ b/tests/pipelines/cogvideo/test_cogvideox_video2video.py @@ -35,6 +35,7 @@ class CogVideoXVideoToVideoPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt", "video"]) + output_shape = (8, 3, 16, 16) # CogVideoX 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"] @@ -129,7 +130,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (8, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.5644, 0.6029, 0.6017, 0.5937, 0.5991, 0.5907, 0.6141, 0.5340, 0.3184, 0.4219, 0.4406, 0.4330, 0.4692, 0.4547, 0.4562, 0.5092]) diff --git a/tests/pipelines/flux/test_pipeline_flux.py b/tests/pipelines/flux/test_pipeline_flux.py index 3965d33e1230..132a0f636181 100644 --- a/tests/pipelines/flux/test_pipeline_flux.py +++ b/tests/pipelines/flux/test_pipeline_flux.py @@ -46,6 +46,7 @@ class FluxPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) diff --git a/tests/pipelines/flux2/test_pipeline_flux2.py b/tests/pipelines/flux2/test_pipeline_flux2.py index 879865e75648..73c480a82809 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2.py +++ b/tests/pipelines/flux2/test_pipeline_flux2.py @@ -23,6 +23,7 @@ class Flux2PipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "height", "width", "guidance_scale", "prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein.py b/tests/pipelines/flux2/test_pipeline_flux2_klein.py index a38eb4de2a90..0d7139b21e16 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein.py +++ b/tests/pipelines/flux2/test_pipeline_flux2_klein.py @@ -35,6 +35,7 @@ class Flux2KleinPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "height", "width", "guidance_scale", "prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein_inpaint.py b/tests/pipelines/flux2/test_pipeline_flux2_klein_inpaint.py index a8385f66db21..ae1924205e80 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein_inpaint.py +++ b/tests/pipelines/flux2/test_pipeline_flux2_klein_inpaint.py @@ -21,6 +21,7 @@ class Flux2KleinInpaintPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "image", "image_reference", "mask_image", "height", "width", "guidance_scale", "prompt_embeds"] ) batch_input_params = frozenset(["prompt", "image", "image_reference", "mask_image"]) + output_shape = (3, 32, 32) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py b/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py index 141814b92b54..29113510c1c8 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py +++ b/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py @@ -23,6 +23,7 @@ class Flux2KleinKVPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = Flux2KleinKVPipeline required_input_params_in_call_signature = frozenset(["prompt", "height", "width", "prompt_embeds", "image"]) batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) @@ -164,7 +165,7 @@ def test_without_image(self): inputs = self.get_dummy_inputs() del inputs["image"] image = pipe(**inputs).images - assert image.shape == (1, 3, 8, 8) + assert image.shape == (1, *self.output_shape) @pytest.mark.skip("Needs to be revisited") def test_encode_prompt_works_in_isolation(self): diff --git a/tests/pipelines/qwenimage/test_qwenimage.py b/tests/pipelines/qwenimage/test_qwenimage.py index bcc9ec771708..989862c52c90 100644 --- a/tests/pipelines/qwenimage/test_qwenimage.py +++ b/tests/pipelines/qwenimage/test_qwenimage.py @@ -36,6 +36,7 @@ class QwenImagePipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "true_cfg_scale", "height", "width", "guidance_scale", "prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self, num_layers: int = 2): torch.manual_seed(0) @@ -129,7 +130,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() image = pipe(**inputs).images generated_image = image[0] - assert generated_image.shape == (3, 32, 32) + assert generated_image.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.5633, 0.6368, 0.6015, 0.5637, 0.5817, 0.5528, 0.5718, 0.6326, 0.4147, 0.3556, 0.5623, 0.4833, 0.4971, 0.5262, 0.4087, 0.5021]) diff --git a/tests/pipelines/qwenimage/test_qwenimage_controlnet.py b/tests/pipelines/qwenimage/test_qwenimage_controlnet.py index 949bd5833a7a..fe7a61dd831a 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_controlnet.py +++ b/tests/pipelines/qwenimage/test_qwenimage_controlnet.py @@ -49,6 +49,7 @@ class QwenImageControlNetPipelineTesterConfig(BasePipelineTesterConfig): ] ) batch_input_params = frozenset(["prompt", "control_image"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -165,7 +166,7 @@ def test_qwen_controlnet(self): inputs = self.get_dummy_inputs() image = pipe(**inputs).images generated_image = image[0] - assert generated_image.shape == (3, 32, 32) + assert generated_image.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4726, 0.5549, 0.6324, 0.6548, 0.4968, 0.4639, 0.4749, 0.4898, 0.4725, 0.4645, 0.4435, 0.3339, 0.3400, 0.4630, 0.3879, 0.4406]) @@ -188,7 +189,7 @@ def test_qwen_controlnet_multicondition(self): image = pipe(**inputs).images generated_image = image[0] - assert generated_image.shape == (3, 32, 32) + assert generated_image.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.6239, 0.6642, 0.5768, 0.6039, 0.5270, 0.5070, 0.5006, 0.5271, 0.4506, 0.3085, 0.3435, 0.5152, 0.5096, 0.5422, 0.4286, 0.5752]) diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit.py b/tests/pipelines/qwenimage/test_qwenimage_edit.py index 3e6c92033169..8d7f8bc914f4 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_edit.py +++ b/tests/pipelines/qwenimage/test_qwenimage_edit.py @@ -38,6 +38,7 @@ class QwenImageEditPipelineTesterConfig(BasePipelineTesterConfig): ["image", "prompt", "negative_prompt", "true_cfg_scale", "height", "width", "guidance_scale", "prompt_embeds"] ) batch_input_params = frozenset(["prompt", "image"]) + output_shape = (3, 32, 32) def get_dummy_components(self): tiny_ckpt_id = "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" @@ -134,7 +135,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() image = pipe(**inputs).images generated_image = image[0] - assert generated_image.shape == (3, 32, 32) + assert generated_image.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.5637, 0.6341, 0.6001, 0.5620, 0.5794, 0.5498, 0.5757, 0.6389, 0.4174, 0.3597, 0.5649, 0.4894, 0.4969, 0.5255, 0.4083, 0.4986]) diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py index 8680792d6767..2eef47169f08 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py +++ b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py @@ -34,6 +34,7 @@ class QwenImageEditPlusPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "true_cfg_scale", "height", "width", "guidance_scale", "prompt_embeds"] ) batch_input_params = frozenset(["prompt", "image"]) + output_shape = (3, 32, 32) def get_dummy_components(self): tiny_ckpt_id = "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" @@ -131,7 +132,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() image = pipe(**inputs).images generated_image = image[0] - assert generated_image.shape == (3, 32, 32) + assert generated_image.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.5640, 0.6339, 0.5997, 0.5607, 0.5799, 0.5496, 0.5760, 0.6393, 0.4172, 0.3595, 0.5655, 0.4896, 0.4971, 0.5255, 0.4088, 0.4987]) diff --git a/tests/pipelines/qwenimage/test_qwenimage_img2img.py b/tests/pipelines/qwenimage/test_qwenimage_img2img.py index 16a2847e730e..27f858e6b512 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_img2img.py +++ b/tests/pipelines/qwenimage/test_qwenimage_img2img.py @@ -38,6 +38,7 @@ class QwenImageImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "image", "height", "width", "guidance_scale", "true_cfg_scale", "strength"] ) batch_input_params = frozenset(["prompt", "image"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -133,7 +134,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() image = pipe(**inputs).images generated_image = image[0] - assert generated_image.shape == (3, 32, 32) + assert generated_image.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.5832, 0.6525, 0.5760, 0.5546, 0.5763, 0.5743, 0.4986, 0.4591, 0.4371, 0.4667, 0.4898, 0.3232, 0.4333, 0.5274, 0.4735, 0.4937]) diff --git a/tests/pipelines/qwenimage/test_qwenimage_inpaint.py b/tests/pipelines/qwenimage/test_qwenimage_inpaint.py index aa113216020b..c5597bf6f460 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_inpaint.py +++ b/tests/pipelines/qwenimage/test_qwenimage_inpaint.py @@ -49,6 +49,7 @@ class QwenImageInpaintPipelineTesterConfig(BasePipelineTesterConfig): ] ) batch_input_params = frozenset(["prompt", "image", "mask_image"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) diff --git a/tests/pipelines/testing_utils/common.py b/tests/pipelines/testing_utils/common.py index 757c7f102288..ac96f0595429 100644 --- a/tests/pipelines/testing_utils/common.py +++ b/tests/pipelines/testing_utils/common.py @@ -112,6 +112,16 @@ def callback_cfg_params(self) -> frozenset: "callback function when dynamically adjusting `guidance_scale`." ) + @property + def output_shape(self) -> tuple: + raise NotImplementedError( + "You need to set the attribute `output_shape` in the child test class. `output_shape` is the expected " + "per-sample shape of the pipeline output for the standard dummy inputs — the shape of a single element " + "of `pipeline(**get_dummy_inputs())[0]` (i.e. with the batch dimension dropped). For an image pipeline " + "requesting `output_type='pt'` it is `(channels, height, width)`; for a video pipeline it is " + "`(num_frames, channels, height, width)`. Analogous to the model-level `BaseModelTesterConfig.output_shape`." + ) + # ==================== Shared helpers ==================== def get_generator(self, seed=0): @@ -212,6 +222,13 @@ def test_save_load_local(self, tmp_path, base_pipe_output, expected_max_differen output_loaded, base_pipe_output, atol=expected_max_difference, msg="Loaded pipeline output changed." ) + def test_output(self, base_pipe_output): + output = base_pipe_output + assert output is not None, "Pipeline output is None." + assert output[0].shape == self.output_shape, ( + f"Output sample shape does not match expected. Expected {self.output_shape}, got {tuple(output[0].shape)}." + ) + def test_pipeline_call_signature(self): assert hasattr(self.pipeline_class, "__call__"), f"{self.pipeline_class} should have a `__call__` method" diff --git a/tests/pipelines/wan/test_wan.py b/tests/pipelines/wan/test_wan.py index e1339cf115e1..1fd2d011971c 100644 --- a/tests/pipelines/wan/test_wan.py +++ b/tests/pipelines/wan/test_wan.py @@ -28,6 +28,7 @@ class WanPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 16, 16) # 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"] @@ -99,7 +100,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (9, 3, 16, 16) + assert generated_video.shape == self.output_shape # 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]) diff --git a/tests/pipelines/wan/test_wan_22.py b/tests/pipelines/wan/test_wan_22.py index 8c0d873bf489..f34f664fe89c 100644 --- a/tests/pipelines/wan/test_wan_22.py +++ b/tests/pipelines/wan/test_wan_22.py @@ -28,6 +28,7 @@ class Wan22PipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 16, 16) # 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"] @@ -115,7 +116,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (9, 3, 16, 16) + assert generated_video.shape == self.output_shape # 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]) @@ -163,6 +164,7 @@ class Wan225BPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 32, 32) # 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"] @@ -243,7 +245,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (9, 3, 32, 32) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4814, 0.4298, 0.5094, 0.4289, 0.5061, 0.4301, 0.5043, 0.4284, 0.5375, 0.5965, 0.5527, 0.6014, 0.5228, 0.6076, 0.6644, 0.5651]) diff --git a/tests/pipelines/wan/test_wan_22_image_to_video.py b/tests/pipelines/wan/test_wan_22_image_to_video.py index b8967594271b..a06884ddc712 100644 --- a/tests/pipelines/wan/test_wan_22_image_to_video.py +++ b/tests/pipelines/wan/test_wan_22_image_to_video.py @@ -30,6 +30,7 @@ class Wan22ImageToVideoPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 16, 16) # 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"] @@ -123,7 +124,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (9, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4527, 0.4526, 0.4498, 0.4539, 0.4521, 0.4524, 0.4533, 0.4535, 0.5154, 0.5353, 0.5200, 0.5174, 0.5434, 0.5301, 0.5199, 0.5216]) @@ -175,6 +176,7 @@ class Wan225BImageToVideoPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 32, 32) # 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"] @@ -261,7 +263,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (9, 3, 32, 32) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4833, 0.4305, 0.5100, 0.4299, 0.5056, 0.4298, 0.5052, 0.4332, 0.5550, 0.6092, 0.5536, 0.5928, 0.5199, 0.5864, 0.6705, 0.5493]) diff --git a/tests/pipelines/wan/test_wan_animate.py b/tests/pipelines/wan/test_wan_animate.py index c76606c81a87..da46417629f5 100644 --- a/tests/pipelines/wan/test_wan_animate.py +++ b/tests/pipelines/wan/test_wan_animate.py @@ -41,6 +41,7 @@ class WanAnimatePipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (17, 3, 16, 16) # 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"] @@ -152,7 +153,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames[0] - assert video.shape == (17, 3, 16, 16) + assert video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4525, 0.4521, 0.4486, 0.4534, 0.4523, 0.4529, 0.454, 0.4533, 0.5055, 0.5203, 0.5363, 0.4827, 0.5057, 0.5176, 0.5117, 0.5139]) @@ -175,7 +176,7 @@ def test_inference_replacement(self): inputs["mask_video"] = [Image.new("L", (height, width))] * num_frames video = pipe(**inputs).frames[0] - assert video.shape == (17, 3, 16, 16) + assert video.shape == self.output_shape @pytest.mark.skip( reason="Setting the Wan Animate latents to zero at the last denoising step does not guarantee that the output" diff --git a/tests/pipelines/wan/test_wan_image_to_video.py b/tests/pipelines/wan/test_wan_image_to_video.py index 0b881f2742bc..a1a7128aab48 100644 --- a/tests/pipelines/wan/test_wan_image_to_video.py +++ b/tests/pipelines/wan/test_wan_image_to_video.py @@ -36,6 +36,7 @@ class WanImageToVideoPipelineTesterConfig(BasePipelineTesterConfig): ["image", "prompt", "negative_prompt", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 16, 16) # 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"] @@ -129,7 +130,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (9, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4528, 0.4525, 0.4493, 0.4537, 0.4521, 0.4532, 0.4543, 0.4536, 0.5084, 0.5252, 0.5211, 0.5120, 0.5419, 0.5355, 0.5169, 0.5213]) @@ -177,6 +178,7 @@ class WanFLFToVideoPipelineTesterConfig(BasePipelineTesterConfig): ["image", "prompt", "negative_prompt", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 16, 16) # 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"] @@ -273,7 +275,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (9, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4525, 0.4525, 0.4497, 0.4537, 0.4520, 0.4529, 0.4540, 0.4535, 0.5157, 0.5449, 0.5201, 0.5192, 0.5398, 0.5374, 0.5162, 0.5112]) diff --git a/tests/pipelines/wan/test_wan_vace.py b/tests/pipelines/wan/test_wan_vace.py index c4d62aaec389..8a6ffecc6c87 100644 --- a/tests/pipelines/wan/test_wan_vace.py +++ b/tests/pipelines/wan/test_wan_vace.py @@ -36,6 +36,7 @@ class WanVACEPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt"]) + output_shape = (17, 3, 16, 16) # 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"] @@ -125,7 +126,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (17, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off 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]) @@ -143,7 +144,7 @@ def test_inference_with_single_reference_image(self): inputs["reference_images"] = Image.new("RGB", (16, 16)) video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (17, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off 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]) @@ -161,7 +162,7 @@ def test_inference_with_multiple_reference_image(self): inputs["reference_images"] = [[Image.new("RGB", (16, 16))] * 2] video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (17, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off 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]) @@ -179,7 +180,7 @@ def test_inference_with_only_transformer(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames[0] - assert video.shape == (17, 3, 16, 16) + assert video.shape == self.output_shape def test_inference_with_only_transformer_2(self): components = self.get_dummy_components() @@ -197,7 +198,7 @@ def test_inference_with_only_transformer_2(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames[0] - assert video.shape == (17, 3, 16, 16) + assert video.shape == self.output_shape 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) diff --git a/tests/pipelines/wan/test_wan_video_to_video.py b/tests/pipelines/wan/test_wan_video_to_video.py index 92a8d46d6bfc..86949488d0b3 100644 --- a/tests/pipelines/wan/test_wan_video_to_video.py +++ b/tests/pipelines/wan/test_wan_video_to_video.py @@ -30,6 +30,7 @@ class WanVideoToVideoPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "negative_prompt", "height", "width", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] ) batch_input_params = frozenset(["prompt", "video"]) + output_shape = (17, 3, 16, 16) # 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"] @@ -100,7 +101,7 @@ def test_inference(self): inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - assert generated_video.shape == (17, 3, 16, 16) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4522, 0.4534, 0.4532, 0.4553, 0.4526, 0.4538, 0.4533, 0.4547, 0.513, 0.5176, 0.5286, 0.4958, 0.4955, 0.5381, 0.5154, 0.5195])