From d2441f8c2f56a902eb84abd7b434e8afaa84c251 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Thu, 23 Jul 2026 04:58:36 -0700 Subject: [PATCH 1/2] refactor cogvideox pipeline tests to the new mixin structure --- tests/pipelines/cogvideo/test_cogvideox.py | 193 +++++++-------------- 1 file changed, 67 insertions(+), 126 deletions(-) diff --git a/tests/pipelines/cogvideo/test_cogvideox.py b/tests/pipelines/cogvideo/test_cogvideox.py index e1797cfaac9c..4745aadac3d1 100644 --- a/tests/pipelines/cogvideo/test_cogvideox.py +++ b/tests/pipelines/cogvideo/test_cogvideox.py @@ -14,9 +14,9 @@ import gc import inspect -import unittest import numpy as np +import pytest import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -24,52 +24,33 @@ from ...testing_utils import ( backend_empty_cache, - enable_full_determinism, numpy_cosine_similarity_distance, require_torch_accelerator, slow, 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 ( +from ..testing_utils import ( + BasePipelineTesterConfig, FasterCacheTesterMixin, FirstBlockCacheTesterMixin, + MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, check_qkv_fusion_matches_attn_procs_length, check_qkv_fusion_processors_exist, - to_np, ) -enable_full_determinism() - - -class CogVideoXPipelineFastTests( - PipelineTesterMixin, - PyramidAttentionBroadcastTesterMixin, - FasterCacheTesterMixin, - FirstBlockCacheTesterMixin, - unittest.TestCase, -): +class CogVideoXPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = CogVideoXPipeline - 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"]) + # 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"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True def get_dummy_components(self, num_layers: int = 1): torch.manual_seed(0) @@ -121,24 +102,19 @@ def get_dummy_components(self, num_layers: int = 1): text_encoder = T5EncoderModel(config) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, "text_encoder": text_encoder, "tokenizer": tokenizer, } - 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": "", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, # Cannot reduce because convolution kernel becomes bigger than sample @@ -146,26 +122,28 @@ def get_dummy_inputs(self, device, seed=0): "width": 16, "num_frames": 8, "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 TestCogVideoXPipeline(CogVideoXPipelineTesterConfig, 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] + assert generated_video.shape == (8, 3, 16, 16) + + # 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]) + # fmt: on - self.assertEqual(generated_video.shape, (8, 3, 16, 16)) - expected_video = torch.randn(8, 3, 16, 16) - max_diff = np.abs(generated_video - expected_video).max() - self.assertLessEqual(max_diff, 1e10) + 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_callback_inputs(self): sig = inspect.signature(self.pipeline_class.__call__) @@ -175,13 +153,9 @@ def test_callback_inputs(self): if not (has_callback_tensor_inputs and has_callback_step_end): return - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - self.assertTrue( - hasattr(pipe, "_callback_tensor_inputs"), - f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs", + pipe = self.get_pipeline().to(torch_device) + assert hasattr(pipe, "_callback_tensor_inputs"), ( + f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs" ) def callback_inputs_subset(pipe, i, t, callback_kwargs): @@ -203,7 +177,7 @@ def callback_inputs_all(pipe, i, t, callback_kwargs): return callback_kwargs - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() # Test passing in a subset inputs["callback_on_step_end"] = callback_inputs_subset @@ -227,56 +201,13 @@ def callback_inputs_change_tensor(pipe, i, t, callback_kwargs): assert output.abs().sum() < 1e10 def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3) - - def test_attention_slicing_forward_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 - ): - if not self.test_attention_slicing: - return - - components = self.get_dummy_components() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].eval() - 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) - output_without_slicing = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=1) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing1 = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=2) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing2 = pipe(**inputs)[0] - - if test_max_difference: - max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max() - max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max() - self.assertLess( - max(max_diff1, max_diff2), - expected_max_diff, - "Attention slicing should not affect the inference results", - ) + super().test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3) def test_vae_tiling(self, expected_diff_max: float = 0.2): - generator_device = "cpu" - components = self.get_dummy_components() - - pipe = self.pipeline_class(**components) - pipe.to("cpu") - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() # Without tiling - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() inputs["height"] = inputs["width"] = 128 output_without_tiling = pipe(**inputs)[0] @@ -287,24 +218,18 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): tile_overlap_factor_height=1 / 12, tile_overlap_factor_width=1 / 12, ) - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() inputs["height"] = inputs["width"] = 128 output_with_tiling = pipe(**inputs)[0] - self.assertLess( - (to_np(output_without_tiling) - to_np(output_with_tiling)).max(), - expected_diff_max, - "VAE tiling should not affect the inference results", + assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, ( + "VAE tiling should not affect the inference results" ) def test_fused_qkv_projections(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() frames = pipe(**inputs).frames # [B, F, C, H, W] original_image_slice = frames[0, -2:, -1, -3:, -3:] @@ -316,12 +241,12 @@ def test_fused_qkv_projections(self): pipe.transformer, pipe.transformer.original_attn_processors ), "Something wrong with the attention processors concerning the fused QKV projections." - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() frames = pipe(**inputs).frames image_slice_fused = frames[0, -2:, -1, -3:, -3:] pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() frames = pipe(**inputs).frames image_slice_disabled = frames[0, -2:, -1, -3:, -3:] @@ -336,18 +261,34 @@ def test_fused_qkv_projections(self): ) +class TestCogVideoXPipelineMemory(CogVideoXPipelineTesterConfig, MemoryTesterMixin): + pass + + +class TestCogVideoXPipelinePyramidAttentionBroadcast( + CogVideoXPipelineTesterConfig, PyramidAttentionBroadcastTesterMixin +): + pass + + +class TestCogVideoXPipelineFasterCache(CogVideoXPipelineTesterConfig, FasterCacheTesterMixin): + pass + + +class TestCogVideoXPipelineFirstBlockCache(CogVideoXPipelineTesterConfig, FirstBlockCacheTesterMixin): + pass + + @slow @require_torch_accelerator -class CogVideoXPipelineIntegrationTests(unittest.TestCase): +class TestCogVideoXPipelineSlow: prompt = "A painting of a squirrel eating a burger." - def setUp(self): - super().setUp() + @pytest.fixture(autouse=True) + def cleanup(self): gc.collect() backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() + yield gc.collect() backend_empty_cache(torch_device) From 71eeec34a609d9719766e9f6522cb823d149ecfc Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Fri, 24 Jul 2026 05:33:52 -0700 Subject: [PATCH 2/2] drop test_callback_inputs override, base mixin covers it --- tests/pipelines/cogvideo/test_cogvideox.py | 56 ---------------------- 1 file changed, 56 deletions(-) diff --git a/tests/pipelines/cogvideo/test_cogvideox.py b/tests/pipelines/cogvideo/test_cogvideox.py index 4745aadac3d1..e81cd2ffc9a9 100644 --- a/tests/pipelines/cogvideo/test_cogvideox.py +++ b/tests/pipelines/cogvideo/test_cogvideox.py @@ -13,7 +13,6 @@ # limitations under the License. import gc -import inspect import numpy as np import pytest @@ -145,61 +144,6 @@ def test_inference(self): generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) assert torch.allclose(generated_slice, expected_slice, atol=1e-3) - def test_callback_inputs(self): - sig = inspect.signature(self.pipeline_class.__call__) - has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters - has_callback_step_end = "callback_on_step_end" in sig.parameters - - if not (has_callback_tensor_inputs and has_callback_step_end): - return - - pipe = self.get_pipeline().to(torch_device) - assert hasattr(pipe, "_callback_tensor_inputs"), ( - f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs" - ) - - def callback_inputs_subset(pipe, i, t, callback_kwargs): - # iterate over callback args - for tensor_name, tensor_value in callback_kwargs.items(): - # check that we're only passing in allowed tensor inputs - assert tensor_name in pipe._callback_tensor_inputs - - return callback_kwargs - - def callback_inputs_all(pipe, i, t, callback_kwargs): - for tensor_name in pipe._callback_tensor_inputs: - assert tensor_name in callback_kwargs - - # iterate over callback args - for tensor_name, tensor_value in callback_kwargs.items(): - # check that we're only passing in allowed tensor inputs - assert tensor_name in pipe._callback_tensor_inputs - - return callback_kwargs - - inputs = self.get_dummy_inputs() - - # Test passing in a subset - inputs["callback_on_step_end"] = callback_inputs_subset - inputs["callback_on_step_end_tensor_inputs"] = ["latents"] - output = pipe(**inputs)[0] - - # Test passing in a everything - inputs["callback_on_step_end"] = callback_inputs_all - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - output = pipe(**inputs)[0] - - def callback_inputs_change_tensor(pipe, i, t, callback_kwargs): - is_last = i == (pipe.num_timesteps - 1) - if is_last: - callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"]) - return callback_kwargs - - inputs["callback_on_step_end"] = callback_inputs_change_tensor - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - output = pipe(**inputs)[0] - assert output.abs().sum() < 1e10 - def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3)