[tests] introduce output_shape property in the pipeline tests. (huggingface#14253)

* introduce output_shape property in the pipeline tests.

* up

* add missing

* up

Author: sayakpaul

.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**, `<Pipeline>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<Pipeline>...`. Add only the mixins that apply:

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])

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])

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])

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])

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)

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)

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)

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)

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):