diff --git a/.ai/references/testing.md b/.ai/references/testing.md index 62e1ca986a07..eb5273907220 100644 --- a/.ai/references/testing.md +++ b/.ai/references/testing.md @@ -34,6 +34,10 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers - `torch.nn.MultiheadAttention` is the common instance: it passes `self.out_proj.weight` straight to `torch.nn.functional.multi_head_attention_forward` instead of calling `self.out_proj`, so the hook on `out_proj` never fires. `SiglipVisionModel`'s attention pooling head wraps one — see `tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py`, whose `image_encoder` is excluded for this reason. - `HunyuanDiTAttentionPool` (`src/diffusers/models/embeddings.py`) shows the same failure without an MHA module: a plain `nn.Module` that hands its `q_proj` / `k_proj` / `v_proj` / `c_proj` weights to `torch.nn.functional.multi_head_attention_forward`, so all four projections stay offloaded rather than just one. `HunyuanDiT2DModel` opts out of group offloading entirely with `_supports_group_offloading = False`. - Before adding a skip or an exclusion, confirm the failure still reproduces — several existing skips are stale, having outlived the upstream cause. +- **A migration that surfaces a `src/` gap marks the test `xfail`, it does not patch the pipeline.** Give the marker a module-level name and a `reason` naming the exact gap (`PNDM_*` in `tests/pipelines/pndm/test_pndm.py` is the worked example), and prefer `strict=True` so the marker reports XPASS — and gets deleted — the day the pipeline is fixed. Use `strict=False` only when one mark covers a group whose members do not all fail. Marking a whole test class keeps the mixin's own marks (`@is_memory`, `@require_accelerator`) intact; overriding individual inherited tests drops the decorators they were declared with, so re-declare those too. +- **`from_pipe` tests** (a pipeline that is a variant of an existing one — PAG, AnimateDiff, ...) compose the shared `FromPipeTesterMixin` (`tests/pipelines/testing_utils/from_pipe.py`, exported from `..testing_utils`) in their own test class. It derives the original pipeline from `pipeline_class.__name__`; set `original_pipeline_repo` on the test class to pull it from a repo other than the default for that class. The unittest-era `PipelineFromPipeTesterMixin` in `tests/pipelines/test_pipelines_common.py` is what it replaces. +- **A hardware gap is a conditional skip, not an xfail.** When a test fails only because the runner's cuDNN build has no kernel for an op — `RuntimeError: GET was unable to find an engine to execute this computation`, as Sana's depthwise `Conv2d` hits in bfloat16 — wrap the call in `skip_if_no_cudnn_engine()` (`tests/testing_utils.py`). It skips on that error and re-raises every other `RuntimeError`, so the test still runs wherever the kernel exists. +- **PAG pipelines** also compose `PAGPipelineTesterMixin` (`tests/pipelines/pag/testing_utils.py`) in place of `PipelineTesterMixin`: it adds `test_pag_disable_enable` and `test_pag_inference` on top, driven by `base_pipeline_class` and the `pag_*` knobs on the test class. Keep `test_pag_applied_layers` per pipeline — which layers PAG resolves to is model-specific. - **IP-Adapter tests** live in their own class decorated with `@is_ip_adapter`, subclassing only the config (not `PipelineTesterMixin`). UNet pipelines that load adapters through the standard `IPAdapterMixin` API compose the shared `IPAdapterTesterMixin` (`tests/pipelines/testing_utils/ip_adapter.py`, exported from `..testing_utils`); pipelines whose IP-Adapter API differs (Flux, for example) keep a bespoke mixin next to their own tests. #### LoRA tests diff --git a/src/diffusers/pipelines/omnigen/pipeline_omnigen.py b/src/diffusers/pipelines/omnigen/pipeline_omnigen.py index 6e5db93d1f35..6564b2a672a0 100644 --- a/src/diffusers/pipelines/omnigen/pipeline_omnigen.py +++ b/src/diffusers/pipelines/omnigen/pipeline_omnigen.py @@ -116,9 +116,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class OmniGenPipeline( - DiffusionPipeline, -): +class OmniGenPipeline(DiffusionPipeline): r""" The OmniGen pipeline for multimodal-to-image generation. diff --git a/tests/pipelines/animatediff/test_animatediff.py b/tests/pipelines/animatediff/test_animatediff.py index 45c13bbd2c11..5baa5efd0aa6 100644 --- a/tests/pipelines/animatediff/test_animatediff.py +++ b/tests/pipelines/animatediff/test_animatediff.py @@ -10,7 +10,6 @@ AutoencoderKL, DDIMScheduler, MotionAdapter, - StableDiffusionPipeline, UNet2DConditionModel, ) @@ -22,8 +21,8 @@ torch_device, ) from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import PipelineFromPipeTesterMixin from ..testing_utils import ( + FromPipeTesterMixin, IPAdapterTesterMixin, LoraMemoryTesterMixin, LoraTesterMixin, @@ -31,7 +30,6 @@ UNetLoraTesterMixin, ) from .testing_utils import ( - FROM_PIPE_SKIP_REASON, FreeInitTesterMixin, FreeNoiseSplitInferenceTesterMixin, MotionPipelineTesterConfig, @@ -128,32 +126,6 @@ class TestAnimateDiffPipeline( FreeInitTesterMixin, FreeNoiseSplitInferenceTesterMixin, ): - def test_from_pipe_consistent_config(self): - original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" - - # create StableDiffusionPipeline - pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False) - - # StableDiffusionPipeline -> AnimateDiffPipeline - pipe_components = self.get_dummy_components() - pipe_additional_components = { - name: component for name, component in pipe_components.items() if name not in pipe_original.components - } - pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components) - - # AnimateDiffPipeline -> StableDiffusionPipeline - original_pipe_additional_components = {} - for name, component in pipe_original.components.items(): - if name not in pipe.components or not isinstance(component, pipe.components[name].__class__): - original_pipe_additional_components[name] = component - - pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components) - - # compare the config - original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")} - original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")} - assert original_config_2 == original_config - def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4): if torch_device == "cpu" and expected_slice is None: # fmt: off @@ -244,10 +216,7 @@ def test_animatediff(self): assert numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice.flatten()) < 1e-3 -@pytest.mark.skip(FROM_PIPE_SKIP_REASON) -class TestAnimateDiffPipelineFromPipe(AnimateDiffPipelineTesterConfig, PipelineFromPipeTesterMixin): - """`from_pipe` forward-pass parity and offload round trip for the AnimateDiff pipeline. +class TestAnimateDiffPipelineFromPipe(AnimateDiffPipelineTesterConfig, FromPipeTesterMixin): + """`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff pipeline.""" - Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above, - but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet. - """ + original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" diff --git a/tests/pipelines/animatediff/test_animatediff_controlnet.py b/tests/pipelines/animatediff/test_animatediff_controlnet.py index 5946a86ace41..fd8ae81f0ec7 100644 --- a/tests/pipelines/animatediff/test_animatediff_controlnet.py +++ b/tests/pipelines/animatediff/test_animatediff_controlnet.py @@ -1,4 +1,3 @@ -import pytest import torch from PIL import Image from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -9,14 +8,13 @@ ControlNetModel, DDIMScheduler, MotionAdapter, - StableDiffusionPipeline, UNet2DConditionModel, ) from ...testing_utils import torch_device from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import PipelineFromPipeTesterMixin from ..testing_utils import ( + FromPipeTesterMixin, IPAdapterTesterMixin, LoraMemoryTesterMixin, LoraTesterMixin, @@ -24,7 +22,6 @@ UNetLoraTesterMixin, ) from .testing_utils import ( - FROM_PIPE_SKIP_REASON, FreeInitTesterMixin, FreeNoiseTesterMixin, MotionPipelineTesterConfig, @@ -142,32 +139,6 @@ def get_free_noise_inputs(self): # `get_dummy_inputs` rather than by overriding `num_frames` on the returned dict. return self.get_dummy_inputs(num_frames=16) - def test_from_pipe_consistent_config(self): - original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" - - # create StableDiffusionPipeline - pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False) - - # StableDiffusionPipeline -> AnimateDiffControlNetPipeline - pipe_components = self.get_dummy_components() - pipe_additional_components = { - name: component for name, component in pipe_components.items() if name not in pipe_original.components - } - pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components) - - # AnimateDiffControlNetPipeline -> StableDiffusionPipeline - original_pipe_additional_components = {} - for name, component in pipe_original.components.items(): - if name not in pipe.components or not isinstance(component, pipe.components[name].__class__): - original_pipe_additional_components[name] = component - - pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components) - - # compare the config - original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")} - original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")} - assert original_config_2 == original_config - def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4): if torch_device == "cpu" and expected_slice is None: # fmt: off @@ -198,12 +169,7 @@ class TestAnimateDiffControlNetPipelineLoRAMemory(AnimateDiffControlNetPipelineT """LoRA x memory-optimization tests (group offload, CPU offload) for the pipeline.""" -@pytest.mark.skip(FROM_PIPE_SKIP_REASON) -class TestAnimateDiffControlNetPipelineFromPipe( - AnimateDiffControlNetPipelineTesterConfig, PipelineFromPipeTesterMixin -): - """`from_pipe` forward-pass parity and offload round trip for the AnimateDiff ControlNet pipeline. +class TestAnimateDiffControlNetPipelineFromPipe(AnimateDiffControlNetPipelineTesterConfig, FromPipeTesterMixin): + """`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff ControlNet pipeline.""" - Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above, - but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet. - """ + original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" diff --git a/tests/pipelines/animatediff/test_animatediff_sparsectrl.py b/tests/pipelines/animatediff/test_animatediff_sparsectrl.py index 967183802bec..7783f8883379 100644 --- a/tests/pipelines/animatediff/test_animatediff_sparsectrl.py +++ b/tests/pipelines/animatediff/test_animatediff_sparsectrl.py @@ -1,4 +1,3 @@ -import pytest import torch from PIL import Image from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -9,14 +8,13 @@ DDIMScheduler, MotionAdapter, SparseControlNetModel, - StableDiffusionPipeline, UNet2DConditionModel, ) from ...testing_utils import assert_tensors_close, torch_device from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import PipelineFromPipeTesterMixin from ..testing_utils import ( + FromPipeTesterMixin, IPAdapterTesterMixin, LoraMemoryTesterMixin, LoraTesterMixin, @@ -24,7 +22,6 @@ UNetLoraTesterMixin, ) from .testing_utils import ( - FROM_PIPE_SKIP_REASON, FreeInitTesterMixin, MotionPipelineTesterConfig, MotionPipelineTesterMixin, @@ -138,32 +135,6 @@ class TestAnimateDiffSparseControlNetPipeline( MotionPipelineTesterMixin, FreeInitTesterMixin, ): - def test_from_pipe_consistent_config(self): - original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" - - # create StableDiffusionPipeline - pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False) - - # StableDiffusionPipeline -> AnimateDiffSparseControlNetPipeline - pipe_components = self.get_dummy_components() - pipe_additional_components = { - name: component for name, component in pipe_components.items() if name not in pipe_original.components - } - pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components) - - # AnimateDiffSparseControlNetPipeline -> StableDiffusionPipeline - original_pipe_additional_components = {} - for name, component in pipe_original.components.items(): - if name not in pipe.components or not isinstance(component, pipe.components[name].__class__): - original_pipe_additional_components[name] = component - - pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components) - - # compare the config - original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")} - original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")} - assert original_config_2 == original_config - def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4): if torch_device == "cpu" and expected_slice is None: # fmt: off @@ -229,12 +200,9 @@ class TestAnimateDiffSparseControlNetPipelineLoRAMemory( """LoRA x memory-optimization tests (group offload, CPU offload) for the pipeline.""" -@pytest.mark.skip(FROM_PIPE_SKIP_REASON) class TestAnimateDiffSparseControlNetPipelineFromPipe( - AnimateDiffSparseControlNetPipelineTesterConfig, PipelineFromPipeTesterMixin + AnimateDiffSparseControlNetPipelineTesterConfig, FromPipeTesterMixin ): - """`from_pipe` forward-pass parity and offload round trip for the AnimateDiff SparseControlNet pipeline. + """`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff SparseControlNet pipeline.""" - Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above, - but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet. - """ + original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" diff --git a/tests/pipelines/animatediff/test_animatediff_video2video.py b/tests/pipelines/animatediff/test_animatediff_video2video.py index 089eb909096a..fd454d1a8182 100644 --- a/tests/pipelines/animatediff/test_animatediff_video2video.py +++ b/tests/pipelines/animatediff/test_animatediff_video2video.py @@ -1,4 +1,3 @@ -import pytest import torch from PIL import Image from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -8,14 +7,13 @@ AutoencoderKL, DDIMScheduler, MotionAdapter, - StableDiffusionPipeline, UNet2DConditionModel, ) from ...testing_utils import torch_device from ..pipeline_params import TEXT_TO_IMAGE_PARAMS, VIDEO_TO_VIDEO_BATCH_PARAMS -from ..test_pipelines_common import PipelineFromPipeTesterMixin from ..testing_utils import ( + FromPipeTesterMixin, IPAdapterTesterMixin, LoraMemoryTesterMixin, LoraTesterMixin, @@ -23,7 +21,6 @@ UNetLoraTesterMixin, ) from .testing_utils import ( - FROM_PIPE_SKIP_REASON, FreeInitTesterMixin, FreeNoiseSplitInferenceTesterMixin, MotionPipelineTesterConfig, @@ -133,32 +130,6 @@ def get_free_noise_inputs(self): inputs["strength"] = 0.5 return inputs - def test_from_pipe_consistent_config(self): - original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" - - # create StableDiffusionPipeline - pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False) - - # StableDiffusionPipeline -> AnimateDiffVideoToVideoPipeline - pipe_components = self.get_dummy_components() - pipe_additional_components = { - name: component for name, component in pipe_components.items() if name not in pipe_original.components - } - pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components) - - # AnimateDiffVideoToVideoPipeline -> StableDiffusionPipeline - original_pipe_additional_components = {} - for name, component in pipe_original.components.items(): - if name not in pipe.components or not isinstance(component, pipe.components[name].__class__): - original_pipe_additional_components[name] = component - - pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components) - - # compare the config - original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")} - original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")} - assert original_config_2 == original_config - def test_latent_inputs(self): pipe = self.get_pipeline().to(torch_device) @@ -191,12 +162,7 @@ class TestAnimateDiffVideoToVideoPipelineLoRAMemory( """LoRA x memory-optimization tests (group offload, CPU offload) for the pipeline.""" -@pytest.mark.skip(FROM_PIPE_SKIP_REASON) -class TestAnimateDiffVideoToVideoPipelineFromPipe( - AnimateDiffVideoToVideoPipelineTesterConfig, PipelineFromPipeTesterMixin -): - """`from_pipe` forward-pass parity and offload round trip for the AnimateDiff video-to-video pipeline. +class TestAnimateDiffVideoToVideoPipelineFromPipe(AnimateDiffVideoToVideoPipelineTesterConfig, FromPipeTesterMixin): + """`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff video-to-video pipeline.""" - Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above, - but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet. - """ + original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" diff --git a/tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py b/tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py index 36b5625972c5..26ad36ebb765 100644 --- a/tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py +++ b/tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py @@ -1,4 +1,3 @@ -import pytest import torch from PIL import Image from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -9,14 +8,13 @@ ControlNetModel, DDIMScheduler, MotionAdapter, - StableDiffusionPipeline, UNet2DConditionModel, ) from ...testing_utils import torch_device from ..pipeline_params import TEXT_TO_IMAGE_PARAMS, VIDEO_TO_VIDEO_BATCH_PARAMS -from ..test_pipelines_common import PipelineFromPipeTesterMixin from ..testing_utils import ( + FromPipeTesterMixin, IPAdapterTesterMixin, LoraMemoryTesterMixin, LoraTesterMixin, @@ -24,7 +22,6 @@ UNetLoraTesterMixin, ) from .testing_utils import ( - FROM_PIPE_SKIP_REASON, FreeInitTesterMixin, FreeNoiseTesterMixin, MotionPipelineTesterConfig, @@ -145,32 +142,6 @@ def get_free_noise_inputs(self): # longer video rather than by passing `num_frames`. return self.get_dummy_inputs(num_frames=16) - def test_from_pipe_consistent_config(self): - original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" - - # create StableDiffusionPipeline - pipe_original = StableDiffusionPipeline.from_pretrained(original_repo, requires_safety_checker=False) - - # StableDiffusionPipeline -> AnimateDiffVideoToVideoControlNetPipeline - pipe_components = self.get_dummy_components() - pipe_additional_components = { - name: component for name, component in pipe_components.items() if name not in pipe_original.components - } - pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components) - - # AnimateDiffVideoToVideoControlNetPipeline -> StableDiffusionPipeline - original_pipe_additional_components = {} - for name, component in pipe_original.components.items(): - if name not in pipe.components or not isinstance(component, pipe.components[name].__class__): - original_pipe_additional_components[name] = component - - pipe_original_2 = StableDiffusionPipeline.from_pipe(pipe, **original_pipe_additional_components) - - # compare the config - original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")} - original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")} - assert original_config_2 == original_config - def test_latent_inputs(self): pipe = self.get_pipeline().to(torch_device) @@ -213,12 +184,9 @@ class TestAnimateDiffVideoToVideoControlNetPipelineLoRAMemory( """LoRA x memory-optimization tests (group offload, CPU offload) for the pipeline.""" -@pytest.mark.skip(FROM_PIPE_SKIP_REASON) class TestAnimateDiffVideoToVideoControlNetPipelineFromPipe( - AnimateDiffVideoToVideoControlNetPipelineTesterConfig, PipelineFromPipeTesterMixin + AnimateDiffVideoToVideoControlNetPipelineTesterConfig, FromPipeTesterMixin ): - """`from_pipe` forward-pass parity and offload round trip for the AnimateDiff video-to-video ControlNet pipeline. + """`from_pipe` round-trip tests against `StableDiffusionPipeline` for the AnimateDiff video-to-video ControlNet pipeline.""" - Parked, not deleted: `test_from_pipe_consistent_config` runs for real as a method on the main test class above, - but the forward-pass checks in `PipelineFromPipeTesterMixin` have no pytest-style equivalent yet. - """ + original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" diff --git a/tests/pipelines/animatediff/testing_utils.py b/tests/pipelines/animatediff/testing_utils.py index 8501ad8de0b9..a3e6fc3d8986 100644 --- a/tests/pipelines/animatediff/testing_utils.py +++ b/tests/pipelines/animatediff/testing_utils.py @@ -24,19 +24,6 @@ from ..testing_utils.common import BasePipelineOutputMixin -# `PipelineFromPipeTesterMixin` (tests/pipelines/test_pipelines_common.py) still covers `from_pipe` forward-pass -# parity and the model-CPU-offload round trip, but it is unittest-era: its tests call `self.get_dummy_inputs(device, -# seed=0)` and `self.assertLess`, neither of which exists on a `BasePipelineTesterConfig` outside a -# `unittest.TestCase`. Un-skipping the parked classes below without porting the mixin first would error, not fail. -# The mixin is still live for the ten `tests/pipelines/pag/` files and `stable_diffusion_adapter`; rewrite it -# pytest-style when those are migrated, then drop these skips. -FROM_PIPE_SKIP_REASON = ( - "`PipelineFromPipeTesterMixin` is still unittest-style and cannot run against `BasePipelineTesterConfig` — " - "these error rather than fail if un-skipped. Port the mixin to pytest (due when `tests/pipelines/pag/` is " - "migrated), then remove this skip." -) - - class MotionPipelineTesterConfig(BasePipelineTesterConfig): """`BasePipelineTesterConfig` for the AnimateDiff pipelines in this directory.""" diff --git a/tests/pipelines/omnigen/test_pipeline_omnigen.py b/tests/pipelines/omnigen/test_pipeline_omnigen.py index 1a758b705042..db4bd0abc879 100644 --- a/tests/pipelines/omnigen/test_pipeline_omnigen.py +++ b/tests/pipelines/omnigen/test_pipeline_omnigen.py @@ -1,7 +1,7 @@ import gc -import unittest import numpy as np +import pytest import torch from transformers import AutoTokenizer @@ -15,15 +15,14 @@ slow, torch_device, ) -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -class OmniGenPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class OmniGenPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = OmniGenPipeline - params = frozenset(["prompt", "guidance_scale"]) - batch_params = frozenset(["prompt"]) - test_xformers_attention = False - test_layerwise_casting = True + required_input_params_in_call_signature = frozenset(["prompt", "guidance_scale"]) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 16, 16) def get_dummy_components(self): torch.manual_seed(0) @@ -54,53 +53,47 @@ def get_dummy_components(self): scheduler = FlowMatchEulerDiscreteScheduler(invert_sigmas=True, num_train_timesteps=1) tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, "tokenizer": tokenizer, } - return components - - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").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": 1, "guidance_scale": 3.0, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "height": 16, "width": 16, } - return inputs - def test_inference(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) - generated_image = pipe(**inputs).images[0] +class TestOmniGenPipeline(OmniGenPipelineTesterConfig, PipelineTesterMixin): + """Core pipeline tests for OmniGen. The old `test_inference` only asserted the output shape, which + `PipelineTesterMixin.test_output` now covers against `output_shape`.""" - self.assertEqual(generated_image.shape, (16, 16, 3)) + +class TestOmniGenPipelineMemory(OmniGenPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the OmniGen pipeline.""" @slow @require_torch_accelerator -class OmniGenPipelineSlowTests(unittest.TestCase): +class TestOmniGenPipelineIntegration: pipeline_class = OmniGenPipeline repo_id = "shitao/OmniGen-v1-diffusers" - 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/ovis_image/test_ovis_image.py b/tests/pipelines/ovis_image/test_ovis_image.py index be5fee50bb1b..bfa951233803 100644 --- a/tests/pipelines/ovis_image/test_ovis_image.py +++ b/tests/pipelines/ovis_image/test_ovis_image.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 Qwen2Tokenizer, Qwen3Config, Qwen3Model @@ -27,29 +24,15 @@ ) from ...testing_utils import 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 PipelineTesterMixin +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -class OvisImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class OvisImagePipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = OvisImagePipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] - ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = 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) @@ -101,47 +84,37 @@ def get_dummy_components(self): "transformer": transformer, } - 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): return { "prompt": "a cat", "negative_prompt": "bad quality", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 2.0, "height": 16, "width": 16, "max_sequence_length": 32, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - 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) - image = pipe(**inputs).images - generated_image = image[0] - self.assertEqual(generated_image.shape, (16, 16, 3)) - self.assertTrue(np.isfinite(image).all()) +class TestOvisImagePipeline(OvisImagePipelineTesterConfig, PipelineTesterMixin): + def test_inference(self, base_pipe_output): + # `test_output` already pins the shape; this one guards against a NaN/inf output. + assert torch.isfinite(base_pipe_output).all() def test_guidance_scale_is_set(self): # The `guidance_scale` property reads `self._guidance_scale`, which `__call__` must initialize. - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() pipe(**inputs) assert pipe.guidance_scale == inputs["guidance_scale"] def test_max_sequence_length_is_used(self): # `max_sequence_length` should bound the encoded prompt length. - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) embeds_16 = pipe.encode_prompt( "a cat", do_classifier_free_guidance=False, device=torch_device, max_sequence_length=16 )[0] @@ -150,3 +123,7 @@ def test_max_sequence_length_is_used(self): )[0] assert embeds_16.shape[1] == 16 assert embeds_32.shape[1] == 32 + + +class TestOvisImagePipelineMemory(OvisImagePipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the OvisImage pipeline.""" diff --git a/tests/pipelines/pag/test_pag_animatediff.py b/tests/pipelines/pag/test_pag_animatediff.py index ba508c88a277..a3ae6d0ee6e6 100644 --- a/tests/pipelines/pag/test_pag_animatediff.py +++ b/tests/pipelines/pag/test_pag_animatediff.py @@ -1,7 +1,4 @@ -import inspect -import unittest - -import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -10,49 +7,28 @@ AnimateDiffPipeline, AutoencoderKL, DDIMScheduler, - DPMSolverMultistepScheduler, - LCMScheduler, MotionAdapter, - StableDiffusionPipeline, UNet2DConditionModel, - UNetMotionModel, ) -from diffusers.models.attention import FreeNoiseTransformerBlock -from diffusers.utils import is_xformers_available -from ...testing_utils import require_accelerator, torch_device -from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import ( - IPAdapterTesterMixin, - PipelineFromPipeTesterMixin, - PipelineTesterMixin, - SDFunctionTesterMixin, +from ...testing_utils import torch_device +from ..animatediff.testing_utils import ( + FreeInitTesterMixin, + FreeNoiseTesterMixin, + MotionPipelineTesterConfig, + MotionPipelineTesterMixin, ) +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..testing_utils import FromPipeTesterMixin, IPAdapterTesterMixin, MemoryTesterMixin +from .testing_utils import PAGPipelineTesterMixin -def to_np(tensor): - if isinstance(tensor, torch.Tensor): - tensor = tensor.detach().cpu().numpy() - - return tensor - - -class AnimateDiffPAGPipelineFastTests( - IPAdapterTesterMixin, SDFunctionTesterMixin, PipelineTesterMixin, PipelineFromPipeTesterMixin, unittest.TestCase -): +class AnimateDiffPAGPipelineTesterConfig(MotionPipelineTesterConfig): pipeline_class = AnimateDiffPAGPipeline - params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - batch_params = TEXT_TO_IMAGE_BATCH_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 = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + # `num_frames` defaults to 16; height/width default to `unet.sample_size * vae_scale_factor` (8 * 2). + output_shape = (16, 3, 16, 16) def get_dummy_components(self): cross_attention_dim = 8 @@ -100,6 +76,7 @@ def get_dummy_components(self): ) text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + torch.manual_seed(0) motion_adapter = MotionAdapter( block_out_channels=block_out_channels, motion_layers_per_block=2, @@ -107,7 +84,7 @@ def get_dummy_components(self): motion_num_attention_heads=4, ) - components = { + return { "unet": unet, "scheduler": scheduler, "vae": vae, @@ -117,375 +94,51 @@ def get_dummy_components(self): "feature_extractor": None, "image_encoder": 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": 7.5, "pag_scale": 3.0, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs - - def test_from_pipe_consistent_config(self): - assert self.original_pipeline_class == StableDiffusionPipeline - original_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" - original_kwargs = {"requires_safety_checker": False} - - # create original_pipeline_class(sd) - pipe_original = self.original_pipeline_class.from_pretrained(original_repo, **original_kwargs) - - # original_pipeline_class(sd) -> pipeline_class - pipe_components = self.get_dummy_components() - pipe_additional_components = {} - for name, component in pipe_components.items(): - if name not in pipe_original.components: - pipe_additional_components[name] = component - - pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components) - - # pipeline_class -> original_pipeline_class(sd) - original_pipe_additional_components = {} - for name, component in pipe_original.components.items(): - if name not in pipe.components or not isinstance(component, pipe.components[name].__class__): - original_pipe_additional_components[name] = component - - pipe_original_2 = self.original_pipeline_class.from_pipe(pipe, **original_pipe_additional_components) - - # compare the config - original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")} - original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")} - assert original_config_2 == original_config - - def test_motion_unet_loading(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - assert isinstance(pipe.unet, UNetMotionModel) - - @unittest.skip("Attention slicing is not enabled in this pipeline") - def test_attention_slicing_forward_pass(self): - pass - - def test_ip_adapter(self): - expected_pipe_slice = None - - if torch_device == "cpu": - expected_pipe_slice = np.array( - [ - 0.5254, - 0.5844, - 0.4705, - 0.4952, - 0.4953, - 0.5932, - 0.5227, - 0.4285, - 0.5430, - 0.4903, - 0.4027, - 0.4942, - 0.3771, - 0.3899, - 0.6014, - 0.4469, - 0.4935, - 0.5543, - 0.5867, - 0.5314, - 0.3727, - 0.5183, - 0.6138, - 0.5170, - 0.5111, - 0.4725, - 0.5741, - ] - ) - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) - - def test_dict_tuple_outputs_equivalent(self): - expected_slice = None - if torch_device == "cpu": - expected_slice = np.array([0.5227, 0.4285, 0.5430, 0.4469, 0.4935, 0.5543, 0.5111, 0.4725, 0.5741]) - return super().test_dict_tuple_outputs_equivalent(expected_slice=expected_slice) - - @require_accelerator - def test_to_device(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - - pipe.to("cpu") - # pipeline creates a new motion UNet under the hood. So we need to check the device from pipe.components - model_devices = [ - component.device.type for component in pipe.components.values() if hasattr(component, "device") - ] - self.assertTrue(all(device == "cpu" for device in model_devices)) - - output_cpu = pipe(**self.get_dummy_inputs("cpu"))[0] - self.assertTrue(np.isnan(output_cpu).sum() == 0) - - pipe.to(torch_device) - model_devices = [ - component.device.type for component in pipe.components.values() if hasattr(component, "device") - ] - self.assertTrue(all(device == torch_device for device in model_devices)) - - output_device = pipe(**self.get_dummy_inputs(torch_device))[0] - self.assertTrue(np.isnan(to_np(output_device)).sum() == 0) - - def test_to_dtype(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - - # pipeline creates a new motion UNet under the hood. So we need to check the dtype from pipe.components - model_dtypes = [component.dtype for component in pipe.components.values() if hasattr(component, "dtype")] - self.assertTrue(all(dtype == torch.float32 for dtype in model_dtypes)) - - pipe.to(dtype=torch.float16) - model_dtypes = [component.dtype for component in pipe.components.values() if hasattr(component, "dtype")] - self.assertTrue(all(dtype == torch.float16 for dtype in model_dtypes)) - - def test_prompt_embeds(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - pipe.to(torch_device) - - inputs = self.get_dummy_inputs(torch_device) - inputs.pop("prompt") - inputs["prompt_embeds"] = torch.randn((1, 4, pipe.text_encoder.config.hidden_size), device=torch_device) - pipe(**inputs) - - def test_free_init(self): - components = self.get_dummy_components() - pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - pipe.to(torch_device) - - inputs_normal = self.get_dummy_inputs(torch_device) - frames_normal = pipe(**inputs_normal).frames[0] - - pipe.enable_free_init( - num_iters=2, - use_fast_sampling=True, - method="butterworth", - order=4, - spatial_stop_frequency=0.25, - temporal_stop_frequency=0.25, - ) - inputs_enable_free_init = self.get_dummy_inputs(torch_device) - frames_enable_free_init = pipe(**inputs_enable_free_init).frames[0] - - pipe.disable_free_init() - inputs_disable_free_init = self.get_dummy_inputs(torch_device) - frames_disable_free_init = pipe(**inputs_disable_free_init).frames[0] - - sum_enabled = np.abs(to_np(frames_normal) - to_np(frames_enable_free_init)).sum() - max_diff_disabled = np.abs(to_np(frames_normal) - to_np(frames_disable_free_init)).max() - self.assertGreater( - sum_enabled, 1e1, "Enabling of FreeInit should lead to results different from the default pipeline results" - ) - self.assertLess( - max_diff_disabled, - 1e-3, - "Disabling of FreeInit should lead to results similar to the default pipeline results", - ) - - def test_free_init_with_schedulers(self): - components = self.get_dummy_components() - pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - pipe.to(torch_device) - - inputs_normal = self.get_dummy_inputs(torch_device) - frames_normal = pipe(**inputs_normal).frames[0] - - schedulers_to_test = [ - DPMSolverMultistepScheduler.from_config( - components["scheduler"].config, - timestep_spacing="linspace", - beta_schedule="linear", - algorithm_type="dpmsolver++", - steps_offset=1, - clip_sample=False, - ), - LCMScheduler.from_config( - components["scheduler"].config, - timestep_spacing="linspace", - beta_schedule="linear", - steps_offset=1, - clip_sample=False, - ), - ] - components.pop("scheduler") - - for scheduler in schedulers_to_test: - components["scheduler"] = scheduler - pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - pipe.to(torch_device) - - pipe.enable_free_init(num_iters=2, use_fast_sampling=False) - - inputs = self.get_dummy_inputs(torch_device) - frames_enable_free_init = pipe(**inputs).frames[0] - sum_enabled = np.abs(to_np(frames_normal) - to_np(frames_enable_free_init)).sum() - - self.assertGreater( - sum_enabled, - 1e1, - "Enabling of FreeInit should lead to results different from the default pipeline results", - ) - - def test_free_noise_blocks(self): - components = self.get_dummy_components() - pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - pipe.to(torch_device) - pipe.enable_free_noise() - for block in pipe.unet.down_blocks: - for motion_module in block.motion_modules: - for transformer_block in motion_module.transformer_blocks: - self.assertTrue( - isinstance(transformer_block, FreeNoiseTransformerBlock), - "Motion module transformer blocks must be an instance of `FreeNoiseTransformerBlock` after enabling FreeNoise.", - ) - pipe.disable_free_noise() - for block in pipe.unet.down_blocks: - for motion_module in block.motion_modules: - for transformer_block in motion_module.transformer_blocks: - self.assertFalse( - isinstance(transformer_block, FreeNoiseTransformerBlock), - "Motion module transformer blocks must not be an instance of `FreeNoiseTransformerBlock` after disabling FreeNoise.", - ) - - def test_free_noise(self): - components = self.get_dummy_components() - pipe: AnimateDiffPAGPipeline = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - pipe.to(torch_device) - - inputs_normal = self.get_dummy_inputs(torch_device) - frames_normal = pipe(**inputs_normal).frames[0] - - for context_length in [8, 9]: - for context_stride in [4, 6]: - pipe.enable_free_noise(context_length, context_stride) - - inputs_enable_free_noise = self.get_dummy_inputs(torch_device) - frames_enable_free_noise = pipe(**inputs_enable_free_noise).frames[0] - - pipe.disable_free_noise() - - inputs_disable_free_noise = self.get_dummy_inputs(torch_device) - frames_disable_free_noise = pipe(**inputs_disable_free_noise).frames[0] - - sum_enabled = np.abs(to_np(frames_normal) - to_np(frames_enable_free_noise)).sum() - max_diff_disabled = np.abs(to_np(frames_normal) - to_np(frames_disable_free_noise)).max() - self.assertGreater( - sum_enabled, - 1e1, - "Enabling of FreeNoise should lead to results different from the default pipeline results", - ) - self.assertLess( - max_diff_disabled, - 1e-4, - "Disabling of FreeNoise should lead to results similar to the default pipeline results", - ) - - @unittest.skipIf( - torch_device != "cuda" or not is_xformers_available(), - reason="XFormers attention is only available with CUDA and `xformers` installed", +class TestAnimateDiffPAGPipeline( + AnimateDiffPAGPipelineTesterConfig, + MotionPipelineTesterMixin, + PAGPipelineTesterMixin, + FreeInitTesterMixin, + FreeNoiseTesterMixin, +): + base_pipeline_class = AnimateDiffPipeline + # AnimateDiff's PAG layers resolve through the motion modules, so the "PAG enabled" leg keeps the pipeline + # default layers and just leaves `pag_scale` at the dummy inputs' 3.0. + pag_enabled_applied_layers = None + + @pytest.mark.skip( + "`AnimateDiffPAGPipeline.check_inputs` rejects a dict `prompt`, so FreeNoise's per-frame prompts cannot be " + "passed to it (the non-PAG `AnimateDiffPipeline` accepts them)." ) - def test_xformers_attention_forwardGenerator_pass(self): - 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) - - inputs = self.get_dummy_inputs(torch_device) - output_without_offload = pipe(**inputs).frames[0] - output_without_offload = ( - output_without_offload.cpu() if torch.is_tensor(output_without_offload) else output_without_offload - ) - - pipe.enable_xformers_memory_efficient_attention() - inputs = self.get_dummy_inputs(torch_device) - output_with_offload = pipe(**inputs).frames[0] - output_with_offload = ( - output_with_offload.cpu() if torch.is_tensor(output_with_offload) else output_without_offload - ) - - max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max() - self.assertLess(max_diff, 1e-4, "XFormers attention should not affect the inference results") - - def test_vae_slicing(self): - return super().test_vae_slicing(image_count=2) - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - components.pop("pag_applied_layers", None) - pipe_sd = AnimateDiffPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) + def test_free_noise_multi_prompt(self): + pass - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." + def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4): + # Re-recorded: `get_dummy_components` now seeds the `MotionAdapter` like every other component, which it + # did not before the migration, so the adapter's weights (and this slice) changed. + if torch_device == "cpu" and expected_slice is None: + # fmt: off + expected_slice = torch.tensor([0.5132, 0.4380, 0.5327, 0.4619, 0.4955, 0.5457, 0.4980, 0.5015, 0.5652]) + # fmt: on + super().test_dict_tuple_outputs_equivalent( + expected_slice=expected_slice, expected_max_difference=expected_max_difference ) - out = pipe_sd(**inputs).frames[0, -3:, -3:, -1] - - components = self.get_dummy_components() - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).frames[0, -3:, -3:, -1] - - # pag enabled - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).frames[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 def test_pag_applied_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline - components.pop("pag_applied_layers", None) - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() # pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers # Note that for motion modules in AnimateDiff, both attn1 and attn2 are self-attention @@ -526,7 +179,7 @@ def test_pag_applied_layers(self): pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["mid_block.attentions.1"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) # pag_applied_layers = "down" should apply to all self-attention layers in down_blocks @@ -551,13 +204,19 @@ def test_pag_applied_layers(self): pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["motion_modules.42"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) - def test_encode_prompt_works_in_isolation(self): - extra_required_param_value_dict = { - "device": torch.device(torch_device).type, - "num_images_per_prompt": 1, - "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, - } - return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) + +class TestAnimateDiffPAGPipelineMemory(AnimateDiffPAGPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the AnimateDiff PAG pipeline.""" + + +class TestAnimateDiffPAGPipelineIPAdapter(AnimateDiffPAGPipelineTesterConfig, IPAdapterTesterMixin): + """IP-Adapter tests for the AnimateDiff PAG pipeline.""" + + +class TestAnimateDiffPAGPipelineFromPipe(AnimateDiffPAGPipelineTesterConfig, FromPipeTesterMixin): + """`from_pipe` round-trip tests against `StableDiffusionPipeline`.""" + + original_pipeline_repo = "hf-internal-testing/tinier-stable-diffusion-pipe" diff --git a/tests/pipelines/pag/test_pag_controlnet_sd.py b/tests/pipelines/pag/test_pag_controlnet_sd.py index c65fadff3bd9..8e551a7d3fbb 100644 --- a/tests/pipelines/pag/test_pag_controlnet_sd.py +++ b/tests/pipelines/pag/test_pag_controlnet_sd.py @@ -13,10 +13,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 from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -30,37 +26,30 @@ ) from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device from ..pipeline_params import ( TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, + FromPipeTesterMixin, IPAdapterTesterMixin, - PipelineFromPipeTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, + MemoryTesterMixin, ) +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionControlNetPAGPipelineFastTests( - PipelineTesterMixin, - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineFromPipeTesterMixin, - unittest.TestCase, -): +class StableDiffusionControlNetPAGPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionControlNetPAGPipeline - params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"}) + output_shape = (3, 64, 64) def get_dummy_components(self, time_cond_proj_dim=None): # Copied from tests.pipelines.controlnet.test_controlnet_sdxl.StableDiffusionXLControlNetPipelineFastTests.get_dummy_components @@ -120,7 +109,7 @@ def get_dummy_components(self, time_cond_proj_dim=None): text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -131,123 +120,88 @@ def get_dummy_components(self, time_cond_proj_dim=None): "feature_extractor": None, "image_encoder": 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) + def get_dummy_inputs(self): + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 + # The conditioning image is drawn from the same generator, which is then handed to the pipeline in the + # state that leaves it — the expected slices below were recorded that way. image = randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device("cpu"), ) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "num_inference_steps": 2, "guidance_scale": 6.0, "pag_scale": 3.0, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "image": image, } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe_sd = StableDiffusionControlNetPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - # pag enabled - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 +class TestStableDiffusionControlNetPAGPipeline( + StableDiffusionControlNetPAGPipelineTesterConfig, PAGPipelineTesterMixin +): + base_pipeline_class = StableDiffusionControlNetPipeline def test_pag_cfg(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array( - [0.4533307, 0.28746086, 0.16543725, 0.786471, 0.5409842, 0.40167668, 0.7050086, 0.69445, 0.42238155] + # Run on CPU: the expected slice below is CPU-specific. + pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"]) + + image = pipe_pag(**self.get_dummy_inputs())[0] + assert image.shape == (1, *self.output_shape), ( + f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}" ) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" + # fmt: off + expected_slice = torch.tensor([0.4533307, 0.28746086, 0.16543725, 0.786471, 0.5409842, 0.40167668, 0.7050086, 0.69445, 0.42238155]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_pag_uncond(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"]) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["guidance_scale"] = 0.0 - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array( - [0.45054412, 0.27958393, 0.15983358, 0.80098593, 0.5432001, 0.4018666, 0.7127423, 0.69932824, 0.42260933] + image = pipe_pag(**inputs)[0] + assert image.shape == (1, *self.output_shape), ( + f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}" ) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" + # fmt: off + expected_slice = torch.tensor([0.45054412, 0.27958393, 0.15983358, 0.80098593, 0.5432001, 0.4018666, 0.7127423, 0.69932824, 0.42260933]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_encode_prompt_works_in_isolation(self): extra_required_param_value_dict = { "device": torch.device(torch_device).type, - "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) + + +class TestStableDiffusionControlNetPAGPipelineMemory( + StableDiffusionControlNetPAGPipelineTesterConfig, MemoryTesterMixin +): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD ControlNet PAG pipeline.""" + + +class TestStableDiffusionControlNetPAGPipelineIPAdapter( + StableDiffusionControlNetPAGPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SD ControlNet PAG pipeline.""" + + +class TestStableDiffusionControlNetPAGPipelineFromPipe( + StableDiffusionControlNetPAGPipelineTesterConfig, FromPipeTesterMixin +): + """`from_pipe` round-trip tests against `StableDiffusionPipeline`.""" diff --git a/tests/pipelines/pag/test_pag_controlnet_sd_inpaint.py b/tests/pipelines/pag/test_pag_controlnet_sd_inpaint.py index a59180cb9540..bb7ccb9eefbc 100644 --- a/tests/pipelines/pag/test_pag_controlnet_sd_inpaint.py +++ b/tests/pipelines/pag/test_pag_controlnet_sd_inpaint.py @@ -15,9 +15,7 @@ # This model implementation is heavily based on: -import inspect import random -import unittest import numpy as np import torch @@ -34,26 +32,24 @@ ) from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, floats_tensor, torch_device +from ...testing_utils import assert_tensors_close, enable_full_determinism, floats_tensor, torch_device from ..pipeline_params import ( TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, ) -from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionControlNetPAGInpaintPipelineFastTests( - PipelineLatentTesterMixin, PipelineKarrasSchedulerTesterMixin, PipelineTesterMixin, unittest.TestCase -): +class StableDiffusionControlNetPAGInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionControlNetPAGInpaintPipeline - params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - image_params = frozenset({"control_image"}) # skip `image` and `mask` for now, only test for control_image - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS + batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS + # The output resolution follows the 64x64 input image. + output_shape = (3, 64, 64) def get_dummy_components(self): # Copied from tests.pipelines.controlnet.test_controlnet_inpaint.ControlNetInpaintPipelineFastTests.get_dummy_components @@ -109,7 +105,7 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -120,130 +116,83 @@ def get_dummy_components(self): "feature_extractor": None, "image_encoder": 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) + def get_dummy_inputs(self): + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 + # The control image is drawn from the same generator, which is then handed to the pipeline in the state + # that leaves it — the expected slices below were recorded that way. control_image = randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device("cpu"), ) - init_image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + init_image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) init_image = init_image.cpu().permute(0, 2, 3, 1)[0] image = Image.fromarray(np.uint8(init_image)).convert("RGB").resize((64, 64)) mask_image = Image.fromarray(np.uint8(init_image + 4)).convert("RGB").resize((64, 64)) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "num_inference_steps": 2, "guidance_scale": 6.0, "pag_scale": 3.0, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "image": image, "mask_image": mask_image, "control_image": control_image, } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe_sd = StableDiffusionControlNetInpaintPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__calss__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - # pag enabled - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 +class TestStableDiffusionControlNetPAGInpaintPipeline( + StableDiffusionControlNetPAGInpaintPipelineTesterConfig, PAGPipelineTesterMixin +): + base_pipeline_class = StableDiffusionControlNetInpaintPipeline def test_pag_cfg(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array( - [0.7277897, 0.61666954, 0.54722667, 0.595576, 0.593909, 0.56389576, 0.41761285, 0.50566983, 0.49766505] + # Run on CPU: the expected slice below is CPU-specific. + pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"]) + + image = pipe_pag(**self.get_dummy_inputs())[0] + assert image.shape == (1, *self.output_shape), ( + f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}" ) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" + # fmt: off + expected_slice = torch.tensor([0.7277897, 0.61666954, 0.54722667, 0.595576, 0.593909, 0.56389576, 0.41761285, 0.50566983, 0.49766505]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_pag_uncond(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"]) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["guidance_scale"] = 0.0 - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array( - [0.7349223, 0.60567534, 0.5428778, 0.6091342, 0.60273147, 0.57611704, 0.42401767, 0.5064247, 0.49535546] + image = pipe_pag(**inputs)[0] + assert image.shape == (1, *self.output_shape), ( + f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}" ) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" + # fmt: off + expected_slice = torch.tensor([0.7349223, 0.60567534, 0.5428778, 0.6091342, 0.60273147, 0.57611704, 0.42401767, 0.5064247, 0.49535546]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_encode_prompt_works_in_isolation(self): extra_required_param_value_dict = { "device": torch.device(torch_device).type, - "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) + + +class TestStableDiffusionControlNetPAGInpaintPipelineMemory( + StableDiffusionControlNetPAGInpaintPipelineTesterConfig, MemoryTesterMixin +): + """Memory tests (CPU offload, group offload, layerwise casting) for the SD ControlNet PAG inpaint pipeline.""" diff --git a/tests/pipelines/pag/test_pag_controlnet_sdxl.py b/tests/pipelines/pag/test_pag_controlnet_sdxl.py index 25d0c77d8627..756b228c0de1 100644 --- a/tests/pipelines/pag/test_pag_controlnet_sdxl.py +++ b/tests/pipelines/pag/test_pag_controlnet_sdxl.py @@ -13,10 +13,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 transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer @@ -30,37 +27,30 @@ ) from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism +from ...testing_utils import assert_tensors_close, enable_full_determinism from ..pipeline_params import ( TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, + FromPipeTesterMixin, IPAdapterTesterMixin, - PipelineFromPipeTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, + MemoryTesterMixin, ) +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionXLControlNetPAGPipelineFastTests( - PipelineTesterMixin, - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineFromPipeTesterMixin, - unittest.TestCase, -): +class StableDiffusionXLControlNetPAGPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLControlNetPAGPipeline - params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"}) + output_shape = (3, 64, 64) def get_dummy_components(self, time_cond_proj_dim=None): # Copied from tests.pipelines.controlnet.test_controlnet_sdxl.StableDiffusionXLControlNetPipelineFastTests.get_dummy_components @@ -151,114 +141,84 @@ def get_dummy_components(self, time_cond_proj_dim=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) + def get_dummy_inputs(self): + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 + # The conditioning image is drawn from the same generator, which is then handed to the pipeline in the + # state that leaves it — the expected slices below were recorded that way. image = randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device("cpu"), ) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "num_inference_steps": 2, "guidance_scale": 6.0, "pag_scale": 3.0, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "image": image, } - return inputs - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() +class TestStableDiffusionXLControlNetPAGPipeline( + StableDiffusionXLControlNetPAGPipelineTesterConfig, PAGPipelineTesterMixin +): + base_pipeline_class = StableDiffusionXLControlNetPipeline - # base pipeline (expect same output when pag is disabled) - pipe_sd = StableDiffusionXLControlNetPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) + def test_pag_cfg(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"]) - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." + image = pipe_pag(**self.get_dummy_inputs())[0] + assert image.shape == (1, *self.output_shape), ( + f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}" ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] + # fmt: off + expected_slice = torch.tensor([0.6864, 0.5436, 0.5644, 0.6136, 0.5541, 0.5910, 0.4519, 0.4634, 0.5252]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) - # pag enabled - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 - - def test_pag_cfg(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) + def test_pag_uncond(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"]) - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] + inputs = self.get_dummy_inputs() + inputs["guidance_scale"] = 0.0 + image = pipe_pag(**inputs)[0] + assert image.shape == (1, *self.output_shape), ( + f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}" + ) - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.6864, 0.5436, 0.5644, 0.6136, 0.5541, 0.5910, 0.4519, 0.4634, 0.5252]) + # fmt: off + expected_slice = torch.tensor([0.6843, 0.5381, 0.5675, 0.6109, 0.5493, 0.5988, 0.4477, 0.4679, 0.5242]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" + @pytest.mark.skip("We test this functionality elsewhere already.") + def test_save_load_optional_components(self): + pass - def test_pag_uncond(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) +class TestStableDiffusionXLControlNetPAGPipelineMemory( + StableDiffusionXLControlNetPAGPipelineTesterConfig, MemoryTesterMixin +): + """Memory tests (CPU offload, group offload, layerwise casting) for the SDXL ControlNet PAG pipeline.""" - inputs = self.get_dummy_inputs(device) - inputs["guidance_scale"] = 0.0 - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.6843, 0.5381, 0.5675, 0.6109, 0.5493, 0.5988, 0.4477, 0.4679, 0.5242]) +class TestStableDiffusionXLControlNetPAGPipelineIPAdapter( + StableDiffusionXLControlNetPAGPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SDXL ControlNet PAG pipeline.""" - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" - @unittest.skip("We test this functionality elsewhere already.") - def test_save_load_optional_components(self): - pass +class TestStableDiffusionXLControlNetPAGPipelineFromPipe( + StableDiffusionXLControlNetPAGPipelineTesterConfig, FromPipeTesterMixin +): + """`from_pipe` round-trip tests against `StableDiffusionXLPipeline`.""" diff --git a/tests/pipelines/pag/test_pag_controlnet_sdxl_img2img.py b/tests/pipelines/pag/test_pag_controlnet_sdxl_img2img.py index e23f92e31119..89e6b1191b0e 100644 --- a/tests/pipelines/pag/test_pag_controlnet_sdxl_img2img.py +++ b/tests/pipelines/pag/test_pag_controlnet_sdxl_img2img.py @@ -13,11 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect import random -import unittest -import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer @@ -30,38 +28,34 @@ UNet2DConditionModel, ) -from ...testing_utils import enable_full_determinism, floats_tensor +from ...testing_utils import assert_tensors_close, enable_full_determinism, floats_tensor, torch_device from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, + FromPipeTesterMixin, IPAdapterTesterMixin, - PipelineFromPipeTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, + MemoryTesterMixin, ) +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionXLControlNetPAGImg2ImgPipelineFastTests( - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, - PipelineFromPipeTesterMixin, - unittest.TestCase, -): +class StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLControlNetPAGImg2ImgPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union( + {"pag_scale", "pag_adaptive_scale"} + ) + batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union( {"add_text_embeds", "add_time_ids", "add_neg_time_ids"} ) + # The output resolution follows the 64x64 conditioning image. + output_shape = (3, 64, 64) # Copied from tests.pipelines.controlnet.test_controlnet_sdxl_img2img.ControlNetPipelineSDXLImg2ImgFastTests.get_dummy_components def get_dummy_components(self, skip_first_text_encoder=False): @@ -153,117 +147,80 @@ def get_dummy_components(self, skip_first_text_encoder=False): # based on tests.pipelines.controlnet.test_controlnet_sdxl_img2img.ControlNetPipelineSDXLImg2ImgFastTests.get_dummy_inputs # add `pag_scale` to the inputs - def get_dummy_inputs(self, device, seed=0): + def get_dummy_inputs(self): controlnet_embedder_scale_factor = 2 image = floats_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), - rng=random.Random(seed), - ).to(device) + rng=random.Random(0), + ).to(torch_device) - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "pag_scale": 3.0, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "image": image, "control_image": image, } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline - pipe_sd = StableDiffusionXLControlNetImg2ImgPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - # pag enable - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 +class TestStableDiffusionXLControlNetPAGImg2ImgPipeline( + StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig, PAGPipelineTesterMixin +): + base_pipeline_class = StableDiffusionXLControlNetImg2ImgPipeline + @pytest.mark.skip("We test this functionality elsewhere already.") def test_save_load_optional_components(self): pass def test_pag_cfg(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() + # Run on CPU: the expected slice below is CPU-specific. + pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"]) + + image = pipe_pag(**self.get_dummy_inputs())[0] + assert image.shape == (1, *self.output_shape), ( + f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}" + ) - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) + # fmt: off + expected_slice = torch.tensor([0.55155665, 0.4650753, 0.46541628, 0.60965055, 0.55995595, 0.49751496, 0.5937391, 0.5700847, 0.44238678]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] + def test_pag_uncond(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe_pag = self.get_pag_pipeline(pag_applied_layers=["mid", "up", "down"]) - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array( - [0.55155665, 0.4650753, 0.46541628, 0.60965055, 0.55995595, 0.49751496, 0.5937391, 0.5700847, 0.44238678] + inputs = self.get_dummy_inputs() + inputs["guidance_scale"] = 0.0 + image = pipe_pag(**inputs)[0] + assert image.shape == (1, *self.output_shape), ( + f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}" ) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" + # fmt: off + expected_slice = torch.tensor([0.549061, 0.46218234, 0.4675981, 0.6109464, 0.5547849, 0.4960261, 0.60211027, 0.5698843, 0.44092298]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) - def test_pag_uncond(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) +class TestStableDiffusionXLControlNetPAGImg2ImgPipelineMemory( + StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig, MemoryTesterMixin +): + """Memory tests (CPU offload, group offload, layerwise casting) for the SDXL ControlNet PAG img2img pipeline.""" - inputs = self.get_dummy_inputs(device) - inputs["guidance_scale"] = 0.0 - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array( - [0.549061, 0.46218234, 0.4675981, 0.6109464, 0.5547849, 0.4960261, 0.60211027, 0.5698843, 0.44092298] - ) +class TestStableDiffusionXLControlNetPAGImg2ImgPipelineIPAdapter( + StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SDXL ControlNet PAG img2img pipeline.""" + - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" +class TestStableDiffusionXLControlNetPAGImg2ImgPipelineFromPipe( + StableDiffusionXLControlNetPAGImg2ImgPipelineTesterConfig, FromPipeTesterMixin +): + """`from_pipe` round-trip tests against `StableDiffusionXLPipeline`.""" diff --git a/tests/pipelines/pag/test_pag_hunyuan_dit.py b/tests/pipelines/pag/test_pag_hunyuan_dit.py index d4dd144a840b..41d2f68aa5c7 100644 --- a/tests/pipelines/pag/test_pag_hunyuan_dit.py +++ b/tests/pipelines/pag/test_pag_hunyuan_dit.py @@ -13,11 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import tempfile -import unittest - -import numpy as np +import pytest import torch from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel @@ -29,22 +25,21 @@ HunyuanDiTPipeline, ) -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 PipelineTesterMixin, to_np +from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class HunyuanDiTPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class HunyuanDiTPAGPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HunyuanDiTPAGPipeline - 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 + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + # `transformer.sample_size` (16) * `vae_scale_factor` (8) / 8 -> the dummy transformer generates 16x16 images. + output_shape = (3, 16, 16) def get_dummy_components(self): torch.manual_seed(0) @@ -72,7 +67,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, @@ -83,159 +78,88 @@ 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", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "use_resolution_binning": False, "pag_scale": 0.0, } - 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) - image = pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] +class TestHunyuanDiTPAGPipeline(HunyuanDiTPAGPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = HunyuanDiTPipeline + # HunyuanDiT's denoiser is a transformer: PAG uses the pipeline's default applied layers, turned on by raising + # `pag_scale` (the dummy inputs disable it with `pag_scale=0.0`). + pag_enabled_applied_layers = None + pag_enabled_scale = 3.0 - 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) + def test_inference(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - @unittest.skip("Not supported.") - def test_sequential_cpu_offload_forward_pass(self): - # TODO(YiYi) need to fix later - pass + image = pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - @unittest.skip("Not supported.") - def test_sequential_offload_forward_pass_twice(self): - # TODO(YiYi) need to fix later - pass + # 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 + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical( - expected_max_diff=1e-3, - ) + super().test_inference_batch_single_identical(expected_max_diff=1e-3) def test_feed_forward_chunking(self): - device = "cpu" + # Run on CPU to keep the device-dependent `torch.Generator` deterministic. + pipe = self.get_pipeline() - 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) - image = pipe(**inputs).images - image_slice_no_chunking = image[0, -3:, -3:, -1] + output_no_chunking = self.run_pipe(pipe) 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] + output_chunking = self.run_pipe(pipe) - 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( + output_chunking, output_no_chunking, atol=1e-4, msg="Forward chunking changed the output." + ) 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 keep the device-dependent `torch.Generator` deterministic. + 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_output = self.run_pipe(pipe) pipe.transformer.fuse_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] + output_fused = self.run_pipe(pipe) 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] + output_disabled = self.run_pipe(pipe) - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( - "Fusion of QKV projections shouldn't affect the outputs." + assert_tensors_close( + output_fused, original_output, atol=1e-2, rtol=1e-2, msg="Fusion of QKV projections changed 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( + output_disabled, + output_fused, + atol=1e-2, + rtol=1e-2, + msg="Outputs changed after the fused QKV projections were 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( + output_disabled, + original_output, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe_sd = HunyuanDiTPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - components = self.get_dummy_components() - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - # pag enabled - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 3.0 - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 - def test_pag_applied_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn1" in k] original_attn_procs = pipe.transformer.attn_processors @@ -265,19 +189,16 @@ def test_pag_applied_layers(self): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert len(pipe.pag_attn_processors) == 2 - @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) + def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4): + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() prompt = inputs["prompt"] generator = inputs["generator"] @@ -325,19 +246,17 @@ def test_save_load_optional_components(self): output = pipe(**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) + inputs = self.get_dummy_inputs() generator = inputs["generator"] num_inference_steps = inputs["num_inference_steps"] @@ -361,5 +280,23 @@ def test_save_load_optional_components(self): output_loaded = pipe_loaded(**inputs)[0] - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, 1e-4) + assert_tensors_close( + output_loaded, + output, + atol=expected_max_difference, + msg="Output changed after dropping optional components.", + ) + + +class TestHunyuanDiTPAGPipelineMemory(HunyuanDiTPAGPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the HunyuanDiT PAG pipeline.""" + + @pytest.mark.skip("Not supported.") + def test_sequential_cpu_offload_forward_pass(self): + # TODO(YiYi) need to fix later + pass + + @pytest.mark.skip("Not supported.") + def test_sequential_offload_forward_pass_twice(self): + # TODO(YiYi) need to fix later + pass diff --git a/tests/pipelines/pag/test_pag_kolors.py b/tests/pipelines/pag/test_pag_kolors.py index dac3f02ca5ef..02341ffcac60 100644 --- a/tests/pipelines/pag/test_pag_kolors.py +++ b/tests/pipelines/pag/test_pag_kolors.py @@ -13,10 +13,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 diffusers import ( @@ -32,29 +29,21 @@ from ..pipeline_params import ( TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS, ) -from ..test_pipelines_common import ( - PipelineFromPipeTesterMixin, - PipelineTesterMixin, -) +from ..testing_utils import BasePipelineTesterConfig, FromPipeTesterMixin, MemoryTesterMixin +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class KolorsPAGPipelineFastTests( - PipelineTesterMixin, - PipelineFromPipeTesterMixin, - unittest.TestCase, -): +class KolorsPAGPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = KolorsPAGPipeline - params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"}) + output_shape = (3, 64, 64) # Copied from tests.pipelines.kolors.test_kolors.KolorsPipelineFastTests.get_dummy_components def get_dummy_components(self, time_cond_proj_dim=None): @@ -112,65 +101,27 @@ def get_dummy_components(self, time_cond_proj_dim=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, "pag_scale": 0.9, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe_sd = KolorsPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - # pag enabled - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 +class TestKolorsPAGPipeline(KolorsPAGPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = KolorsPipeline + # fmt: off + expected_pag_slice = torch.tensor([0.26030684, 0.43192005, 0.4042826, 0.4189067, 0.5181305, 0.3832534, 0.472135, 0.4145031, 0.43726248]) + # fmt: on def test_pag_applied_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() # pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers all_self_attn_layers = [k for k in pipe.unet.attn_processors.keys() if "attn1" in k] @@ -201,7 +152,7 @@ def test_pag_applied_layers(self): # pag_applied_layers = ["mid.block_0.attentions_1"] does not exist in the model pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["mid_block.attentions.1"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) # pag_applied_layers = "down" should apply to all self-attention layers in down_blocks @@ -212,7 +163,7 @@ def test_pag_applied_layers(self): pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["down_blocks.0"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) pipe.unet.set_attn_processor(original_attn_procs.copy()) @@ -225,33 +176,16 @@ def test_pag_applied_layers(self): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert len(pipe.pag_attn_processors) == 2 - def test_pag_inference(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() + def test_inference_batch_single_identical(self): + super().test_inference_batch_single_identical(expected_max_diff=3e-3) - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) + def test_encode_prompt_works_in_isolation(self): + return super().test_encode_prompt_works_in_isolation(atol=1e-3, rtol=1e-3) - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array( - [0.26030684, 0.43192005, 0.4042826, 0.4189067, 0.5181305, 0.3832534, 0.472135, 0.4145031, 0.43726248] - ) +class TestKolorsPAGPipelineMemory(KolorsPAGPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kolors PAG pipeline.""" - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=3e-3) - - def test_encode_prompt_works_in_isolation(self): - return super().test_encode_prompt_works_in_isolation(atol=1e-3, rtol=1e-3) +class TestKolorsPAGPipelineFromPipe(KolorsPAGPipelineTesterConfig, FromPipeTesterMixin): + """`from_pipe` round-trip tests against `KolorsPipeline`.""" diff --git a/tests/pipelines/pag/test_pag_pixart_sigma.py b/tests/pipelines/pag/test_pag_pixart_sigma.py index 2c0b89d0cfe0..607d3ec82644 100644 --- a/tests/pipelines/pag/test_pag_pixart_sigma.py +++ b/tests/pipelines/pag/test_pag_pixart_sigma.py @@ -13,15 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import tempfile -import unittest - -import numpy as np +import pytest import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel -import diffusers from diffusers import ( AutoencoderKL, DDIMScheduler, @@ -29,33 +24,30 @@ PixArtSigmaPipeline, PixArtTransformer2DModel, ) -from diffusers.utils import logging from ...testing_utils import ( - CaptureLogger, + assert_tensors_close, 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 PipelineTesterMixin, assert_mean_pixel_difference, to_np +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class PixArtSigmaPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class PixArtSigmaPAGPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = PixArtSigmaPAGPipeline - params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - params = set(params) - params.remove("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 + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - { + "cross_attention_kwargs" + } + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + # `transformer.sample_size` (8) * `vae_scale_factor` (8) / 8 -> the dummy transformer generates 8x8 images. + output_shape = (3, 8, 8) def get_dummy_components(self): torch.manual_seed(0) @@ -85,76 +77,40 @@ def get_dummy_components(self): tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler, "text_encoder": text_encoder, "tokenizer": tokenizer, } - return components - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + def get_dummy_inputs(self): + return { "prompt": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 1.0, "pag_scale": 3.0, "use_resolution_binning": False, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe = PixArtSigmaPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe.__class__.__name__}." - ) - out = pipe(**inputs).images[0, -3:, -3:, -1] - - # pag disabled with pag_scale=0.0 - components["pag_applied_layers"] = ["blocks.1"] - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - # pag enabled - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 +class TestPixArtSigmaPAGPipeline(PixArtSigmaPAGPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = PixArtSigmaPipeline + # PixArt's denoiser is a transformer, so PAG resolves per transformer block rather than the UNet's mid/up/down. + pag_enabled_applied_layers = ["blocks.1"] + # `test_pag_inference` builds the pipeline with the class default (`blocks.1`), as it did before the migration. + pag_inference_applied_layers = None + # fmt: off + expected_pag_slice = torch.tensor([0.6499, 0.3250, 0.3572, 0.6780, 0.4453, 0.4582, 0.2770, 0.5168, 0.4594]) + # fmt: on def test_pag_applied_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() # "attn1" should apply to all self-attention layers. all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn1" in k] @@ -162,189 +118,38 @@ def test_pag_applied_layers(self): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert set(pipe.pag_attn_processors) == set(all_self_attn_layers) - def test_pag_inference(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == ( - 1, - 8, - 8, - 3, - ), f"the shape of the output image should be (1, 8, 8, 3) but got {image.shape}" - expected_slice = np.array([0.6499, 0.3250, 0.3572, 0.6780, 0.4453, 0.4582, 0.2770, 0.5168, 0.4594]) - - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3): + # Run on CPU: sliced attention is compared against a full-attention run of the same pipeline. + pipe = self.get_pipeline() - # Because the PAG PixArt Sigma has `pag_applied_layers`. - # Also, we shouldn't be doing `set_default_attn_processor()` after loading - # the pipeline with `pag_applied_layers`. - def test_save_load_local(self, expected_max_difference=1e-4): - 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) - output = pipe(**inputs)[0] - - logger = logging.get_logger("diffusers.pipelines.pipeline_utils") - logger.setLevel(diffusers.logging.INFO) - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, safe_serialization=False) - - with CaptureLogger(logger) as cap_logger: - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, pag_applied_layers=["blocks.1"]) - - for name in pipe_loaded.components.keys(): - if name not in pipe_loaded._optional_components: - assert name in str(cap_logger) - - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output_loaded = pipe_loaded(**inputs)[0] - - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, expected_max_difference) - - # We shouldn't be setting `set_default_attn_processor` here. - def test_attention_slicing_forward_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 - ): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - 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] + output_without_slicing = self.run_pipe(pipe) pipe.enable_attention_slicing(slice_size=1) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing1 = pipe(**inputs)[0] + output_with_slicing_1 = self.run_pipe(pipe) 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", - ) + output_with_slicing_2 = self.run_pipe(pipe) - if test_mean_pixel_difference: - assert_mean_pixel_difference(to_np(output_with_slicing1[0]), to_np(output_without_slicing[0])) - assert_mean_pixel_difference(to_np(output_with_slicing2[0]), to_np(output_without_slicing[0])) - - # Because we have `pag_applied_layers` we cannot directly apply - # `set_default_attn_processor` - def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - generator_device = "cpu" - if expected_slice is None: - output = pipe(**self.get_dummy_inputs(generator_device))[0] - else: - output = expected_slice - - output_tuple = pipe(**self.get_dummy_inputs(generator_device), return_dict=False)[0] - - if expected_slice is None: - max_diff = np.abs(to_np(output) - to_np(output_tuple)).max() - else: - if output_tuple.ndim != 5: - max_diff = np.abs(to_np(output) - to_np(output_tuple)[0, -3:, -3:, -1].flatten()).max() - else: - max_diff = np.abs(to_np(output) - to_np(output_tuple)[0, -3:, -3:, -1, -1].flatten()).max() - - self.assertLess(max_diff, expected_max_difference) - - # Same reason as above - def test_inference_batch_single_identical( - self, - batch_size=2, - expected_max_diff=1e-4, - additional_params_copy_to_batched_inputs=["num_inference_steps"], - ): - 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) - # Reset generator in case it is has been used in self.get_dummy_inputs - inputs["generator"] = self.get_generator(0) - - logger = logging.get_logger(pipe.__module__) - logger.setLevel(level=diffusers.logging.FATAL) - - # batchify inputs - batched_inputs = {} - batched_inputs.update(inputs) - - for name in self.batch_params: - if name not in inputs: - continue - - value = inputs[name] - if name == "prompt": - len_prompt = len(value) - batched_inputs[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)] - batched_inputs[name][-1] = 100 * "very long" - - else: - batched_inputs[name] = batch_size * [value] - - if "generator" in inputs: - batched_inputs["generator"] = [self.get_generator(i) for i in range(batch_size)] - - if "batch_size" in inputs: - batched_inputs["batch_size"] = batch_size - - for arg in additional_params_copy_to_batched_inputs: - batched_inputs[arg] = inputs[arg] - - output = pipe(**inputs) - output_batch = pipe(**batched_inputs) - - assert output_batch[0].shape[0] == batch_size - - max_diff = np.abs(to_np(output_batch[0][0]) - to_np(output[0][0])).max() - assert max_diff < expected_max_diff - - # Because we're passing `pag_applied_layers` (type of List) in the components as well. - def test_components_function(self): - init_components = self.get_dummy_components() - init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float, list))} - - pipe = self.pipeline_class(**init_components) + assert_tensors_close( + output_with_slicing_1, + output_without_slicing, + atol=expected_max_diff, + msg="Attention slicing (slice_size=1) changed the output.", + ) + assert_tensors_close( + output_with_slicing_2, + output_without_slicing, + atol=expected_max_diff, + msg="Attention slicing (slice_size=2) changed the output.", + ) - self.assertTrue(hasattr(pipe, "components")) - self.assertTrue(set(pipe.components.keys()) == set(init_components.keys())) + def test_inference_batch_single_identical(self): + super().test_inference_batch_single_identical(batch_size=2) - @unittest.skip("Test is already covered through encode_prompt isolation.") + @pytest.mark.skip("Test is already covered through encode_prompt isolation.") def test_save_load_optional_components(self): pass + + +class TestPixArtSigmaPAGPipelineMemory(PixArtSigmaPAGPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the PixArt-sigma PAG pipeline.""" diff --git a/tests/pipelines/pag/test_pag_sana.py b/tests/pipelines/pag/test_pag_sana.py index d5c5d0824af7..2c35e8c1cfa5 100644 --- a/tests/pipelines/pag/test_pag_sana.py +++ b/tests/pipelines/pag/test_pag_sana.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 transformers import Gemma2Config, Gemma2ForCausalLM, GemmaTokenizer @@ -27,31 +24,26 @@ SanaTransformer2DModel, ) -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 PipelineTesterMixin, to_np +from ...testing_utils import ( + assert_tensors_close, + enable_full_determinism, + require_accelerator, + skip_if_no_cudnn_engine, + torch_device, +) +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class SanaPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class SanaPAGPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = SanaPAGPipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] - ) - test_xformers_attention = False + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -115,183 +107,64 @@ def get_dummy_components(self): text_encoder = Gemma2ForCausalLM(text_encoder_config) tokenizer = GemmaTokenizer.from_pretrained("hf-internal-testing/dummy-gemma") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, "text_encoder": text_encoder, "tokenizer": tokenizer, } - return components - - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + + def get_dummy_inputs(self): + return { "prompt": "", "negative_prompt": "", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "pag_scale": 3.0, "height": 32, "width": 32, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). "output_type": "pt", "complex_human_instruction": None, } - 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) - image = pipe(**inputs)[0] - generated_image = image[0] - - self.assertEqual(generated_image.shape, (3, 32, 32)) - expected_image = torch.randn(3, 32, 32) - max_diff = np.abs(generated_image - expected_image).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 - - 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] +class TestSanaPAGPipeline(SanaPAGPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = SanaPipeline + # Only the "PAG off reproduces the base pipeline" leg was asserted before the migration. + check_pag_changes_output = False - 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_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe_sd = SanaPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] + def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3): + # Run on CPU: sliced attention is compared against a full-attention run of the same pipeline. + pipe = self.get_pipeline() - components = self.get_dummy_components() + output_without_slicing = self.run_pipe(pipe) - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) + pipe.enable_attention_slicing(slice_size=1) + output_with_slicing_1 = self.run_pipe(pipe) - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] + pipe.enable_attention_slicing(slice_size=2) + output_with_slicing_2 = self.run_pipe(pipe) - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 + assert_tensors_close( + output_with_slicing_1, + output_without_slicing, + atol=expected_max_diff, + msg="Attention slicing (slice_size=1) changed the output.", + ) + assert_tensors_close( + output_with_slicing_2, + output_without_slicing, + atol=expected_max_diff, + msg="Attention slicing (slice_size=2) changed the output.", + ) def test_pag_applied_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn1" in k] original_attn_procs = pipe.transformer.attn_processors @@ -322,18 +195,33 @@ def test_pag_applied_layers(self): assert len(pipe.pag_attn_processors) == 2 # 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 - def test_float16_inference(self): - # Requires higher tolerance as model seems very sensitive to dtype - super().test_float16_inference(expected_max_diff=0.08) + # Sana's multiscale linear attention runs a depthwise `Conv2d`, which some cuDNN builds have no bfloat16 + # engine for. The decorators below repeat the ones the base method is declared with: overriding a test drops + # the marks it inherited. + @pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="half-precision inference requires CUDA or XPU") + @require_accelerator + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=str) + def test_half_precision_inference_no_nan(self, dtype): + with skip_if_no_cudnn_engine(): + super().test_half_precision_inference_no_nan(dtype) + + +class TestSanaPAGPipelineMemory(SanaPAGPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Sana PAG pipeline.""" + + # Layerwise casting computes in bfloat16, which lands on the same missing depthwise-conv engine as above. + def test_layerwise_casting_inference(self): + with skip_if_no_cudnn_engine(): + super().test_layerwise_casting_inference() diff --git a/tests/pipelines/pag/test_pag_sd.py b/tests/pipelines/pag/test_pag_sd.py index 1dd3ef298fd0..a009fbc2d41f 100644 --- a/tests/pipelines/pag/test_pag_sd.py +++ b/tests/pipelines/pag/test_pag_sd.py @@ -14,10 +14,9 @@ # limitations under the License. import gc -import inspect -import unittest import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -40,33 +39,26 @@ from ..pipeline_params import ( TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, + FromPipeTesterMixin, IPAdapterTesterMixin, - PipelineFromPipeTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, + MemoryTesterMixin, ) +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionPAGPipelineFastTests( - PipelineTesterMixin, - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineFromPipeTesterMixin, - unittest.TestCase, -): +class StableDiffusionPAGPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionPAGPipeline - params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"}) + output_shape = (3, 64, 64) def get_dummy_components(self, time_cond_proj_dim=None): cross_attention_dim = 8 @@ -116,7 +108,7 @@ def get_dummy_components(self, time_cond_proj_dim=None): text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "scheduler": scheduler, "vae": vae, @@ -126,67 +118,28 @@ def get_dummy_components(self, time_cond_proj_dim=None): "feature_extractor": None, "image_encoder": 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, "pag_scale": 0.9, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe_sd = StableDiffusionPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - # pag enabled - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 +class TestStableDiffusionPAGPipeline(StableDiffusionPAGPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = StableDiffusionPipeline + # fmt: off + expected_pag_slice = torch.tensor([0.23171297, 0.44669262, 0.48407662, 0.29981518, 0.36721927, 0.46788025, 0.46333545, 0.3314417, 0.42078137]) + # fmt: on def test_pag_applied_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() # pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers all_self_attn_layers = [k for k in pipe.unet.attn_processors.keys() if "attn1" in k] @@ -224,7 +177,7 @@ def test_pag_applied_layers(self): # pag_applied_layers = ["mid.block_0.attentions_1"] does not exist in the model pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["mid_block.attentions.1"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) # pag_applied_layers = "down" should apply to all self-attention layers in down_blocks @@ -239,7 +192,7 @@ def test_pag_applied_layers(self): pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["down_blocks.0"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) pipe.unet.set_attn_processor(original_attn_procs.copy()) @@ -252,52 +205,37 @@ def test_pag_applied_layers(self): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert len(pipe.pag_attn_processors) == 1 - def test_pag_inference(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - - expected_slice = np.array( - [0.23171297, 0.44669262, 0.48407662, 0.29981518, 0.36721927, 0.46788025, 0.46333545, 0.3314417, 0.42078137] - ) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) - def test_encode_prompt_works_in_isolation(self): extra_required_param_value_dict = { "device": torch.device(torch_device).type, - "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) +class TestStableDiffusionPAGPipelineMemory(StableDiffusionPAGPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD PAG pipeline.""" + + +class TestStableDiffusionPAGPipelineIPAdapter(StableDiffusionPAGPipelineTesterConfig, IPAdapterTesterMixin): + """IP-Adapter tests for the SD PAG pipeline.""" + + +class TestStableDiffusionPAGPipelineFromPipe(StableDiffusionPAGPipelineTesterConfig, FromPipeTesterMixin): + """`from_pipe` round-trip tests against `StableDiffusionPipeline`.""" + + @slow @require_torch_accelerator -class StableDiffusionPAGPipelineIntegrationTests(unittest.TestCase): +class TestStableDiffusionPAGPipelineIntegration: pipeline_class = StableDiffusionPAGPipeline repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - 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/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index 7f755ea8e170..6958985c954a 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -1,7 +1,3 @@ -import inspect -import unittest - -import numpy as np import torch from transformers import ( AutoConfig, @@ -20,19 +16,19 @@ StableDiffusion3Pipeline, ) -from ...testing_utils import ( - torch_device, -) -from ..test_pipelines_common import ( - PipelineTesterMixin, +from ...testing_utils import assert_tensors_close, torch_device +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, check_qkv_fusion_matches_attn_procs_length, check_qkv_fusion_processors_exist, ) +from .testing_utils import PAGPipelineTesterMixin -class StableDiffusion3PAGPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class StableDiffusion3PAGPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusion3PAGPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "prompt", "height", @@ -43,8 +39,8 @@ class StableDiffusion3PAGPipelineFastTests(unittest.TestCase, PipelineTesterMixi "negative_prompt_embeds", ] ) - batch_params = frozenset(["prompt", "negative_prompt"]) - test_xformers_attention = False + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -117,64 +113,50 @@ def get_dummy_components(self): "vae": vae, } - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").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`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "pag_scale": 0.0, } - return inputs - def test_stable_diffusion_3_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) - output_same_prompt = pipe(**inputs).images[0] +class TestStableDiffusion3PAGPipeline(StableDiffusion3PAGPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = StableDiffusion3Pipeline + # The dummy inputs disable PAG (`pag_scale=0.0`); only the "PAG off reproduces the base pipeline" leg is + # asserted here, as it was before the migration. + check_pag_changes_output = False - inputs = self.get_dummy_inputs(torch_device) - inputs["prompt_2"] = "a different prompt" - inputs["prompt_3"] = "another different prompt" - output_different_prompts = pipe(**inputs).images[0] + def test_stable_diffusion_3_different_prompts(self): + pipe = self.get_pipeline().to(torch_device) - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + output_same_prompt = self.run_pipe(pipe)[0] + output_different_prompts = self.run_pipe( + pipe, prompt_2="a different prompt", prompt_3="another different prompt" + )[0] # Outputs should be different here - assert max_diff > 1e-2 + assert (output_same_prompt - output_different_prompts).abs().max() > 1e-2 def test_stable_diffusion_3_different_negative_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - - inputs = self.get_dummy_inputs(torch_device) - output_same_prompt = pipe(**inputs).images[0] - - inputs = self.get_dummy_inputs(torch_device) - inputs["negative_prompt_2"] = "deformed" - inputs["negative_prompt_3"] = "blurry" - output_different_prompts = pipe(**inputs).images[0] + pipe = self.get_pipeline().to(torch_device) - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + output_same_prompt = self.run_pipe(pipe)[0] + output_different_prompts = self.run_pipe(pipe, negative_prompt_2="deformed", negative_prompt_3="blurry")[0] # Outputs should be different here - assert max_diff > 1e-2 + assert (output_same_prompt - output_different_prompts).abs().max() > 1e-2 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 keep the device-dependent `torch.Generator` deterministic. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - original_image_slice = image[0, -3:, -3:, -1] + original_output = self.run_pipe(pipe) # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added # to the pipeline level. @@ -186,62 +168,31 @@ 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) - image = pipe(**inputs).images - image_slice_fused = image[0, -3:, -3:, -1] + output_fused = self.run_pipe(pipe) pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_disabled = image[0, -3:, -3:, -1] + output_disabled = self.run_pipe(pipe) - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( - "Fusion of QKV projections shouldn't affect the outputs." + assert_tensors_close( + output_fused, original_output, atol=1e-3, rtol=1e-3, msg="Fusion of QKV projections changed the outputs." ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert_tensors_close( + output_disabled, + output_fused, + atol=1e-3, + rtol=1e-3, + msg="Outputs changed after the fused QKV projections were 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( + output_disabled, + original_output, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe_sd = StableDiffusion3Pipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - components = self.get_dummy_components() - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - def test_pag_applied_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() all_self_attn_layers = [k for k in pipe.transformer.attn_processors.keys() if "attn" in k] original_attn_procs = pipe.transformer.attn_processors @@ -270,3 +221,7 @@ def test_pag_applied_layers(self): pag_layers = ["blocks.0", r"blocks\.1"] pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert len(pipe.pag_attn_processors) == 2 + + +class TestStableDiffusion3PAGPipelineMemory(StableDiffusion3PAGPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 PAG pipeline.""" diff --git a/tests/pipelines/pag/test_pag_sd3_img2img.py b/tests/pipelines/pag/test_pag_sd3_img2img.py index ede146915c55..85018e96c6a2 100644 --- a/tests/pipelines/pag/test_pag_sd3_img2img.py +++ b/tests/pipelines/pag/test_pag_sd3_img2img.py @@ -1,9 +1,8 @@ import gc -import inspect import random -import unittest import numpy as np +import pytest import torch from transformers import ( AutoConfig, @@ -33,29 +32,30 @@ torch_device, ) from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( - PipelineTesterMixin, -) +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusion3PAGImg2ImgPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class StableDiffusion3PAGImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusion3PAGImg2ImgPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - {"height", "width"} - required_optional_params = PipelineTesterMixin.required_optional_params - {"latents"} - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - image_latens_params = IMAGE_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union( + {"pag_scale", "pag_adaptive_scale"} + ) - {"height", "width"} + batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS - - test_xformers_attention = False + # Img2img derives the latents from the input image, so `__call__` takes no `latents`. + optional_input_params = frozenset( + ["num_inference_steps", "num_images_per_prompt", "generator", "output_type", "return_dict"] + ) + # The output resolution follows the 32x32 input image. + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -128,93 +128,49 @@ def get_dummy_components(self): "vae": vae, } - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) image = image / 2 + 0.5 - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "image": image, - "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`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "pag_scale": 0.7, } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - # base pipeline (expect same output when pag is disabled) - pipe_sd = StableDiffusion3Img2ImgPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - components = self.get_dummy_components() +class TestStableDiffusion3PAGImg2ImgPipeline(StableDiffusion3PAGImg2ImgPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = StableDiffusion3Img2ImgPipeline + # SD3's denoiser is a transformer, so PAG resolves per transformer block rather than the UNet's mid/up/down. + pag_inference_applied_layers = ["blocks.0"] + # Only the "PAG off reproduces the base pipeline" leg was asserted before the migration. + check_pag_changes_output = False + # fmt: off + expected_pag_slice = torch.tensor([0.741577, 0.5491905, 0.59911674, 0.7702221, 0.65531653, 0.60989463, 0.491042, 0.5380116, 0.5592475]) + # fmt: on - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - - def test_pag_inference(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["blocks.0"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == ( - 1, - 32, - 32, - 3, - ), f"the shape of the output image should be (1, 32, 32, 3) but got {image.shape}" - - expected_slice = np.array( - [0.741577, 0.5491905, 0.59911674, 0.7702221, 0.65531653, 0.60989463, 0.491042, 0.5380116, 0.5592475] - ) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) +class TestStableDiffusion3PAGImg2ImgPipelineMemory(StableDiffusion3PAGImg2ImgPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 PAG img2img pipeline.""" @slow @require_torch_accelerator -class StableDiffusion3PAGImg2ImgPipelineIntegrationTests(unittest.TestCase): +class TestStableDiffusion3PAGImg2ImgPipelineIntegration: pipeline_class = StableDiffusion3PAGImg2ImgPipeline repo_id = "stabilityai/stable-diffusion-3-medium-diffusers" - 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/pag/test_pag_sd_img2img.py b/tests/pipelines/pag/test_pag_sd_img2img.py index 1e9b3c24c9ac..4906255017fa 100644 --- a/tests/pipelines/pag/test_pag_sd_img2img.py +++ b/tests/pipelines/pag/test_pag_sd_img2img.py @@ -14,17 +14,15 @@ # limitations under the License. import gc -import inspect import random -import unittest import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer from diffusers import ( AutoencoderKL, - AutoencoderTiny, AutoPipelineForImage2Image, EulerDiscreteScheduler, StableDiffusionImg2ImgPipeline, @@ -42,36 +40,34 @@ torch_device, ) from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, IPAdapterTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, + MemoryTesterMixin, ) +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionPAGImg2ImgPipelineFastTests( - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineTesterMixin, - unittest.TestCase, -): +class StableDiffusionPAGImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionPAGImg2ImgPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - {"height", "width"} - required_optional_params = PipelineTesterMixin.required_optional_params - {"latents"} - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union( + {"pag_scale", "pag_adaptive_scale"} + ) - {"height", "width"} + batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS + # Img2img derives the latents from the input image, so `__call__` takes no `latents`. + optional_input_params = frozenset( + ["num_inference_steps", "num_images_per_prompt", "generator", "output_type", "return_dict"] + ) + # The output resolution follows the 32x32 input image. + output_shape = (3, 32, 32) def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) @@ -117,7 +113,7 @@ def get_dummy_components(self, time_cond_proj_dim=None): text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "scheduler": scheduler, "vae": vae, @@ -127,111 +123,58 @@ def get_dummy_components(self, time_cond_proj_dim=None): "feature_extractor": None, "image_encoder": None, } - return components - def get_dummy_tiny_autoencoder(self): - return AutoencoderTiny(in_channels=3, out_channels=3, latent_channels=4) - - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) image = image / 2 + 0.5 - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "image": image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "pag_scale": 0.9, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe_sd = StableDiffusionImg2ImgPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - # pag enabled - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 - - def test_pag_inference(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == ( - 1, - 32, - 32, - 3, - ), f"the shape of the output image should be (1, 32, 32, 3) but got {image.shape}" - - expected_slice = np.array( - [0.4508819, 0.49191576, 0.42664427, 0.66448534, 0.5606137, 0.43760118, 0.58251137, 0.5944448, 0.51642907] - ) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) +class TestStableDiffusionPAGImg2ImgPipeline(StableDiffusionPAGImg2ImgPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = StableDiffusionImg2ImgPipeline + # fmt: off + expected_pag_slice = torch.tensor([0.4508819, 0.49191576, 0.42664427, 0.66448534, 0.5606137, 0.43760118, 0.58251137, 0.5944448, 0.51642907]) + # fmt: on def test_encode_prompt_works_in_isolation(self): extra_required_param_value_dict = { "device": torch.device(torch_device).type, - "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) +class TestStableDiffusionPAGImg2ImgPipelineMemory(StableDiffusionPAGImg2ImgPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD PAG img2img pipeline.""" + + +class TestStableDiffusionPAGImg2ImgPipelineIPAdapter( + StableDiffusionPAGImg2ImgPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SD PAG img2img pipeline.""" + + @slow @require_torch_accelerator -class StableDiffusionPAGImg2ImgPipelineIntegrationTests(unittest.TestCase): +class TestStableDiffusionPAGImg2ImgPipelineIntegration: pipeline_class = StableDiffusionPAGImg2ImgPipeline repo_id = "Jiali/stable-diffusion-1.5" - 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/pag/test_pag_sd_inpaint.py b/tests/pipelines/pag/test_pag_sd_inpaint.py index 5a78ac6ade12..c13f91053a48 100644 --- a/tests/pipelines/pag/test_pag_sd_inpaint.py +++ b/tests/pipelines/pag/test_pag_sd_inpaint.py @@ -15,9 +15,9 @@ import gc import random -import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -26,6 +26,7 @@ AutoencoderKL, AutoPipelineForInpainting, PNDMScheduler, + StableDiffusionInpaintPipeline, StableDiffusionPAGInpaintPipeline, UNet2DConditionModel, ) @@ -44,32 +45,29 @@ TEXT_GUIDED_IMAGE_INPAINTING_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, + FromPipeTesterMixin, IPAdapterTesterMixin, - PipelineFromPipeTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, + MemoryTesterMixin, ) +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionPAGInpaintPipelineFastTests( - PipelineTesterMixin, - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineFromPipeTesterMixin, - unittest.TestCase, -): +class StableDiffusionPAGInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionPAGInpaintPipeline - params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - image_params = frozenset([]) - image_latents_params = frozenset([]) + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS.union( + {"pag_scale", "pag_adaptive_scale"} + ) + batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union( {"add_text_embeds", "add_time_ids", "mask", "masked_image_latents"} ) + # The output resolution follows the 64x64 input image. + output_shape = (3, 64, 64) def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) @@ -109,7 +107,7 @@ def get_dummy_components(self, time_cond_proj_dim=None): text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "scheduler": scheduler, "vae": vae, @@ -119,42 +117,39 @@ def get_dummy_components(self, time_cond_proj_dim=None): "feature_extractor": None, "image_encoder": None, } - return components - def get_dummy_inputs(self, device, seed=0): + def get_dummy_inputs(self): # TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) image = image.cpu().permute(0, 2, 3, 1)[0] init_image = Image.fromarray(np.uint8(image)).convert("RGB").resize((64, 64)) # create mask image[8:, 8:, :] = 255 mask_image = Image.fromarray(np.uint8(image)).convert("L").resize((64, 64)) - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "image": init_image, "mask_image": mask_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "strength": 1.0, "pag_scale": 0.9, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - def test_pag_applied_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - # base pipeline - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) +class TestStableDiffusionPAGInpaintPipeline(StableDiffusionPAGInpaintPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = StableDiffusionInpaintPipeline + # fmt: off + expected_pag_slice = torch.tensor([0.7173, 0.5821, 0.6031, 0.5765, 0.6412, 0.6558, 0.5803, 0.5675, 0.5246]) + # fmt: on + + def test_pag_applied_layers(self): + pipe = self.get_pipeline() # pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers all_self_attn_layers = [k for k in pipe.unet.attn_processors.keys() if "attn1" in k] @@ -192,7 +187,7 @@ def test_pag_applied_layers(self): # pag_applied_layers = ["mid.block_0.attentions_1"] does not exist in the model pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["mid_block.attentions.1"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) # pag_applied_layers = "down" should apply to all self-attention layers in down_blocks @@ -207,7 +202,7 @@ def test_pag_applied_layers(self): pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["down_blocks.0"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) pipe.unet.set_attn_processor(original_attn_procs.copy()) @@ -220,50 +215,41 @@ def test_pag_applied_layers(self): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert len(pipe.pag_attn_processors) == 1 - def test_pag_inference(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - - expected_slice = np.array([0.7173, 0.5821, 0.6031, 0.5765, 0.6412, 0.6558, 0.5803, 0.5675, 0.5246]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" - def test_encode_prompt_works_in_isolation(self): extra_required_param_value_dict = { "device": torch.device(torch_device).type, - "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict, atol=1e-3, rtol=1e-3) +class TestStableDiffusionPAGInpaintPipelineMemory(StableDiffusionPAGInpaintPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD PAG inpaint pipeline.""" + + +class TestStableDiffusionPAGInpaintPipelineIPAdapter( + StableDiffusionPAGInpaintPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SD PAG inpaint pipeline.""" + + +class TestStableDiffusionPAGInpaintPipelineFromPipe( + StableDiffusionPAGInpaintPipelineTesterConfig, FromPipeTesterMixin +): + """`from_pipe` round-trip tests against `StableDiffusionPipeline`.""" + + @slow @require_torch_accelerator -class StableDiffusionPAGPipelineIntegrationTests(unittest.TestCase): +class TestStableDiffusionPAGInpaintPipelineIntegration: pipeline_class = StableDiffusionPAGInpaintPipeline repo_id = "stable-diffusion-v1-5/stable-diffusion-v1-5" - 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/pag/test_pag_sdxl.py b/tests/pipelines/pag/test_pag_sdxl.py index 99700b926802..24fbc7c20c93 100644 --- a/tests/pipelines/pag/test_pag_sdxl.py +++ b/tests/pipelines/pag/test_pag_sdxl.py @@ -14,10 +14,9 @@ # limitations under the License. import gc -import inspect -import unittest import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer @@ -41,33 +40,26 @@ from ..pipeline_params import ( TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, + FromPipeTesterMixin, IPAdapterTesterMixin, - PipelineFromPipeTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, + MemoryTesterMixin, ) +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionXLPAGPipelineFastTests( - PipelineTesterMixin, - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineFromPipeTesterMixin, - unittest.TestCase, -): +class StableDiffusionXLPAGPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLPAGPipeline - params = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"add_text_embeds", "add_time_ids"}) + output_shape = (3, 64, 64) def get_dummy_components(self, time_cond_proj_dim=None): # Copied from tests.pipelines.stable_diffusion_xl.test_stable_diffusion_xl.StableDiffusionXLPipelineTesterConfig.get_dummy_components @@ -129,7 +121,7 @@ def get_dummy_components(self, time_cond_proj_dim=None): text_encoder_2 = CLIPTextModelWithProjection(text_encoder_config) tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "scheduler": scheduler, "vae": vae, @@ -140,67 +132,28 @@ def get_dummy_components(self, time_cond_proj_dim=None): "image_encoder": 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, "pag_scale": 0.9, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline (expect same output when pag is disabled) - pipe_sd = StableDiffusionXLPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - # pag enabled - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 +class TestStableDiffusionXLPAGPipeline(StableDiffusionXLPAGPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = StableDiffusionXLPipeline + # fmt: off + expected_pag_slice = torch.tensor([0.5565, 0.5305, 0.4652, 0.4330, 0.4823, 0.4640, 0.5191, 0.4983, 0.4684]) + # fmt: on def test_pag_applied_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # base pipeline - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline() # pag_applied_layers = ["mid","up","down"] should apply to all self-attention layers all_self_attn_layers = [k for k in pipe.unet.attn_processors.keys() if "attn1" in k] @@ -234,7 +187,7 @@ def test_pag_applied_layers(self): # pag_applied_layers = ["mid.block_0.attentions_1"] does not exist in the model pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["mid_block.attentions.1"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) # pag_applied_layers = "down" should apply to all self-attention layers in down_blocks @@ -249,7 +202,7 @@ def test_pag_applied_layers(self): pipe.unet.set_attn_processor(original_attn_procs.copy()) pag_layers = ["down_blocks.0"] - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) pipe.unet.set_attn_processor(original_attn_procs.copy()) @@ -262,47 +215,34 @@ def test_pag_applied_layers(self): pipe._set_pag_attn_processor(pag_applied_layers=pag_layers, do_classifier_free_guidance=False) assert len(pipe.pag_attn_processors) == 2 - def test_pag_inference(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() + @pytest.mark.skip("We test this functionality elsewhere already.") + def test_save_load_optional_components(self): + pass - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] +class TestStableDiffusionXLPAGPipelineMemory(StableDiffusionXLPAGPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL PAG pipeline.""" - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.5565, 0.5305, 0.4652, 0.4330, 0.4823, 0.4640, 0.5191, 0.4983, 0.4684]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) +class TestStableDiffusionXLPAGPipelineIPAdapter(StableDiffusionXLPAGPipelineTesterConfig, IPAdapterTesterMixin): + """IP-Adapter tests for the SDXL PAG pipeline.""" - @unittest.skip("We test this functionality elsewhere already.") - def test_save_load_optional_components(self): - pass + +class TestStableDiffusionXLPAGPipelineFromPipe(StableDiffusionXLPAGPipelineTesterConfig, FromPipeTesterMixin): + """`from_pipe` round-trip tests against `StableDiffusionXLPipeline`.""" @slow @require_torch_accelerator -class StableDiffusionXLPAGPipelineIntegrationTests(unittest.TestCase): +class TestStableDiffusionXLPAGPipelineIntegration: pipeline_class = StableDiffusionXLPAGPipeline repo_id = "stabilityai/stable-diffusion-xl-base-1.0" - 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/pag/test_pag_sdxl_img2img.py b/tests/pipelines/pag/test_pag_sdxl_img2img.py index ca427ac8285c..1f8497ada573 100644 --- a/tests/pipelines/pag/test_pag_sdxl_img2img.py +++ b/tests/pipelines/pag/test_pag_sdxl_img2img.py @@ -14,11 +14,10 @@ # limitations under the License. import gc -import inspect import random -import unittest import numpy as np +import pytest import torch from transformers import ( CLIPImageProcessor, @@ -49,37 +48,33 @@ torch_device, ) from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, + FromPipeTesterMixin, IPAdapterTesterMixin, - PipelineFromPipeTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, + MemoryTesterMixin, ) +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionXLPAGImg2ImgPipelineFastTests( - PipelineTesterMixin, - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineFromPipeTesterMixin, - unittest.TestCase, -): +class StableDiffusionXLPAGImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLPAGImg2ImgPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - {"height", "width"} - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union( + {"pag_scale", "pag_adaptive_scale"} + ) - {"height", "width"} + batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union( {"add_text_embeds", "add_time_ids", "add_neg_time_ids"} ) + # The output resolution follows the 32x32 input image. + output_shape = (3, 32, 32) # based on tests.pipelines.stable_diffusion_xl.test_stable_diffusion_xl_img2img_pipeline.get_dummy_components def get_dummy_components( @@ -183,101 +178,64 @@ def get_dummy_components( # based on tests.pipelines.stable_diffusion_xl.test_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipelineTesterConfig # add `pag_scale` to the inputs - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) image = image / 2 + 0.5 - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "image": image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "pag_scale": 3.0, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "strength": 0.8, } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(requires_aesthetics_score=True) - - # base pipeline - pipe_sd = StableDiffusionXLImg2ImgPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] +class TestStableDiffusionXLPAGImg2ImgPipeline(StableDiffusionXLPAGImg2ImgPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = StableDiffusionXLImg2ImgPipeline + # The expected slice below was recorded against the aesthetics-score configuration. + pag_component_kwargs = {"requires_aesthetics_score": True} + # fmt: off + expected_pag_slice = torch.tensor([0.4566, 0.4907, 0.4374, 0.6633, 0.5626, 0.4494, 0.5771, 0.6011, 0.5245]) + # fmt: on - # pag enabled - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 + @pytest.mark.skip("We test this functionality elsewhere already.") + def test_save_load_optional_components(self): + pass - def test_pag_inference(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(requires_aesthetics_score=True) - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) +class TestStableDiffusionXLPAGImg2ImgPipelineMemory( + StableDiffusionXLPAGImg2ImgPipelineTesterConfig, MemoryTesterMixin +): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL PAG img2img pipeline.""" - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - assert image.shape == ( - 1, - 32, - 32, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.4566, 0.4907, 0.4374, 0.6633, 0.5626, 0.4494, 0.5771, 0.6011, 0.5245]) +class TestStableDiffusionXLPAGImg2ImgPipelineIPAdapter( + StableDiffusionXLPAGImg2ImgPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SDXL PAG img2img pipeline.""" - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" - @unittest.skip("We test this functionality elsewhere already.") - def test_save_load_optional_components(self): - pass +class TestStableDiffusionXLPAGImg2ImgPipelineFromPipe( + StableDiffusionXLPAGImg2ImgPipelineTesterConfig, FromPipeTesterMixin +): + """`from_pipe` round-trip tests against `StableDiffusionXLPipeline`.""" @slow @require_torch_accelerator -class StableDiffusionXLPAGImg2ImgPipelineIntegrationTests(unittest.TestCase): +class TestStableDiffusionXLPAGImg2ImgPipelineIntegration: repo_id = "stabilityai/stable-diffusion-xl-base-1.0" - 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/pag/test_pag_sdxl_inpaint.py b/tests/pipelines/pag/test_pag_sdxl_inpaint.py index cf22090b60f5..14e39716572b 100644 --- a/tests/pipelines/pag/test_pag_sdxl_inpaint.py +++ b/tests/pipelines/pag/test_pag_sdxl_inpaint.py @@ -14,11 +14,10 @@ # limitations under the License. import gc -import inspect import random -import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import ( @@ -54,32 +53,29 @@ TEXT_GUIDED_IMAGE_INPAINTING_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, + FromPipeTesterMixin, IPAdapterTesterMixin, - PipelineFromPipeTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, + MemoryTesterMixin, ) +from .testing_utils import PAGPipelineTesterMixin enable_full_determinism() -class StableDiffusionXLPAGInpaintPipelineFastTests( - PipelineTesterMixin, - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineFromPipeTesterMixin, - unittest.TestCase, -): +class StableDiffusionXLPAGInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLPAGInpaintPipeline - params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS.union({"pag_scale", "pag_adaptive_scale"}) - batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - image_params = frozenset([]) - image_latents_params = frozenset([]) + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS.union( + {"pag_scale", "pag_adaptive_scale"} + ) + batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union( {"add_text_embeds", "add_time_ids", "mask", "masked_image_latents"} ) + # The output resolution follows the 64x64 input image. + output_shape = (3, 64, 64) # based on tests.pipelines.stable_diffusion_xl.test_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipelineTesterConfig.get_dummy_components def get_dummy_components( @@ -181,108 +177,71 @@ def get_dummy_components( } return components - def get_dummy_inputs(self, device, seed=0): + def get_dummy_inputs(self): # TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) image = image.cpu().permute(0, 2, 3, 1)[0] init_image = Image.fromarray(np.uint8(image)).convert("RGB").resize((64, 64)) # create mask image[8:, 8:, :] = 255 mask_image = Image.fromarray(np.uint8(image)).convert("L").resize((64, 64)) - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "image": init_image, "mask_image": mask_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "strength": 1.0, "pag_scale": 0.9, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - - def test_pag_disable_enable(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(requires_aesthetics_score=True) - - # base pipeline - pipe_sd = StableDiffusionXLInpaintPipeline(**components) - pipe_sd = pipe_sd.to(device) - pipe_sd.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["pag_scale"] - assert "pag_scale" not in inspect.signature(pipe_sd.__call__).parameters, ( - f"`pag_scale` should not be a call parameter of the base pipeline {pipe_sd.__class__.__name__}." - ) - out = pipe_sd(**inputs).images[0, -3:, -3:, -1] - # pag disabled with pag_scale=0.0 - pipe_pag = self.pipeline_class(**components) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - inputs["pag_scale"] = 0.0 - out_pag_disabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] +class TestStableDiffusionXLPAGInpaintPipeline(StableDiffusionXLPAGInpaintPipelineTesterConfig, PAGPipelineTesterMixin): + base_pipeline_class = StableDiffusionXLInpaintPipeline + # The expected slice below was recorded against the aesthetics-score configuration. + pag_component_kwargs = {"requires_aesthetics_score": True} + # fmt: off + expected_pag_slice = torch.tensor([0.7893, 0.5446, 0.5826, 0.6441, 0.6660, 0.7566, 0.6605, 0.5838, 0.5160]) + # fmt: on - # pag enabled - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - out_pag_enabled = pipe_pag(**inputs).images[0, -3:, -3:, -1] - - assert np.abs(out.flatten() - out_pag_disabled.flatten()).max() < 1e-3 - assert np.abs(out.flatten() - out_pag_enabled.flatten()).max() > 1e-3 + @pytest.mark.skip("We test this functionality elsewhere already.") + def test_save_load_optional_components(self): + pass - def test_pag_inference(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(requires_aesthetics_score=True) - pipe_pag = self.pipeline_class(**components, pag_applied_layers=["mid", "up", "down"]) - pipe_pag = pipe_pag.to(device) - pipe_pag.set_progress_bar_config(disable=None) +class TestStableDiffusionXLPAGInpaintPipelineMemory( + StableDiffusionXLPAGInpaintPipelineTesterConfig, MemoryTesterMixin +): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL PAG inpaint pipeline.""" - inputs = self.get_dummy_inputs(device) - image = pipe_pag(**inputs).images - image_slice = image[0, -3:, -3:, -1] - assert image.shape == ( - 1, - 64, - 64, - 3, - ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.7893, 0.5446, 0.5826, 0.6441, 0.6660, 0.7566, 0.6605, 0.5838, 0.5160]) +class TestStableDiffusionXLPAGInpaintPipelineIPAdapter( + StableDiffusionXLPAGInpaintPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SDXL PAG inpaint pipeline.""" - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" - @unittest.skip("We test this functionality elsewhere already.") - def test_save_load_optional_components(self): - pass +class TestStableDiffusionXLPAGInpaintPipelineFromPipe( + StableDiffusionXLPAGInpaintPipelineTesterConfig, FromPipeTesterMixin +): + """`from_pipe` round-trip tests against `StableDiffusionXLPipeline`.""" @slow @require_torch_accelerator -class StableDiffusionXLPAGInpaintPipelineIntegrationTests(unittest.TestCase): +class TestStableDiffusionXLPAGInpaintPipelineIntegration: repo_id = "stabilityai/stable-diffusion-xl-base-1.0" - 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/pag/testing_utils.py b/tests/pipelines/pag/testing_utils.py new file mode 100644 index 000000000000..7e19d4cd73d8 --- /dev/null +++ b/tests/pipelines/pag/testing_utils.py @@ -0,0 +1,121 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect + +import pytest + +from ...testing_utils import assert_tensors_close, torch_device +from ..testing_utils import PipelineTesterMixin + + +class PAGPipelineTesterMixin(PipelineTesterMixin): + """`PipelineTesterMixin` plus the two tests every PAG pipeline shares. + + A PAG pipeline is an existing pipeline with perturbed-attention guidance layered on, so the tests all take the + same shape: run the pipeline it derives from, then check that PAG at `pag_scale=0.0` reproduces that output and + that enabling PAG moves it. Subclasses supply the knobs below; anything pipeline-specific (which layers PAG + resolves to, for instance) stays a method on the concrete test class. + """ + + # The non-PAG pipeline this one derives from. Required. + base_pipeline_class = None + + # `pag_applied_layers` for the "PAG enabled" leg of `test_pag_disable_enable`. `None` keeps the pipeline default. + pag_enabled_applied_layers = ["mid", "up", "down"] + + # `pag_scale` for the "PAG enabled" leg. `None` keeps the value from `get_dummy_inputs()`. + pag_enabled_scale = None + + # Some PAG pipelines only assert the disabled leg: their dummy denoiser is small enough that PAG's effect is not + # reliably above the tolerance. Set to `False` there. + check_pag_changes_output = True + + # `pag_applied_layers` the `test_pag_inference` pipeline is built with. `None` keeps the pipeline default. + pag_inference_applied_layers = ["mid", "up", "down"] + + # CPU-specific expected slice for `test_pag_inference`, laid out as the flattened `output[0, -1, -3:, -3:]` + # corner of the `"pt"` output. `None` skips the test. + expected_pag_slice = None + + # Extra kwargs for the `get_dummy_components()` call the two tests below build their pipelines from, when the + # PAG comparison needs a configuration other than the default one — the SDXL img2img and inpaint testers pin + # `requires_aesthetics_score=True`, which is what their expected slices were recorded against. + pag_component_kwargs = {} + + def get_pag_components(self): + return self.get_dummy_components(**self.pag_component_kwargs) + + def get_pag_pipeline(self, components=None, **pag_kwargs): + """Build the pipeline under test with explicit PAG constructor kwargs (`pag_applied_layers`, ...).""" + components = components if components is not None else self.get_dummy_components() + pipe = self.pipeline_class(**components, **pag_kwargs) + pipe.set_progress_bar_config(disable=None) + return pipe + + def test_pag_disable_enable(self): + # Run on CPU to keep the device-dependent `torch.Generator` deterministic. + components = self.get_pag_components() + + # base pipeline (expect same output when pag is disabled) + pipe_base = self.base_pipeline_class(**components) + pipe_base.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs() + del inputs["pag_scale"] + assert "pag_scale" not in inspect.signature(pipe_base.__call__).parameters, ( + f"`pag_scale` should not be a call parameter of the base pipeline {pipe_base.__class__.__name__}." + ) + out = pipe_base(**inputs)[0] + + # pag disabled with pag_scale=0.0 + pipe_pag = self.get_pipeline(**self.get_pag_components()) + out_pag_disabled = self.run_pipe(pipe_pag, pag_scale=0.0) + + assert_tensors_close(out_pag_disabled, out, atol=1e-3, msg="PAG at `pag_scale=0.0` changed the output.") + + if not self.check_pag_changes_output: + return + + # pag enabled + pag_kwargs = {} + if self.pag_enabled_applied_layers is not None: + pag_kwargs["pag_applied_layers"] = self.pag_enabled_applied_layers + pipe_pag = self.get_pag_pipeline(self.get_pag_components(), **pag_kwargs) + + extra_inputs = {} if self.pag_enabled_scale is None else {"pag_scale": self.pag_enabled_scale} + out_pag_enabled = self.run_pipe(pipe_pag, **extra_inputs) + + assert (out - out_pag_enabled).abs().max() > 1e-3, "Enabling PAG should change the output." + + def test_pag_inference(self): + if self.expected_pag_slice is None: + pytest.skip(f"No CPU expected slice pinned for {self.pipeline_class.__name__}.") + if torch_device != "cpu": + pytest.skip("The expected slice is CPU-specific.") + + pag_kwargs = ( + {} + if self.pag_inference_applied_layers is None + else {"pag_applied_layers": self.pag_inference_applied_layers} + ) + pipe_pag = self.get_pag_pipeline(self.get_pag_components(), **pag_kwargs) + + image = pipe_pag(**self.get_dummy_inputs())[0] + assert image.shape == (1, *self.output_shape), ( + f"the shape of the output image should be {(1, *self.output_shape)} but got {tuple(image.shape)}" + ) + + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), self.expected_pag_slice, atol=1e-3) diff --git a/tests/pipelines/pixart_alpha/test_pixart.py b/tests/pipelines/pixart_alpha/test_pixart.py index b68841159b4d..091cc3dfa07d 100644 --- a/tests/pipelines/pixart_alpha/test_pixart.py +++ b/tests/pipelines/pixart_alpha/test_pixart.py @@ -14,10 +14,9 @@ # limitations under the License. import gc -import tempfile -import unittest import numpy as np +import pytest import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -30,6 +29,7 @@ from diffusers.utils.import_utils import is_torch_neuronx_available from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, backend_synchronize, enable_full_determinism, @@ -39,23 +39,19 @@ 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 PipelineTesterMixin, to_np +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class PixArtAlphaPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class PixArtAlphaPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = PixArtAlphaPipeline - 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 - test_group_offloading = True + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + # `transformer.sample_size` (8) * `vae_scale_factor` (8) / 8 -> the dummy transformer generates 8x8 images. + output_shape = (3, 8, 8) def get_dummy_components(self): torch.manual_seed(0) @@ -87,98 +83,74 @@ def get_dummy_components(self): tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler, "text_encoder": text_encoder, "tokenizer": tokenizer, } - return components - - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + + def get_dummy_inputs(self): + return { "prompt": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "use_resolution_binning": False, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - @unittest.skip("Not supported.") - def test_sequential_cpu_offload_forward_pass(self): - # TODO(PVP, Sayak) need to fix later - return +class TestPixArtAlphaPipeline(PixArtAlphaPipelineTesterConfig, PipelineTesterMixin): 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) + # 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_slice = image[0, -3:, -3:, -1] + image = pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - self.assertEqual(image.shape, (1, 8, 8, 3)) - expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + # fmt: off + expected_slice = torch.tensor([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_inference_non_square_images(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() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs, height=32, width=48).images - image_slice = image[0, -3:, -3:, -1] - self.assertEqual(image.shape, (1, 32, 48, 3)) + image = pipe(**self.get_dummy_inputs(), height=32, width=48).images + assert image.shape == (1, 3, 32, 48) - expected_slice = np.array([0.6493, 0.537, 0.4081, 0.4762, 0.3695, 0.4711, 0.3026, 0.5218, 0.5263]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + # fmt: off + expected_slice = torch.tensor([0.6493, 0.537, 0.4081, 0.4762, 0.3695, 0.4711, 0.3026, 0.5218, 0.5263]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) - @unittest.skip("Test is already covered through encode_prompt isolation.") + @pytest.mark.skip("Test is already covered through encode_prompt isolation.") def test_save_load_optional_components(self): pass - def test_inference_with_embeddings_and_multiple_images(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) - - prompt = inputs["prompt"] - generator = inputs["generator"] - num_inference_steps = inputs["num_inference_steps"] - output_type = inputs["output_type"] + def test_inference_with_embeddings_and_multiple_images(self, tmp_path): + pipe = self.get_pipeline().to(torch_device) - prompt_embeds, prompt_attn_mask, negative_prompt_embeds, neg_prompt_attn_mask = pipe.encode_prompt(prompt) + inputs = self.get_dummy_inputs() + prompt_embeds, prompt_attn_mask, negative_prompt_embeds, neg_prompt_attn_mask = pipe.encode_prompt( + inputs["prompt"] + ) # inputs with prompt converted to embeddings - inputs = { + embedding_inputs = { "prompt_embeds": prompt_embeds, "prompt_attention_mask": prompt_attn_mask, "negative_prompt": None, "negative_prompt_embeds": negative_prompt_embeds, "negative_prompt_attention_mask": neg_prompt_attn_mask, - "generator": generator, - "num_inference_steps": num_inference_steps, - "output_type": output_type, + "generator": inputs["generator"], + "num_inference_steps": inputs["num_inference_steps"], + "output_type": inputs["output_type"], "num_images_per_prompt": 2, "use_resolution_binning": False, } @@ -187,97 +159,76 @@ def test_inference_with_embeddings_and_multiple_images(self): for optional_component in pipe._optional_components: setattr(pipe, optional_component, None) - output = pipe(**inputs)[0] + output = pipe(**embedding_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"] + embedding_inputs["generator"] = self.get_generator(0) + output_loaded = pipe_loaded(**embedding_inputs)[0] - # inputs with prompt converted to embeddings - inputs = { - "prompt_embeds": prompt_embeds, - "prompt_attention_mask": prompt_attn_mask, - "negative_prompt": None, - "negative_prompt_embeds": negative_prompt_embeds, - "negative_prompt_attention_mask": neg_prompt_attn_mask, - "generator": generator, - "num_inference_steps": num_inference_steps, - "output_type": output_type, - "num_images_per_prompt": 2, - "use_resolution_binning": False, - } - - output_loaded = pipe_loaded(**inputs)[0] - - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, 1e-4) + assert_tensors_close( + output_loaded, output, atol=1e-4, msg="Output changed after dropping optional components." + ) def test_inference_with_multiple_images_per_prompt(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() - inputs = self.get_dummy_inputs(device) - inputs["num_images_per_prompt"] = 2 - image = pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] + image = pipe(**self.get_dummy_inputs(), num_images_per_prompt=2).images + assert image.shape == (2, *self.output_shape) - self.assertEqual(image.shape, (2, 8, 8, 3)) - expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + # fmt: off + expected_slice = torch.tensor([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.483, 0.2583, 0.5331, 0.4852]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_raises_warning_for_mask_feature(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) + inputs = self.get_dummy_inputs() inputs.update({"mask_feature": True}) - with self.assertWarns(FutureWarning) as warning_ctx: + with pytest.warns(FutureWarning, match="mask_feature"): _ = pipe(**inputs).images - assert "mask_feature" in str(warning_ctx.warning) - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=1e-3) + super().test_inference_batch_single_identical(expected_max_diff=1e-3) + + +class TestPixArtAlphaPipelineMemory(PixArtAlphaPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the PixArt-alpha pipeline.""" + + @pytest.mark.skip("Not supported.") + def test_sequential_cpu_offload_forward_pass(self): + # TODO(PVP, Sayak) need to fix later + pass + + @pytest.mark.skip("Not supported.") + def test_sequential_offload_forward_pass_twice(self): + # TODO(PVP, Sayak) need to fix later + pass @slow @require_torch_accelerator -class PixArtAlphaPipelineIntegrationTests(unittest.TestCase): +class TestPixArtAlphaPipelineIntegration: ckpt_id_1024 = "PixArt-alpha/PixArt-XL-2-1024-MS" ckpt_id_512 = "PixArt-alpha/PixArt-XL-2-512x512" prompt = "A small cactus with a happy face in the Sahara desert." - 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) @@ -296,7 +247,7 @@ def test_pixart_1024(self): max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice) # Neuron uses bfloat16 internally which has lower precision than float16 on CUDA atol = 1e-2 if is_torch_neuronx_available() else 1e-4 - self.assertLessEqual(max_diff, atol) + assert max_diff <= atol def test_pixart_512(self): generator = torch.Generator("cpu").manual_seed(0) @@ -314,7 +265,7 @@ def test_pixart_512(self): max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice) # Neuron uses bfloat16 internally which has lower precision than float16 on CUDA atol = 1e-2 if is_torch_neuronx_available() else 1e-4 - self.assertLessEqual(max_diff, atol) + assert max_diff <= atol def test_pixart_1024_without_resolution_binning(self): generator = torch.manual_seed(0) @@ -419,9 +370,6 @@ def test_pixart_512_neuron_compile(self): output_type="np", ).images - self.assertEqual(image.shape, (1, 512, 512, 3)) - self.assertFalse(np.isnan(image).any(), "Output contains NaN values") - self.assertTrue( - (image >= 0.0).all() and (image <= 1.0).all(), - "Output pixel values outside [0, 1]", - ) + assert image.shape == (1, 512, 512, 3) + assert not np.isnan(image).any(), "Output contains NaN values" + assert (image >= 0.0).all() and (image <= 1.0).all(), "Output pixel values outside [0, 1]" diff --git a/tests/pipelines/pixart_sigma/test_pixart.py b/tests/pipelines/pixart_sigma/test_pixart.py index a6f73ff50f00..e65867dc11bd 100644 --- a/tests/pipelines/pixart_sigma/test_pixart.py +++ b/tests/pipelines/pixart_sigma/test_pixart.py @@ -14,10 +14,9 @@ # limitations under the License. import gc -import tempfile -import unittest import numpy as np +import pytest import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -30,6 +29,7 @@ from ...testing_utils import ( Expectations, + assert_tensors_close, backend_empty_cache, enable_full_determinism, numpy_cosine_similarity_distance, @@ -37,28 +37,25 @@ 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 PixArtSigmaPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class PixArtSigmaPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = PixArtSigmaPipeline - 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 - test_group_offloading = True + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + # `transformer.sample_size` (8) * `vae_scale_factor` (8) / 8 -> the dummy transformer generates 8x8 images. + output_shape = (3, 8, 8) def get_dummy_components(self): torch.manual_seed(0) @@ -90,94 +87,70 @@ def get_dummy_components(self): tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler, "text_encoder": text_encoder, "tokenizer": tokenizer, } - return components - - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) - inputs = { + + def get_dummy_inputs(self): + return { "prompt": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "use_resolution_binning": False, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } - return inputs - @unittest.skip("Not supported.") - def test_sequential_cpu_offload_forward_pass(self): - # TODO(PVP, Sayak) need to fix later - return +class TestPixArtSigmaPipeline(PixArtSigmaPipelineTesterConfig, 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 + assert image.shape == (1, *self.output_shape) - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - self.assertEqual(image.shape, (1, 8, 8, 3)) - expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.4830, 0.2583, 0.5331, 0.4852]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + # fmt: off + expected_slice = torch.tensor([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.4830, 0.2583, 0.5331, 0.4852]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_inference_non_square_images(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) - image = pipe(**inputs, height=32, width=48).images - image_slice = image[0, -3:, -3:, -1] - self.assertEqual(image.shape, (1, 32, 48, 3)) + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - expected_slice = np.array([0.6493, 0.5370, 0.4081, 0.4762, 0.3695, 0.4711, 0.3026, 0.5218, 0.5263]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + image = pipe(**self.get_dummy_inputs(), height=32, width=48).images + assert image.shape == (1, 3, 32, 48) - def test_inference_with_embeddings_and_multiple_images(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + # fmt: off + expected_slice = torch.tensor([0.6493, 0.5370, 0.4081, 0.4762, 0.3695, 0.4711, 0.3026, 0.5218, 0.5263]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) - inputs = self.get_dummy_inputs(torch_device) + def test_inference_with_embeddings_and_multiple_images(self, tmp_path): + 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"] - - prompt_embeds, prompt_attn_mask, negative_prompt_embeds, neg_prompt_attn_mask = pipe.encode_prompt(prompt) + inputs = self.get_dummy_inputs() + prompt_embeds, prompt_attn_mask, negative_prompt_embeds, neg_prompt_attn_mask = pipe.encode_prompt( + inputs["prompt"] + ) # inputs with prompt converted to embeddings - inputs = { + embedding_inputs = { "prompt_embeds": prompt_embeds, "prompt_attention_mask": prompt_attn_mask, "negative_prompt": None, "negative_prompt_embeds": negative_prompt_embeds, "negative_prompt_attention_mask": neg_prompt_attn_mask, - "generator": generator, - "num_inference_steps": num_inference_steps, - "output_type": output_type, + "generator": inputs["generator"], + "num_inference_steps": inputs["num_inference_steps"], + "output_type": inputs["output_type"], "num_images_per_prompt": 2, "use_resolution_binning": False, } @@ -186,80 +159,49 @@ def test_inference_with_embeddings_and_multiple_images(self): for optional_component in pipe._optional_components: setattr(pipe, optional_component, None) - output = pipe(**inputs)[0] + output = pipe(**embedding_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_attn_mask, - "negative_prompt": None, - "negative_prompt_embeds": negative_prompt_embeds, - "negative_prompt_attention_mask": neg_prompt_attn_mask, - "generator": generator, - "num_inference_steps": num_inference_steps, - "output_type": output_type, - "num_images_per_prompt": 2, - "use_resolution_binning": False, - } - - output_loaded = pipe_loaded(**inputs)[0] + embedding_inputs["generator"] = self.get_generator(0) + output_loaded = pipe_loaded(**embedding_inputs)[0] - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, 1e-4) + assert_tensors_close( + output_loaded, output, atol=1e-4, msg="Output changed after dropping optional components." + ) def test_inference_with_multiple_images_per_prompt(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() - inputs = self.get_dummy_inputs(device) - inputs["num_images_per_prompt"] = 2 - image = pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] + image = pipe(**self.get_dummy_inputs(), num_images_per_prompt=2).images + assert image.shape == (2, *self.output_shape) - self.assertEqual(image.shape, (2, 8, 8, 3)) - expected_slice = np.array([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.4830, 0.2583, 0.5331, 0.4852]) - max_diff = np.abs(image_slice.flatten() - expected_slice).max() - self.assertLessEqual(max_diff, 1e-3) + # fmt: off + expected_slice = torch.tensor([0.6319, 0.3526, 0.3806, 0.6327, 0.4639, 0.4830, 0.2583, 0.5331, 0.4852]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) - @unittest.skip("Test is already covered through encode_prompt isolation.") + @pytest.mark.skip("Test is already covered through encode_prompt isolation.") def test_save_load_optional_components(self): pass def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=1e-3) + super().test_inference_batch_single_identical(expected_max_diff=1e-3) 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 keep the device-dependent `torch.Generator` deterministic. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - original_image_slice = image[0, -3:, -3:, -1] + original_output = self.run_pipe(pipe) # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added # to the pipeline level. @@ -271,40 +213,56 @@ 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) - image = pipe(**inputs).images - image_slice_fused = image[0, -3:, -3:, -1] + output_fused = self.run_pipe(pipe) pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_disabled = image[0, -3:, -3:, -1] + output_disabled = self.run_pipe(pipe) - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( - "Fusion of QKV projections shouldn't affect the outputs." + assert_tensors_close( + output_fused, original_output, atol=1e-3, rtol=1e-3, msg="Fusion of QKV projections changed the outputs." ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert_tensors_close( + output_disabled, + output_fused, + atol=1e-3, + rtol=1e-3, + msg="Outputs changed after the fused QKV projections were 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( + output_disabled, + original_output, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) +class TestPixArtSigmaPipelineMemory(PixArtSigmaPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the PixArt-sigma pipeline.""" + + @pytest.mark.skip("Not supported.") + def test_sequential_cpu_offload_forward_pass(self): + # TODO(PVP, Sayak) need to fix later + pass + + @pytest.mark.skip("Not supported.") + def test_sequential_offload_forward_pass_twice(self): + # TODO(PVP, Sayak) need to fix later + pass + + @slow @require_torch_accelerator -class PixArtSigmaPipelineIntegrationTests(unittest.TestCase): +class TestPixArtSigmaPipelineIntegration: ckpt_id_1024 = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" ckpt_id_512 = "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" prompt = "A small cactus with a happy face in the Sahara desert." - 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) @@ -321,7 +279,7 @@ def test_pixart_1024(self): expected_slice = np.array([0.4517, 0.4446, 0.4375, 0.449, 0.4399, 0.4365, 0.4583, 0.4629, 0.4473]) max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice) - self.assertLessEqual(max_diff, 1e-4) + assert max_diff <= 1e-4 def test_pixart_512(self): generator = torch.Generator("cpu").manual_seed(0) @@ -349,7 +307,7 @@ def test_pixart_512(self): expected_slice = expected_slices.get_expectation() max_diff = numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice) - self.assertLessEqual(max_diff, 1e-4) + assert max_diff <= 1e-4 def test_pixart_1024_without_resolution_binning(self): generator = torch.manual_seed(0) diff --git a/tests/pipelines/pndm/test_pndm.py b/tests/pipelines/pndm/test_pndm.py index 5bca4bdede03..520a79ffb1e2 100644 --- a/tests/pipelines/pndm/test_pndm.py +++ b/tests/pipelines/pndm/test_pndm.py @@ -13,61 +13,159 @@ # 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 diffusers import PNDMPipeline, PNDMScheduler, UNet2DModel -from ...testing_utils import enable_full_determinism, nightly, require_torch, torch_device +from ...testing_utils import ( + enable_full_determinism, + nightly, + require_accelerator, + require_torch, + torch_device, +) +from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class PNDMPipelineFastTests(unittest.TestCase): - @property - def dummy_uncond_unet(self): +# `PNDMPipeline.__call__` postprocesses unconditionally — it always runs `image.cpu().permute(0, 2, 3, 1).numpy()` +# and only branches afterwards to decide whether to wrap the result in PIL. There is no `output_type="pt"` path, so +# `get_dummy_inputs` below has to ask for `"np"`, and every shared test that compares outputs with +# `assert_tensors_close` (torch-only) or calls `torch.isnan` on them fails on the numpy array it gets back. +# +# The sibling unconditional pipelines already have a `"pt"` path (`DDIMPipeline`, `DDPMPipeline`); adding one here +# is a `src/` change and out of scope for this test migration, so the affected tests are marked `xfail` rather than +# skipped: whoever adds the `"pt"` branch will see them XPASS and can drop these markers. +NO_PT_OUTPUT = pytest.mark.xfail( + reason="`PNDMPipeline` has no `output_type='pt'` path and always returns a numpy array.", + strict=True, +) + +# `PNDMPipeline` samples its initial noise with `randn_tensor(..., device=self.device)`: no `dtype=`, so the noise +# stays float32 and disagrees with a half-precision or layerwise-cast UNet. Same story as above — a `src/` gap, so +# the accelerator-only tests it breaks are marked `xfail`. +UNSUPPORTED_DTYPE = pytest.mark.xfail( + reason="`PNDMPipeline` samples its initial noise without `dtype=self.unet.dtype`, so it stays float32.", + strict=True, +) + +# The memory mixin trips over three separate `src/` gaps at once, so its tests carry one marker between them: +# - the numpy output above (`test_group_offloading_inference`, `test_pipeline_level_group_offloading_inference`, +# `test_pipeline_with_accelerator_device_map`), +# - the float32 noise above (`test_layerwise_casting_inference`), +# - `randn_tensor(..., device=self.device)` reading the *pipeline's* device rather than `self._execution_device`, +# which under sequential offload is `meta` (`test_sequential_cpu_offload_forward_pass`, +# `test_sequential_offload_forward_pass_twice`), +# - and no `model_cpu_offload_seq`, which `enable_model_cpu_offload` requires +# (`test_model_cpu_offload_forward_pass`, `test_cpu_offload_forward_pass_twice`). +# `strict=False` because `test_pipeline_level_group_offloading_sanity_checks` never runs the pipeline and so passes +# — it reports XPASS. The class-level marker is what keeps `MemoryTesterMixin`'s own `@is_memory` / +# `@require_accelerator` marks intact; overriding the eight failing tests individually would drop the +# `@require_accelerate_version_greater` gates they are declared with. +UNSUPPORTED_MEMORY_OPTIMIZATIONS = pytest.mark.xfail( + reason=( + "`PNDMPipeline` returns numpy, samples noise without `dtype=`/`_execution_device`, and declares no " + "`model_cpu_offload_seq`." + ), + strict=False, +) + + +class PNDMPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = PNDMPipeline + required_input_params_in_call_signature = UNCONDITIONAL_IMAGE_GENERATION_PARAMS + batch_input_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS + # PNDM is unconditional and samples its own noise: there is no prompt to repeat + # (`num_images_per_prompt`) and no user-suppliable `latents`. + optional_input_params = frozenset(["num_inference_steps", "generator", "output_type", "return_dict"]) + # `(height, width, channels)` — the numpy layout, not the `(channels, height, width)` the other configs get + # from `output_type="pt"` (see `NO_PT_OUTPUT` above). + output_shape = (8, 8, 3) + + def get_dummy_components(self): torch.manual_seed(0) - model = UNet2DModel( - block_out_channels=(32, 64), - layers_per_block=2, - sample_size=32, + unet = UNet2DModel( + block_out_channels=(4, 8), + layers_per_block=1, + norm_num_groups=4, + sample_size=8, in_channels=3, out_channels=3, down_block_types=("DownBlock2D", "AttnDownBlock2D"), up_block_types=("AttnUpBlock2D", "UpBlock2D"), ) - return model + scheduler = PNDMScheduler() + return {"unet": unet, "scheduler": scheduler} + + def get_dummy_inputs(self): + return { + "batch_size": 1, + "generator": self.get_generator(0), + # `PNDMScheduler` runs Runge-Kutta warm-up steps, so it needs at least 4 inference steps. + "num_inference_steps": 4, + # `"np"` rather than the usual `"pt"` — see `NO_PT_OUTPUT` above. + "output_type": "np", + } + +class TestPNDMPipeline(PNDMPipelineTesterConfig, PipelineTesterMixin): def test_inference(self): - unet = self.dummy_uncond_unet - scheduler = PNDMScheduler() + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - pndm = PNDMPipeline(unet=unet, scheduler=scheduler) - pndm.to(torch_device) - pndm.set_progress_bar_config(disable=None) + image = pipe(**self.get_dummy_inputs()).images + generated_image = image[0] + assert generated_image.shape == self.output_shape - generator = torch.manual_seed(0) - image = pndm(generator=generator, num_inference_steps=20, output_type="np").images + expected_slice = np.array([0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0]) + assert np.abs(generated_image[-3:, -3:, -1].flatten() - expected_slice).max() < 1e-2 - generator = torch.manual_seed(0) - image_from_tuple = pndm(generator=generator, num_inference_steps=20, output_type="np", return_dict=False)[0] + @NO_PT_OUTPUT + def test_save_load_local(self, tmp_path, base_pipe_output, expected_max_difference=5e-4): + super().test_save_load_local(tmp_path, base_pipe_output, expected_max_difference) - image_slice = image[0, -3:, -3:, -1] - image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1] + @NO_PT_OUTPUT + def test_inference_batch_single_identical(self): + super().test_inference_batch_single_identical() - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0]) + @NO_PT_OUTPUT + def test_dict_tuple_outputs_equivalent(self): + super().test_dict_tuple_outputs_equivalent() - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 + # The three overrides below re-declare the base methods' skip decorators: overriding a test drops the marks + # the base declared it with, and without them these would run (and xfail for the wrong reason) on CPU. + @NO_PT_OUTPUT + @require_accelerator + def test_to_device(self): + super().test_to_device() + + @UNSUPPORTED_DTYPE + @pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="half-precision inference requires CUDA or XPU") + @require_accelerator + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=str) + def test_half_precision_inference_no_nan(self, dtype): + super().test_half_precision_inference_no_nan(dtype) + + @UNSUPPORTED_DTYPE + @pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU") + @require_accelerator + def test_save_load_float16(self, tmp_path, expected_max_diff=1e-2): + super().test_save_load_float16(tmp_path, expected_max_diff) + + +@UNSUPPORTED_MEMORY_OPTIMIZATIONS +class TestPNDMPipelineMemory(PNDMPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the PNDM pipeline.""" @nightly @require_torch -class PNDMPipelineIntegrationTests(unittest.TestCase): +class TestPNDMPipelineIntegration: def test_inference_cifar10(self): model_id = "google/ddpm-cifar10-32" diff --git a/tests/pipelines/prx/test_pipeline_prx.py b/tests/pipelines/prx/test_pipeline_prx.py index 104f7b0b7553..f4e4c17621b0 100644 --- a/tests/pipelines/prx/test_pipeline_prx.py +++ b/tests/pipelines/prx/test_pipeline_prx.py @@ -1,6 +1,4 @@ -import unittest - -import numpy as np +import pytest import torch from transformers import AutoTokenizer from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig, T5GemmaModuleConfig @@ -11,26 +9,32 @@ from diffusers.pipelines.prx.pipeline_prx import PRXPipeline from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from ...testing_utils import assert_tensors_close from ..pipeline_params import TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +# `T5GemmaEncoder` is instantiated here from a hand-built config rather than loaded from a repo, and transformers v5 +# cannot round-trip that through `save_pretrained`/`from_pretrained`, so every test that reloads the pipeline from +# disk is skipped. +T5GEMMA_SERIALIZATION_SKIP_REASON = "Custom T5GemmaEncoder not compatible with transformers v5." +# Both PRX pipelines read `callback_on_step_end`'s inputs out of `locals()` but throw away what the callback +# returns, so a callback that rewrites `latents` (or `prompt_embeds`) has no effect on the denoising loop. Every +# other diffusers pipeline pops those keys back off the returned dict. Fixing it is a `src/` change and out of +# scope for this test migration, so the one shared test that exercises the write-back is marked `xfail`: whoever +# adds the pop-back will see it XPASS and can drop this marker. +CALLBACK_OUTPUTS_IGNORED = pytest.mark.xfail( + reason="`PRX` pipelines discard the dict `callback_on_step_end` returns, so callback edits to `latents` are lost.", + strict=True, +) -class PRXPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + +class PRXPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = PRXPipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"]) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True - - @classmethod - def setUpClass(cls): - # Ensure PRXPipeline has an _execution_device property expected by __call__ - if not isinstance(getattr(PRXPipeline, "_execution_device", None), property): - try: - setattr(PRXPipeline, "_execution_device", property(lambda self: torch.device("cpu"))) - except Exception: - pass + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_input_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -102,116 +106,54 @@ def get_dummy_components(self): "tokenizer": tokenizer, } - 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): return { "prompt": "", "negative_prompt": "", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 1.0, "height": 32, "width": 32, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). "output_type": "pt", "use_resolution_binning": False, } - def test_inference(self): - device = "cpu" - components = self.get_dummy_components() - pipe = PRXPipeline(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) - try: - pipe.register_to_config(_execution_device="cpu") - except Exception: - pass - - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs)[0] - generated_image = image[0] - - self.assertEqual(generated_image.shape, (3, 32, 32)) - expected_image = torch.zeros(3, 32, 32) - max_diff = np.abs(generated_image - expected_image).max() - self.assertLessEqual(max_diff, 1e10) +class TestPRXPipeline(PRXPipelineTesterConfig, PipelineTesterMixin): + @CALLBACK_OUTPUTS_IGNORED def test_callback_inputs(self): - components = self.get_dummy_components() - pipe = PRXPipeline(**components) - pipe = pipe.to("cpu") - pipe.set_progress_bar_config(disable=None) - try: - pipe.register_to_config(_execution_device="cpu") - except Exception: - pass - self.assertTrue( - hasattr(pipe, "_callback_tensor_inputs"), - f" {PRXPipeline} 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): - for tensor_name in callback_kwargs.keys(): - 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 - for tensor_name in callback_kwargs.keys(): - assert tensor_name in pipe._callback_tensor_inputs - return callback_kwargs - - inputs = self.get_dummy_inputs("cpu") - - inputs["callback_on_step_end"] = callback_inputs_subset - inputs["callback_on_step_end_tensor_inputs"] = ["latents"] - _ = pipe(**inputs)[0] - - inputs["callback_on_step_end"] = callback_inputs_all - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - _ = pipe(**inputs)[0] + super().test_callback_inputs() def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3): - if not self.test_attention_slicing: - return + # Run on CPU: sliced attention is compared against a full-attention run of the same pipeline. + pipe = self.get_pipeline() - 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("cpu") - pipe.set_progress_bar_config(disable=None) - - def to_np_local(tensor): - if isinstance(tensor, torch.Tensor): - return tensor.detach().cpu().numpy() - return tensor - - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - output_without_slicing = pipe(**inputs)[0] + output_without_slicing = self.run_pipe(pipe) pipe.enable_attention_slicing(slice_size=1) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing1 = pipe(**inputs)[0] + output_with_slicing_1 = self.run_pipe(pipe) pipe.enable_attention_slicing(slice_size=2) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing2 = pipe(**inputs)[0] + output_with_slicing_2 = self.run_pipe(pipe) - max_diff1 = np.abs(to_np_local(output_with_slicing1) - to_np_local(output_without_slicing)).max() - max_diff2 = np.abs(to_np_local(output_with_slicing2) - to_np_local(output_without_slicing)).max() - self.assertLess(max(max_diff1, max_diff2), expected_max_diff) + assert_tensors_close( + output_with_slicing_1, + output_without_slicing, + atol=expected_max_diff, + msg="Attention slicing (slice_size=1) changed the output.", + ) + assert_tensors_close( + output_with_slicing_2, + output_without_slicing, + atol=expected_max_diff, + msg="Attention slicing (slice_size=2) changed the output.", + ) def test_inference_with_autoencoder_dc(self): - """Test PRXPipeline with AutoencoderDC (DCAE) instead of AutoencoderKL.""" - device = "cpu" - + """PRXPipeline should also work with an `AutoencoderDC` (DCAE) in place of the `AutoencoderKL`.""" components = self.get_dummy_components() torch.manual_seed(0) @@ -240,43 +182,38 @@ def test_inference_with_autoencoder_dc(self): ).eval() components["vae"] = vae_dc + pipe = self.get_pipeline(**components) - pipe = PRXPipeline(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - expected_scale_factor = vae_dc.spatial_compression_ratio - self.assertEqual(pipe.vae_scale_factor, expected_scale_factor) - - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs)[0] - generated_image = image[0] + assert pipe.vae_scale_factor == vae_dc.spatial_compression_ratio - self.assertEqual(generated_image.shape, (3, 32, 32)) - expected_image = torch.zeros(3, 32, 32) - max_diff = np.abs(generated_image - expected_image).max() - self.assertLessEqual(max_diff, 1e10) + output = self.run_pipe(pipe) + assert output[0].shape == self.output_shape + assert torch.isfinite(output).all() - @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") + @pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON) def test_loading_with_variants(self): pass - @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") - def test_pipeline_with_accelerator_device_map(self): + @pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON) + def test_save_load_local(self): pass - @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") - def test_save_load_local(self): + @pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON) + def test_save_load_float16(self): pass - @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") + @pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON) def test_save_load_optional_components(self): pass - @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") - def test_dtype_dict(self): + @pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON) + def test_torch_dtype_dict(self): pass - @unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.") - def test_dtype_alias(self): + +class TestPRXPipelineMemory(PRXPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the PRX pipeline.""" + + @pytest.mark.skip(T5GEMMA_SERIALIZATION_SKIP_REASON) + def test_pipeline_with_accelerator_device_map(self): pass diff --git a/tests/pipelines/prx/test_pipeline_prx_pixel.py b/tests/pipelines/prx/test_pipeline_prx_pixel.py index 5ecb489a8143..187a3d8a01d7 100644 --- a/tests/pipelines/prx/test_pipeline_prx_pixel.py +++ b/tests/pipelines/prx/test_pipeline_prx_pixel.py @@ -1,6 +1,4 @@ -import unittest - -import numpy as np +import pytest import torch from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model @@ -8,19 +6,29 @@ from diffusers.pipelines.prx.pipeline_prx_pixel import PRXPixelPipeline from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from ...testing_utils import assert_tensors_close, torch_device from ..pipeline_params import TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +# Both PRX pipelines read `callback_on_step_end`'s inputs out of `locals()` but throw away what the callback +# returns, so a callback that rewrites `latents` (or `prompt_embeds`) has no effect on the denoising loop. Every +# other diffusers pipeline pops those keys back off the returned dict. Fixing it is a `src/` change and out of +# scope for this test migration, so the one shared test that exercises the write-back is marked `xfail`: whoever +# adds the pop-back will see it XPASS and can drop this marker. +CALLBACK_OUTPUTS_IGNORED = pytest.mark.xfail( + reason="`PRX` pipelines discard the dict `callback_on_step_end` returns, so callback edits to `latents` are lost.", + strict=True, +) -class PRXPixelPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class PRXPixelPipelineTesterConfig(BasePipelineTesterConfig): # PRXPixelPipeline is standalone: it inherits from DiffusionPipeline (not PRXPipeline) and always has its own # image_processor, so it denoises raw RGB in pixel space and supports output_type="pil"/"np" without a VAE. pipeline_class = PRXPixelPipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"]) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_input_params = frozenset(["prompt", "negative_prompt", "num_images_per_prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -64,108 +72,75 @@ def get_dummy_components(self): "prompt_max_tokens": 16, } - 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): return { "prompt": "", "negative_prompt": "", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 1.0, "height": 32, "width": 32, # Pixel-space PRX has no VAE and returns raw (C, H, W) tensors for output_type="pt". The generic # PipelineTesterMixin tests compare these tensors directly, so default to "pt" here; the PIL/np default - # path is exercised explicitly in test_inference and test_inference_pil_and_np_output. + # path is exercised explicitly in test_inference_pil_and_np_output. "output_type": "pt", # 32px is not in the 1024 aspect-ratio bins, so binning must be disabled for these tiny fast tests. "use_resolution_binning": False, } - def _build_pipe(self, device="cpu"): - components = self.get_dummy_components() - pipe = PRXPixelPipeline(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) - return pipe - def test_inference(self): - device = "cpu" - pipe = self._build_pipe(device) +class TestPRXPixelPipeline(PRXPixelPipelineTesterConfig, PipelineTesterMixin): + @CALLBACK_OUTPUTS_IGNORED + def test_callback_inputs(self): + super().test_callback_inputs() + def test_pixel_space_has_no_vae(self): # Pixel space: vae_scale_factor is always 1, and the pipeline always carries an image processor # so postprocessing (and the default output_type="pil") works without any VAE. - self.assertEqual(pipe.vae_scale_factor, 1) - self.assertIsNotNone(pipe.image_processor) - - # Default output is PIL (no VAE needed: the image processor denormalizes the denoised pixels directly). - inputs = self.get_dummy_inputs(device) - inputs.pop("output_type") # default is "pil" - images = pipe(**inputs).images - self.assertEqual(len(images), 1) - self.assertEqual(images[0].size, (32, 32)) + pipe = self.get_pipeline() - # Raw "pt" output is the denoised RGB tensor at the requested resolution. - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs)[0] - generated_image = image[0] - self.assertEqual(generated_image.shape, (3, 32, 32)) - expected_image = torch.zeros(3, 32, 32) - max_diff = np.abs(generated_image.cpu().numpy() - expected_image.numpy()).max() - self.assertLessEqual(max_diff, 1e10) + assert pipe.vae_scale_factor == 1 + assert pipe.image_processor is not None def test_inference_batch(self): - device = "cpu" - pipe = self._build_pipe(device) + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = ["", ""] - inputs["negative_prompt"] = ["", ""] - image = pipe(**inputs)[0] + image = self.run_pipe(pipe, prompt=["", ""], negative_prompt=["", ""]) - self.assertEqual(image.shape[0], 2) - self.assertEqual(tuple(image.shape[1:]), (3, 32, 32)) + assert image.shape[0] == 2 + assert tuple(image.shape[1:]) == self.output_shape def test_inference_with_cfg(self): - device = "cpu" - pipe = self._build_pipe(device) + pipe = self.get_pipeline() # CFG off. - inputs = self.get_dummy_inputs(device) - inputs["guidance_scale"] = 1.0 - out_no_cfg = pipe(**inputs)[0] - self.assertFalse(pipe.do_classifier_free_guidance) - self.assertEqual(out_no_cfg[0].shape, (3, 32, 32)) + out_no_cfg = self.run_pipe(pipe, guidance_scale=1.0) + assert not pipe.do_classifier_free_guidance + assert out_no_cfg[0].shape == self.output_shape # CFG on. - inputs = self.get_dummy_inputs(device) - inputs["guidance_scale"] = 5.0 - out_cfg = pipe(**inputs)[0] - self.assertTrue(pipe.do_classifier_free_guidance) - self.assertEqual(out_cfg[0].shape, (3, 32, 32)) + out_cfg = self.run_pipe(pipe, guidance_scale=5.0) + assert pipe.do_classifier_free_guidance + assert out_cfg[0].shape == self.output_shape # Guidance should actually change the output. - max_diff = np.abs(out_no_cfg.cpu().numpy() - out_cfg.cpu().numpy()).max() - self.assertGreater(max_diff, 0.0) + assert not torch.allclose(out_no_cfg, out_cfg) def test_inference_with_prompt_embeds(self): - device = "cpu" - pipe = self._build_pipe(device) + pipe = self.get_pipeline() # Precompute embeddings via the public encode_prompt API (CFG on so we get negatives too). prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = ( pipe.encode_prompt( prompt="a prompt", - device=device, + device=torch.device("cpu"), do_classifier_free_guidance=True, negative_prompt="", ) ) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs.pop("prompt") inputs.pop("negative_prompt") inputs["guidance_scale"] = 5.0 @@ -175,31 +150,29 @@ def test_inference_with_prompt_embeds(self): inputs["negative_prompt_attention_mask"] = negative_prompt_attention_mask image = pipe(**inputs)[0] - self.assertEqual(image[0].shape, (3, 32, 32)) + assert image[0].shape == self.output_shape def test_inference_pil_and_np_output(self): # The default output_type="pil" must work without a VAE: the denoised pixels are denormalized # directly by the image processor instead of being decoded. - device = "cpu" - pipe = self._build_pipe(device) + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs.pop("output_type") # default is "pil" images = pipe(**inputs).images - self.assertEqual(len(images), 1) - self.assertEqual(images[0].size, (32, 32)) + assert len(images) == 1 + assert images[0].size == (32, 32) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["output_type"] = "np" images = pipe(**inputs).images - self.assertEqual(images.shape, (1, 32, 32, 3)) - self.assertGreaterEqual(images.min(), 0.0) - self.assertLessEqual(images.max(), 1.0) + assert images.shape == (1, 32, 32, 3) + assert images.min() >= 0.0 + assert images.max() <= 1.0 def test_non_multiple_size_raises(self): # height/width must be divisible by vae_scale_factor * transformer patch_size; check_inputs must raise # a clear ValueError instead of letting the transformer fail on an invalid reshape mid-denoising. - device = "cpu" components = self.get_dummy_components() torch.manual_seed(0) components["transformer"] = PRXTransformer2DModel( @@ -214,82 +187,45 @@ def test_non_multiple_size_raises(self): bottleneck_size=8, resolution_embeds=True, ) - pipe = PRXPixelPipeline(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline(**components) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["height"] = 31 # vae_scale_factor (1) * patch_size (2) = 2; 31 is not a multiple - with self.assertRaisesRegex(ValueError, "divisible"): + with pytest.raises(ValueError, match="divisible"): pipe(**inputs) - def test_callback_inputs(self): - device = "cpu" - pipe = self._build_pipe(device) - self.assertTrue( - hasattr(pipe, "_callback_tensor_inputs"), - f" {PRXPixelPipeline} 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): - for tensor_name in callback_kwargs.keys(): - 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 - for tensor_name in callback_kwargs.keys(): - assert tensor_name in pipe._callback_tensor_inputs - return callback_kwargs - - inputs = self.get_dummy_inputs(device) - inputs["callback_on_step_end"] = callback_inputs_subset - inputs["callback_on_step_end_tensor_inputs"] = ["latents"] - _ = pipe(**inputs)[0] - - inputs = self.get_dummy_inputs(device) - inputs["callback_on_step_end"] = callback_inputs_all - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - _ = pipe(**inputs)[0] - def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3): - # Overridden: the mixin version calls assert_mean_pixel_difference, which assumes HWC image - # arrays. Pixel-space PRX has no VAE; compare raw (C, H, W) tensors directly ("pt") instead of - # going through PIL. - if not self.test_attention_slicing: - return + # Run on CPU: sliced attention is compared against a full-attention run of the same pipeline. + pipe = self.get_pipeline() - 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("cpu") - pipe.set_progress_bar_config(disable=None) - - def to_np_local(tensor): - if isinstance(tensor, torch.Tensor): - return tensor.detach().cpu().numpy() - return tensor - - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - output_without_slicing = pipe(**inputs)[0] + output_without_slicing = self.run_pipe(pipe) pipe.enable_attention_slicing(slice_size=1) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing1 = pipe(**inputs)[0] + output_with_slicing_1 = self.run_pipe(pipe) pipe.enable_attention_slicing(slice_size=2) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing2 = pipe(**inputs)[0] + output_with_slicing_2 = self.run_pipe(pipe) + + assert_tensors_close( + output_with_slicing_1, + output_without_slicing, + atol=expected_max_diff, + msg="Attention slicing (slice_size=1) changed the output.", + ) + assert_tensors_close( + output_with_slicing_2, + output_without_slicing, + atol=expected_max_diff, + msg="Attention slicing (slice_size=2) changed the output.", + ) + + def test_encode_prompt_works_in_isolation(self): + extra_required_param_value_dict = { + "device": torch.device(torch_device).type, + "do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0, + } + return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) - max_diff1 = np.abs(to_np_local(output_with_slicing1) - to_np_local(output_without_slicing)).max() - max_diff2 = np.abs(to_np_local(output_with_slicing2) - to_np_local(output_without_slicing)).max() - self.assertLess(max(max_diff1, max_diff2), expected_max_diff) - @unittest.skip("Slow original-vs-diffusers parity test is optional and intentionally skipped for fast CI.") - def test_prx_pixel_original_parity(self): - pass +class TestPRXPixelPipelineMemory(PRXPixelPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the PRXPixel pipeline.""" diff --git a/tests/pipelines/testing_utils/__init__.py b/tests/pipelines/testing_utils/__init__.py index 94204f246e08..9b756ec64693 100644 --- a/tests/pipelines/testing_utils/__init__.py +++ b/tests/pipelines/testing_utils/__init__.py @@ -7,6 +7,7 @@ TaylorSeerCacheTesterMixin, ) from .common import BasePipelineTesterConfig, PipelineTesterMixin +from .from_pipe import FromPipeTesterMixin from .ip_adapter import IPAdapterTesterMixin from .lora import LoraMemoryTesterMixin, LoraTesterMixin, UNetLoraTesterMixin from .memory import ( @@ -26,6 +27,7 @@ __all__ = [ "BasePipelineTesterConfig", "PipelineTesterMixin", + "FromPipeTesterMixin", "IPAdapterTesterMixin", "LoraTesterMixin", "LoraMemoryTesterMixin", diff --git a/tests/pipelines/testing_utils/from_pipe.py b/tests/pipelines/testing_utils/from_pipe.py new file mode 100644 index 000000000000..d76027d1cff9 --- /dev/null +++ b/tests/pipelines/testing_utils/from_pipe.py @@ -0,0 +1,208 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect + +from diffusers import KolorsPipeline, StableDiffusionPipeline, StableDiffusionXLPipeline + +from ...testing_utils import ( + assert_tensors_close, + require_accelerate_version_greater, + require_accelerator, + torch_device, +) +from .common import BasePipelineOutputMixin + + +# The repo and constructor kwargs used to build the "original" pipeline a variant is derived from, keyed by that +# original pipeline class. +ORIGINAL_PIPELINE_REPOS = { + StableDiffusionPipeline: ("hf-internal-testing/tiny-stable-diffusion-torch", {"requires_safety_checker": False}), + StableDiffusionXLPipeline: ( + "hf-internal-testing/tiny-stable-diffusion-xl-pipe", + {"requires_aesthetics_score": True, "force_zeros_for_empty_prompt": False}, + ), + KolorsPipeline: ("hf-internal-testing/tiny-kolors-pipe", {"force_zeros_for_empty_prompt": False}), +} + + +class FromPipeTesterMixin(BasePipelineOutputMixin): + """`DiffusionPipeline.from_pipe` tests for pipelines that are variants of an existing one. + + Composed with `BasePipelineTesterConfig`, which supplies `pipeline_class`, `get_dummy_components()` and + `get_dummy_inputs()`. The pytest-style successor of `PipelineFromPipeTesterMixin` in + `tests/pipelines/test_pipelines_common.py`. + """ + + # Set on the test class to pull the original pipeline from a repo other than the default for its class. + original_pipeline_repo = None + + @property + def original_pipeline_class(self): + """The pipeline this one is a variant of — the source `from_pipe` is expected to round-trip through.""" + name = self.pipeline_class.__name__.lower() + if "xl" in name: + return StableDiffusionXLPipeline + elif "kolors" in name: + return KolorsPipeline + return StableDiffusionPipeline + + def get_dummy_inputs_pipe(self): + inputs = self.get_dummy_inputs() + inputs["return_dict"] = False + return inputs + + def get_dummy_inputs_for_pipe_original(self): + """The dummy inputs, restricted to the parameters the original pipeline's `__call__` accepts.""" + original_call_params = set(inspect.signature(self.original_pipeline_class.__call__).parameters.keys()) + return {k: v for k, v in self.get_dummy_inputs_pipe().items() if k in original_call_params} + + def _split_components(self, components): + """Split the given components into what the original pipeline expects and what only this one does.""" + original_expected_modules, _ = self.original_pipeline_class._get_signature_keys(self.original_pipeline_class) + + # components of this pipeline that the original one also expects + original_pipe_components = {} + # components this pipeline doesn't have, but the original one expects + original_pipe_additional_components = {} + # components this pipeline has, but the original one doesn't expect + current_pipe_additional_components = {} + + for name, component in components.items(): + if name in original_expected_modules: + original_pipe_components[name] = component + else: + current_pipe_additional_components[name] = component + + for name in original_expected_modules: + if name not in original_pipe_components: + if name in self.original_pipeline_class._optional_components: + original_pipe_additional_components[name] = None + else: + raise ValueError(f"missing required module for {self.original_pipeline_class.__name__}: {name}") + + return ( + {**original_pipe_components, **original_pipe_additional_components}, + current_pipe_additional_components, + ) + + def _build_original_pipeline(self, components): + """Build the original pipeline out of `components`, plus the components only this pipeline has. + + Both pipelines are built from the *same* component instances — rebuilding them per pipeline would make the + comparison depend on `get_dummy_components()` being bit-identical across calls. + """ + original_components, current_pipe_additional_components = self._split_components(components) + pipe_original = self.original_pipeline_class(**original_components) + pipe_original.set_progress_bar_config(disable=None) + return pipe_original, current_pipe_additional_components + + def test_from_pipe_consistent_config(self): + original_repo, original_kwargs = ORIGINAL_PIPELINE_REPOS[self.original_pipeline_class] + original_repo = self.original_pipeline_repo or original_repo + + # create original_pipeline_class(sd/sdxl/kolors) + pipe_original = self.original_pipeline_class.from_pretrained(original_repo, **original_kwargs) + + # original_pipeline_class -> pipeline_class + pipe_additional_components = { + name: component + for name, component in self.get_dummy_components().items() + if name not in pipe_original.components + } + pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components) + + # pipeline_class -> original_pipeline_class + original_pipe_additional_components = {} + for name, component in pipe_original.components.items(): + if name not in pipe.components or not isinstance(component, pipe.components[name].__class__): + original_pipe_additional_components[name] = component + + pipe_original_2 = self.original_pipeline_class.from_pipe(pipe, **original_pipe_additional_components) + + # compare the config + original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")} + original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")} + assert original_config_2 == original_config + + def test_from_pipe_consistent_forward_pass(self, expected_max_diff=1e-3): + components = self.get_dummy_components() + pipe_original, current_pipe_additional_components = self._build_original_pipeline(components) + pipe_original.to(torch_device) + + output_original = pipe_original(**self.get_dummy_inputs_for_pipe_original())[0] + + # `from_pipe` must not repurpose the original pipeline's attention processors — PAG and friends install + # their own on the derived pipeline only. + original_attn_processor_types = { + name: {k: type(v) for k, v in component.attn_processors.items()} + for name, component in pipe_original.components.items() + if hasattr(component, "attn_processors") + } + + pipe = self.get_pipeline(**components).to(torch_device) + output = pipe(**self.get_dummy_inputs_pipe())[0] + + pipe_from_original = self.pipeline_class.from_pipe(pipe_original, **current_pipe_additional_components) + pipe_from_original.to(torch_device) + pipe_from_original.set_progress_bar_config(disable=None) + output_from_original = pipe_from_original(**self.get_dummy_inputs_pipe())[0] + + assert_tensors_close( + output_from_original, + output, + atol=expected_max_diff, + msg="The outputs of the pipelines created with `from_pipe` and `__init__` are different.", + ) + + output_original_2 = pipe_original(**self.get_dummy_inputs_for_pipe_original())[0] + assert_tensors_close( + output_original_2, + output_original, + atol=expected_max_diff, + msg="`from_pipe` should not change the output of original pipeline.", + ) + + for name, expected_types in original_attn_processor_types.items(): + component = pipe_original.components[name] + assert {k: type(v) for k, v in component.attn_processors.items()} == expected_types, ( + f"`from_pipe` changed the attention processors of `{name}` in the original pipeline." + ) + + @require_accelerator + @require_accelerate_version_greater("0.14.0") + def test_from_pipe_consistent_forward_pass_cpu_offload(self, expected_max_diff=1e-3): + components = self.get_dummy_components() + + # Build the original pipeline before running anything. Both pipelines share one scheduler object, and some + # `__init__`s edit that object in place: `StableDiffusionPipeline`, for one, rewrites a `steps_offset` of 0 + # to 1. Building it later would put the two forward passes below on different schedules. + pipe_original, current_pipe_additional_components = self._build_original_pipeline(components) + + pipe = self.get_pipeline(**components) + pipe.enable_model_cpu_offload(device=torch_device) + output = pipe(**self.get_dummy_inputs_pipe())[0] + + pipe_from_original = self.pipeline_class.from_pipe(pipe_original, **current_pipe_additional_components) + pipe_from_original.set_progress_bar_config(disable=None) + pipe_from_original.enable_model_cpu_offload(device=torch_device) + output_from_original = pipe_from_original(**self.get_dummy_inputs_pipe())[0] + + assert_tensors_close( + output_from_original, + output, + atol=expected_max_diff, + msg="The outputs of the pipelines created with `from_pipe` and `__init__` are different.", + ) diff --git a/tests/testing_utils.py b/tests/testing_utils.py index c35f975285c4..30f633032a2a 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -1018,6 +1018,25 @@ def export_to_gif(image: list[PIL.Image.Image], output_gif_path: str = None) -> return output_gif_path +@contextmanager +def skip_if_no_cudnn_engine(): + """ + Skip the enclosing test when cuDNN has no kernel for an op the pipeline runs. + + cuDNN does not ship an engine for every (op, dtype, layout) combination, and the set it covers differs between + cuDNN builds and GPU architectures. When none applies it raises `RuntimeError: GET was unable to find an engine + to execute this computation` — for instance for Sana's depthwise `Conv2d` in bfloat16. That is a property of the + runner, not of the code under test, so tests that can hit it are skipped there rather than failed. Any other + `RuntimeError` propagates untouched. + """ + try: + yield + except RuntimeError as e: + if "unable to find an engine" not in str(e): + raise + pytest.skip(f"cuDNN has no engine for this computation on {torch_device}: {e}") + + @contextmanager def buffered_writer(raw_f): f = io.BufferedWriter(raw_f)