From a0f96e21b2435f8e9163ad6c1dfa18de0f0f2e8d Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Fri, 24 Jul 2026 11:29:02 -0700 Subject: [PATCH] refactor cogvideox fun control pipeline tests to the new mixin structure --- .../cogvideo/test_cogvideox_fun_control.py | 212 ++++-------------- 1 file changed, 43 insertions(+), 169 deletions(-) diff --git a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py index 30eab652c011..f79c1e71225b 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py +++ b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import unittest import numpy as np import torch @@ -22,41 +20,25 @@ from diffusers import AutoencoderKLCogVideoX, CogVideoXFunControlPipeline, CogVideoXTransformer3DModel, DDIMScheduler -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 ( +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 CogVideoXFunControlPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class CogVideoXFunControlPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = CogVideoXFunControlPipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"control_video"}) - 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", "control_video"]) + # 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): torch.manual_seed(0) @@ -108,169 +90,63 @@ def get_dummy_components(self): 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: int = 0, num_frames: int = 8): - 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 because convolution kernel becomes bigger than sample height = 16 width = 16 - control_video = [Image.new("RGB", (width, height))] * num_frames + control_video = [Image.new("RGB", (width, height))] * 8 - inputs = { + return { "prompt": "dance monkey", "negative_prompt": "", "control_video": control_video, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "height": height, "width": width, "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 TestCogVideoXFunControlPipeline(CogVideoXFunControlPipelineTesterConfig, 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 - - 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 + # 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]) + # fmt: on - 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.5): # NOTE(aryan): This requires a higher expected_max_diff than other CogVideoX pipelines - 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] @@ -281,24 +157,18 @@ def test_vae_tiling(self, expected_diff_max: float = 0.5): 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).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:] @@ -310,12 +180,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:] @@ -328,3 +198,7 @@ def test_fused_qkv_projections(self): assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( "Original outputs should match when fused QKV projections are disabled." ) + + +class TestCogVideoXFunControlPipelineMemory(CogVideoXFunControlPipelineTesterConfig, MemoryTesterMixin): + pass