From e48bfdd3863c222a022842f4cef645aa8f037a83 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Fri, 24 Jul 2026 07:27:39 -0700 Subject: [PATCH] refactor cogvideox image to video pipeline tests to the new mixin structure --- .../cogvideo/test_cogvideox_image2video.py | 213 +++++------------- 1 file changed, 52 insertions(+), 161 deletions(-) diff --git a/tests/pipelines/cogvideo/test_cogvideox_image2video.py b/tests/pipelines/cogvideo/test_cogvideox_image2video.py index fc9e35c51b59..f00fd30e7928 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_image2video.py +++ b/tests/pipelines/cogvideo/test_cogvideox_image2video.py @@ -13,10 +13,9 @@ # limitations under the License. import gc -import inspect -import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -26,41 +25,39 @@ 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, + MemoryTesterMixin, PipelineTesterMixin, check_qkv_fusion_matches_attn_procs_length, check_qkv_fusion_processors_exist, - to_np, ) -enable_full_determinism() - - -class CogVideoXImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class CogVideoXImageToVideoPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = CogVideoXImageToVideoPipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"image"}) - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( + required_input_params_in_call_signature = frozenset( [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", + "image", + "prompt", + "negative_prompt", + "height", + "width", + "guidance_scale", + "prompt_embeds", + "negative_prompt_embeds", ] ) - test_xformers_attention = False + batch_input_params = frozenset(["prompt", "image"]) + # 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"] + ) def get_dummy_components(self): torch.manual_seed(0) @@ -126,12 +123,7 @@ def get_dummy_components(self): } 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): # Cannot reduce below 16 because convolution kernel becomes bigger than sample # Cannot reduce below 32 because 3D RoPE errors out image_height = 16 @@ -141,137 +133,42 @@ def get_dummy_inputs(self, device, seed=0): "image": image, "prompt": "dance monkey", "negative_prompt": "", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "height": image_height, "width": image_width, "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 TestCogVideoXImageToVideoPipeline(CogVideoXImageToVideoPipelineTesterConfig, 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) - 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) - - 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 - - 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", - ) - - 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 + # 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]) + # fmt: on - 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(torch_device) - - # 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 + 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_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.3): # Note(aryan): Investigate why this needs a bit higher tolerance - generator_device = "cpu" components = self.get_dummy_components() # The reason to modify it this way is because I2V Transformer limits the generation to resolutions used during initialization. @@ -284,12 +181,10 @@ def test_vae_tiling(self, expected_diff_max: float = 0.3): sample_width=16, ) - pipe = self.pipeline_class(**components) - pipe.to("cpu") - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline(**components) # 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] @@ -300,24 +195,18 @@ def test_vae_tiling(self, expected_diff_max: float = 0.3): 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:] @@ -329,12 +218,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:] @@ -349,18 +238,20 @@ def test_fused_qkv_projections(self): ) +class TestCogVideoXImageToVideoPipelineMemory(CogVideoXImageToVideoPipelineTesterConfig, MemoryTesterMixin): + pass + + @slow @require_torch_accelerator -class CogVideoXImageToVideoPipelineIntegrationTests(unittest.TestCase): +class TestCogVideoXImageToVideoPipelineSlow: 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)