diff --git a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py index d723c7941ae9..959034c89fdd 100644 --- a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py +++ b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py @@ -951,6 +951,7 @@ def __call__( encoder_hidden_states_t5=prompt_embeds_t5, encoder_hidden_states_llama3=prompt_embeds_llama3, pooled_embeds=pooled_prompt_embeds, + attention_kwargs=self.attention_kwargs, return_dict=False, )[0] noise_pred = -noise_pred diff --git a/tests/pipelines/hidream_image/test_pipeline_hidream.py b/tests/pipelines/hidream_image/test_pipeline_hidream.py index e98858f2e59a..fc4b93daee29 100644 --- a/tests/pipelines/hidream_image/test_pipeline_hidream.py +++ b/tests/pipelines/hidream_image/test_pipeline_hidream.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - -import numpy as np import torch from transformers import ( AutoConfig, @@ -34,23 +31,26 @@ HiDreamImageTransformer2DModel, ) -from ...testing_utils import enable_full_determinism -from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import PipelineTesterMixin +from ...testing_utils import assert_tensors_close, enable_full_determinism +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() -class HiDreamImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class HiDreamImagePipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HiDreamImagePipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs", "prompt_embeds", "negative_prompt_embeds"} - 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 = PipelineTesterMixin.required_optional_params - test_xformers_attention = False - test_layerwise_casting = True + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt"] + ) + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 128, 128) def get_dummy_components(self): torch.manual_seed(0) @@ -107,7 +107,7 @@ def get_dummy_components(self): scheduler = FlowMatchEulerDiscreteScheduler() - components = { + return { "scheduler": scheduler, "vae": vae, "text_encoder": text_encoder, @@ -120,42 +120,46 @@ def get_dummy_components(self): "tokenizer_4": tokenizer_4, "transformer": transformer, } - 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": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, - "output_type": "np", + # 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 TestHiDreamImagePipeline(HiDreamImagePipelineTesterConfig, 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) - image = pipe(**inputs)[0] + image = pipe(**self.get_dummy_inputs())[0] generated_image = image[0] - self.assertEqual(generated_image.shape, (128, 128, 3)) + assert generated_image.shape == self.output_shape # fmt: off - expected_slice = np.array([0.3841, 0.5881, 0.3958, 0.5585, 0.4646, 0.4385, 0.4529, 0.4178, 0.6595, 0.2871, 0.7331, 0.5203, 0.5525, 0.5296, 0.4695, 0.5524]) + expected_slice = torch.tensor([0.3841, 0.5585, 0.4529, 0.5123, 0.4314, 0.5122, 0.4885, 0.4554, 0.4653, 0.3369, 0.1799, 0.6126, 0.6202, 0.2871, 0.5525, 0.5524]) # fmt: on generated_slice = generated_image.flatten() - generated_slice = np.concatenate([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(np.allclose(generated_slice, expected_slice, atol=5e-3)) + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + assert_tensors_close(generated_slice, expected_slice, atol=5e-3) + + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=3e-4): + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) + + +class TestHiDreamImagePipelineMemory(HiDreamImagePipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the HiDream Image pipeline.""" + + +class TestHiDreamImagePipelineLoRA(HiDreamImagePipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the HiDream Image pipeline.""" + - def test_inference_batch_single_identical(self): - super().test_inference_batch_single_identical(expected_max_diff=3e-4) +class TestHiDreamImagePipelineLoRAMemory(HiDreamImagePipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the HiDream Image pipeline.""" diff --git a/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py b/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py index 7d8a244da166..8021186000d9 100644 --- a/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py +++ b/tests/pipelines/hunyuan_image_21/test_hunyuanimage.py @@ -12,9 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - -import numpy as np +import pytest import torch from transformers import ( ByT5Tokenizer, @@ -33,36 +31,23 @@ HunyuanImageTransformer2DModel, ) -from ...testing_utils import enable_full_determinism -from ..test_pipelines_common import FirstBlockCacheTesterMixin, PipelineTesterMixin, to_np +from ...testing_utils import assert_tensors_close, enable_full_determinism +from ..testing_utils import ( + BasePipelineTesterConfig, + FirstBlockCacheTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() -class HunyuanImagePipelineFastTests( - PipelineTesterMixin, - FirstBlockCacheTesterMixin, - unittest.TestCase, -): +class HunyuanImagePipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HunyuanImagePipeline - params = frozenset(["prompt", "height", "width"]) - batch_params = frozenset(["prompt", "negative_prompt"]) - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] - ) - - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True - test_attention_slicing = False + required_input_params_in_call_signature = frozenset(["prompt", "height", "width"]) + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 16, 16) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1, guidance_embeds: bool = False): torch.manual_seed(0) @@ -153,7 +138,7 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1, text_encoder_2 = T5EncoderModel(t5_config) tokenizer_2 = ByT5Tokenizer() - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -164,122 +149,93 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1, "guider": guider, "ocr_guider": ocr_guider, } - 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": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 5, "height": 16, "width": 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 TestHunyuanImagePipeline(HunyuanImagePipelineTesterConfig, 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) - image = pipe(**inputs).images + image = pipe(**self.get_dummy_inputs()).images generated_image = image[0] - self.assertEqual(generated_image.shape, (3, 16, 16)) + assert generated_image.shape == self.output_shape - expected_slice_np = np.array( - [0.6252659, 0.51482046, 0.60799813, 0.59267783, 0.488082, 0.5857634, 0.523781, 0.58028054, 0.5674121] - ) - output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy() + # fmt: off + expected_slice = torch.tensor([0.6252659, 0.51482046, 0.60799813, 0.59267783, 0.488082, 0.5857634, 0.523781, 0.58028054, 0.5674121]) + # fmt: on - self.assertTrue( - np.abs(output_slice - expected_slice_np).max() < 1e-3, - f"output_slice: {output_slice}, expected_slice_np: {expected_slice_np}", - ) + generated_slice = generated_image[0, -3:, -3:].flatten() + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) def test_inference_guider(self): - device = "cpu" - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() pipe.guider = pipe.guider.new(guidance_scale=1000) pipe.ocr_guider = pipe.ocr_guider.new(guidance_scale=1000) - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images + image = pipe(**self.get_dummy_inputs()).images generated_image = image[0] - self.assertEqual(generated_image.shape, (3, 16, 16)) + assert generated_image.shape == self.output_shape - expected_slice_np = np.array( - [0.6068114, 0.48716035, 0.5984431, 0.60241306, 0.48849544, 0.5624479, 0.53696984, 0.58964247, 0.54248774] - ) - output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy() + # fmt: off + expected_slice = torch.tensor([0.6068114, 0.48716035, 0.5984431, 0.60241306, 0.48849544, 0.5624479, 0.53696984, 0.58964247, 0.54248774]) + # fmt: on - self.assertTrue( - np.abs(output_slice - expected_slice_np).max() < 1e-3, - f"output_slice: {output_slice}, expected_slice_np: {expected_slice_np}", - ) + generated_slice = generated_image[0, -3:, -3:].flatten() + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) def test_inference_with_distilled_guidance(self): - device = "cpu" + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline(**self.get_dummy_components(guidance_embeds=True)) - components = self.get_dummy_components(guidance_embeds=True) - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["distilled_guidance_scale"] = 3.5 - image = pipe(**inputs).images + image = pipe(**self.get_dummy_inputs(), distilled_guidance_scale=3.5).images generated_image = image[0] - self.assertEqual(generated_image.shape, (3, 16, 16)) + assert generated_image.shape == self.output_shape - expected_slice_np = np.array( - [0.63667065, 0.5187377, 0.66757566, 0.6320319, 0.4913387, 0.54813194, 0.5335031, 0.5736143, 0.5461346] - ) - output_slice = generated_image[0, -3:, -3:].flatten().cpu().numpy() + # fmt: off + expected_slice = torch.tensor([0.63667065, 0.5187377, 0.66757566, 0.6320319, 0.4913387, 0.54813194, 0.5335031, 0.5736143, 0.5461346]) + # fmt: on - self.assertTrue( - np.abs(output_slice - expected_slice_np).max() < 1e-3, - f"output_slice: {output_slice}, expected_slice_np: {expected_slice_np}", - ) + generated_slice = generated_image[0, -3:, -3:].flatten() + assert_tensors_close(generated_slice, expected_slice, atol=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["height"] = inputs["width"] = 128 - output_without_tiling = pipe(**inputs)[0] + output_without_tiling = self.run_pipe(pipe, height=128, width=128) # With tiling pipe.vae.enable_tiling(tile_sample_min_size=96) - inputs = self.get_dummy_inputs(generator_device) - 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", + output_with_tiling = self.run_pipe(pipe, height=128, width=128) + + assert_tensors_close( + output_with_tiling, + output_without_tiling, + atol=expected_diff_max, + msg="VAE tiling should not affect the inference results.", ) - @unittest.skip("TODO: Test not supported for now because needs to be adjusted to work with guiders.") + @pytest.mark.skip("TODO: Test not supported for now because needs to be adjusted to work with guiders.") def test_encode_prompt_works_in_isolation(self): pass + + +class TestHunyuanImagePipelineMemory(HunyuanImagePipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the HunyuanImage pipeline.""" + + +class TestHunyuanImagePipelineFirstBlockCache(HunyuanImagePipelineTesterConfig, FirstBlockCacheTesterMixin): + """First Block Cache tests for the HunyuanImage pipeline.""" diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py b/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py index 7bd845f8baef..3596c5ca6941 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py @@ -12,10 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import unittest - -import numpy as np +import pytest import torch from PIL import Image from transformers import ( @@ -37,37 +34,33 @@ HunyuanVideoTransformer3DModel, ) -from ...testing_utils import enable_full_determinism, torch_device -from ..test_pipelines_common import PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, to_np +from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, + PyramidAttentionBroadcastTesterMixin, +) enable_full_determinism() -class HunyuanVideoImageToVideoPipelineFastTests( - PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, unittest.TestCase -): +class HunyuanVideoImageToVideoPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HunyuanVideoImageToVideoPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( ["image", "prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] ) - batch_params = frozenset(["prompt", "image"]) - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + batch_input_params = frozenset(["prompt", "image"]) + # NOTE: The generated video has 4 fewer frames than requested because they are dropped in the pipeline. + output_shape = (5, 3, 16, 16) + # HunyuanVideo is a video pipeline: it exposes `num_videos_per_prompt`, not `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) - # there is no xformers processor for Flux - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True - def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = HunyuanVideoTransformer3DModel( @@ -176,7 +169,7 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): size=224, ) - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -186,18 +179,12 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): "tokenizer_2": tokenizer_2, "image_processor": image_processor, } - 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): image_height = 16 image_width = 16 image = Image.new("RGB", (image_width, image_height)) - inputs = { + return { "image": image, "prompt": "dance monkey", "prompt_template": { @@ -207,30 +194,26 @@ def get_dummy_inputs(self, device, seed=0): "image_emb_start": 5, "image_emb_end": 54, }, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 4.5, "height": image_height, "width": image_width, "num_frames": 9, "max_sequence_length": 64, + # 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 TestHunyuanVideoImageToVideoPipeline(HunyuanVideoImageToVideoPipelineTesterConfig, 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) - video = pipe(**inputs).frames + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - # NOTE: The expected video has 4 lesser frames because they are dropped in the pipeline - self.assertEqual(generated_video.shape, (5, 3, 16, 16)) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4477, 0.4781, 0.4478, 0.5687, 0.3446, 0.1606, 0.2699, 0.3613, 0.5592, 0.6789, 0.6793, 0.5311, 0.5175, 0.3748, 0.4228, 0.4149]) @@ -238,119 +221,16 @@ def test_inference(self): generated_slice = generated_video.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue( - torch.allclose(generated_slice, expected_slice, atol=1e-3), - "The generated video does not match the expected slice.", - ) - - 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", + assert_tensors_close( + generated_slice, expected_slice, atol=1e-3, msg="The generated video does not match the expected slice." ) - 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(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 - - 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() - 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", - ) - - def test_vae_tiling(self, expected_diff_max: float = 0.2): + def test_vae_tiling(self, expected_diff_max: float = 0.6): # Seems to require higher tolerance than the other tests - expected_diff_max = 0.6 - 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().to(torch_device) # Without tiling - inputs = self.get_dummy_inputs(generator_device) - inputs["height"] = inputs["width"] = 128 - output_without_tiling = pipe(**inputs)[0] + output_without_tiling = self.run_pipe(pipe, height=128, width=128) # With tiling pipe.vae.enable_tiling( @@ -359,31 +239,47 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): tile_sample_stride_height=64, tile_sample_stride_width=64, ) - inputs = self.get_dummy_inputs(generator_device) - 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", + output_with_tiling = self.run_pipe(pipe, height=128, width=128) + + assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, ( + "VAE tiling should not affect the inference results." ) # TODO(aryan): Create a dummy gemma model with smol vocab size - @unittest.skip( + @pytest.mark.skip( "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass - @unittest.skip( + @pytest.mark.skip( "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass - @unittest.skip( + @pytest.mark.skip( "Encode prompt currently does not work in isolation because of requiring image embeddings from image processor. The test does not handle this case, or we need to rewrite encode_prompt." ) def test_encode_prompt_works_in_isolation(self): pass + + +class TestHunyuanVideoImageToVideoPipelineMemory(HunyuanVideoImageToVideoPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the HunyuanVideo I2V pipeline.""" + + +class TestHunyuanVideoImageToVideoPipelinePyramidAttentionBroadcast( + HunyuanVideoImageToVideoPipelineTesterConfig, PyramidAttentionBroadcastTesterMixin +): + """Pyramid Attention Broadcast cache tests for the HunyuanVideo I2V pipeline.""" + + +class TestHunyuanVideoImageToVideoPipelineLoRA(HunyuanVideoImageToVideoPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the HunyuanVideo I2V pipeline.""" + + +class TestHunyuanVideoImageToVideoPipelineLoRAMemory( + HunyuanVideoImageToVideoPipelineTesterConfig, LoraMemoryTesterMixin +): + """LoRA x memory-optimization tests (group offload, CPU offload) for the HunyuanVideo I2V pipeline.""" diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py b/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py index dde3ba718ba5..63299cda1901 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py @@ -12,10 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import unittest - -import numpy as np +import pytest import torch from PIL import Image from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer, LlamaConfig, LlamaModel, LlamaTokenizer @@ -27,37 +24,32 @@ HunyuanVideoTransformer3DModel, ) -from ...testing_utils import enable_full_determinism, torch_device -from ..test_pipelines_common import PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, to_np +from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, + PyramidAttentionBroadcastTesterMixin, +) enable_full_determinism() -class HunyuanSkyreelsImageToVideoPipelineFastTests( - PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, unittest.TestCase -): +class HunyuanSkyreelsImageToVideoPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HunyuanSkyreelsImageToVideoPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( ["image", "prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] ) - batch_params = frozenset(["prompt", "image"]) - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + batch_input_params = frozenset(["prompt", "image"]) + output_shape = (9, 3, 16, 16) + # HunyuanVideo is a video pipeline: it exposes `num_videos_per_prompt`, not `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) - # there is no xformers processor for Flux - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True - def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = HunyuanVideoTransformer3DModel( @@ -141,7 +133,7 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): text_encoder_2 = CLIPTextModel(clip_text_encoder_config) tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -150,48 +142,39 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): "tokenizer": tokenizer, "tokenizer_2": tokenizer_2, } - 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): image_height = 16 image_width = 16 image = Image.new("RGB", (image_width, image_height)) - inputs = { + return { "image": image, "prompt": "dance monkey", "prompt_template": { "template": "{}", "crop_start": 0, }, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 4.5, - "height": 16, - "width": 16, + "height": image_height, + "width": image_width, # 4 * k + 1 is the recommendation "num_frames": 9, "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 TestHunyuanSkyreelsImageToVideoPipeline(HunyuanSkyreelsImageToVideoPipelineTesterConfig, 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) - video = pipe(**inputs).frames + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (9, 3, 16, 16)) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.5979, 0.5689, 0.5049, 0.4954, 0.4626, 0.5027, 0.4998, 0.5639, 0.5746, 0.5710, 0.5034, 0.5987, 0.6288, 0.5199, 0.5518, 0.5783]) @@ -199,119 +182,16 @@ def test_inference(self): generated_slice = generated_video.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue( - torch.allclose(generated_slice, expected_slice, atol=1e-3), - "The generated video does not match the expected slice.", - ) - - 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", + assert_tensors_close( + generated_slice, expected_slice, atol=1e-3, msg="The generated video does not match the expected slice." ) - 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(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 - - 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() - 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", - ) - - def test_vae_tiling(self, expected_diff_max: float = 0.2): + def test_vae_tiling(self, expected_diff_max: float = 0.6): # Seems to require higher tolerance than the other tests - expected_diff_max = 0.6 - 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().to(torch_device) # Without tiling - inputs = self.get_dummy_inputs(generator_device) - inputs["height"] = inputs["width"] = 128 - output_without_tiling = pipe(**inputs)[0] + output_without_tiling = self.run_pipe(pipe, height=128, width=128) # With tiling pipe.vae.enable_tiling( @@ -320,25 +200,43 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): tile_sample_stride_height=64, tile_sample_stride_width=64, ) - inputs = self.get_dummy_inputs(generator_device) - 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", + output_with_tiling = self.run_pipe(pipe, height=128, width=128) + + assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, ( + "VAE tiling should not affect the inference results." ) # TODO(aryan): Create a dummy gemma model with smol vocab size - @unittest.skip( + @pytest.mark.skip( "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass - @unittest.skip( + @pytest.mark.skip( "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass + + +class TestHunyuanSkyreelsImageToVideoPipelineMemory( + HunyuanSkyreelsImageToVideoPipelineTesterConfig, MemoryTesterMixin +): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Skyreels I2V pipeline.""" + + +class TestHunyuanSkyreelsImageToVideoPipelinePyramidAttentionBroadcast( + HunyuanSkyreelsImageToVideoPipelineTesterConfig, PyramidAttentionBroadcastTesterMixin +): + """Pyramid Attention Broadcast cache tests for the Skyreels I2V pipeline.""" + + +class TestHunyuanSkyreelsImageToVideoPipelineLoRA(HunyuanSkyreelsImageToVideoPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Skyreels I2V pipeline.""" + + +class TestHunyuanSkyreelsImageToVideoPipelineLoRAMemory( + HunyuanSkyreelsImageToVideoPipelineTesterConfig, LoraMemoryTesterMixin +): + """LoRA x memory-optimization tests (group offload, CPU offload) for the Skyreels I2V pipeline.""" diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py b/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py index 843d1843857c..7a0fcc731cd1 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py @@ -12,10 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import unittest - -import numpy as np +import pytest import torch from PIL import Image from transformers import ( @@ -31,49 +28,42 @@ from diffusers import ( AutoencoderKLHunyuanVideo, - FasterCacheConfig, FlowMatchEulerDiscreteScheduler, HunyuanVideoFramepackPipeline, HunyuanVideoFramepackTransformer3DModel, ) from ...testing_utils import ( + assert_tensors_close, enable_full_determinism, torch_device, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, FasterCacheTesterMixin, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, - to_np, ) enable_full_determinism() -class HunyuanVideoFramepackPipelineFastTests( - PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, FasterCacheTesterMixin, unittest.TestCase -): +class HunyuanVideoFramepackPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HunyuanVideoFramepackPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( ["image", "prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] ) - batch_params = frozenset(["image", "prompt"]) - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + batch_input_params = frozenset(["image", "prompt"]) + output_shape = (13, 3, 32, 32) + # Framepack is a video pipeline (`num_videos_per_prompt`) and takes `image_latents` rather than `latents`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "output_type", "return_dict"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True - # `image_encoder` is a `SiglipVisionModel`, whose attention pooling head # (`SiglipMultiheadAttentionPoolingHead`) wraps a `torch.nn.MultiheadAttention`. That hands # `self.out_proj.weight` to `torch.nn.functional.multi_head_attention_forward` instead of calling @@ -83,14 +73,6 @@ class HunyuanVideoFramepackPipelineFastTests( # so exclude just this one rather than skipping the test. group_offloading_leaf_level_exclude_modules = ["image_encoder"] - faster_cache_config = FasterCacheConfig( - spatial_attention_block_skip_range=2, - spatial_attention_timestep_skip_range=(-1, 901), - unconditional_batch_skip_range=2, - attention_weight_callback=lambda _: 0.5, - is_guidance_distilled=True, - ) - def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) transformer = HunyuanVideoFramepackTransformer3DModel( @@ -183,7 +165,7 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): ) image_encoder = SiglipVisionModel.from_pretrained("hf-internal-testing/tiny-random-SiglipVisionModel") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -194,25 +176,19 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): "feature_extractor": feature_extractor, "image_encoder": image_encoder, } - 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): image_height = 32 image_width = 32 image = Image.new("RGB", (image_width, image_height)) - inputs = { + return { "image": image, "prompt": "dance monkey", "prompt_template": { "template": "{}", "crop_start": 0, }, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 4.5, "height": image_height, @@ -220,22 +196,19 @@ def get_dummy_inputs(self, device, seed=0): "num_frames": 9, "latent_window_size": 3, "max_sequence_length": 256, + # 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 TestHunyuanVideoFramepackPipeline(HunyuanVideoFramepackPipelineTesterConfig, 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) - video = pipe(**inputs).frames + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (13, 3, 32, 32)) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.3628, 0.3380, 0.3421, 0.3505, 0.3362, 0.3268, 0.4167, 0.4063, 0.5225, 0.4693, 0.4827, 0.4583, 0.4144, 0.3983, 0.4089, 0.4587]) @@ -243,58 +216,68 @@ def test_inference(self): generated_slice = generated_video.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue( - torch.allclose(generated_slice, expected_slice, atol=1e-3), - "The generated video does not match the expected slice.", + assert_tensors_close( + generated_slice, expected_slice, atol=1e-3, msg="The generated video does not match the expected slice." + ) + + def test_vae_tiling(self, expected_diff_max: float = 0.6): + # Seems to require higher tolerance than the other tests + pipe = self.get_pipeline().to(torch_device) + + # Without tiling + output_without_tiling = self.run_pipe(pipe, height=128, width=128) + + # With tiling + pipe.vae.enable_tiling( + tile_sample_min_height=96, + tile_sample_min_width=96, + tile_sample_stride_height=64, + tile_sample_stride_width=64, + ) + output_with_tiling = self.run_pipe(pipe, height=128, width=128) + + assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, ( + "VAE tiling should not affect the inference results." ) 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", + # Framepack does not fit the shared version of this test: with `output_type="latent"` it returns the + # accumulated history as a one-element list rather than a tensor, and that history keeps the leading + # image-conditioning latent frame the callback never sees. So zeroing `latents` in the callback zeroes + # every generated frame but not that first one — assert exactly that instead. + 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(): + for tensor_name in callback_kwargs: # 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(): + for tensor_name in callback_kwargs: # 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) + inputs = self.get_dummy_inputs() + inputs["output_type"] = "latent" # 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] + pipe(**inputs) - # Test passing in a everything + # Test passing in everything inputs["callback_on_step_end"] = callback_inputs_all inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - output = pipe(**inputs)[0] + pipe(**inputs) def callback_inputs_change_tensor(pipe, i, t, callback_kwargs): is_last = i == (pipe.num_timesteps - 1) @@ -304,83 +287,28 @@ def callback_inputs_change_tensor(pipe, i, t, 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_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() - 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", - ) - - def test_vae_tiling(self, expected_diff_max: float = 0.2): - # Seems to require higher tolerance than the other tests - expected_diff_max = 0.6 - generator_device = "cpu" - components = self.get_dummy_components() + latents = pipe(**inputs)[0][0] + # Frame 0 is the image-conditioning latent; the rest are what the denoising loop produced. + assert latents[:, :, 1:].abs().sum() == 0 - pipe = self.pipeline_class(**components) - pipe.to("cpu") - pipe.set_progress_bar_config(disable=None) + # TODO(aryan): Create a dummy gemma model with smol vocab size + @pytest.mark.skip( + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + ) + def test_inference_batch_consistent(self): + pass - # Without tiling - inputs = self.get_dummy_inputs(generator_device) - inputs["height"] = inputs["width"] = 128 - output_without_tiling = pipe(**inputs)[0] + @pytest.mark.skip( + "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." + ) + def test_inference_batch_single_identical(self): + pass - # With tiling - pipe.vae.enable_tiling( - tile_sample_min_height=96, - tile_sample_min_width=96, - tile_sample_stride_height=64, - tile_sample_stride_width=64, - ) - inputs = self.get_dummy_inputs(generator_device) - 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", - ) - def test_float16_inference(self, expected_max_diff=0.2): - # NOTE: this test needs a higher tolerance because of multiple forwards through - # the model, which compounds the overall fp32 vs fp16 numerical differences. It - # shouldn't be expected that the results are the same, so we bump the tolerance. - return super().test_float16_inference(expected_max_diff) +class TestHunyuanVideoFramepackPipelineMemory(HunyuanVideoFramepackPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Framepack pipeline.""" - @unittest.skip("The image_encoder uses SiglipVisionModel, which does not support sequential CPU offloading.") + @pytest.mark.skip("The image_encoder uses SiglipVisionModel, which does not support sequential CPU offloading.") def test_sequential_cpu_offload_forward_pass(self): # https://github.com/huggingface/transformers/blob/21cb353b7b4f77c6f5f5c3341d660f86ff416d04/src/transformers/models/siglip/modeling_siglip.py#L803 # This is because it instantiates it's attention layer from torch.nn.MultiheadAttention, which calls to @@ -389,7 +317,7 @@ def test_sequential_cpu_offload_forward_pass(self): # this test because of MHA (example: HunyuanDiT because of AttentionPooling layer). pass - @unittest.skip("The image_encoder uses SiglipVisionModel, which does not support sequential CPU offloading.") + @pytest.mark.skip("The image_encoder uses SiglipVisionModel, which does not support sequential CPU offloading.") def test_sequential_offload_forward_pass_twice(self): # https://github.com/huggingface/transformers/blob/21cb353b7b4f77c6f5f5c3341d660f86ff416d04/src/transformers/models/siglip/modeling_siglip.py#L803 # This is because it instantiates it's attention layer from torch.nn.MultiheadAttention, which calls to @@ -398,23 +326,29 @@ def test_sequential_offload_forward_pass_twice(self): # this test because of MHA (example: HunyuanDiT because of AttentionPooling layer). pass - @unittest.skip("The image_encoder uses SiglipVisionModel, which does not support group offloading.") - def test_pipeline_level_group_offloading_inference(self): - # Same root cause as the sequential CPU offloading skips above: the attention layer is a - # torch.nn.MultiheadAttention, which passes `self.out_proj.weight` to - # `torch.nn.functional.multi_head_attention_forward` instead of calling `self.out_proj`. The leaf-level - # onload hook on `out_proj` is therefore never triggered and its weights stay on the CPU. - pass - # TODO(aryan): Create a dummy gemma model with smol vocab size - @unittest.skip( - "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." - ) - def test_inference_batch_consistent(self): - pass +class TestHunyuanVideoFramepackPipelinePyramidAttentionBroadcast( + HunyuanVideoFramepackPipelineTesterConfig, PyramidAttentionBroadcastTesterMixin +): + """Pyramid Attention Broadcast cache tests for the Framepack pipeline.""" - @unittest.skip( - "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." - ) - def test_inference_batch_single_identical(self): - pass + +class TestHunyuanVideoFramepackPipelineFasterCache(HunyuanVideoFramepackPipelineTesterConfig, FasterCacheTesterMixin): + """FasterCache tests for the Framepack pipeline.""" + + # Framepack is guidance-distilled, so the FasterCache tester must skip the low/high-frequency-delta checks. + FASTER_CACHE_CONFIG = { + "spatial_attention_block_skip_range": 2, + "spatial_attention_timestep_skip_range": (-1, 901), + "unconditional_batch_skip_range": 2, + "attention_weight_callback": lambda _: 0.5, + "is_guidance_distilled": True, + } + + +class TestHunyuanVideoFramepackPipelineLoRA(HunyuanVideoFramepackPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the Framepack pipeline.""" + + +class TestHunyuanVideoFramepackPipelineLoRAMemory(HunyuanVideoFramepackPipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the Framepack pipeline.""" diff --git a/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py b/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py index cefe48d01d9a..6f823252d792 100644 --- a/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py +++ b/tests/pipelines/hunyuan_video1_5/test_hunyuan_1_5.py @@ -12,8 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - +import pytest import torch from transformers import ( AutoConfig, @@ -32,16 +31,16 @@ ) from diffusers.guiders import ClassifierFreeGuidance -from ...testing_utils import enable_full_determinism -from ..test_pipelines_common import PipelineTesterMixin +from ...testing_utils import assert_tensors_close, enable_full_determinism +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class HunyuanVideo15PipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class HunyuanVideo15PipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HunyuanVideo15Pipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "prompt", "negative_prompt", @@ -57,12 +56,12 @@ class HunyuanVideo15PipelineFastTests(PipelineTesterMixin, unittest.TestCase): "negative_prompt_embeds_mask_2", ] ) - batch_params = ["prompt", "negative_prompt"] - required_optional_params = frozenset(["num_inference_steps", "generator", "latents", "return_dict"]) - test_attention_slicing = False - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = False + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (9, 3, 16, 16) + # HunyuanVideo 1.5 is a video pipeline: it exposes `num_videos_per_prompt`, not `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, num_layers: int = 1): torch.manual_seed(0) @@ -126,7 +125,7 @@ def get_dummy_components(self, num_layers: int = 1): guider = ClassifierFreeGuidance(guidance_scale=1.0) - components = { + return { "transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler, @@ -136,59 +135,49 @@ def get_dummy_components(self, num_layers: int = 1): "tokenizer_2": tokenizer_2, "guider": guider, } - 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": "monkey", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "height": 16, "width": 16, "num_frames": 9, + # 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) - inputs = self.get_dummy_inputs(device) - result = pipe(**inputs) - video = result.frames +class TestHunyuanVideo15Pipeline(HunyuanVideo15PipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (9, 3, 16, 16)) - generated_slice = generated_video.flatten() - generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4296, 0.5549, 0.3088, 0.9115, 0.5049, 0.7926, 0.5549, 0.8618, 0.5091, 0.5075, 0.7117, 0.5292, 0.7053, 0.4864, 0.5206, 0.3878]) # fmt: on - self.assertTrue( - torch.abs(generated_slice - expected_slice).max() < 1e-3, - f"output_slice: {generated_slice}, expected_slice: {expected_slice}", - ) + generated_slice = generated_video.flatten() + generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) - @unittest.skip("TODO: Test not supported for now because needs to be adjusted to work with guiders.") + @pytest.mark.skip("TODO: Test not supported for now because needs to be adjusted to work with guiders.") def test_encode_prompt_works_in_isolation(self): pass - @unittest.skip("Needs to be revisited.") + @pytest.mark.skip("Needs to be revisited.") def test_inference_batch_consistent(self): - super().test_inference_batch_consistent() + pass - @unittest.skip("Needs to be revisited.") + @pytest.mark.skip("Needs to be revisited.") def test_inference_batch_single_identical(self): - super().test_inference_batch_single_identical() + pass + + +class TestHunyuanVideo15PipelineMemory(HunyuanVideo15PipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the HunyuanVideo 1.5 pipeline.""" diff --git a/tests/pipelines/hunyuandit/test_hunyuan_dit.py b/tests/pipelines/hunyuandit/test_hunyuan_dit.py index 9b17ac606481..80cd76a2bcf8 100644 --- a/tests/pipelines/hunyuandit/test_hunyuan_dit.py +++ b/tests/pipelines/hunyuandit/test_hunyuan_dit.py @@ -14,16 +14,16 @@ # limitations under the License. import gc -import tempfile -import unittest import numpy as np +import pytest import torch from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel from diffusers import AutoencoderKL, DDPMScheduler, HunyuanDiT2DModel, HunyuanDiTPipeline from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, enable_full_determinism, numpy_cosine_similarity_distance, @@ -31,27 +31,24 @@ 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 ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +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 HunyuanDiTPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class HunyuanDiTPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HunyuanDiTPipeline - 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 = PipelineTesterMixin.required_optional_params - test_layerwise_casting = True + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + output_shape = (3, 16, 16) def get_dummy_components(self): torch.manual_seed(0) @@ -79,7 +76,7 @@ def get_dummy_components(self): text_encoder_2 = T5EncoderModel(config) tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler, @@ -90,43 +87,36 @@ def get_dummy_components(self): "safety_checker": None, "feature_extractor": None, } - 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": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, - "output_type": "np", "use_resolution_binning": False, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", } - return inputs + +class TestHunyuanDiTPipeline(HunyuanDiTPipelineTesterConfig, PipelineTesterMixin): def test_inference(self): - device = "cpu" + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + image = pipe(**self.get_dummy_inputs()).images + generated_image = image[0] + assert generated_image.shape == self.output_shape - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] + # fmt: off + expected_slice = torch.tensor([0.56939435, 0.34541583, 0.35915792, 0.46489206, 0.38775963, 0.45004836, 0.5957267, 0.59481275, 0.33287364]) + # fmt: on - self.assertEqual(image.shape, (1, 16, 16, 3)) - expected_slice = np.array( - [0.56939435, 0.34541583, 0.35915792, 0.46489206, 0.38775963, 0.45004836, 0.5957267, 0.59481275, 0.33287364] - ) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + generated_slice = generated_image[-1, -3:, -3:].flatten() + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) - @unittest.skip("The HunyuanDiT Attention pooling layer does not support sequential CPU offloading.") + @pytest.mark.skip("The HunyuanDiT Attention pooling layer does not support sequential CPU offloading.") def test_sequential_cpu_offload_forward_pass(self): # TODO(YiYi) need to fix later # This is because it instantiates it's attention layer from torch.nn.MultiheadAttention, which calls to @@ -135,7 +125,7 @@ def test_sequential_cpu_offload_forward_pass(self): # this test because of MHA (example: HunyuanVideo Framepack) pass - @unittest.skip("The HunyuanDiT Attention pooling layer does not support sequential CPU offloading.") + @pytest.mark.skip("The HunyuanDiT Attention pooling layer does not support sequential CPU offloading.") def test_sequential_offload_forward_pass_twice(self): # TODO(YiYi) need to fix later # This is because it instantiates it's attention layer from torch.nn.MultiheadAttention, which calls to @@ -144,42 +134,26 @@ def test_sequential_offload_forward_pass_twice(self): # this test because of MHA (example: HunyuanVideo Framepack) pass - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical( - expected_max_diff=1e-3, - ) + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-3): + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) def test_feed_forward_chunking(self): - device = "cpu" - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_no_chunking = image[0, -3:, -3:, -1] + image_no_chunking = pipe(**self.get_dummy_inputs()).images pipe.transformer.enable_forward_chunking(chunk_size=1, dim=0) - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_chunking = image[0, -3:, -3:, -1] + image_chunking = pipe(**self.get_dummy_inputs()).images - max_diff = np.abs(to_np(image_slice_no_chunking) - to_np(image_slice_chunking)).max() - self.assertLess(max_diff, 1e-4) + assert_tensors_close( + image_chunking, image_no_chunking, atol=1e-4, msg="Feed forward chunking should not affect the outputs." + ) 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) + # Run on CPU to ensure determinism for the device-dependent torch.Generator. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - inputs["return_dict"] = False - image = pipe(**inputs)[0] - original_image_slice = image[0, -3:, -3:, -1] + original_image = pipe(**self.get_dummy_inputs(), return_dict=False)[0] pipe.transformer.fuse_qkv_projections() # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added @@ -192,52 +166,50 @@ 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["return_dict"] = False - image_fused = pipe(**inputs)[0] - image_slice_fused = image_fused[0, -3:, -3:, -1] + image_fused = pipe(**self.get_dummy_inputs(), return_dict=False)[0] pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - inputs["return_dict"] = False - image_disabled = pipe(**inputs)[0] - image_slice_disabled = image_disabled[0, -3:, -3:, -1] - - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( - "Fusion of QKV projections shouldn't affect the outputs." + image_disabled = pipe(**self.get_dummy_inputs(), return_dict=False)[0] + + assert_tensors_close( + image_fused, + original_image, + atol=1e-2, + rtol=1e-2, + msg="Fusion of QKV projections shouldn't affect the outputs.", ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert_tensors_close( + image_disabled, + image_fused, + atol=1e-2, + rtol=1e-2, + msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", ) - 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." + assert_tensors_close( + image_disabled, + original_image, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) - @unittest.skip( + @pytest.mark.skip( "Test not supported as `encode_prompt` is called two times separately which deivates from about 99% of the pipelines we have." ) def test_encode_prompt_works_in_isolation(self): pass - def test_save_load_optional_components(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) + def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4): + pipe = self.get_pipeline().to(torch_device) - prompt = inputs["prompt"] - generator = inputs["generator"] - num_inference_steps = inputs["num_inference_steps"] - output_type = inputs["output_type"] + inputs = self.get_dummy_inputs() ( prompt_embeds, negative_prompt_embeds, prompt_attention_mask, negative_prompt_attention_mask, - ) = pipe.encode_prompt(prompt, device=torch_device, dtype=torch.float32, text_encoder_index=0) + ) = pipe.encode_prompt(inputs["prompt"], device=torch_device, dtype=torch.float32, text_encoder_index=0) ( prompt_embeds_2, @@ -245,86 +217,74 @@ def test_save_load_optional_components(self): prompt_attention_mask_2, negative_prompt_attention_mask_2, ) = pipe.encode_prompt( - prompt, + inputs["prompt"], device=torch_device, dtype=torch.float32, text_encoder_index=1, ) - # inputs with prompt converted to embeddings - inputs = { - "prompt_embeds": prompt_embeds, - "prompt_attention_mask": prompt_attention_mask, - "negative_prompt_embeds": negative_prompt_embeds, - "negative_prompt_attention_mask": negative_prompt_attention_mask, - "prompt_embeds_2": prompt_embeds_2, - "prompt_attention_mask_2": prompt_attention_mask_2, - "negative_prompt_embeds_2": negative_prompt_embeds_2, - "negative_prompt_attention_mask_2": negative_prompt_attention_mask_2, - "generator": generator, - "num_inference_steps": num_inference_steps, - "output_type": output_type, - "use_resolution_binning": False, - } + def embedded_inputs(): + # Inputs with the prompt already converted to embeddings. + return { + "prompt_embeds": prompt_embeds, + "prompt_attention_mask": prompt_attention_mask, + "negative_prompt_embeds": negative_prompt_embeds, + "negative_prompt_attention_mask": negative_prompt_attention_mask, + "prompt_embeds_2": prompt_embeds_2, + "prompt_attention_mask_2": prompt_attention_mask_2, + "negative_prompt_embeds_2": negative_prompt_embeds_2, + "negative_prompt_attention_mask_2": negative_prompt_attention_mask_2, + "generator": self.get_generator(0), + "num_inference_steps": inputs["num_inference_steps"], + "output_type": inputs["output_type"], + "use_resolution_binning": False, + } # set all optional components to None for optional_component in pipe._optional_components: setattr(pipe, optional_component, None) - output = pipe(**inputs)[0] + output = pipe(**embedded_inputs())[0] - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir) - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir) - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) + pipe.save_pretrained(tmp_path) + pipe_loaded = self.pipeline_class.from_pretrained(tmp_path) + pipe_loaded.to(torch_device) + pipe_loaded.set_progress_bar_config(disable=None) for optional_component in pipe._optional_components: - self.assertTrue( - getattr(pipe_loaded, optional_component) is None, - f"`{optional_component}` did not stay set to None after loading.", + assert getattr(pipe_loaded, optional_component) is None, ( + f"`{optional_component}` did not stay set to None after loading." ) - inputs = self.get_dummy_inputs(torch_device) - - generator = inputs["generator"] - num_inference_steps = inputs["num_inference_steps"] - output_type = inputs["output_type"] - - # inputs with prompt converted to embeddings - inputs = { - "prompt_embeds": prompt_embeds, - "prompt_attention_mask": prompt_attention_mask, - "negative_prompt_embeds": negative_prompt_embeds, - "negative_prompt_attention_mask": negative_prompt_attention_mask, - "prompt_embeds_2": prompt_embeds_2, - "prompt_attention_mask_2": prompt_attention_mask_2, - "negative_prompt_embeds_2": negative_prompt_embeds_2, - "negative_prompt_attention_mask_2": negative_prompt_attention_mask_2, - "generator": generator, - "num_inference_steps": num_inference_steps, - "output_type": output_type, - "use_resolution_binning": False, - } + output_loaded = pipe_loaded(**embedded_inputs())[0] + + assert_tensors_close( + output_loaded, output, atol=expected_max_difference, msg="Reloaded pipeline output differs." + ) + - output_loaded = pipe_loaded(**inputs)[0] +class TestHunyuanDiTPipelineMemory(HunyuanDiTPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the HunyuanDiT pipeline.""" - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, 1e-4) + @pytest.mark.skip("The HunyuanDiT Attention pooling layer does not support sequential CPU offloading.") + def test_sequential_cpu_offload_forward_pass(self): + pass + + @pytest.mark.skip("The HunyuanDiT Attention pooling layer does not support sequential CPU offloading.") + def test_sequential_offload_forward_pass_twice(self): + pass @slow @require_torch_accelerator -class HunyuanDiTPipelineIntegrationTests(unittest.TestCase): +class TestHunyuanDiTPipelineIntegration: prompt = "一个宇航员在骑马" - 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) diff --git a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py index 86519fea0d30..e571e565e461 100644 --- a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py +++ b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py @@ -14,10 +14,10 @@ # limitations under the License. import gc -import unittest from unittest import mock import numpy as np +import pytest import torch from transformers import ( CLIPImageProcessor, @@ -51,7 +51,7 @@ enable_full_determinism() -class IPAdapterNightlyTestsMixin(unittest.TestCase): +class IPAdapterNightlyTestsMixin: dtype = torch.float16 _SD_PIPELINE_RANDN_TENSOR_TARGETS = { @@ -63,15 +63,12 @@ class IPAdapterNightlyTestsMixin(unittest.TestCase): StableDiffusionXLInpaintPipeline: "diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.randn_tensor", } - def setUp(self): - # clean up the VRAM before each test - super().setUp() + @pytest.fixture(autouse=True) + def cleanup(self): + # clean up the VRAM before and after each test gc.collect() backend_empty_cache(torch_device) - - def tearDown(self): - # clean up the VRAM after each test - super().tearDown() + yield gc.collect() backend_empty_cache(torch_device) @@ -82,14 +79,14 @@ def get_fixed_randn_tensor_patch(self, pipeline, shape=(1, 4, 64, 64), seed=33): fixed_noise = self.get_fixed_noise(shape=shape, seed=seed) def fake_randn_tensor(requested_shape, generator=None, device=None, dtype=None, layout=None): - self.assertEqual(tuple(requested_shape), tuple(fixed_noise.shape)) + assert tuple(requested_shape) == tuple(fixed_noise.shape) return fixed_noise.to(device=device, dtype=dtype) for pipeline_cls, target in self._SD_PIPELINE_RANDN_TENSOR_TARGETS.items(): if isinstance(pipeline, pipeline_cls): return mock.patch(target, side_effect=fake_randn_tensor) - self.fail(f"No fixed randn_tensor patch target configured for pipeline type {type(pipeline)}") + raise AssertionError(f"No fixed randn_tensor patch target configured for pipeline type {type(pipeline)}") def get_image_encoder(self, repo_id, subfolder): image_encoder = CLIPVisionModelWithProjection.from_pretrained( @@ -195,7 +192,7 @@ def get_dummy_inputs( @slow @require_torch_accelerator -class IPAdapterSDIntegrationTests(IPAdapterNightlyTestsMixin): +class TestIPAdapterSDIntegration(IPAdapterNightlyTestsMixin): def test_text_to_image(self): image_encoder = self.get_image_encoder(repo_id="h94/IP-Adapter", subfolder="models/image_encoder") pipeline = StableDiffusionPipeline.from_pretrained( @@ -325,16 +322,15 @@ def test_text_to_image_model_cpu_offload(self): with self.get_fixed_randn_tensor_patch(pipeline): output_with_offload = pipeline(**inputs).images max_diff = np.abs(output_with_offload - output_without_offload).max() - self.assertLess(max_diff, 1e-3, "CPU offloading should not affect the inference results") + assert max_diff < 1e-3, "CPU offloading should not affect the inference results" offloaded_modules = [ v for k, v in pipeline.components.items() if isinstance(v, torch.nn.Module) and k not in pipeline._exclude_from_cpu_offload ] - ( - self.assertTrue(all(v.device.type == "cpu" for v in offloaded_modules)), - f"Not offloaded: {[v for v in offloaded_modules if v.device.type != 'cpu']}", + assert all(v.device.type == "cpu" for v in offloaded_modules), ( + f"Not offloaded: {[v for v in offloaded_modules if v.device.type != 'cpu']}" ) def test_text_to_image_full_face(self): @@ -436,7 +432,7 @@ def test_text_to_image_face_id(self): @slow @require_torch_accelerator -class IPAdapterSDXLIntegrationTests(IPAdapterNightlyTestsMixin): +class TestIPAdapterSDXLIntegration(IPAdapterNightlyTestsMixin): def test_text_to_image_sdxl(self): image_encoder = self.get_image_encoder(repo_id="h94/IP-Adapter", subfolder="sdxl_models/image_encoder") feature_extractor = self.get_image_processor("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k") diff --git a/tests/pipelines/testing_utils/memory.py b/tests/pipelines/testing_utils/memory.py index 6c7986b1bb5c..6c97e0323035 100644 --- a/tests/pipelines/testing_utils/memory.py +++ b/tests/pipelines/testing_utils/memory.py @@ -287,6 +287,7 @@ def enable_group_offload_on_component(pipe, group_offloading_kwargs): "text_encoder", "text_encoder_2", "text_encoder_3", + "text_encoder_4", "transformer", "transformer_2", "unet",