From b4e6819fbef62e6ed4bdc88807b246973beb5f72 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 1 Sep 2026 04:42:39 +0000 Subject: [PATCH 1/6] migrate imports --- .../pipelines/chroma/test_pipeline_chroma.py | 2 +- .../chroma/test_pipeline_chroma_img2img.py | 2 +- .../controlnet_flux/test_controlnet_flux.py | 8 +- tests/pipelines/deepfloyd_if/test_if.py | 2 +- .../pipelines/deepfloyd_if/test_if_img2img.py | 2 +- .../test_if_img2img_superresolution.py | 2 +- .../deepfloyd_if/test_if_inpainting.py | 2 +- .../test_if_inpainting_superresolution.py | 2 +- .../deepfloyd_if/test_if_superresolution.py | 2 +- tests/pipelines/flux/test_pipeline_flux.py | 2 +- .../flux/test_pipeline_flux_img2img.py | 8 +- .../flux/test_pipeline_flux_inpaint.py | 8 +- .../flux/test_pipeline_flux_kontext.py | 2 +- .../test_pipeline_flux_kontext_inpaint.py | 2 +- tests/pipelines/flux/testing_utils.py | 162 ------------------ tests/pipelines/kandinsky/test_kandinsky.py | 2 +- .../kandinsky/test_kandinsky_img2img.py | 2 +- .../kandinsky/test_kandinsky_inpaint.py | 2 +- tests/pipelines/shap_e/test_shap_e.py | 2 +- tests/pipelines/shap_e/test_shap_e_img2img.py | 2 +- .../test_stable_diffusion_adapter.py | 3 +- .../stable_unclip/test_stable_unclip.py | 2 +- .../test_stable_unclip_img2img.py | 2 +- tests/pipelines/test_pipelines_common.py | 11 +- tests/pipelines/testing_utils/__init__.py | 5 +- tests/pipelines/testing_utils/ip_adapter.py | 145 +++++++++++++++- tests/pipelines/testing_utils/utils.py | 13 ++ 27 files changed, 198 insertions(+), 201 deletions(-) delete mode 100644 tests/pipelines/flux/testing_utils.py diff --git a/tests/pipelines/chroma/test_pipeline_chroma.py b/tests/pipelines/chroma/test_pipeline_chroma.py index ecaf056bf716..ba05f28fc345 100644 --- a/tests/pipelines/chroma/test_pipeline_chroma.py +++ b/tests/pipelines/chroma/test_pipeline_chroma.py @@ -4,9 +4,9 @@ from diffusers import AutoencoderKL, ChromaPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler from ...testing_utils import assert_tensors_close, torch_device -from ..flux.testing_utils import FluxIPAdapterTesterMixin from ..testing_utils import ( BasePipelineTesterConfig, + FluxIPAdapterTesterMixin, MemoryTesterMixin, PipelineTesterMixin, check_qkv_fused_layers_exist, diff --git a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py index 224d2bd6c98f..09aef4c62c6f 100644 --- a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py +++ b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py @@ -6,9 +6,9 @@ from diffusers import AutoencoderKL, ChromaImg2ImgPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler from ...testing_utils import assert_tensors_close, floats_tensor, torch_device -from ..flux.testing_utils import FluxIPAdapterTesterMixin from ..testing_utils import ( BasePipelineTesterConfig, + FluxIPAdapterTesterMixin, MemoryTesterMixin, PipelineTesterMixin, check_qkv_fused_layers_exist, diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux.py b/tests/pipelines/controlnet_flux/test_controlnet_flux.py index 90c5e617f61b..40f426af1214 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux.py @@ -40,8 +40,12 @@ require_big_accelerator, torch_device, ) -from ..flux.testing_utils import FluxIPAdapterTesterMixin -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + FluxIPAdapterTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() diff --git a/tests/pipelines/deepfloyd_if/test_if.py b/tests/pipelines/deepfloyd_if/test_if.py index 96bf699c4bc5..5d60dcb62e35 100644 --- a/tests/pipelines/deepfloyd_if/test_if.py +++ b/tests/pipelines/deepfloyd_if/test_if.py @@ -36,11 +36,11 @@ slow, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, PipelineOffloadTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/deepfloyd_if/test_if_img2img.py b/tests/pipelines/deepfloyd_if/test_if_img2img.py index 8dd0cea76e37..a79bef01e99b 100644 --- a/tests/pipelines/deepfloyd_if/test_if_img2img.py +++ b/tests/pipelines/deepfloyd_if/test_if_img2img.py @@ -37,11 +37,11 @@ slow, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, PipelineOffloadTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py b/tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py index 19d3607aa541..f12f163d3eee 100644 --- a/tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py +++ b/tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py @@ -37,11 +37,11 @@ slow, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, PipelineOffloadTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/deepfloyd_if/test_if_inpainting.py b/tests/pipelines/deepfloyd_if/test_if_inpainting.py index 9b89c5fe9ddc..9728c743b362 100644 --- a/tests/pipelines/deepfloyd_if/test_if_inpainting.py +++ b/tests/pipelines/deepfloyd_if/test_if_inpainting.py @@ -37,11 +37,11 @@ slow, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, PipelineOffloadTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py b/tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py index 6040042737b5..1d49208eabcf 100644 --- a/tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py +++ b/tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py @@ -37,11 +37,11 @@ slow, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, PipelineOffloadTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/deepfloyd_if/test_if_superresolution.py b/tests/pipelines/deepfloyd_if/test_if_superresolution.py index 6443c2478357..270e96b3e9a2 100644 --- a/tests/pipelines/deepfloyd_if/test_if_superresolution.py +++ b/tests/pipelines/deepfloyd_if/test_if_superresolution.py @@ -37,11 +37,11 @@ slow, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, PipelineOffloadTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/flux/test_pipeline_flux.py b/tests/pipelines/flux/test_pipeline_flux.py index 70cd88a48c9b..46dbdc68f74e 100644 --- a/tests/pipelines/flux/test_pipeline_flux.py +++ b/tests/pipelines/flux/test_pipeline_flux.py @@ -31,6 +31,7 @@ BasePipelineTesterConfig, FasterCacheTesterMixin, FirstBlockCacheTesterMixin, + FluxIPAdapterTesterMixin, LoraMemoryTesterMixin, LoraTesterMixin, MagCacheTesterMixin, @@ -39,7 +40,6 @@ PyramidAttentionBroadcastTesterMixin, TaylorSeerCacheTesterMixin, ) -from .testing_utils import FluxIPAdapterTesterMixin if is_peft_available(): diff --git a/tests/pipelines/flux/test_pipeline_flux_img2img.py b/tests/pipelines/flux/test_pipeline_flux_img2img.py index 93736df73ff5..f54aaba31cb0 100644 --- a/tests/pipelines/flux/test_pipeline_flux_img2img.py +++ b/tests/pipelines/flux/test_pipeline_flux_img2img.py @@ -6,8 +6,12 @@ from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxImg2ImgPipeline, FluxTransformer2DModel from ...testing_utils import floats_tensor, torch_device -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -from .testing_utils import FluxIPAdapterTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + FluxIPAdapterTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) class FluxImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): diff --git a/tests/pipelines/flux/test_pipeline_flux_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_inpaint.py index 87e778083134..98d0b0359207 100644 --- a/tests/pipelines/flux/test_pipeline_flux_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_inpaint.py @@ -6,8 +6,12 @@ from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxInpaintPipeline, FluxTransformer2DModel from ...testing_utils import floats_tensor, torch_device -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -from .testing_utils import FluxIPAdapterTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + FluxIPAdapterTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) class FluxInpaintPipelineTesterConfig(BasePipelineTesterConfig): diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext.py b/tests/pipelines/flux/test_pipeline_flux_kontext.py index 1fb5a52f7c84..d2c1645d9e1e 100644 --- a/tests/pipelines/flux/test_pipeline_flux_kontext.py +++ b/tests/pipelines/flux/test_pipeline_flux_kontext.py @@ -13,11 +13,11 @@ from ..testing_utils import ( BasePipelineTesterConfig, FasterCacheTesterMixin, + FluxIPAdapterTesterMixin, MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, ) -from .testing_utils import FluxIPAdapterTesterMixin class FluxKontextPipelineTesterConfig(BasePipelineTesterConfig): diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py index cfae5b4a2a38..234d95b25ebf 100644 --- a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py @@ -14,11 +14,11 @@ from ..testing_utils import ( BasePipelineTesterConfig, FasterCacheTesterMixin, + FluxIPAdapterTesterMixin, MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, ) -from .testing_utils import FluxIPAdapterTesterMixin class FluxKontextInpaintPipelineTesterConfig(BasePipelineTesterConfig): diff --git a/tests/pipelines/flux/testing_utils.py b/tests/pipelines/flux/testing_utils.py deleted file mode 100644 index f17c3916109f..000000000000 --- a/tests/pipelines/flux/testing_utils.py +++ /dev/null @@ -1,162 +0,0 @@ -# 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 typing import Any - -import torch - -from diffusers.loaders import FluxIPAdapterMixin - -from ...testing_utils import assert_tensors_close, is_ip_adapter, torch_device -from ..testing_utils.common import BasePipelineOutputMixin - - -@is_ip_adapter -class FluxIPAdapterTesterMixin(BasePipelineOutputMixin): - """IP-Adapter tests shared by the Flux pipelines in this directory and by the other pipelines built on the - Flux IP-Adapter API (Chroma, Flux ControlNet). - - Flux has its own IP-Adapter API (`FluxIPAdapterMixin`, image embeddings sized after the transformer's - `pooled_projection_dim`), so it doesn't reuse the Stable Diffusion `IPAdapterTesterMixin`. Compose it with a - `BasePipelineTesterConfig` subclass in its own test class, separate from the `PipelineTesterMixin` one. - """ - - def test_pipeline_signature(self): - parameters = inspect.signature(self.pipeline_class.__call__).parameters - - assert issubclass(self.pipeline_class, FluxIPAdapterMixin) - assert "ip_adapter_image" in parameters, ( - "`ip_adapter_image` argument must be supported by the `__call__` method" - ) - assert "ip_adapter_image_embeds" in parameters, ( - "`ip_adapter_image_embeds` argument must be supported by the `__call__` method" - ) - - def _get_dummy_image_embeds(self, image_embed_dim: int = 768): - return torch.randn((1, 1, image_embed_dim), device=torch_device) - - def _modify_inputs_for_ip_adapter_test(self, inputs: dict[str, Any]): - inputs["negative_prompt"] = "" - if "true_cfg_scale" in inspect.signature(self.pipeline_class.__call__).parameters: - inputs["true_cfg_scale"] = 4.0 - # Request torch outputs so comparisons run on torch tensors directly (see `BasePipelineTesterConfig`). - inputs["output_type"] = "pt" - inputs["return_dict"] = False - return inputs - - def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=None): - r"""Tests for IP-Adapter. - - The following scenarios are tested: - - Single IP-Adapter with scale=0 should produce same output as no IP-Adapter. - - Multi IP-Adapter with scale=0 should produce same output as no IP-Adapter. - - Single IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. - - Multi IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. - """ - # The state dict builder is imported here rather than at module scope: it lives in a model test module - # that calls `enable_full_determinism()` on import, which would otherwise flip that global for every test - # collected alongside this directory. - from ...models.transformers.test_models_transformer_flux import create_flux_ip_adapter_state_dict - - # Raising the tolerance for this test when it's run on a CPU because we compare against static slices and - # that can be shaky (with a VVVV low probability). - expected_max_diff = 9e-4 if torch_device == "cpu" else expected_max_diff - - components = self.get_dummy_components() - pipe = self.get_pipeline(**components).to(torch_device) - image_embed_dim = ( - pipe.transformer.config.pooled_projection_dim - if hasattr(pipe.transformer.config, "pooled_projection_dim") - else 768 - ) - - # forward pass without ip adapter - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - if expected_pipe_slice is None: - output_without_adapter = pipe(**inputs)[0] - else: - output_without_adapter = expected_pipe_slice - - # 1. Single IP-Adapter test cases - adapter_state_dict = create_flux_ip_adapter_state_dict(pipe.transformer) - # Load through the pipeline's public IP-Adapter API. `image_encoder_pretrained_model_name_or_path=None` - # skips fetching a CLIP image encoder since we feed pre-computed `ip_adapter_image_embeds` directly. - pipe.load_ip_adapter(adapter_state_dict, weight_name="", image_encoder_pretrained_model_name_or_path=None) - - # forward pass with single ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - pipe.set_ip_adapter_scale(0.0) - output_without_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with single ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - pipe.set_ip_adapter_scale(42.0) - output_with_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten() - - assert_tensors_close( - output_without_adapter_scale, - output_without_adapter, - atol=expected_max_diff, - msg="Output without ip-adapter must be same as normal inference", - ) - max_diff_with_adapter_scale = (output_with_adapter_scale - output_without_adapter).abs().max() - assert max_diff_with_adapter_scale > 1e-2, "Output with ip-adapter must be different from normal inference" - - # 2. Multi IP-Adapter test cases - adapter_state_dict_1 = create_flux_ip_adapter_state_dict(pipe.transformer) - adapter_state_dict_2 = create_flux_ip_adapter_state_dict(pipe.transformer) - pipe.load_ip_adapter( - [adapter_state_dict_1, adapter_state_dict_2], - weight_name=["", ""], - image_encoder_pretrained_model_name_or_path=None, - ) - - # forward pass with multi ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - pipe.set_ip_adapter_scale([0.0, 0.0]) - output_without_multi_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_without_multi_adapter_scale = output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with multi ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - pipe.set_ip_adapter_scale([42.0, 42.0]) - output_with_multi_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_with_multi_adapter_scale = output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten() - - assert_tensors_close( - output_without_multi_adapter_scale, - output_without_adapter, - atol=expected_max_diff, - msg="Output without multi-ip-adapter must be same as normal inference", - ) - max_diff_with_multi_adapter_scale = (output_with_multi_adapter_scale - output_without_adapter).abs().max() - assert max_diff_with_multi_adapter_scale > 1e-2, ( - "Output with multi-ip-adapter scale must be different from normal inference" - ) diff --git a/tests/pipelines/kandinsky/test_kandinsky.py b/tests/pipelines/kandinsky/test_kandinsky.py index 3e0e6f7d7f77..e63cca5153cf 100644 --- a/tests/pipelines/kandinsky/test_kandinsky.py +++ b/tests/pipelines/kandinsky/test_kandinsky.py @@ -31,11 +31,11 @@ slow, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/kandinsky/test_kandinsky_img2img.py b/tests/pipelines/kandinsky/test_kandinsky_img2img.py index 7f095b485701..1f72c6c35069 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_img2img.py +++ b/tests/pipelines/kandinsky/test_kandinsky_img2img.py @@ -44,11 +44,11 @@ slow, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/kandinsky/test_kandinsky_inpaint.py b/tests/pipelines/kandinsky/test_kandinsky_inpaint.py index 1232678cebb0..43eda3965157 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_inpaint.py +++ b/tests/pipelines/kandinsky/test_kandinsky_inpaint.py @@ -36,11 +36,11 @@ require_torch_accelerator, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/shap_e/test_shap_e.py b/tests/pipelines/shap_e/test_shap_e.py index a41ae265c797..aa8cc23aa8e3 100644 --- a/tests/pipelines/shap_e/test_shap_e.py +++ b/tests/pipelines/shap_e/test_shap_e.py @@ -29,11 +29,11 @@ require_torch_accelerator, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/shap_e/test_shap_e_img2img.py b/tests/pipelines/shap_e/test_shap_e_img2img.py index 9c572fbdf4da..bab08826d170 100644 --- a/tests/pipelines/shap_e/test_shap_e_img2img.py +++ b/tests/pipelines/shap_e/test_shap_e_img2img.py @@ -32,11 +32,11 @@ require_torch_accelerator, torch_device, ) -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py b/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py index 6f4346c0520b..039722e1ff36 100644 --- a/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py +++ b/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py @@ -49,7 +49,8 @@ torch_device, ) from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS -from ..test_pipelines_common import PipelineFromPipeTesterMixin, PipelineTesterMixin, assert_mean_pixel_difference +from ..test_pipelines_common import PipelineFromPipeTesterMixin, PipelineTesterMixin +from ..testing_utils import assert_mean_pixel_difference enable_full_determinism() diff --git a/tests/pipelines/stable_unclip/test_stable_unclip.py b/tests/pipelines/stable_unclip/test_stable_unclip.py index ebdf9403cebf..befd4e2bf192 100644 --- a/tests/pipelines/stable_unclip/test_stable_unclip.py +++ b/tests/pipelines/stable_unclip/test_stable_unclip.py @@ -26,11 +26,11 @@ torch_device, ) from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py b/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py index 5a0bc5cc3238..5f9ee5a5eaf6 100644 --- a/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py +++ b/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py @@ -32,11 +32,11 @@ torch_device, ) from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS -from ..test_pipelines_common import assert_mean_pixel_difference from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, + assert_mean_pixel_difference, ) diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index 106ba55cf149..58019ecf931a 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -69,6 +69,7 @@ skip_mps, torch_device, ) +from .testing_utils import assert_mean_pixel_difference def to_np(tensor): @@ -2934,13 +2935,3 @@ def run_forward(pipe): assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-4), ( "Outputs after disabling cache should match original inference exactly." ) - - -# Some models (e.g. unCLIP) are extremely likely to significantly deviate depending on which hardware is used. -# This helper function is used to check that the image doesn't deviate on average more than 10 pixels from a -# reference image. -def assert_mean_pixel_difference(image, expected_image, expected_max_diff=10): - image = np.asarray(DiffusionPipeline.numpy_to_pil(image)[0], dtype=np.float32) - expected_image = np.asarray(DiffusionPipeline.numpy_to_pil(expected_image)[0], dtype=np.float32) - avg_diff = np.abs(image - expected_image).mean() - assert avg_diff < expected_max_diff, f"Error image deviates {avg_diff} pixels on average" diff --git a/tests/pipelines/testing_utils/__init__.py b/tests/pipelines/testing_utils/__init__.py index 9b756ec64693..bfb93aaec03e 100644 --- a/tests/pipelines/testing_utils/__init__.py +++ b/tests/pipelines/testing_utils/__init__.py @@ -8,7 +8,7 @@ ) from .common import BasePipelineTesterConfig, PipelineTesterMixin from .from_pipe import FromPipeTesterMixin -from .ip_adapter import IPAdapterTesterMixin +from .ip_adapter import FluxIPAdapterTesterMixin, IPAdapterTesterMixin from .lora import LoraMemoryTesterMixin, LoraTesterMixin, UNetLoraTesterMixin from .memory import ( GroupOffloadTesterMixin, @@ -17,6 +17,7 @@ PipelineOffloadTesterMixin, ) from .utils import ( + assert_mean_pixel_difference, check_qkv_fused_layers_exist, check_qkv_fusion_matches_attn_procs_length, check_qkv_fusion_processors_exist, @@ -29,6 +30,7 @@ "PipelineTesterMixin", "FromPipeTesterMixin", "IPAdapterTesterMixin", + "FluxIPAdapterTesterMixin", "LoraTesterMixin", "LoraMemoryTesterMixin", "UNetLoraTesterMixin", @@ -42,6 +44,7 @@ "FirstBlockCacheTesterMixin", "TaylorSeerCacheTesterMixin", "MagCacheTesterMixin", + "assert_mean_pixel_difference", "check_qkv_fused_layers_exist", "check_qkv_fusion_matches_attn_procs_length", "check_qkv_fusion_processors_exist", diff --git a/tests/pipelines/testing_utils/ip_adapter.py b/tests/pipelines/testing_utils/ip_adapter.py index 5548b7468482..1354f5e3bfd2 100644 --- a/tests/pipelines/testing_utils/ip_adapter.py +++ b/tests/pipelines/testing_utils/ip_adapter.py @@ -14,11 +14,12 @@ # limitations under the License. import inspect +from typing import Any import pytest import torch -from diffusers.loaders import IPAdapterMixin +from diffusers.loaders import FluxIPAdapterMixin, IPAdapterMixin from ...testing_utils import assert_tensors_close, is_ip_adapter, torch_device from .common import BasePipelineOutputMixin @@ -29,8 +30,8 @@ class IPAdapterTesterMixin(BasePipelineOutputMixin): """IP-Adapter tests for UNet pipelines that load adapters through the standard `IPAdapterMixin` API. Compose it with a `BasePipelineTesterConfig` subclass in its own test class, separate from the - `PipelineTesterMixin` one. Pipelines whose IP-Adapter API differs (Flux, for example) keep their tests in - their own test module instead. + `PipelineTesterMixin` one. Pipelines whose IP-Adapter API differs use a mixin of their own; see + `FluxIPAdapterTesterMixin` below. """ def _get_dummy_image_embeds(self, cross_attention_dim: int = 32): @@ -239,3 +240,141 @@ def test_ip_adapter_faceid(self, expected_max_diff: float = 1e-4): ) max_diff_with_adapter_scale = (output_with_adapter_scale - output_without_adapter).abs().max() assert max_diff_with_adapter_scale > 1e-3, "Output with ip-adapter must be different from normal inference" + + +@is_ip_adapter +class FluxIPAdapterTesterMixin(BasePipelineOutputMixin): + """IP-Adapter tests shared by the Flux pipelines and by the other pipelines built on the Flux IP-Adapter + API (Chroma, Flux ControlNet). + + Flux has its own IP-Adapter API (`FluxIPAdapterMixin`, image embeddings sized after the transformer's + `pooled_projection_dim`), so it doesn't reuse the Stable Diffusion `IPAdapterTesterMixin`. Compose it with a + `BasePipelineTesterConfig` subclass in its own test class, separate from the `PipelineTesterMixin` one. + """ + + def test_pipeline_signature(self): + parameters = inspect.signature(self.pipeline_class.__call__).parameters + + assert issubclass(self.pipeline_class, FluxIPAdapterMixin) + assert "ip_adapter_image" in parameters, ( + "`ip_adapter_image` argument must be supported by the `__call__` method" + ) + assert "ip_adapter_image_embeds" in parameters, ( + "`ip_adapter_image_embeds` argument must be supported by the `__call__` method" + ) + + def _get_dummy_image_embeds(self, image_embed_dim: int = 768): + return torch.randn((1, 1, image_embed_dim), device=torch_device) + + def _modify_inputs_for_ip_adapter_test(self, inputs: dict[str, Any]): + inputs["negative_prompt"] = "" + if "true_cfg_scale" in inspect.signature(self.pipeline_class.__call__).parameters: + inputs["true_cfg_scale"] = 4.0 + # Request torch outputs so comparisons run on torch tensors directly (see `BasePipelineTesterConfig`). + inputs["output_type"] = "pt" + inputs["return_dict"] = False + return inputs + + def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=None): + r"""Tests for IP-Adapter. + + The following scenarios are tested: + - Single IP-Adapter with scale=0 should produce same output as no IP-Adapter. + - Multi IP-Adapter with scale=0 should produce same output as no IP-Adapter. + - Single IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. + - Multi IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. + """ + # The state dict builder is imported here rather than at module scope: it lives in a model test module + # that calls `enable_full_determinism()` on import, which would otherwise flip that global for every test + # collected alongside this directory. + from ...models.transformers.test_models_transformer_flux import create_flux_ip_adapter_state_dict + + # Raising the tolerance for this test when it's run on a CPU because we compare against static slices and + # that can be shaky (with a VVVV low probability). + expected_max_diff = 9e-4 if torch_device == "cpu" else expected_max_diff + + components = self.get_dummy_components() + pipe = self.get_pipeline(**components).to(torch_device) + image_embed_dim = ( + pipe.transformer.config.pooled_projection_dim + if hasattr(pipe.transformer.config, "pooled_projection_dim") + else 768 + ) + + # forward pass without ip adapter + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + if expected_pipe_slice is None: + output_without_adapter = pipe(**inputs)[0] + else: + output_without_adapter = expected_pipe_slice + + # 1. Single IP-Adapter test cases + adapter_state_dict = create_flux_ip_adapter_state_dict(pipe.transformer) + # Load through the pipeline's public IP-Adapter API. `image_encoder_pretrained_model_name_or_path=None` + # skips fetching a CLIP image encoder since we feed pre-computed `ip_adapter_image_embeds` directly. + pipe.load_ip_adapter(adapter_state_dict, weight_name="", image_encoder_pretrained_model_name_or_path=None) + + # forward pass with single ip adapter, but scale=0 which should have no effect + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] + inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] + pipe.set_ip_adapter_scale(0.0) + output_without_adapter_scale = pipe(**inputs)[0] + if expected_pipe_slice is not None: + output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten() + + # forward pass with single ip adapter, but with scale of adapter weights + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] + inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] + pipe.set_ip_adapter_scale(42.0) + output_with_adapter_scale = pipe(**inputs)[0] + if expected_pipe_slice is not None: + output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten() + + assert_tensors_close( + output_without_adapter_scale, + output_without_adapter, + atol=expected_max_diff, + msg="Output without ip-adapter must be same as normal inference", + ) + max_diff_with_adapter_scale = (output_with_adapter_scale - output_without_adapter).abs().max() + assert max_diff_with_adapter_scale > 1e-2, "Output with ip-adapter must be different from normal inference" + + # 2. Multi IP-Adapter test cases + adapter_state_dict_1 = create_flux_ip_adapter_state_dict(pipe.transformer) + adapter_state_dict_2 = create_flux_ip_adapter_state_dict(pipe.transformer) + pipe.load_ip_adapter( + [adapter_state_dict_1, adapter_state_dict_2], + weight_name=["", ""], + image_encoder_pretrained_model_name_or_path=None, + ) + + # forward pass with multi ip adapter, but scale=0 which should have no effect + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 + inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 + pipe.set_ip_adapter_scale([0.0, 0.0]) + output_without_multi_adapter_scale = pipe(**inputs)[0] + if expected_pipe_slice is not None: + output_without_multi_adapter_scale = output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten() + + # forward pass with multi ip adapter, but with scale of adapter weights + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 + inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 + pipe.set_ip_adapter_scale([42.0, 42.0]) + output_with_multi_adapter_scale = pipe(**inputs)[0] + if expected_pipe_slice is not None: + output_with_multi_adapter_scale = output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten() + + assert_tensors_close( + output_without_multi_adapter_scale, + output_without_adapter, + atol=expected_max_diff, + msg="Output without multi-ip-adapter must be same as normal inference", + ) + max_diff_with_multi_adapter_scale = (output_with_multi_adapter_scale - output_without_adapter).abs().max() + assert max_diff_with_multi_adapter_scale > 1e-2, ( + "Output with multi-ip-adapter scale must be different from normal inference" + ) diff --git a/tests/pipelines/testing_utils/utils.py b/tests/pipelines/testing_utils/utils.py index a20c0ca9becb..58bc15f2ab8b 100644 --- a/tests/pipelines/testing_utils/utils.py +++ b/tests/pipelines/testing_utils/utils.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import numpy as np + +from diffusers import DiffusionPipeline from diffusers.models.attention import AttentionModuleMixin @@ -52,3 +55,13 @@ def check_qkv_fused_layers_exist(model, layer_names): is_fused = is_fused_attribute_set and is_fused_layer is_fused_submodules.append(is_fused) return all(is_fused_submodules) + + +# Some models (e.g. unCLIP) are extremely likely to significantly deviate depending on which hardware is used. +# This helper function is used to check that the image doesn't deviate on average more than 10 pixels from a +# reference image. +def assert_mean_pixel_difference(image, expected_image, expected_max_diff=10): + image = np.asarray(DiffusionPipeline.numpy_to_pil(image)[0], dtype=np.float32) + expected_image = np.asarray(DiffusionPipeline.numpy_to_pil(expected_image)[0], dtype=np.float32) + avg_diff = np.abs(image - expected_image).mean() + assert avg_diff < expected_max_diff, f"Error image deviates {avg_diff} pixels on average" From c02fb4c42d4f29ef9bc4add09de1fe511030fa32 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 1 Sep 2026 05:42:51 +0000 Subject: [PATCH 2/6] migrate stable diffusion adapter tests --- .../test_stable_diffusion_adapter.py | 659 +++++++----------- tests/pipelines/testing_utils/common.py | 16 +- 2 files changed, 262 insertions(+), 413 deletions(-) diff --git a/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py b/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py index 039722e1ff36..a740ea1fe58c 100644 --- a/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py +++ b/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py @@ -15,14 +15,12 @@ import gc import random -import unittest import numpy as np +import pytest import torch -from parameterized import parameterized from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer -import diffusers from diffusers import ( AutoencoderKL, LCMScheduler, @@ -32,8 +30,6 @@ T2IAdapter, UNet2DConditionModel, ) -from diffusers.utils import logging -from diffusers.utils.import_utils import is_xformers_available from ...testing_utils import ( backend_empty_cache, @@ -49,19 +45,107 @@ torch_device, ) from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS -from ..test_pipelines_common import PipelineFromPipeTesterMixin, PipelineTesterMixin -from ..testing_utils import assert_mean_pixel_difference +from ..testing_utils import ( + BasePipelineTesterConfig, + FromPipeTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() -class AdapterTests: +# `StableDiffusionAdapterPipeline.__call__` has no `output_type="pt"` path: it branches on `"latent"` and `"pil"` +# and lets everything else fall through to the deprecated `decode_latents()`, which always returns a numpy array in +# `(batch, height, width, channels)`. So `get_dummy_inputs()` below has to ask for `"np"`, and every shared test +# that compares outputs with `assert_tensors_close` (torch-only) fails on the numpy array it gets back. +# +# The sibling SD pipelines route their postprocessing through `self.image_processor.postprocess(...)` and do +# support `"pt"`; adding that 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="`StableDiffusionAdapterPipeline` has no `output_type='pt'` path and always returns a numpy array.", + strict=True, +) + +# Same gap, applied to a whole class. Every `MemoryTesterMixin` test that runs the pipeline and compares its output +# hits the numpy array above (`torch.allclose`/`torch.isnan` reject it), but +# `test_pipeline_level_group_offloading_sanity_checks` never runs the pipeline and so passes — hence `strict=False`, +# which lets it report XPASS. Marking the class rather than each test keeps `MemoryTesterMixin`'s own `@is_memory` / +# `@require_accelerator` marks intact. +NO_PT_OUTPUT_NON_STRICT = pytest.mark.xfail( + reason="`StableDiffusionAdapterPipeline` has no `output_type='pt'` path and always returns a numpy array.", + strict=False, +) + + +class AdapterPipelineTesterConfig(BasePipelineTesterConfig): + """Shared testing contract for the three T2I-Adapter variants. + + Each variant subclass sets `adapter_type`; the multi-adapter one also raises `num_conditioning_images`, since + `MultiAdapter` takes one conditioning image per adapter. + """ + pipeline_class = StableDiffusionAdapterPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS + batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS + # `(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 = (64, 64, 3) + + # One of "full_adapter", "light_adapter" or "multi_adapter", set by the variant subclass. + adapter_type = None + # Number of conditioning images `get_dummy_inputs()` builds — one per adapter. + num_conditioning_images = 1 + + def get_adapter(self, channels, downscale_factor): + """Build this variant's adapter over `channels`, downscaling the conditioning image by `downscale_factor`.""" + if self.adapter_type in ("full_adapter", "light_adapter"): + return T2IAdapter( + in_channels=3, + channels=channels, + num_res_blocks=2, + downscale_factor=downscale_factor, + adapter_type=self.adapter_type, + ) + elif self.adapter_type == "multi_adapter": + return MultiAdapter( + [ + T2IAdapter( + in_channels=3, + channels=channels, + num_res_blocks=2, + downscale_factor=downscale_factor, + adapter_type="full_adapter", + ) + for _ in range(2) + ] + ) + raise ValueError( + f"Unknown adapter type: {self.adapter_type}, must be one of 'full_adapter', 'light_adapter', or " + "'multi_adapter'" + ) + + def get_text_encoder_and_tokenizer(self): + torch.manual_seed(0) + text_encoder_config = CLIPTextConfig( + bos_token_id=0, + eos_token_id=2, + hidden_size=32, + intermediate_size=37, + layer_norm_eps=1e-05, + num_attention_heads=4, + num_hidden_layers=5, + pad_token_id=1, + vocab_size=1000, + ) + text_encoder = CLIPTextModel(text_encoder_config) + tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + return text_encoder, tokenizer - def get_dummy_components(self, adapter_type, time_cond_proj_dim=None): + def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) unet = UNet2DConditionModel( block_out_channels=(32, 64), @@ -84,56 +168,12 @@ def get_dummy_components(self, adapter_type, time_cond_proj_dim=None): up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], latent_channels=4, ) - torch.manual_seed(0) - text_encoder_config = CLIPTextConfig( - bos_token_id=0, - eos_token_id=2, - hidden_size=32, - intermediate_size=37, - layer_norm_eps=1e-05, - num_attention_heads=4, - num_hidden_layers=5, - pad_token_id=1, - vocab_size=1000, - ) - text_encoder = CLIPTextModel(text_encoder_config) - tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + text_encoder, tokenizer = self.get_text_encoder_and_tokenizer() torch.manual_seed(0) + adapter = self.get_adapter(channels=[32, 64], downscale_factor=2) - if adapter_type == "full_adapter" or adapter_type == "light_adapter": - adapter = T2IAdapter( - in_channels=3, - channels=[32, 64], - num_res_blocks=2, - downscale_factor=2, - adapter_type=adapter_type, - ) - elif adapter_type == "multi_adapter": - adapter = MultiAdapter( - [ - T2IAdapter( - in_channels=3, - channels=[32, 64], - num_res_blocks=2, - downscale_factor=2, - adapter_type="full_adapter", - ), - T2IAdapter( - in_channels=3, - channels=[32, 64], - num_res_blocks=2, - downscale_factor=2, - adapter_type="full_adapter", - ), - ] - ) - else: - raise ValueError( - f"Unknown adapter type: {adapter_type}, must be one of 'full_adapter', 'light_adapter', or 'multi_adapter''" - ) - - components = { + return { "adapter": adapter, "unet": unet, "scheduler": scheduler, @@ -143,12 +183,11 @@ def get_dummy_components(self, adapter_type, time_cond_proj_dim=None): "safety_checker": None, "feature_extractor": None, } - return components - def get_dummy_components_with_full_downscaling(self, adapter_type): - """Get dummy components with x8 VAE downscaling and 4 UNet down blocks. - These dummy components are intended to fully-exercise the T2I-Adapter - downscaling behavior. + def get_dummy_components_with_full_downscaling(self): + """Dummy components with x8 VAE downscaling and 4 UNet down blocks. + + These dummy components are intended to fully-exercise the T2I-Adapter downscaling behavior. """ torch.manual_seed(0) unet = UNet2DConditionModel( @@ -171,56 +210,12 @@ def get_dummy_components_with_full_downscaling(self, adapter_type): up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"], latent_channels=4, ) - torch.manual_seed(0) - text_encoder_config = CLIPTextConfig( - bos_token_id=0, - eos_token_id=2, - hidden_size=32, - intermediate_size=37, - layer_norm_eps=1e-05, - num_attention_heads=4, - num_hidden_layers=5, - pad_token_id=1, - vocab_size=1000, - ) - text_encoder = CLIPTextModel(text_encoder_config) - tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + text_encoder, tokenizer = self.get_text_encoder_and_tokenizer() torch.manual_seed(0) + adapter = self.get_adapter(channels=[32, 32, 32, 64], downscale_factor=8) - if adapter_type == "full_adapter" or adapter_type == "light_adapter": - adapter = T2IAdapter( - in_channels=3, - channels=[32, 32, 32, 64], - num_res_blocks=2, - downscale_factor=8, - adapter_type=adapter_type, - ) - elif adapter_type == "multi_adapter": - adapter = MultiAdapter( - [ - T2IAdapter( - in_channels=3, - channels=[32, 32, 32, 64], - num_res_blocks=2, - downscale_factor=8, - adapter_type="full_adapter", - ), - T2IAdapter( - in_channels=3, - channels=[32, 32, 32, 64], - num_res_blocks=2, - downscale_factor=8, - adapter_type="full_adapter", - ), - ] - ) - else: - raise ValueError( - f"Unknown adapter type: {adapter_type}, must be one of 'full_adapter', 'light_adapter', or 'multi_adapter''" - ) - - components = { + return { "adapter": adapter, "unet": unet, "scheduler": scheduler, @@ -230,54 +225,40 @@ def get_dummy_components_with_full_downscaling(self, adapter_type): "safety_checker": None, "feature_extractor": None, } - return components - - def get_dummy_inputs(self, device, seed=0, height=64, width=64, num_images=1): - if num_images == 1: - image = floats_tensor((1, 3, height, width), rng=random.Random(seed)).to(device) - else: - image = [ - floats_tensor((1, 3, height, width), rng=random.Random(seed)).to(device) for _ in range(num_images) - ] - - 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, height=64, width=64): + # Every conditioning image is drawn from the same seed, so the adapters all see the same input. + images = [ + floats_tensor((1, 3, height, width), rng=random.Random(0)).to(torch_device) + for _ in range(self.num_conditioning_images) + ] + + return { "prompt": "A painting of a squirrel eating a burger", - "image": image, - "generator": generator, + "image": images[0] if self.num_conditioning_images == 1 else images, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, + # `"np"` rather than the usual `"pt"` — see `NO_PT_OUTPUT` above. "output_type": "np", } - return inputs - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - @unittest.skipIf( - torch_device != "cuda" or not is_xformers_available(), - reason="XFormers attention is only available with CUDA and `xformers` installed", - ) - def test_xformers_attention_forwardGenerator_pass(self): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) +class AdapterPipelineTesterMixin(PipelineTesterMixin): + """Core pipeline tests plus the adapter-specific ones shared by all three variants.""" - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) - - @parameterized.expand( + @pytest.mark.parametrize( + "dim", [ # (dim=264) The internal feature map will be 33x33 after initial pixel unshuffling (downscaled x8). - (((4 * 8 + 1) * 8),), + ((4 * 8 + 1) * 8), # (dim=272) The internal feature map will be 17x17 after the first T2I down block (downscaled x16). - (((4 * 4 + 1) * 16),), + ((4 * 4 + 1) * 16), # (dim=288) The internal feature map will be 9x9 after the second T2I down block (downscaled x32). - (((4 * 2 + 1) * 32),), + ((4 * 2 + 1) * 32), # (dim=320) The internal feature map will be 5x5 after the third T2I down block (downscaled x64). - (((4 * 1 + 1) * 64),), - ] + ((4 * 1 + 1) * 64), + ], ) def test_multiple_image_dimensions(self, dim): """Test that the T2I-Adapter pipeline supports any input dimension that @@ -288,335 +269,193 @@ def test_multiple_image_dimensions(self, dim): Note that we have selected `dim` values to produce odd resolutions at each downscaling level. """ - components = self.get_dummy_components_with_full_downscaling() - sd_pipe = StableDiffusionAdapterPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline(**self.get_dummy_components_with_full_downscaling()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device, height=dim, width=dim) - image = sd_pipe(**inputs).images + image = sd_pipe(**self.get_dummy_inputs(height=dim, width=dim)).images assert image.shape == (1, dim, dim, 3) def test_adapter_lcm(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionAdapterPipeline(**components) + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) sd_pipe.scheduler = LCMScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - output = sd_pipe(**inputs) - image = output.images - image_slice = image[0, -3:, -3:, -1] + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - assert image.shape == (1, 64, 64, 3) + # fmt: off expected_slice = np.array([0.4532, 0.5410, 0.4295, 0.5327, 0.6015, 0.4396, 0.5432, 0.4957, 0.4827]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # fmt: on + assert np.abs(image[0, -3:, -3:, -1].flatten() - expected_slice).max() < 1e-2 def test_adapter_lcm_custom_timesteps(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionAdapterPipeline(**components) + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) sd_pipe.scheduler = LCMScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() del inputs["num_inference_steps"] inputs["timesteps"] = [999, 499] - output = sd_pipe(**inputs) - image = output.images - - image_slice = image[0, -3:, -3:, -1] + image = sd_pipe(**inputs).images + assert image.shape == (1, *self.output_shape) - assert image.shape == (1, 64, 64, 3) + # Custom timesteps matching the default schedule reproduce `test_adapter_lcm`'s output. + # fmt: off expected_slice = np.array([0.4532, 0.5410, 0.4295, 0.5327, 0.6015, 0.4396, 0.5432, 0.4957, 0.4827]) + # fmt: on + assert np.abs(image[0, -3:, -3:, -1].flatten() - expected_slice).max() < 1e-2 - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - - 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, - } - return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) - - -class StableDiffusionFullAdapterPipelineFastTests( - AdapterTests, PipelineTesterMixin, PipelineFromPipeTesterMixin, unittest.TestCase -): - def get_dummy_components(self, time_cond_proj_dim=None): - return super().get_dummy_components("full_adapter", time_cond_proj_dim=time_cond_proj_dim) - - def get_dummy_components_with_full_downscaling(self): - return super().get_dummy_components_with_full_downscaling("full_adapter") + # The four overrides below only attach `NO_PT_OUTPUT`; none of these base methods carry decorators of their own. + @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) - def test_stable_diffusion_adapter_default_case(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionAdapterPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + @NO_PT_OUTPUT + def test_dict_tuple_outputs_equivalent(self): + super().test_dict_tuple_outputs_equivalent() - inputs = self.get_dummy_inputs(device) - image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] + @NO_PT_OUTPUT + def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4): + super().test_save_load_optional_components(tmp_path, expected_max_difference) - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.5248, 0.5794, 0.4504, 0.4649, 0.6327, 0.4491, 0.4922, 0.5155, 0.4938]) - assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-3 + @NO_PT_OUTPUT + def test_inference_batch_single_identical(self): + super().test_inference_batch_single_identical(expected_max_diff=2e-3) + # These three re-declare the base methods' decorators, which overriding would otherwise drop. + @NO_PT_OUTPUT + @pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="half-precision inference requires CUDA or XPU") @require_accelerator - @require_accelerate_version_greater("0.14.0") - def test_from_pipe_consistent_forward_pass_cpu_offload(self): - super().test_from_pipe_consistent_forward_pass_cpu_offload(expected_max_diff=6e-3) + @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) + @NO_PT_OUTPUT + @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) -class StableDiffusionLightAdapterPipelineFastTests(AdapterTests, PipelineTesterMixin, unittest.TestCase): - def get_dummy_components(self, time_cond_proj_dim=None): - return super().get_dummy_components("light_adapter", time_cond_proj_dim=time_cond_proj_dim) + @NO_PT_OUTPUT + @require_accelerator + def test_to_device(self): + super().test_to_device() - def get_dummy_components_with_full_downscaling(self): - return super().get_dummy_components_with_full_downscaling("light_adapter") + @NO_PT_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) - def test_stable_diffusion_adapter_default_case(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionAdapterPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device) - image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] +class FullAdapterPipelineTesterConfig(AdapterPipelineTesterConfig): + adapter_type = "full_adapter" - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.5463, 0.5897, 0.4547, 0.4751, 0.6357, 0.4527, 0.4924, 0.5190, 0.4969]) - assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-3 +class LightAdapterPipelineTesterConfig(AdapterPipelineTesterConfig): + adapter_type = "light_adapter" -class StableDiffusionMultiAdapterPipelineFastTests(AdapterTests, PipelineTesterMixin, unittest.TestCase): - def get_dummy_components(self, time_cond_proj_dim=None): - return super().get_dummy_components("multi_adapter", time_cond_proj_dim=time_cond_proj_dim) - def get_dummy_components_with_full_downscaling(self): - return super().get_dummy_components_with_full_downscaling("multi_adapter") +class MultiAdapterPipelineTesterConfig(AdapterPipelineTesterConfig): + adapter_type = "multi_adapter" + num_conditioning_images = 2 - def get_dummy_inputs(self, device, height=64, width=64, seed=0): - inputs = super().get_dummy_inputs(device, seed, height=height, width=width, num_images=2) + def get_dummy_inputs(self, height=64, width=64): + inputs = super().get_dummy_inputs(height=height, width=width) inputs["adapter_conditioning_scale"] = [0.5, 0.5] return inputs - def test_stable_diffusion_adapter_default_case(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionAdapterPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + def batch_input(self, name, value, batch_size): + # `image` holds one conditioning image per adapter, and the pipeline sizes the adapter state off each + # adapter's own batch (it never expands that state to the prompt's batch size). So the batch dimension is + # the inner list — one batch per adapter — not the outer one. + if name == "image": + return [batch_size * [image] for image in value] + return super().batch_input(name, value, batch_size) - inputs = self.get_dummy_inputs(device) - image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.5368, 0.5864, 0.4573, 0.4682, 0.6317, 0.4550, 0.4931, 0.5175, 0.4986]) - assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-3 - - def test_inference_batch_consistent( - self, batch_sizes=[2, 4, 13], 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) +class TestStableDiffusionFullAdapterPipeline(FullAdapterPipelineTesterConfig, AdapterPipelineTesterMixin): + def test_stable_diffusion_adapter_default_case(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(torch_device) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - logger = logging.get_logger(pipe.__module__) - logger.setLevel(level=diffusers.logging.FATAL) + # fmt: off + expected_slice = np.array([0.5248, 0.5794, 0.4504, 0.4649, 0.6327, 0.4491, 0.4922, 0.5155, 0.4938]) + # fmt: on + assert np.abs(image[0, -3:, -3:, -1].flatten() - expected_slice).max() < 5e-3 - # batchify inputs - for batch_size in batch_sizes: - batched_inputs = {} - for name, value in inputs.items(): - if name in self.batch_params: - # prompt is string - if name == "prompt": - len_prompt = len(value) - # make unequal batch sizes - batched_inputs[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)] - # make last batch super long - batched_inputs[name][-1] = 100 * "very long" - elif name == "image": - batched_images = [] +class TestStableDiffusionLightAdapterPipeline(LightAdapterPipelineTesterConfig, AdapterPipelineTesterMixin): + def test_stable_diffusion_adapter_default_case(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - for image in value: - batched_images.append(batch_size * [image]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - batched_inputs[name] = batched_images - else: - batched_inputs[name] = batch_size * [value] + # fmt: off + expected_slice = np.array([0.5463, 0.5897, 0.4547, 0.4751, 0.6357, 0.4527, 0.4924, 0.5190, 0.4969]) + # fmt: on + assert np.abs(image[0, -3:, -3:, -1].flatten() - expected_slice).max() < 5e-3 - elif name == "batch_size": - batched_inputs[name] = batch_size - else: - batched_inputs[name] = value - for arg in additional_params_copy_to_batched_inputs: - batched_inputs[arg] = inputs[arg] +class TestStableDiffusionMultiAdapterPipeline(MultiAdapterPipelineTesterConfig, AdapterPipelineTesterMixin): + def test_stable_diffusion_adapter_default_case(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - batched_inputs["output_type"] = "np" + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - if self.pipeline_class.__name__ == "DanceDiffusionPipeline": - batched_inputs.pop("output_type") + # fmt: off + expected_slice = np.array([0.5368, 0.5864, 0.4573, 0.4682, 0.6317, 0.4550, 0.4931, 0.5175, 0.4986]) + # fmt: on + assert np.abs(image[0, -3:, -3:, -1].flatten() - expected_slice).max() < 5e-3 - output = pipe(**batched_inputs) - assert len(output[0]) == batch_size +@NO_PT_OUTPUT_NON_STRICT +class TestStableDiffusionFullAdapterPipelineMemory(FullAdapterPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the full adapter.""" - batched_inputs["output_type"] = "np" - if self.pipeline_class.__name__ == "DanceDiffusionPipeline": - batched_inputs.pop("output_type") +@NO_PT_OUTPUT_NON_STRICT +class TestStableDiffusionLightAdapterPipelineMemory(LightAdapterPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the light adapter.""" - output = pipe(**batched_inputs)[0] - assert output.shape[0] == batch_size +@NO_PT_OUTPUT_NON_STRICT +class TestStableDiffusionMultiAdapterPipelineMemory(MultiAdapterPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the multi adapter.""" - logger.setLevel(level=diffusers.logging.WARNING) - def test_num_images_per_prompt(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) +class TestStableDiffusionFullAdapterPipelineFromPipe(FullAdapterPipelineTesterConfig, FromPipeTesterMixin): + """`from_pipe` round-trip tests against `StableDiffusionPipeline` for the full adapter.""" - batch_sizes = [1, 2] - num_images_per_prompts = [1, 2] - - for batch_size in batch_sizes: - for num_images_per_prompt in num_images_per_prompts: - inputs = self.get_dummy_inputs(torch_device) - - for key in inputs.keys(): - if key in self.batch_params: - if key == "image": - batched_images = [] - - for image in inputs[key]: - batched_images.append(batch_size * [image]) - - inputs[key] = batched_images - else: - inputs[key] = batch_size * [inputs[key]] - - images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt)[0] - - assert images.shape[0] == batch_size * num_images_per_prompt - - def test_inference_batch_single_identical( - self, - batch_size=3, - test_max_difference=None, - test_mean_pixel_difference=None, - relax_max_difference=False, - expected_max_diff=2e-3, - additional_params_copy_to_batched_inputs=["num_inference_steps"], - ): - if test_max_difference is None: - # TODO(Pedro) - not sure why, but not at all reproducible at the moment it seems - # make sure that batched and non-batched is identical - test_max_difference = torch_device != "mps" - - if test_mean_pixel_difference is None: - # TODO same as above - test_mean_pixel_difference = torch_device != "mps" - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + @NO_PT_OUTPUT + def test_from_pipe_consistent_forward_pass(self, expected_max_diff=1e-3): + super().test_from_pipe_consistent_forward_pass(expected_max_diff) - inputs = self.get_dummy_inputs(torch_device) - - logger = logging.get_logger(pipe.__module__) - logger.setLevel(level=diffusers.logging.FATAL) - - # batchify inputs - batched_inputs = {} - for name, value in inputs.items(): - if name in self.batch_params: - # prompt is string - if name == "prompt": - len_prompt = len(value) - # make unequal batch sizes - batched_inputs[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)] - - # make last batch super long - batched_inputs[name][-1] = 100 * "very long" - elif name == "image": - batched_images = [] - - for image in value: - batched_images.append(batch_size * [image]) - - batched_inputs[name] = batched_images - else: - batched_inputs[name] = batch_size * [value] - elif name == "batch_size": - batched_inputs[name] = batch_size - elif name == "generator": - batched_inputs[name] = [self.get_generator(i) for i in range(batch_size)] - else: - batched_inputs[name] = value - - for arg in additional_params_copy_to_batched_inputs: - batched_inputs[arg] = inputs[arg] - - if self.pipeline_class.__name__ != "DanceDiffusionPipeline": - batched_inputs["output_type"] = "np" - - output_batch = pipe(**batched_inputs) - assert output_batch[0].shape[0] == batch_size - - inputs["generator"] = self.get_generator(0) - - output = pipe(**inputs) - - logger.setLevel(level=diffusers.logging.WARNING) - if test_max_difference: - if relax_max_difference: - # Taking the median of the largest differences - # is resilient to outliers - diff = np.abs(output_batch[0][0] - output[0][0]) - diff = diff.flatten() - diff.sort() - max_diff = np.median(diff[-5:]) - else: - max_diff = np.abs(output_batch[0][0] - output[0][0]).max() - assert max_diff < expected_max_diff - - if test_mean_pixel_difference: - assert_mean_pixel_difference(output_batch[0][0], output[0][0]) + # Re-declared because overriding an inherited test drops the decorators it was declared with. + @NO_PT_OUTPUT + @require_accelerator + @require_accelerate_version_greater("0.14.0") + def test_from_pipe_consistent_forward_pass_cpu_offload(self): + super().test_from_pipe_consistent_forward_pass_cpu_offload(expected_max_diff=6e-3) @slow @require_torch_accelerator -class StableDiffusionAdapterPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestStableDiffusionAdapterPipelineIntegration: + @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/testing_utils/common.py b/tests/pipelines/testing_utils/common.py index 9efc428d4f69..771951341ded 100644 --- a/tests/pipelines/testing_utils/common.py +++ b/tests/pipelines/testing_utils/common.py @@ -183,6 +183,16 @@ def output_shape(self) -> tuple: def is_text_stack_component(self, name: str) -> bool: return any(key in name for key in self.text_stack_component_names) + def batch_input(self, name, value, batch_size): + """Expand one `batch_input_params` entry into a batch of `batch_size`. + + Defaults to repeating the value, which is what a single tensor/string input needs. Override it for an input + whose batch dimension isn't the outer list — a list holding one conditioning image per adapter, say, where + each adapter takes its own batch (see the multi-adapter tests in + `tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py`). + """ + return batch_size * [value] + def get_generator(self, seed=0): # Always build the generator on CPU: a CPU generator works with a pipeline placed on any device (the tensor # is created on CPU and moved), whereas an accelerator generator cannot seed a CPU tensor (see @@ -350,7 +360,7 @@ def test_inference_batch_consistent(self, batch_sizes=[2], batch_generator=True) # make last batch super long batched_input[name][-1] = 100 * "very long" else: - batched_input[name] = batch_size * [value] + batched_input[name] = self.batch_input(name, value, batch_size) if batch_generator and "generator" in inputs: batched_input["generator"] = [self.get_generator(i) for i in range(batch_size)] @@ -391,7 +401,7 @@ def test_inference_batch_single_identical( 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] + batched_inputs[name] = self.batch_input(name, value, batch_size) if "generator" in inputs: batched_inputs["generator"] = [self.get_generator(i) for i in range(batch_size)] @@ -586,7 +596,7 @@ def test_num_images_per_prompt(self): for key in inputs.keys(): if key in self.batch_input_params: - inputs[key] = batch_size * [inputs[key]] + inputs[key] = self.batch_input(key, inputs[key], batch_size) images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt)[0] From 5b1597b37e005426f4323f4fe199328bf56180df Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 1 Sep 2026 06:42:41 +0000 Subject: [PATCH 3/6] up --- tests/pipelines/kandinsky/test_kandinsky_prior.py | 13 ------------- .../pipelines/kandinsky2_2/test_kandinsky_prior.py | 13 ------------- .../kandinsky2_2/test_kandinsky_prior_emb2emb.py | 13 ------------- tests/pipelines/testing_utils/memory.py | 11 ++++++++--- 4 files changed, 8 insertions(+), 42 deletions(-) diff --git a/tests/pipelines/kandinsky/test_kandinsky_prior.py b/tests/pipelines/kandinsky/test_kandinsky_prior.py index 2dc0aa3c15ab..a06b82551a7a 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_prior.py +++ b/tests/pipelines/kandinsky/test_kandinsky_prior.py @@ -46,15 +46,6 @@ "`clip_mean`, `clip_std`), which group offloading never onloads." ) -# A second, independent gap: the component-scoped test only offloads the denoiser under the names -# `transformer`/`unet`/`controlnet`/`adapter`, and only puts `vae`/`vqvae`/`image_encoder` back on the accelerator. -# A prior pipeline's denoiser is called `prior`, so it matches neither list and is left on CPU while the text -# encoder is onloaded. Fixing this means widening the mixin's component lists, not changing the pipeline. -COMPONENT_GROUP_OFFLOAD_XFAIL_REASON = ( - "`GroupOffloadTesterMixin.test_group_offloading_inference` neither offloads nor places a component named " - "`prior`, so it stays on CPU while the onloaded text encoder runs on the accelerator." -) - class KandinskyPriorPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = KandinskyPriorPipeline @@ -214,10 +205,6 @@ def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff= class TestKandinskyPriorPipelineMemory(KandinskyPriorPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky prior pipeline.""" - @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - def test_group_offloading_inference(self): - super().test_group_offloading_inference() - @pytest.mark.xfail(condition=True, reason=PIPELINE_GROUP_OFFLOAD_XFAIL_REASON, strict=True) def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expected_max_difference=1e-4): super().test_pipeline_level_group_offloading_inference( diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py b/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py index 38455288b377..5ba60c49c1bc 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py @@ -48,15 +48,6 @@ "`clip_mean`, `clip_std`), which group offloading never onloads." ) -# A second, independent gap: the component-scoped test only offloads the denoiser under the names -# `transformer`/`unet`/`controlnet`/`adapter`, and only puts `vae`/`vqvae`/`image_encoder` back on the accelerator. -# A prior pipeline's denoiser is called `prior`, so it matches neither list and is left on CPU while the text -# encoder is onloaded. Fixing this means widening the mixin's component lists, not changing the pipeline. -COMPONENT_GROUP_OFFLOAD_XFAIL_REASON = ( - "`GroupOffloadTesterMixin.test_group_offloading_inference` neither offloads nor places a component named " - "`prior`, so it stays on CPU while the onloaded text encoder runs on the accelerator." -) - class KandinskyV22PriorPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = KandinskyV22PriorPipeline @@ -249,10 +240,6 @@ class TestKandinskyV22PriorPipelineMemory(KandinskyV22PriorPipelineTesterConfig, """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky 2.2 prior pipeline.""" - @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - def test_group_offloading_inference(self): - super().test_group_offloading_inference() - @pytest.mark.xfail(condition=True, reason=PIPELINE_GROUP_OFFLOAD_XFAIL_REASON, strict=True) def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expected_max_difference=1e-4): super().test_pipeline_level_group_offloading_inference( diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py b/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py index e2fd4858bd3d..06f2d3de9eb0 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py @@ -54,15 +54,6 @@ "`clip_mean`, `clip_std`), which group offloading never onloads." ) -# A second, independent gap: the component-scoped test only offloads the denoiser under the names -# `transformer`/`unet`/`controlnet`/`adapter`, and only puts `vae`/`vqvae`/`image_encoder` back on the accelerator. -# A prior pipeline's denoiser is called `prior`, so it matches neither list and is left on CPU while the text -# encoder is onloaded. Fixing this means widening the mixin's component lists, not changing the pipeline. -COMPONENT_GROUP_OFFLOAD_XFAIL_REASON = ( - "`GroupOffloadTesterMixin.test_group_offloading_inference` neither offloads nor places a component named " - "`prior`, so it stays on CPU while the onloaded text encoder runs on the accelerator." -) - class KandinskyV22PriorEmb2EmbPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = KandinskyV22PriorEmb2EmbPipeline @@ -233,10 +224,6 @@ class TestKandinskyV22PriorEmb2EmbPipelineMemory(KandinskyV22PriorEmb2EmbPipelin """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky 2.2 prior emb2emb pipeline.""" - @pytest.mark.xfail(condition=True, reason=COMPONENT_GROUP_OFFLOAD_XFAIL_REASON, strict=True) - def test_group_offloading_inference(self): - super().test_group_offloading_inference() - @pytest.mark.xfail(condition=True, reason=PIPELINE_GROUP_OFFLOAD_XFAIL_REASON, strict=True) def test_pipeline_level_group_offloading_inference(self, base_pipe_output, expected_max_difference=1e-4): super().test_pipeline_level_group_offloading_inference( diff --git a/tests/pipelines/testing_utils/memory.py b/tests/pipelines/testing_utils/memory.py index 981b717c33df..b067e51417bf 100644 --- a/tests/pipelines/testing_utils/memory.py +++ b/tests/pipelines/testing_utils/memory.py @@ -283,6 +283,7 @@ def create_pipe(): return self.get_pipeline() def enable_group_offload_on_component(pipe, group_offloading_kwargs): + offloaded_component_names = set() # We intentionally don't test VAE's here. This is because some tests enable tiling on the VAE. If # tiling is enabled and a forward pass is run, when accelerator streams are used, the execution order # of the layers is not traced correctly. This causes errors. For apply group offloading to VAE, a @@ -318,9 +319,13 @@ def enable_group_offload_on_component(pipe, group_offloading_kwargs): for module in component.modules() if hasattr(module, "_diffusers_hook") ) - for component_name in ["vae", "vqvae", "image_encoder"]: - component = getattr(pipe, component_name, None) - if isinstance(component, torch.nn.Module): + offloaded_component_names.add(component_name) + # Everything not group-offloaded above still has to be moved to the accelerator: the pipeline is built + # on CPU, so any module component left behind would meet accelerator inputs with CPU weights. Covers + # the VAE skipped above as well as components no pipeline-agnostic list would name — the depth + # estimator of `StableDiffusionDepth2ImgPipeline`, for one. + for component_name, component in pipe.components.items(): + if component_name not in offloaded_component_names and isinstance(component, torch.nn.Module): component.to(torch_device) def run_forward(pipe): From a622a8735545f15a771048d0a96c761924f07650 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 1 Sep 2026 07:30:00 +0000 Subject: [PATCH 4/6] move last remaining bits. --- .ai/references/testing.md | 2 +- .github/workflows/nightly_tests.yml | 1 - .github/workflows/pr_tests_gpu.yml | 2 +- .github/workflows/release_tests_fast.yml | 1 - .../test_models_transformer_flux.py | 2 +- tests/pipelines/test_pipelines_common.py | 2937 ----------------- tests/pipelines/testing_utils/from_pipe.py | 3 +- utils/extract_tests_from_mixin.py | 6 +- 8 files changed, 8 insertions(+), 2946 deletions(-) delete mode 100644 tests/pipelines/test_pipelines_common.py diff --git a/.ai/references/testing.md b/.ai/references/testing.md index 5c93d37b1db5..437ed6a081a0 100644 --- a/.ai/references/testing.md +++ b/.ai/references/testing.md @@ -35,7 +35,7 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers - `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. +- **`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. It replaces the unittest-era `PipelineFromPipeTesterMixin`, since removed. - **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. - **`encode_prompt` reading a component that isn't a text encoder or tokenizer?** `test_encode_prompt_works_in_isolation` rebuilds the pipeline with only the components whose names contain `text` or `tokenizer`. When `encode_prompt` also needs another one — a `processor` used for chat templating, say — list it in `text_stack_component_names` on the config class rather than re-implementing the test. diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index 0fedb5350602..46d0f6c3cda5 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.github/workflows/nightly_tests.yml @@ -311,7 +311,6 @@ jobs: pytest -n 1 --max-worker-restart=0 --dist=loadfile \ --make-reports=tests_torch_minimum_version_cuda \ tests/models/test_modeling_common.py \ - tests/pipelines/test_pipelines_common.py \ tests/pipelines/test_pipeline_utils.py \ tests/pipelines/test_pipelines.py \ tests/pipelines/test_pipelines_auto.py \ diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml index 24b08d2b5727..4527b96ebc74 100644 --- a/.github/workflows/pr_tests_gpu.yml +++ b/.github/workflows/pr_tests_gpu.yml @@ -14,7 +14,7 @@ on: - "src/diffusers/loaders/lora_base.py" - "src/diffusers/loaders/lora_pipeline.py" - "src/diffusers/loaders/peft.py" - - "tests/pipelines/test_pipelines_common.py" + - "tests/pipelines/testing_utils/**.py" - "tests/models/test_modeling_common.py" - "examples/**/*.py" workflow_dispatch: diff --git a/.github/workflows/release_tests_fast.yml b/.github/workflows/release_tests_fast.yml index f83f2abda7b8..3350c57bb132 100644 --- a/.github/workflows/release_tests_fast.yml +++ b/.github/workflows/release_tests_fast.yml @@ -205,7 +205,6 @@ jobs: pytest -n 1 --max-worker-restart=0 --dist=loadfile \ --make-reports=tests_torch_minimum_cuda \ tests/models/test_modeling_common.py \ - tests/pipelines/test_pipelines_common.py \ tests/pipelines/test_pipeline_utils.py \ tests/pipelines/test_pipelines.py \ tests/pipelines/test_pipelines_auto.py \ diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index 53af9eedc50c..c7b1ce1dca21 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -64,7 +64,7 @@ # TODO: This standalone function maintains backward compatibility with pipeline tests -# (tests/pipelines/test_pipelines_common.py) and will be refactored. +# (tests/pipelines/testing_utils/ip_adapter.py) and will be refactored. def create_flux_ip_adapter_state_dict(model) -> dict[str, dict[str, Any]]: """Create a dummy IP Adapter state dict for Flux transformer testing.""" ip_cross_attn_state_dict = {} diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py deleted file mode 100644 index 58019ecf931a..000000000000 --- a/tests/pipelines/test_pipelines_common.py +++ /dev/null @@ -1,2937 +0,0 @@ -import gc -import inspect -import json -import os -import tempfile -import unittest -from typing import Any, Callable, Dict - -import numpy as np -import PIL.Image -import pytest -import torch -import torch.nn as nn - -import diffusers -from diffusers import ( - AsymmetricAutoencoderKL, - AutoencoderKL, - AutoencoderTiny, - ConsistencyDecoderVAE, - DiffusionPipeline, - FasterCacheConfig, - KolorsPipeline, - PyramidAttentionBroadcastConfig, - StableDiffusionPipeline, - StableDiffusionXLPipeline, - UNet2DConditionModel, - apply_faster_cache, -) -from diffusers.hooks import apply_group_offloading -from diffusers.hooks.faster_cache import FasterCacheBlockHook, FasterCacheDenoiserHook -from diffusers.hooks.first_block_cache import FirstBlockCacheConfig -from diffusers.hooks.mag_cache import MagCacheConfig -from diffusers.hooks.pyramid_attention_broadcast import PyramidAttentionBroadcastHook -from diffusers.hooks.taylorseer_cache import TaylorSeerCacheConfig -from diffusers.image_processor import VaeImageProcessor -from diffusers.loaders import FluxIPAdapterMixin, IPAdapterMixin -from diffusers.models.attention import AttentionModuleMixin -from diffusers.models.attention_processor import AttnProcessor -from diffusers.models.controlnets.controlnet_xs import UNetControlNetXSModel -from diffusers.models.unets.unet_3d_condition import UNet3DConditionModel -from diffusers.models.unets.unet_i2vgen_xl import I2VGenXLUNet -from diffusers.models.unets.unet_motion_model import UNetMotionModel -from diffusers.pipelines.pipeline_utils import StableDiffusionMixin -from diffusers.schedulers import KarrasDiffusionSchedulers -from diffusers.utils import logging -from diffusers.utils.import_utils import is_xformers_available -from diffusers.utils.source_code_parsing_utils import ReturnNameVisitor - -from ..models.autoencoders.vae import ( - get_asym_autoencoder_kl_config, - get_autoencoder_kl_config, - get_autoencoder_tiny_config, - get_consistency_vae_config, -) -from ..models.transformers.test_models_transformer_flux import create_flux_ip_adapter_state_dict -from ..models.unets.test_models_unet_2d_condition import ( - create_ip_adapter_faceid_state_dict, - create_ip_adapter_state_dict, -) -from ..testing_utils import ( - CaptureLogger, - backend_empty_cache, - numpy_cosine_similarity_distance, - require_accelerate_version_greater, - require_accelerator, - require_torch, - require_torch_accelerator, - skip_mps, - torch_device, -) -from .testing_utils import assert_mean_pixel_difference - - -def to_np(tensor): - if isinstance(tensor, torch.Tensor): - tensor = tensor.detach().cpu().numpy() - - return tensor - - -def check_same_shape(tensor_list): - shapes = [tensor.shape for tensor in tensor_list] - return all(shape == shapes[0] for shape in shapes[1:]) - - -def check_qkv_fusion_matches_attn_procs_length(model, original_attn_processors): - current_attn_processors = model.attn_processors - return len(current_attn_processors) == len(original_attn_processors) - - -def check_qkv_fusion_processors_exist(model): - current_attn_processors = model.attn_processors - proc_names = [v.__class__.__name__ for _, v in current_attn_processors.items()] - return all(p.startswith("Fused") for p in proc_names) - - -def check_qkv_fused_layers_exist(model, layer_names): - is_fused_submodules = [] - for submodule in model.modules(): - if not isinstance(submodule, AttentionModuleMixin) or not submodule._supports_qkv_fusion: - continue - is_fused_attribute_set = submodule.fused_projections - is_fused_layer = True - for layer in layer_names: - is_fused_layer = is_fused_layer and getattr(submodule, layer, None) is not None - is_fused = is_fused_attribute_set and is_fused_layer - is_fused_submodules.append(is_fused) - return all(is_fused_submodules) - - -class SDFunctionTesterMixin: - """ - This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes. - It provides a set of common tests for PyTorch pipeline that inherit from StableDiffusionMixin, e.g. vae_slicing, vae_tiling, freeu, etc. - """ - - def test_vae_slicing(self, image_count=4): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - # components["scheduler"] = LMSDiscreteScheduler.from_config(components["scheduler"].config) - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * image_count - if "image" in inputs: # fix batch size mismatch in I2V_Gen pipeline - inputs["image"] = [inputs["image"]] * image_count - output_1 = pipe(**inputs) - - # make sure sliced vae decode yields the same result - pipe.vae.enable_slicing() - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * image_count - if "image" in inputs: - inputs["image"] = [inputs["image"]] * image_count - inputs["return_dict"] = False - output_2 = pipe(**inputs) - - assert np.abs(output_2[0].flatten() - output_1[0].flatten()).max() < 1e-2 - - def test_vae_tiling(self): - components = self.get_dummy_components() - - # make sure here that pndm scheduler skips prk - if "safety_checker" in components: - components["safety_checker"] = None - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - inputs["return_dict"] = False - - # Test that tiled decode at 512x512 yields the same result as the non-tiled decode - output_1 = pipe(**inputs)[0] - - # make sure tiled vae decode yields the same result - pipe.vae.enable_tiling() - inputs = self.get_dummy_inputs(torch_device) - inputs["return_dict"] = False - output_2 = pipe(**inputs)[0] - - assert np.abs(to_np(output_2) - to_np(output_1)).max() < 5e-1 - - # test that tiled decode works with various shapes - shapes = [(1, 4, 73, 97), (1, 4, 65, 49)] - with torch.no_grad(): - for shape in shapes: - zeros = torch.zeros(shape).to(torch_device) - pipe.vae.decode(zeros) - - # MPS currently doesn't support ComplexFloats, which are required for FreeU - see https://github.com/huggingface/diffusers/issues/7569. - @skip_mps - def test_freeu(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - # Normal inference - inputs = self.get_dummy_inputs(torch_device) - inputs["return_dict"] = False - inputs["output_type"] = "np" - output = pipe(**inputs)[0] - - # FreeU-enabled inference - pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4) - inputs = self.get_dummy_inputs(torch_device) - inputs["return_dict"] = False - inputs["output_type"] = "np" - output_freeu = pipe(**inputs)[0] - - # FreeU-disabled inference - pipe.disable_freeu() - freeu_keys = {"s1", "s2", "b1", "b2"} - for upsample_block in pipe.unet.up_blocks: - for key in freeu_keys: - assert getattr(upsample_block, key) is None, f"Disabling of FreeU should have set {key} to None." - - inputs = self.get_dummy_inputs(torch_device) - inputs["return_dict"] = False - inputs["output_type"] = "np" - output_no_freeu = pipe(**inputs)[0] - - assert not np.allclose(output[0, -3:, -3:, -1], output_freeu[0, -3:, -3:, -1]), ( - "Enabling of FreeU should lead to different results." - ) - assert np.allclose(output, output_no_freeu, atol=1e-2), ( - f"Disabling of FreeU should lead to results similar to the default pipeline results but Max Abs Error={np.abs(output_no_freeu - output).max()}." - ) - - 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) - - inputs = self.get_dummy_inputs(device) - inputs["return_dict"] = False - image = pipe(**inputs)[0] - original_image_slice = image[0, -3:, -3:, -1] - - pipe.fuse_qkv_projections() - for _, component in pipe.components.items(): - if ( - isinstance(component, nn.Module) - and hasattr(component, "original_attn_processors") - and component.original_attn_processors is not None - ): - assert check_qkv_fusion_processors_exist(component), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length(component, component.original_attn_processors), ( - "Something wrong with the attention processors concerning the fused QKV projections." - ) - - inputs = self.get_dummy_inputs(device) - inputs["return_dict"] = False - image_fused = pipe(**inputs)[0] - image_slice_fused = image_fused[0, -3:, -3:, -1] - - pipe.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - inputs["return_dict"] = False - image_disabled = pipe(**inputs)[0] - image_slice_disabled = image_disabled[0, -3:, -3:, -1] - - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( - "Fusion of QKV projections shouldn't affect the outputs." - ) - 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 np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." - ) - - -class IPAdapterTesterMixin: - """ - This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes. - It provides a set of common tests for pipelines that support IP Adapters. - """ - - def test_pipeline_signature(self): - parameters = inspect.signature(self.pipeline_class.__call__).parameters - - assert issubclass(self.pipeline_class, IPAdapterMixin) - self.assertIn( - "ip_adapter_image", - parameters, - "`ip_adapter_image` argument must be supported by the `__call__` method", - ) - self.assertIn( - "ip_adapter_image_embeds", - parameters, - "`ip_adapter_image_embeds` argument must be supported by the `__call__` method", - ) - - def _get_dummy_image_embeds(self, cross_attention_dim: int = 32): - return torch.randn((2, 1, cross_attention_dim), device=torch_device) - - def _get_dummy_faceid_image_embeds(self, cross_attention_dim: int = 32): - return torch.randn((2, 1, 1, cross_attention_dim), device=torch_device) - - def _get_dummy_masks(self, input_size: int = 64): - _masks = torch.zeros((1, 1, input_size, input_size), device=torch_device) - _masks[0, :, :, : int(input_size / 2)] = 1 - return _masks - - def _modify_inputs_for_ip_adapter_test(self, inputs: Dict[str, Any]): - parameters = inspect.signature(self.pipeline_class.__call__).parameters - if "image" in parameters.keys() and "strength" in parameters.keys(): - inputs["num_inference_steps"] = 4 - - inputs["output_type"] = "np" - inputs["return_dict"] = False - return inputs - - def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=None): - r"""Tests for IP-Adapter. - - The following scenarios are tested: - - Single IP-Adapter with scale=0 should produce same output as no IP-Adapter. - - Multi IP-Adapter with scale=0 should produce same output as no IP-Adapter. - - Single IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. - - Multi IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. - """ - # Raising the tolerance for this test when it's run on a CPU because we - # compare against static slices and that can be shaky (with a VVVV low probability). - expected_max_diff = 9e-4 if torch_device == "cpu" else expected_max_diff - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) - cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32) - - # forward pass without ip adapter - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - if expected_pipe_slice is None: - output_without_adapter = pipe(**inputs)[0] - else: - output_without_adapter = expected_pipe_slice - - # 1. Single IP-Adapter test cases - adapter_state_dict = create_ip_adapter_state_dict(pipe.unet) - pipe.unet._load_ip_adapter_weights(adapter_state_dict) - - # forward pass with single ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] - pipe.set_ip_adapter_scale(0.0) - output_without_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with single ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] - pipe.set_ip_adapter_scale(42.0) - output_with_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten() - - max_diff_without_adapter_scale = np.abs(output_without_adapter_scale - output_without_adapter).max() - max_diff_with_adapter_scale = np.abs(output_with_adapter_scale - output_without_adapter).max() - - self.assertLess( - max_diff_without_adapter_scale, - expected_max_diff, - "Output without ip-adapter must be same as normal inference", - ) - self.assertGreater( - max_diff_with_adapter_scale, 1e-2, "Output with ip-adapter must be different from normal inference" - ) - - # 2. Multi IP-Adapter test cases - adapter_state_dict_1 = create_ip_adapter_state_dict(pipe.unet) - adapter_state_dict_2 = create_ip_adapter_state_dict(pipe.unet) - pipe.unet._load_ip_adapter_weights([adapter_state_dict_1, adapter_state_dict_2]) - - # forward pass with multi ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] * 2 - pipe.set_ip_adapter_scale([0.0, 0.0]) - output_without_multi_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_without_multi_adapter_scale = output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with multi ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] * 2 - pipe.set_ip_adapter_scale([42.0, 42.0]) - output_with_multi_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_with_multi_adapter_scale = output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten() - - max_diff_without_multi_adapter_scale = np.abs( - output_without_multi_adapter_scale - output_without_adapter - ).max() - max_diff_with_multi_adapter_scale = np.abs(output_with_multi_adapter_scale - output_without_adapter).max() - self.assertLess( - max_diff_without_multi_adapter_scale, - expected_max_diff, - "Output without multi-ip-adapter must be same as normal inference", - ) - self.assertGreater( - max_diff_with_multi_adapter_scale, - 1e-2, - "Output with multi-ip-adapter scale must be different from normal inference", - ) - - def test_ip_adapter_cfg(self, expected_max_diff: float = 1e-4): - parameters = inspect.signature(self.pipeline_class.__call__).parameters - - if "guidance_scale" not in parameters: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) - cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32) - - adapter_state_dict = create_ip_adapter_state_dict(pipe.unet) - pipe.unet._load_ip_adapter_weights(adapter_state_dict) - pipe.set_ip_adapter_scale(1.0) - - # forward pass with CFG not applied - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)[0].unsqueeze(0)] - inputs["guidance_scale"] = 1.0 - out_no_cfg = pipe(**inputs)[0] - - # forward pass with CFG applied - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] - inputs["guidance_scale"] = 7.5 - out_cfg = pipe(**inputs)[0] - - assert out_cfg.shape == out_no_cfg.shape - - def test_ip_adapter_masks(self, expected_max_diff: float = 1e-4): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) - cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32) - sample_size = pipe.unet.config.get("sample_size", 32) - block_out_channels = pipe.vae.config.get("block_out_channels", [128, 256, 512, 512]) - input_size = sample_size * (2 ** (len(block_out_channels) - 1)) - - # forward pass without ip adapter - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - output_without_adapter = pipe(**inputs)[0] - output_without_adapter = output_without_adapter[0, -3:, -3:, -1].flatten() - - adapter_state_dict = create_ip_adapter_state_dict(pipe.unet) - pipe.unet._load_ip_adapter_weights(adapter_state_dict) - - # forward pass with single ip adapter and masks, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] - inputs["cross_attention_kwargs"] = {"ip_adapter_masks": [self._get_dummy_masks(input_size)]} - pipe.set_ip_adapter_scale(0.0) - output_without_adapter_scale = pipe(**inputs)[0] - output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with single ip adapter and masks, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] - inputs["cross_attention_kwargs"] = {"ip_adapter_masks": [self._get_dummy_masks(input_size)]} - pipe.set_ip_adapter_scale(42.0) - output_with_adapter_scale = pipe(**inputs)[0] - output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten() - - max_diff_without_adapter_scale = np.abs(output_without_adapter_scale - output_without_adapter).max() - max_diff_with_adapter_scale = np.abs(output_with_adapter_scale - output_without_adapter).max() - - self.assertLess( - max_diff_without_adapter_scale, - expected_max_diff, - "Output without ip-adapter must be same as normal inference", - ) - self.assertGreater( - max_diff_with_adapter_scale, 1e-3, "Output with ip-adapter must be different from normal inference" - ) - - def test_ip_adapter_faceid(self, expected_max_diff: float = 1e-4): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) - cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32) - - # forward pass without ip adapter - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - output_without_adapter = pipe(**inputs)[0] - output_without_adapter = output_without_adapter[0, -3:, -3:, -1].flatten() - - adapter_state_dict = create_ip_adapter_faceid_state_dict(pipe.unet) - pipe.unet._load_ip_adapter_weights(adapter_state_dict) - - # forward pass with single ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_faceid_image_embeds(cross_attention_dim)] - pipe.set_ip_adapter_scale(0.0) - output_without_adapter_scale = pipe(**inputs)[0] - output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with single ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_faceid_image_embeds(cross_attention_dim)] - pipe.set_ip_adapter_scale(42.0) - output_with_adapter_scale = pipe(**inputs)[0] - output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten() - - max_diff_without_adapter_scale = np.abs(output_without_adapter_scale - output_without_adapter).max() - max_diff_with_adapter_scale = np.abs(output_with_adapter_scale - output_without_adapter).max() - - self.assertLess( - max_diff_without_adapter_scale, - expected_max_diff, - "Output without ip-adapter must be same as normal inference", - ) - self.assertGreater( - max_diff_with_adapter_scale, 1e-3, "Output with ip-adapter must be different from normal inference" - ) - - -class FluxIPAdapterTesterMixin: - """ - This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes. - It provides a set of common tests for pipelines that support IP Adapters. - """ - - def test_pipeline_signature(self): - parameters = inspect.signature(self.pipeline_class.__call__).parameters - - assert issubclass(self.pipeline_class, FluxIPAdapterMixin) - self.assertIn( - "ip_adapter_image", - parameters, - "`ip_adapter_image` argument must be supported by the `__call__` method", - ) - self.assertIn( - "ip_adapter_image_embeds", - parameters, - "`ip_adapter_image_embeds` argument must be supported by the `__call__` method", - ) - - def _get_dummy_image_embeds(self, image_embed_dim: int = 768): - return torch.randn((1, 1, image_embed_dim), device=torch_device) - - def _modify_inputs_for_ip_adapter_test(self, inputs: Dict[str, Any]): - inputs["negative_prompt"] = "" - if "true_cfg_scale" in inspect.signature(self.pipeline_class.__call__).parameters: - inputs["true_cfg_scale"] = 4.0 - inputs["output_type"] = "np" - inputs["return_dict"] = False - return inputs - - def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=None): - r"""Tests for IP-Adapter. - - The following scenarios are tested: - - Single IP-Adapter with scale=0 should produce same output as no IP-Adapter. - - Multi IP-Adapter with scale=0 should produce same output as no IP-Adapter. - - Single IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. - - Multi IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. - """ - # Raising the tolerance for this test when it's run on a CPU because we - # compare against static slices and that can be shaky (with a VVVV low probability). - expected_max_diff = 9e-4 if torch_device == "cpu" else expected_max_diff - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) - image_embed_dim = ( - pipe.transformer.config.pooled_projection_dim - if hasattr(pipe.transformer.config, "pooled_projection_dim") - else 768 - ) - - # forward pass without ip adapter - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - if expected_pipe_slice is None: - output_without_adapter = pipe(**inputs)[0] - else: - output_without_adapter = expected_pipe_slice - - # 1. Single IP-Adapter test cases - adapter_state_dict = create_flux_ip_adapter_state_dict(pipe.transformer) - pipe.transformer._load_ip_adapter_weights(adapter_state_dict) - - # forward pass with single ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - pipe.set_ip_adapter_scale(0.0) - output_without_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with single ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - pipe.set_ip_adapter_scale(42.0) - output_with_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten() - - max_diff_without_adapter_scale = np.abs(output_without_adapter_scale - output_without_adapter).max() - max_diff_with_adapter_scale = np.abs(output_with_adapter_scale - output_without_adapter).max() - - self.assertLess( - max_diff_without_adapter_scale, - expected_max_diff, - "Output without ip-adapter must be same as normal inference", - ) - self.assertGreater( - max_diff_with_adapter_scale, 1e-2, "Output with ip-adapter must be different from normal inference" - ) - - # 2. Multi IP-Adapter test cases - adapter_state_dict_1 = create_flux_ip_adapter_state_dict(pipe.transformer) - adapter_state_dict_2 = create_flux_ip_adapter_state_dict(pipe.transformer) - pipe.transformer._load_ip_adapter_weights([adapter_state_dict_1, adapter_state_dict_2]) - - # forward pass with multi ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - pipe.set_ip_adapter_scale([0.0, 0.0]) - output_without_multi_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_without_multi_adapter_scale = output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with multi ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device)) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - pipe.set_ip_adapter_scale([42.0, 42.0]) - output_with_multi_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_with_multi_adapter_scale = output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten() - - max_diff_without_multi_adapter_scale = np.abs( - output_without_multi_adapter_scale - output_without_adapter - ).max() - max_diff_with_multi_adapter_scale = np.abs(output_with_multi_adapter_scale - output_without_adapter).max() - self.assertLess( - max_diff_without_multi_adapter_scale, - expected_max_diff, - "Output without multi-ip-adapter must be same as normal inference", - ) - self.assertGreater( - max_diff_with_multi_adapter_scale, - 1e-2, - "Output with multi-ip-adapter scale must be different from normal inference", - ) - - -class PipelineLatentTesterMixin: - """ - This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes. - It provides a set of common tests for PyTorch pipeline that has vae, e.g. - equivalence of different input and output types, etc. - """ - - @property - def image_params(self) -> frozenset: - raise NotImplementedError( - "You need to set the attribute `image_params` in the child test class. " - "`image_params` are tested for if all accepted input image types (i.e. `pt`,`pil`,`np`) are producing same results" - ) - - @property - def image_latents_params(self) -> frozenset: - raise NotImplementedError( - "You need to set the attribute `image_latents_params` in the child test class. " - "`image_latents_params` are tested for if passing latents directly are producing same results" - ) - - def get_dummy_inputs_by_type(self, device, seed=0, input_image_type="pt", output_type="np"): - inputs = self.get_dummy_inputs(device, seed) - - def convert_to_pt(image): - if isinstance(image, torch.Tensor): - input_image = image - elif isinstance(image, np.ndarray): - input_image = VaeImageProcessor.numpy_to_pt(image) - elif isinstance(image, PIL.Image.Image): - input_image = VaeImageProcessor.pil_to_numpy(image) - input_image = VaeImageProcessor.numpy_to_pt(input_image) - else: - raise ValueError(f"unsupported input_image_type {type(image)}") - return input_image - - def convert_pt_to_type(image, input_image_type): - if input_image_type == "pt": - input_image = image - elif input_image_type == "np": - input_image = VaeImageProcessor.pt_to_numpy(image) - elif input_image_type == "pil": - input_image = VaeImageProcessor.pt_to_numpy(image) - input_image = VaeImageProcessor.numpy_to_pil(input_image) - else: - raise ValueError(f"unsupported input_image_type {input_image_type}.") - return input_image - - for image_param in self.image_params: - if image_param in inputs.keys(): - inputs[image_param] = convert_pt_to_type( - convert_to_pt(inputs[image_param]).to(device), input_image_type - ) - - inputs["output_type"] = output_type - - return inputs - - def test_pt_np_pil_outputs_equivalent(self, expected_max_diff=1e-4): - self._test_pt_np_pil_outputs_equivalent(expected_max_diff=expected_max_diff) - - def _test_pt_np_pil_outputs_equivalent(self, expected_max_diff=1e-4, input_image_type="pt"): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - output_pt = pipe( - **self.get_dummy_inputs_by_type(torch_device, input_image_type=input_image_type, output_type="pt") - )[0] - output_np = pipe( - **self.get_dummy_inputs_by_type(torch_device, input_image_type=input_image_type, output_type="np") - )[0] - output_pil = pipe( - **self.get_dummy_inputs_by_type(torch_device, input_image_type=input_image_type, output_type="pil") - )[0] - - max_diff = np.abs(output_pt.cpu().numpy().transpose(0, 2, 3, 1) - output_np).max() - self.assertLess( - max_diff, expected_max_diff, "`output_type=='pt'` generate different results from `output_type=='np'`" - ) - - max_diff = np.abs(np.array(output_pil[0]) - (output_np * 255).round()).max() - self.assertLess(max_diff, 2.0, "`output_type=='pil'` generate different results from `output_type=='np'`") - - def test_pt_np_pil_inputs_equivalent(self): - if len(self.image_params) == 0: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - out_input_pt = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="pt"))[0] - out_input_np = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="np"))[0] - out_input_pil = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="pil"))[0] - - max_diff = np.abs(out_input_pt - out_input_np).max() - self.assertLess(max_diff, 1e-4, "`input_type=='pt'` generate different result from `input_type=='np'`") - max_diff = np.abs(out_input_pil - out_input_np).max() - self.assertLess(max_diff, 1e-2, "`input_type=='pt'` generate different result from `input_type=='np'`") - - def test_latents_input(self): - if len(self.image_latents_params) == 0: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.image_processor = VaeImageProcessor(do_resize=False, do_normalize=False) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - out = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="pt"))[0] - - vae = components["vae"] - inputs = self.get_dummy_inputs_by_type(torch_device, input_image_type="pt") - generator = inputs["generator"] - for image_param in self.image_latents_params: - if image_param in inputs.keys(): - inputs[image_param] = ( - vae.encode(inputs[image_param]).latent_dist.sample(generator) * vae.config.scaling_factor - ) - out_latents_inputs = pipe(**inputs)[0] - - max_diff = np.abs(out - out_latents_inputs).max() - self.assertLess(max_diff, 1e-4, "passing latents as image input generate different result from passing image") - - def test_multi_vae(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - block_out_channels = pipe.vae.config.block_out_channels - norm_num_groups = pipe.vae.config.norm_num_groups - - vae_classes = [AutoencoderKL, AsymmetricAutoencoderKL, ConsistencyDecoderVAE, AutoencoderTiny] - configs = [ - get_autoencoder_kl_config(block_out_channels, norm_num_groups), - get_asym_autoencoder_kl_config(block_out_channels, norm_num_groups), - get_consistency_vae_config(block_out_channels, norm_num_groups), - get_autoencoder_tiny_config(block_out_channels), - ] - - out_np = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="np"))[0] - - for vae_cls, config in zip(vae_classes, configs): - vae = vae_cls(**config) - vae = vae.to(torch_device) - components["vae"] = vae - vae_pipe = self.pipeline_class(**components) - out_vae_np = vae_pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="np"))[0] - - assert out_vae_np.shape == out_np.shape - - -@require_torch -class PipelineFromPipeTesterMixin: - @property - def original_pipeline_class(self): - if "xl" in self.pipeline_class.__name__.lower(): - original_pipeline_class = StableDiffusionXLPipeline - elif "kolors" in self.pipeline_class.__name__.lower(): - original_pipeline_class = KolorsPipeline - else: - original_pipeline_class = StableDiffusionPipeline - - return original_pipeline_class - - def get_dummy_inputs_pipe(self, device, seed=0): - inputs = self.get_dummy_inputs(device, seed=seed) - inputs["output_type"] = "np" - inputs["return_dict"] = False - return inputs - - def get_dummy_inputs_for_pipe_original(self, device, seed=0): - inputs = {} - for k, v in self.get_dummy_inputs_pipe(device, seed=seed).items(): - if k in set(inspect.signature(self.original_pipeline_class.__call__).parameters.keys()): - inputs[k] = v - return inputs - - def test_from_pipe_consistent_config(self): - if self.original_pipeline_class == StableDiffusionPipeline: - original_repo = "hf-internal-testing/tiny-stable-diffusion-torch" - original_kwargs = {"requires_safety_checker": False} - elif self.original_pipeline_class == StableDiffusionXLPipeline: - original_repo = "hf-internal-testing/tiny-stable-diffusion-xl-pipe" - original_kwargs = {"requires_aesthetics_score": True, "force_zeros_for_empty_prompt": False} - elif self.original_pipeline_class == KolorsPipeline: - original_repo = "hf-internal-testing/tiny-kolors-pipe" - original_kwargs = {"force_zeros_for_empty_prompt": False} - else: - raise ValueError( - "original_pipeline_class must be either StableDiffusionPipeline or StableDiffusionXLPipeline" - ) - - # create original_pipeline_class(sd/sdxl) - pipe_original = self.original_pipeline_class.from_pretrained(original_repo, **original_kwargs) - - # original_pipeline_class(sd/sdxl) -> 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/sdxl) - 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() - original_expected_modules, _ = self.original_pipeline_class._get_signature_keys(self.original_pipeline_class) - - # pipeline components that are also expected to be in the original pipeline - original_pipe_components = {} - # additional components that are not in the pipeline, but expected in the original pipeline - original_pipe_additional_components = {} - # additional components that are in the pipeline, but not expected in the original pipeline - 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.__class__}: {name}") - - pipe_original = self.original_pipeline_class(**original_pipe_components, **original_pipe_additional_components) - for component in pipe_original.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_original.to(torch_device) - pipe_original.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs_for_pipe_original(torch_device) - output_original = pipe_original(**inputs)[0] - - 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_pipe(torch_device) - output = pipe(**inputs)[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) - inputs = self.get_dummy_inputs_pipe(torch_device) - output_from_original = pipe_from_original(**inputs)[0] - - max_diff = np.abs(to_np(output) - to_np(output_from_original)).max() - self.assertLess( - max_diff, - expected_max_diff, - "The outputs of the pipelines created with `from_pipe` and `__init__` are different.", - ) - - inputs = self.get_dummy_inputs_for_pipe_original(torch_device) - output_original_2 = pipe_original(**inputs)[0] - - max_diff = np.abs(to_np(output_original) - to_np(output_original_2)).max() - self.assertLess(max_diff, expected_max_diff, "`from_pipe` should not change the output of original pipeline.") - - for component in pipe_original.components.values(): - if hasattr(component, "attn_processors"): - assert all(type(proc) == AttnProcessor for proc in component.attn_processors.values()), ( - "`from_pipe` changed the attention processor in 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() - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe.enable_model_cpu_offload(device=torch_device) - pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs_pipe(torch_device) - output = pipe(**inputs)[0] - - original_expected_modules, _ = self.original_pipeline_class._get_signature_keys(self.original_pipeline_class) - # pipeline components that are also expected to be in the original pipeline - original_pipe_components = {} - # additional components that are not in the pipeline, but expected in the original pipeline - original_pipe_additional_components = {} - # additional components that are in the pipeline, but not expected in the original pipeline - 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.__class__}: {name}") - - pipe_original = self.original_pipeline_class(**original_pipe_components, **original_pipe_additional_components) - for component in pipe_original.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_original.set_progress_bar_config(disable=None) - - pipe_from_original = self.pipeline_class.from_pipe(pipe_original, **current_pipe_additional_components) - for component in pipe_from_original.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - - pipe_from_original.enable_model_cpu_offload(device=torch_device) - pipe_from_original.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs_pipe(torch_device) - output_from_original = pipe_from_original(**inputs)[0] - - max_diff = np.abs(to_np(output) - to_np(output_from_original)).max() - self.assertLess( - max_diff, - expected_max_diff, - "The outputs of the pipelines created with `from_pipe` and `__init__` are different.", - ) - - -@require_torch -class PipelineKarrasSchedulerTesterMixin: - """ - This mixin is designed to be used with unittest.TestCase classes. - It provides a set of common tests for each PyTorch pipeline that makes use of KarrasDiffusionSchedulers - equivalence of dict and tuple outputs, etc. - """ - - def test_karras_schedulers_shape( - self, num_inference_steps_for_strength=4, num_inference_steps_for_strength_for_iterations=5 - ): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - # make sure that PNDM does not need warm-up - pipe.scheduler.register_to_config(skip_prk_steps=True) - - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = 2 - - if "strength" in inputs: - inputs["num_inference_steps"] = num_inference_steps_for_strength - inputs["strength"] = 0.5 - - outputs = [] - for scheduler_enum in KarrasDiffusionSchedulers: - if "KDPM2" in scheduler_enum.name: - inputs["num_inference_steps"] = num_inference_steps_for_strength_for_iterations - - scheduler_cls = getattr(diffusers, scheduler_enum.name) - pipe.scheduler = scheduler_cls.from_config(pipe.scheduler.config) - output = pipe(**inputs)[0] - outputs.append(output) - - if "KDPM2" in scheduler_enum.name: - inputs["num_inference_steps"] = 2 - - assert check_same_shape(outputs) - - -@require_torch -class PipelineTesterMixin: - """ - This mixin is designed to be used with unittest.TestCase classes. - It provides a set of common tests for each PyTorch pipeline, e.g. saving and loading the pipeline, - equivalence of dict and tuple outputs, etc. - """ - - # Canonical parameters that are passed to `__call__` regardless - # of the type of pipeline. They are always optional and have common - # sense default values. - required_optional_params = frozenset( - [ - "num_inference_steps", - "num_images_per_prompt", - "generator", - "latents", - "output_type", - "return_dict", - ] - ) - - # set these parameters to False in the child class if the pipeline does not support the corresponding functionality - test_attention_slicing = True - - test_xformers_attention = True - test_layerwise_casting = False - test_group_offloading = False - - # Components that cannot be offloaded at leaf level. See `BasePipelineTesterConfig` in - # `tests/pipelines/testing_utils/common.py`, which declares the same attribute, for the rationale. - group_offloading_leaf_level_exclude_modules = [] - - def get_generator(self, seed): - device = torch_device if torch_device != "mps" else "cpu" - generator = torch.Generator(device).manual_seed(seed) - return generator - - @property - def pipeline_class(self) -> Callable | DiffusionPipeline: - raise NotImplementedError( - "You need to set the attribute `pipeline_class = ClassNameOfPipeline` in the child test class. " - "See existing pipeline tests for reference." - ) - - def get_dummy_components(self): - raise NotImplementedError( - "You need to implement `get_dummy_components(self)` in the child test class. " - "See existing pipeline tests for reference." - ) - - def get_dummy_inputs(self, device, seed=0): - raise NotImplementedError( - "You need to implement `get_dummy_inputs(self, device, seed)` in the child test class. " - "See existing pipeline tests for reference." - ) - - @property - def params(self) -> frozenset: - raise NotImplementedError( - "You need to set the attribute `params` in the child test class. " - "`params` are checked for if all values are present in `__call__`'s signature." - " You can set `params` using one of the common set of parameters defined in `pipeline_params.py`" - " e.g., `TEXT_TO_IMAGE_PARAMS` defines the common parameters used in text to " - "image pipelines, including prompts and prompt embedding overrides." - "If your pipeline's set of arguments has minor changes from one of the common sets of arguments, " - "do not make modifications to the existing common sets of arguments. I.e. a text to image pipeline " - "with non-configurable height and width arguments should set the attribute as " - "`params = TEXT_TO_IMAGE_PARAMS - {'height', 'width'}`. " - "See existing pipeline tests for reference." - ) - - @property - def batch_params(self) -> frozenset: - raise NotImplementedError( - "You need to set the attribute `batch_params` in the child test class. " - "`batch_params` are the parameters required to be batched when passed to the pipeline's " - "`__call__` method. `pipeline_params.py` provides some common sets of parameters such as " - "`TEXT_TO_IMAGE_BATCH_PARAMS`, `IMAGE_VARIATION_BATCH_PARAMS`, etc... If your pipeline's " - "set of batch arguments has minor changes from one of the common sets of batch arguments, " - "do not make modifications to the existing common sets of batch arguments. I.e. a text to " - "image pipeline `negative_prompt` is not batched should set the attribute as " - "`batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - {'negative_prompt'}`. " - "See existing pipeline tests for reference." - ) - - @property - def callback_cfg_params(self) -> frozenset: - raise NotImplementedError( - "You need to set the attribute `callback_cfg_params` in the child test class that requires to run test_callback_cfg. " - "`callback_cfg_params` are the parameters that needs to be passed to the pipeline's callback " - "function when dynamically adjusting `guidance_scale`. They are variables that require special" - "treatment when `do_classifier_free_guidance` is `True`. `pipeline_params.py` provides some common" - " sets of parameters such as `TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS`. If your pipeline's " - "set of cfg arguments has minor changes from one of the common sets of cfg arguments, " - "do not make modifications to the existing common sets of cfg arguments. I.e. for inpaint pipeline, you " - " need to adjust batch size of `mask` and `masked_image_latents` so should set the attribute as" - "`callback_cfg_params = TEXT_TO_IMAGE_CFG_PARAMS.union({'mask', 'masked_image_latents'})`" - ) - - def setUp(self): - # clean up the VRAM before each test - super().setUp() - torch.compiler.reset() - gc.collect() - backend_empty_cache(torch_device) - - # Skip tests for pipelines that inherit from DeprecatedPipelineMixin - from diffusers.pipelines.pipeline_utils import DeprecatedPipelineMixin - - if hasattr(self, "pipeline_class") and issubclass(self.pipeline_class, DeprecatedPipelineMixin): - import pytest - - pytest.skip(reason=f"Deprecated Pipeline: {self.pipeline_class.__name__}") - - def tearDown(self): - # clean up the VRAM after each test in case of CUDA runtime errors - super().tearDown() - torch.compiler.reset() - gc.collect() - backend_empty_cache(torch_device) - - def test_save_load_local(self, expected_max_difference=5e-4): - components = self.get_dummy_components() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].eval() - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - 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) - - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - - 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) - - def test_pipeline_call_signature(self): - self.assertTrue( - hasattr(self.pipeline_class, "__call__"), f"{self.pipeline_class} should have a `__call__` method" - ) - - parameters = inspect.signature(self.pipeline_class.__call__).parameters - - optional_parameters = set() - - for k, v in parameters.items(): - if v.default != inspect._empty: - optional_parameters.add(k) - - parameters = set(parameters.keys()) - parameters.remove("self") - parameters.discard("kwargs") # kwargs can be added if arguments of pipeline call function are deprecated - - remaining_required_parameters = set() - - for param in self.params: - if param not in parameters: - remaining_required_parameters.add(param) - - self.assertTrue( - len(remaining_required_parameters) == 0, - f"Required parameters not present: {remaining_required_parameters}", - ) - - remaining_required_optional_parameters = set() - - for param in self.required_optional_params: - if param not in optional_parameters: - remaining_required_optional_parameters.add(param) - - self.assertTrue( - len(remaining_required_optional_parameters) == 0, - f"Required optional parameters not present: {remaining_required_optional_parameters}", - ) - - def test_inference_batch_consistent(self, batch_sizes=[2]): - self._test_inference_batch_consistent(batch_sizes=batch_sizes) - - def _test_inference_batch_consistent( - self, batch_sizes=[2], additional_params_copy_to_batched_inputs=["num_inference_steps"], batch_generator=True - ): - 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) - inputs["generator"] = self.get_generator(0) - - logger = logging.get_logger(pipe.__module__) - logger.setLevel(level=diffusers.logging.FATAL) - - # prepare batched inputs - batched_inputs = [] - for batch_size in batch_sizes: - batched_input = {} - batched_input.update(inputs) - - for name in self.batch_params: - if name not in inputs: - continue - - value = inputs[name] - if name == "prompt": - len_prompt = len(value) - # make unequal batch sizes - batched_input[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)] - - # make last batch super long - batched_input[name][-1] = 100 * "very long" - - else: - batched_input[name] = batch_size * [value] - - if batch_generator and "generator" in inputs: - batched_input["generator"] = [self.get_generator(i) for i in range(batch_size)] - - if "batch_size" in inputs: - batched_input["batch_size"] = batch_size - - batched_inputs.append(batched_input) - - logger.setLevel(level=diffusers.logging.WARNING) - for batch_size, batched_input in zip(batch_sizes, batched_inputs): - output = pipe(**batched_input) - assert len(output[0]) == batch_size - - def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-4): - self._test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) - - 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() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].eval() - pipe = self.pipeline_class(**components) - for components in pipe.components.values(): - if hasattr(components, "set_default_attn_processor"): - components.set_default_attn_processor() - - 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 - - def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4): - components = self.get_dummy_components() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].eval() - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - generator_device = "cpu" - 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) - - 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))} - - pipe = self.pipeline_class(**init_components) - - self.assertTrue(hasattr(pipe, "components")) - self.assertTrue(set(pipe.components.keys()) == set(init_components.keys())) - - @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU") - @require_accelerator - def test_float16_inference(self, expected_max_diff=5e-2): - 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) - - components = self.get_dummy_components() - pipe_fp16 = self.pipeline_class(**components) - for component in pipe_fp16.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_fp16.to(torch_device, torch.float16) - pipe_fp16.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - # Reset generator in case it is used inside dummy inputs - if "generator" in inputs: - inputs["generator"] = self.get_generator(0) - output = pipe(**inputs)[0] - - fp16_inputs = self.get_dummy_inputs(torch_device) - # Reset generator in case it is used inside dummy inputs - if "generator" in fp16_inputs: - fp16_inputs["generator"] = self.get_generator(0) - output_fp16 = pipe_fp16(**fp16_inputs)[0] - - if isinstance(output, torch.Tensor): - output = output.cpu() - output_fp16 = output_fp16.cpu() - - max_diff = numpy_cosine_similarity_distance(output.flatten(), output_fp16.flatten()) - assert max_diff < expected_max_diff - - @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU") - @require_accelerator - def test_save_load_float16(self, expected_max_diff=1e-2): - components = self.get_dummy_components() - for name, module in components.items(): - # Account for components with _keep_in_fp32_modules - if hasattr(module, "_keep_in_fp32_modules") and module._keep_in_fp32_modules is not None: - for name, param in module.named_parameters(): - if any( - module_to_keep_in_fp32 in name.split(".") - for module_to_keep_in_fp32 in module._keep_in_fp32_modules - ): - param.data = param.data.to(torch_device).to(torch.float32) - else: - param.data = param.data.to(torch_device).to(torch.float16) - for name, buf in module.named_buffers(): - if not buf.is_floating_point(): - buf.data = buf.data.to(torch_device) - elif any( - module_to_keep_in_fp32 in name.split(".") - for module_to_keep_in_fp32 in module._keep_in_fp32_modules - ): - buf.data = buf.data.to(torch_device).to(torch.float32) - else: - buf.data = buf.data.to(torch_device).to(torch.float16) - - elif hasattr(module, "half"): - components[name] = module.to(torch_device).half() - - for key, component in components.items(): - if hasattr(component, "eval"): - component.eval() - - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs)[0] - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir) - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, dtype=torch.float16) - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - for name, component in pipe_loaded.components.items(): - if hasattr(component, "dtype"): - self.assertTrue( - component.dtype == torch.float16, - f"`{name}.dtype` switched from `float16` to {component.dtype} after loading.", - ) - - 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_diff, "The output of the fp16 pipeline changed after saving and loading." - ) - - def test_save_load_optional_components(self, expected_max_difference=1e-4): - if not hasattr(self.pipeline_class, "_optional_components"): - return - if not self.pipeline_class._optional_components: - return - components = self.get_dummy_components() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].eval() - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - # set all optional components to None - for optional_component in pipe._optional_components: - setattr(pipe, optional_component, None) - - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - torch.manual_seed(0) - output = pipe(**inputs)[0] - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, safe_serialization=False) - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir) - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - 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.", - ) - - inputs = self.get_dummy_inputs(generator_device) - torch.manual_seed(0) - 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) - - @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") - model_devices = [ - component.device.type for component in components.values() if getattr(component, "device", None) - ] - 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 components.values() if getattr(component, "device", None) - ] - 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) - - model_dtypes = [component.dtype for component in components.values() if getattr(component, "dtype", None)] - self.assertTrue(all(dtype == torch.float32 for dtype in model_dtypes)) - - pipe.to(dtype=torch.float16) - model_dtypes = [component.dtype for component in components.values() if getattr(component, "dtype", None)] - self.assertTrue(all(dtype == torch.float16 for dtype in model_dtypes)) - - def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3): - self._test_attention_slicing_forward_pass(expected_max_diff=expected_max_diff) - - def _test_attention_slicing_forward_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 - ): - if not self.test_attention_slicing: - return - - components = self.get_dummy_components() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].eval() - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - output_without_slicing = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=1) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing1 = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=2) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing2 = pipe(**inputs)[0] - - if test_max_difference: - max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max() - max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max() - self.assertLess( - max(max_diff1, max_diff2), - expected_max_diff, - "Attention slicing should not affect the inference results", - ) - - 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])) - - @require_accelerator - @require_accelerate_version_greater("0.14.0") - def test_sequential_cpu_offload_forward_pass(self, expected_max_diff=1e-4): - import accelerate - - 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) - torch.manual_seed(0) - output_without_offload = pipe(**inputs)[0] - - pipe.enable_sequential_cpu_offload(device=torch_device) - assert pipe._execution_device.type == torch_device - - inputs = self.get_dummy_inputs(generator_device) - torch.manual_seed(0) - output_with_offload = pipe(**inputs)[0] - - max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max() - self.assertLess(max_diff, expected_max_diff, "CPU offloading should not affect the inference results") - - # make sure all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are offloaded correctly - offloaded_modules = { - k: v - for k, v in pipe.components.items() - if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload - } - # 1. all offloaded modules should be saved to cpu and moved to meta device - self.assertTrue( - all(v.device.type == "meta" for v in offloaded_modules.values()), - f"Not offloaded: {[k for k, v in offloaded_modules.items() if v.device.type != 'meta']}", - ) - # 2. all offloaded modules should have hook installed - self.assertTrue( - all(hasattr(v, "_hf_hook") for k, v in offloaded_modules.items()), - f"No hook attached: {[k for k, v in offloaded_modules.items() if not hasattr(v, '_hf_hook')]}", - ) - # 3. all offloaded modules should have correct hooks installed, should be either one of these two - # - `AlignDevicesHook` - # - a SequentialHook` that contains `AlignDevicesHook` - offloaded_modules_with_incorrect_hooks = {} - for k, v in offloaded_modules.items(): - if hasattr(v, "_hf_hook"): - if isinstance(v._hf_hook, accelerate.hooks.SequentialHook): - # if it is a `SequentialHook`, we loop through its `hooks` attribute to check if it only contains `AlignDevicesHook` - for hook in v._hf_hook.hooks: - if not isinstance(hook, accelerate.hooks.AlignDevicesHook): - offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook.hooks[0]) - elif not isinstance(v._hf_hook, accelerate.hooks.AlignDevicesHook): - offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook) - - self.assertTrue( - len(offloaded_modules_with_incorrect_hooks) == 0, - f"Not installed correct hook: {offloaded_modules_with_incorrect_hooks}", - ) - - @require_accelerator - @require_accelerate_version_greater("0.17.0") - def test_model_cpu_offload_forward_pass(self, expected_max_diff=2e-4): - import accelerate - - generator_device = "cpu" - 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 = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(generator_device) - torch.manual_seed(0) - output_without_offload = pipe(**inputs)[0] - - pipe.enable_model_cpu_offload(device=torch_device) - assert pipe._execution_device.type == torch_device - - inputs = self.get_dummy_inputs(generator_device) - torch.manual_seed(0) - output_with_offload = pipe(**inputs)[0] - - max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max() - self.assertLess(max_diff, expected_max_diff, "CPU offloading should not affect the inference results") - - # make sure all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are offloaded correctly - offloaded_modules = { - k: v - for k, v in pipe.components.items() - if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload - } - # 1. check if all offloaded modules are saved to cpu - self.assertTrue( - all(v.device.type == "cpu" for v in offloaded_modules.values()), - f"Not offloaded: {[k for k, v in offloaded_modules.items() if v.device.type != 'cpu']}", - ) - # 2. check if all offloaded modules have hooks installed - self.assertTrue( - all(hasattr(v, "_hf_hook") for k, v in offloaded_modules.items()), - f"No hook attached: {[k for k, v in offloaded_modules.items() if not hasattr(v, '_hf_hook')]}", - ) - # 3. check if all offloaded modules have correct type of hooks installed, should be `CpuOffload` - offloaded_modules_with_incorrect_hooks = {} - for k, v in offloaded_modules.items(): - if hasattr(v, "_hf_hook") and not isinstance(v._hf_hook, accelerate.hooks.CpuOffload): - offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook) - - self.assertTrue( - len(offloaded_modules_with_incorrect_hooks) == 0, - f"Not installed correct hook: {offloaded_modules_with_incorrect_hooks}", - ) - - @require_accelerator - @require_accelerate_version_greater("0.17.0") - def test_cpu_offload_forward_pass_twice(self, expected_max_diff=2e-4): - import accelerate - - generator_device = "cpu" - 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.set_progress_bar_config(disable=None) - - pipe.enable_model_cpu_offload() - inputs = self.get_dummy_inputs(generator_device) - output_with_offload = pipe(**inputs)[0] - - pipe.enable_model_cpu_offload() - inputs = self.get_dummy_inputs(generator_device) - output_with_offload_twice = pipe(**inputs)[0] - - max_diff = np.abs(to_np(output_with_offload) - to_np(output_with_offload_twice)).max() - self.assertLess( - max_diff, expected_max_diff, "running CPU offloading 2nd time should not affect the inference results" - ) - - # make sure all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are offloaded correctly - offloaded_modules = { - k: v - for k, v in pipe.components.items() - if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload - } - # 1. check if all offloaded modules are saved to cpu - self.assertTrue( - all(v.device.type == "cpu" for v in offloaded_modules.values()), - f"Not offloaded: {[k for k, v in offloaded_modules.items() if v.device.type != 'cpu']}", - ) - # 2. check if all offloaded modules have hooks installed - self.assertTrue( - all(hasattr(v, "_hf_hook") for k, v in offloaded_modules.items()), - f"No hook attached: {[k for k, v in offloaded_modules.items() if not hasattr(v, '_hf_hook')]}", - ) - # 3. check if all offloaded modules have correct type of hooks installed, should be `CpuOffload` - offloaded_modules_with_incorrect_hooks = {} - for k, v in offloaded_modules.items(): - if hasattr(v, "_hf_hook") and not isinstance(v._hf_hook, accelerate.hooks.CpuOffload): - offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook) - - self.assertTrue( - len(offloaded_modules_with_incorrect_hooks) == 0, - f"Not installed correct hook: {offloaded_modules_with_incorrect_hooks}", - ) - - @require_accelerator - @require_accelerate_version_greater("0.14.0") - def test_sequential_offload_forward_pass_twice(self, expected_max_diff=2e-4): - import accelerate - - generator_device = "cpu" - 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.set_progress_bar_config(disable=None) - - pipe.enable_sequential_cpu_offload(device=torch_device) - inputs = self.get_dummy_inputs(generator_device) - output_with_offload = pipe(**inputs)[0] - - pipe.enable_sequential_cpu_offload(device=torch_device) - inputs = self.get_dummy_inputs(generator_device) - output_with_offload_twice = pipe(**inputs)[0] - - max_diff = np.abs(to_np(output_with_offload) - to_np(output_with_offload_twice)).max() - self.assertLess( - max_diff, expected_max_diff, "running sequential offloading second time should have the inference results" - ) - - # make sure all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are offloaded correctly - offloaded_modules = { - k: v - for k, v in pipe.components.items() - if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload - } - # 1. check if all offloaded modules are moved to meta device - self.assertTrue( - all(v.device.type == "meta" for v in offloaded_modules.values()), - f"Not offloaded: {[k for k, v in offloaded_modules.items() if v.device.type != 'meta']}", - ) - # 2. check if all offloaded modules have hook installed - self.assertTrue( - all(hasattr(v, "_hf_hook") for k, v in offloaded_modules.items()), - f"No hook attached: {[k for k, v in offloaded_modules.items() if not hasattr(v, '_hf_hook')]}", - ) - # 3. check if all offloaded modules have correct hooks installed, should be either one of these two - # - `AlignDevicesHook` - # - a SequentialHook` that contains `AlignDevicesHook` - offloaded_modules_with_incorrect_hooks = {} - for k, v in offloaded_modules.items(): - if hasattr(v, "_hf_hook"): - if isinstance(v._hf_hook, accelerate.hooks.SequentialHook): - # if it is a `SequentialHook`, we loop through its `hooks` attribute to check if it only contains `AlignDevicesHook` - for hook in v._hf_hook.hooks: - if not isinstance(hook, accelerate.hooks.AlignDevicesHook): - offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook.hooks[0]) - elif not isinstance(v._hf_hook, accelerate.hooks.AlignDevicesHook): - offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook) - - self.assertTrue( - len(offloaded_modules_with_incorrect_hooks) == 0, - f"Not installed correct hook: {offloaded_modules_with_incorrect_hooks}", - ) - - @unittest.skipIf( - torch_device != "cuda" or not is_xformers_available(), - reason="XFormers attention is only available with CUDA and `xformers` installed", - ) - def test_xformers_attention_forwardGenerator_pass(self): - self._test_xformers_attention_forwardGenerator_pass() - - def _test_xformers_attention_forwardGenerator_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-4 - ): - if not self.test_xformers_attention: - 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) - - inputs = self.get_dummy_inputs(torch_device) - output_without_offload = pipe(**inputs)[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)[0] - output_with_offload = ( - output_with_offload.cpu() if torch.is_tensor(output_with_offload) else output_without_offload - ) - - if test_max_difference: - max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max() - self.assertLess(max_diff, expected_max_diff, "XFormers attention should not affect the inference results") - - if test_mean_pixel_difference: - assert_mean_pixel_difference(output_with_offload[0], output_without_offload[0]) - - def test_num_images_per_prompt(self): - sig = inspect.signature(self.pipeline_class.__call__) - - if "num_images_per_prompt" not in sig.parameters: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - batch_sizes = [1, 2] - num_images_per_prompts = [1, 2] - - for batch_size in batch_sizes: - for num_images_per_prompt in num_images_per_prompts: - inputs = self.get_dummy_inputs(torch_device) - - for key in inputs.keys(): - if key in self.batch_params: - inputs[key] = batch_size * [inputs[key]] - - images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt)[0] - - assert images.shape[0] == batch_size * num_images_per_prompt - - def test_cfg(self): - sig = inspect.signature(self.pipeline_class.__call__) - - if "guidance_scale" not in sig.parameters: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - - inputs["guidance_scale"] = 1.0 - out_no_cfg = pipe(**inputs)[0] - - inputs["guidance_scale"] = 7.5 - out_cfg = pipe(**inputs)[0] - - assert out_cfg.shape == out_no_cfg.shape - - 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"] - inputs["output_type"] = "latent" - 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 - inputs["output_type"] = "latent" - 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 - inputs["output_type"] = "latent" - output = pipe(**inputs)[0] - assert output.abs().sum() == 0 - - def test_callback_cfg(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 - - if "guidance_scale" not in sig.parameters: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - 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_increase_guidance(pipe, i, t, callback_kwargs): - pipe._guidance_scale += 1.0 - - return callback_kwargs - - inputs = self.get_dummy_inputs(torch_device) - - # use cfg guidance because some pipelines modify the shape of the latents - # outside of the denoising loop - inputs["guidance_scale"] = 2.0 - inputs["callback_on_step_end"] = callback_increase_guidance - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - _ = pipe(**inputs)[0] - - # we increase the guidance scale by 1.0 at every step - # check that the guidance scale is increased by the number of scheduler timesteps - # accounts for models that modify the number of inference steps based on strength - assert pipe.guidance_scale == (inputs["guidance_scale"] + pipe.num_timesteps) - - def test_serialization_with_variants(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - model_components = [ - component_name for component_name, component in pipe.components.items() if isinstance(component, nn.Module) - ] - variant = "fp16" - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False) - - with open(f"{tmpdir}/model_index.json", "r") as f: - config = json.load(f) - - for subfolder in os.listdir(tmpdir): - if not os.path.isfile(subfolder) and subfolder in model_components: - folder_path = os.path.join(tmpdir, subfolder) - is_folder = os.path.isdir(folder_path) and subfolder in config - assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - - def test_loading_with_variants(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - variant = "fp16" - - def is_nan(tensor): - if tensor.ndimension() == 0: - has_nan = torch.isnan(tensor).item() - else: - has_nan = torch.isnan(tensor).any() - return has_nan - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False) - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, variant=variant) - - model_components_pipe = { - component_name: component - for component_name, component in pipe.components.items() - if isinstance(component, nn.Module) - } - model_components_pipe_loaded = { - component_name: component - for component_name, component in pipe_loaded.components.items() - if isinstance(component, nn.Module) - } - for component_name in model_components_pipe: - pipe_component = model_components_pipe[component_name] - pipe_loaded_component = model_components_pipe_loaded[component_name] - - model_loaded_params = dict(pipe_loaded_component.named_parameters()) - model_original_params = dict(pipe_component.named_parameters()) - - for name, p1 in model_original_params.items(): - # Skip tied weights that aren't saved with variants (transformers v5 behavior) - if name not in model_loaded_params: - continue - - p2 = model_loaded_params[name] - # nan check for luminanext (mps). - if not (is_nan(p1) and is_nan(p2)): - self.assertTrue(torch.equal(p1, p2)) - - def test_loading_with_incorrect_variants_raises_error(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - variant = "fp16" - - with tempfile.TemporaryDirectory() as tmpdir: - # Don't save with variants. - pipe.save_pretrained(tmpdir, safe_serialization=False) - - with self.assertRaises(ValueError) as error: - _ = self.pipeline_class.from_pretrained(tmpdir, variant=variant) - - assert f"You are trying to load the model files of the `variant={variant}`" in str(error.exception) - - def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4): - if not hasattr(self.pipeline_class, "encode_prompt"): - return - - components = self.get_dummy_components() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].eval() - - # We initialize the pipeline with only text encoders and tokenizers, - # mimicking a real-world scenario. - components_with_text_encoders = {} - for k in components: - if "text" in k or "tokenizer" in k: - components_with_text_encoders[k] = components[k] - else: - components_with_text_encoders[k] = None - pipe_with_just_text_encoder = self.pipeline_class(**components_with_text_encoders) - pipe_with_just_text_encoder = pipe_with_just_text_encoder.to(torch_device) - - # Get inputs and also the args of `encode_prompts`. - inputs = self.get_dummy_inputs(torch_device) - encode_prompt_signature = inspect.signature(pipe_with_just_text_encoder.encode_prompt) - encode_prompt_parameters = list(encode_prompt_signature.parameters.values()) - - # Required args in encode_prompt with those with no default. - required_params = [] - for param in encode_prompt_parameters: - if param.name == "self" or param.name == "kwargs": - continue - if param.default is inspect.Parameter.empty: - required_params.append(param.name) - - # Craft inputs for the `encode_prompt()` method to run in isolation. - encode_prompt_param_names = [p.name for p in encode_prompt_parameters if p.name != "self"] - input_keys = list(inputs.keys()) - encode_prompt_inputs = {k: inputs.pop(k) for k in input_keys if k in encode_prompt_param_names} - - pipe_call_signature = inspect.signature(pipe_with_just_text_encoder.__call__) - pipe_call_parameters = pipe_call_signature.parameters - - # For each required arg in encode_prompt, check if it's missing - # in encode_prompt_inputs. If so, see if __call__ has a default - # for that arg and use it if available. - for required_param_name in required_params: - if required_param_name not in encode_prompt_inputs: - pipe_call_param = pipe_call_parameters.get(required_param_name, None) - if pipe_call_param is not None and pipe_call_param.default is not inspect.Parameter.empty: - # Use the default from pipe.__call__ - encode_prompt_inputs[required_param_name] = pipe_call_param.default - elif extra_required_param_value_dict is not None and isinstance(extra_required_param_value_dict, dict): - encode_prompt_inputs[required_param_name] = extra_required_param_value_dict[required_param_name] - else: - raise ValueError( - f"Required parameter '{required_param_name}' in " - f"encode_prompt has no default in either encode_prompt or __call__." - ) - - # Compute `encode_prompt()`. - with torch.no_grad(): - encoded_prompt_outputs = pipe_with_just_text_encoder.encode_prompt(**encode_prompt_inputs) - - # Programmatically determine the return names of `encode_prompt.` - ast_visitor = ReturnNameVisitor() - encode_prompt_tree = ast_visitor.get_ast_tree(cls=self.pipeline_class) - ast_visitor.visit(encode_prompt_tree) - prompt_embed_kwargs = ast_visitor.return_names - prompt_embeds_kwargs = dict(zip(prompt_embed_kwargs, encoded_prompt_outputs)) - - # Pack the outputs of `encode_prompt`. - adapted_prompt_embeds_kwargs = { - k: prompt_embeds_kwargs.pop(k) for k in list(prompt_embeds_kwargs.keys()) if k in pipe_call_parameters - } - - # now initialize a pipeline without text encoders and compute outputs with the - # `encode_prompt()` outputs and other relevant inputs. - components_with_text_encoders = {} - for k in components: - if "text" in k or "tokenizer" in k: - components_with_text_encoders[k] = None - else: - components_with_text_encoders[k] = components[k] - pipe_without_text_encoders = self.pipeline_class(**components_with_text_encoders).to(torch_device) - - # Set `negative_prompt` to None as we have already calculated its embeds - # if it was present in `inputs`. This is because otherwise we will interfere wrongly - # for non-None `negative_prompt` values as defaults (PixArt for example). - pipe_without_tes_inputs = {**inputs, **adapted_prompt_embeds_kwargs} - if ( - pipe_call_parameters.get("negative_prompt", None) is not None - and pipe_call_parameters.get("negative_prompt").default is not None - ): - pipe_without_tes_inputs.update({"negative_prompt": None}) - - # Pipelines like attend and excite have `prompt` as a required argument. - if ( - pipe_call_parameters.get("prompt", None) is not None - and pipe_call_parameters.get("prompt").default is inspect.Parameter.empty - and pipe_call_parameters.get("prompt_embeds", None) is not None - and pipe_call_parameters.get("prompt_embeds").default is None - ): - pipe_without_tes_inputs.update({"prompt": None}) - - pipe_out = pipe_without_text_encoders(**pipe_without_tes_inputs)[0] - - # Compare against regular pipeline outputs. - full_pipe = self.pipeline_class(**components).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) - pipe_out_2 = full_pipe(**inputs)[0] - - if isinstance(pipe_out, np.ndarray) and isinstance(pipe_out_2, np.ndarray): - self.assertTrue(np.allclose(pipe_out, pipe_out_2, atol=atol, rtol=rtol)) - elif isinstance(pipe_out, torch.Tensor) and isinstance(pipe_out_2, torch.Tensor): - self.assertTrue(torch.allclose(pipe_out, pipe_out_2, atol=atol, rtol=rtol)) - - def test_StableDiffusionMixin_component(self): - """Any pipeline that have LDMFuncMixin should have vae and unet components.""" - if not issubclass(self.pipeline_class, StableDiffusionMixin): - return - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - self.assertTrue(hasattr(pipe, "vae") and isinstance(pipe.vae, (AutoencoderKL, AutoencoderTiny))) - self.assertTrue( - hasattr(pipe, "unet") - and isinstance( - pipe.unet, - (UNet2DConditionModel, UNet3DConditionModel, I2VGenXLUNet, UNetMotionModel, UNetControlNetXSModel), - ) - ) - - @pytest.mark.xfail( - condition=torch_device == "mps", - reason="MPS does not support float8 casting.", - strict=True, - ) - def test_layerwise_casting_inference(self): - if not self.test_layerwise_casting: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device, dtype=torch.bfloat16) - pipe.set_progress_bar_config(disable=None) - - denoiser = pipe.transformer if hasattr(pipe, "transformer") else pipe.unet - denoiser.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16) - - inputs = self.get_dummy_inputs(torch_device) - _ = pipe(**inputs)[0] - - @require_torch_accelerator - def test_group_offloading_inference(self): - if not self.test_group_offloading: - return - - def create_pipe(): - torch.manual_seed(0) - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - return pipe - - def enable_group_offload_on_component(pipe, group_offloading_kwargs): - # We intentionally don't test VAE's here. This is because some tests enable tiling on the VAE. If - # tiling is enabled and a forward pass is run, when accelerator streams are used, the execution order of - # the layers is not traced correctly. This causes errors. For apply group offloading to VAE, a - # warmup forward pass (even with dummy small inputs) is recommended. - for component_name in [ - "text_encoder", - "text_encoder_2", - "text_encoder_3", - "transformer", - "unet", - "controlnet", - ]: - if not hasattr(pipe, component_name): - continue - component = getattr(pipe, component_name) - if not getattr(component, "_supports_group_offloading", True): - continue - if hasattr(component, "enable_group_offload"): - # For diffusers ModelMixin implementations - component.enable_group_offload(torch.device(torch_device), **group_offloading_kwargs) - else: - # For other models not part of diffusers - apply_group_offloading( - component, onload_device=torch.device(torch_device), **group_offloading_kwargs - ) - self.assertTrue( - all( - module._diffusers_hook.get_hook("group_offloading") is not None - for module in component.modules() - if hasattr(module, "_diffusers_hook") - ) - ) - for component_name in ["vae", "vqvae", "image_encoder"]: - component = getattr(pipe, component_name, None) - if isinstance(component, torch.nn.Module): - component.to(torch_device) - - def run_forward(pipe): - torch.manual_seed(0) - inputs = self.get_dummy_inputs(torch_device) - return pipe(**inputs)[0] - - pipe = create_pipe().to(torch_device) - output_without_group_offloading = run_forward(pipe) - - pipe = create_pipe() - enable_group_offload_on_component(pipe, {"offload_type": "block_level", "num_blocks_per_group": 1}) - output_with_group_offloading1 = run_forward(pipe) - - pipe = create_pipe() - enable_group_offload_on_component(pipe, {"offload_type": "leaf_level"}) - output_with_group_offloading2 = run_forward(pipe) - - if torch.is_tensor(output_without_group_offloading): - output_without_group_offloading = output_without_group_offloading.detach().cpu().numpy() - output_with_group_offloading1 = output_with_group_offloading1.detach().cpu().numpy() - output_with_group_offloading2 = output_with_group_offloading2.detach().cpu().numpy() - - self.assertTrue(np.allclose(output_without_group_offloading, output_with_group_offloading1, atol=1e-4)) - self.assertTrue(np.allclose(output_without_group_offloading, output_with_group_offloading2, atol=1e-4)) - - def test_dtype_dict(self): - components = self.get_dummy_components() - if not components: - self.skipTest("No dummy components defined.") - - pipe = self.pipeline_class(**components) - specified_key = next(iter(components.keys())) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, dtype=dtype_dict) - - for name, component in loaded_pipe.components.items(): - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) - self.assertEqual( - component.dtype, - expected_dtype, - f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}", - ) - - def test_dtype_alias(self): - # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - loaded_pipe = self.pipeline_class.from_pretrained(tmpdirname, dtype=torch.float16) - - for name, component in loaded_pipe.components.items(): - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - self.assertEqual( - component.dtype, - torch.float16, - f"Component '{name}' has dtype {component.dtype} but expected {torch.float16}", - ) - - def test_pipeline_with_accelerator_device_map(self, expected_max_difference=1e-4): - components = self.get_dummy_components() - # Set text encoders to eval mode to match from_pretrained behavior - # This ensures deterministic outputs when models are loaded with device_map - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].eval() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - torch.manual_seed(0) - inputs = self.get_dummy_inputs(torch_device) - inputs["generator"] = torch.manual_seed(0) - out = pipe(**inputs)[0] - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir) - loaded_pipe = self.pipeline_class.from_pretrained(tmpdir, device_map=torch_device) - for component in loaded_pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - inputs["generator"] = torch.manual_seed(0) - loaded_out = loaded_pipe(**inputs)[0] - max_diff = np.abs(to_np(out) - to_np(loaded_out)).max() - self.assertLess(max_diff, expected_max_difference) - - @require_torch_accelerator - def test_pipeline_level_group_offloading_sanity_checks(self): - components = self.get_dummy_components() - pipe: DiffusionPipeline = self.pipeline_class(**components) - - for name, component in pipe.components.items(): - if hasattr(component, "_supports_group_offloading"): - if not component._supports_group_offloading: - pytest.skip(f"{self.pipeline_class.__name__} is not suitable for this test.") - - module_names = sorted( - [name for name, component in pipe.components.items() if isinstance(component, torch.nn.Module)] - ) - exclude_module_name = module_names[0] - offload_device = "cpu" - pipe.enable_group_offload( - onload_device=torch_device, - offload_device=offload_device, - offload_type="leaf_level", - exclude_modules=exclude_module_name, - ) - excluded_module = getattr(pipe, exclude_module_name) - self.assertTrue(torch.device(excluded_module.device).type == torch.device(torch_device).type) - - for name, component in pipe.components.items(): - if name not in [exclude_module_name] and isinstance(component, torch.nn.Module): - # `component.device` prints the `onload_device` type. We should probably override the - # `device` property in `ModelMixin`. - # Skip modules with no parameters (e.g., dummy safety checkers with only buffers) - params = list(component.parameters()) - if not params: - continue - component_device = params[0].device - self.assertTrue(torch.device(component_device).type == torch.device(offload_device).type) - - @require_torch_accelerator - def test_pipeline_level_group_offloading_inference(self, expected_max_difference=1e-4): - components = self.get_dummy_components() - pipe: DiffusionPipeline = self.pipeline_class(**components) - - for name, component in pipe.components.items(): - if hasattr(component, "_supports_group_offloading"): - if not component._supports_group_offloading: - pytest.skip(f"{self.pipeline_class.__name__} is not suitable for this test.") - - # Regular inference. - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - torch.manual_seed(0) - inputs = self.get_dummy_inputs(torch_device) - inputs["generator"] = torch.manual_seed(0) - out = pipe(**inputs)[0] - - pipe.to("cpu") - del pipe - - # Inference with offloading - pipe: DiffusionPipeline = self.pipeline_class(**components) - offload_device = "cpu" - pipe.enable_group_offload( - onload_device=torch_device, - offload_device=offload_device, - offload_type="leaf_level", - exclude_modules=self.group_offloading_leaf_level_exclude_modules, - ) - pipe.set_progress_bar_config(disable=None) - inputs["generator"] = torch.manual_seed(0) - out_offload = pipe(**inputs)[0] - - max_diff = np.abs(to_np(out) - to_np(out_offload)).max() - self.assertLess(max_diff, expected_max_difference) - - -class PyramidAttentionBroadcastTesterMixin: - pab_config = PyramidAttentionBroadcastConfig( - spatial_attention_block_skip_range=2, - spatial_attention_timestep_skip_range=(100, 800), - spatial_attention_block_identifiers=["transformer_blocks"], - ) - - def test_pyramid_attention_broadcast_layers(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - num_layers = 0 - num_single_layers = 0 - dummy_component_kwargs = {} - dummy_component_parameters = inspect.signature(self.get_dummy_components).parameters - if "num_layers" in dummy_component_parameters: - num_layers = 2 - dummy_component_kwargs["num_layers"] = num_layers - if "num_single_layers" in dummy_component_parameters: - num_single_layers = 2 - dummy_component_kwargs["num_single_layers"] = num_single_layers - - components = self.get_dummy_components(**dummy_component_kwargs) - pipe = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - - self.pab_config.current_timestep_callback = lambda: pipe.current_timestep - denoiser = pipe.transformer if hasattr(pipe, "transformer") else pipe.unet - denoiser.enable_cache(self.pab_config) - - expected_hooks = 0 - if self.pab_config.spatial_attention_block_skip_range is not None: - expected_hooks += num_layers + num_single_layers - if self.pab_config.temporal_attention_block_skip_range is not None: - expected_hooks += num_layers + num_single_layers - if self.pab_config.cross_attention_block_skip_range is not None: - expected_hooks += num_layers + num_single_layers - - denoiser = pipe.transformer if hasattr(pipe, "transformer") else pipe.unet - count = 0 - for module in denoiser.modules(): - if hasattr(module, "_diffusers_hook"): - hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast") - if hook is None: - continue - count += 1 - self.assertTrue( - isinstance(hook, PyramidAttentionBroadcastHook), - "Hook should be of type PyramidAttentionBroadcastHook.", - ) - self.assertTrue(hook.state.cache is None, "Cache should be None at initialization.") - self.assertEqual(count, expected_hooks, "Number of hooks should match the expected number.") - - # Perform dummy inference step to ensure state is updated - def pab_state_check_callback(pipe, i, t, kwargs): - for module in denoiser.modules(): - if hasattr(module, "_diffusers_hook"): - hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast") - if hook is None: - continue - self.assertTrue( - hook.state.cache is not None, - "Cache should have updated during inference.", - ) - self.assertTrue( - hook.state.iteration == i + 1, - "Hook iteration state should have updated during inference.", - ) - return {} - - inputs = self.get_dummy_inputs(device) - inputs["num_inference_steps"] = 2 - inputs["callback_on_step_end"] = pab_state_check_callback - pipe(**inputs)[0] - - # After inference, reset_stateful_hooks is called within the pipeline, which should have reset the states - for module in denoiser.modules(): - if hasattr(module, "_diffusers_hook"): - hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast") - if hook is None: - continue - self.assertTrue( - hook.state.cache is None, - "Cache should be reset to None after inference.", - ) - self.assertTrue( - hook.state.iteration == 0, - "Iteration should be reset to 0 after inference.", - ) - - def test_pyramid_attention_broadcast_inference(self, expected_atol: float = 0.2): - # We need to use higher tolerance because we are using a random model. With a converged/trained - # model, the tolerance can be lower. - - device = "cpu" # ensure determinism for the device-dependent torch.Generator - num_layers = 2 - components = self.get_dummy_components(num_layers=num_layers) - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].eval() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - # Run inference without PAB - inputs = self.get_dummy_inputs(device) - inputs["num_inference_steps"] = 4 - output = pipe(**inputs)[0] - original_image_slice = output.flatten() - original_image_slice = np.concatenate((original_image_slice[:8], original_image_slice[-8:])) - - # Run inference with PAB enabled - self.pab_config.current_timestep_callback = lambda: pipe.current_timestep - denoiser = pipe.transformer if hasattr(pipe, "transformer") else pipe.unet - denoiser.enable_cache(self.pab_config) - - inputs = self.get_dummy_inputs(device) - inputs["num_inference_steps"] = 4 - output = pipe(**inputs)[0] - image_slice_pab_enabled = output.flatten() - image_slice_pab_enabled = np.concatenate((image_slice_pab_enabled[:8], image_slice_pab_enabled[-8:])) - - # Run inference with PAB disabled - denoiser.disable_cache() - - inputs = self.get_dummy_inputs(device) - inputs["num_inference_steps"] = 4 - output = pipe(**inputs)[0] - image_slice_pab_disabled = output.flatten() - image_slice_pab_disabled = np.concatenate((image_slice_pab_disabled[:8], image_slice_pab_disabled[-8:])) - - assert np.allclose(original_image_slice, image_slice_pab_enabled, atol=expected_atol), ( - "PAB outputs should not differ much in specified timestep range." - ) - assert np.allclose(original_image_slice, image_slice_pab_disabled, atol=1e-4), ( - "Outputs from normal inference and after disabling cache should not differ." - ) - - -class FasterCacheTesterMixin: - faster_cache_config = FasterCacheConfig( - spatial_attention_block_skip_range=2, - spatial_attention_timestep_skip_range=(-1, 901), - unconditional_batch_skip_range=2, - attention_weight_callback=lambda _: 0.5, - ) - - def test_faster_cache_basic_warning_or_errors_raised(self): - components = self.get_dummy_components() - - logger = logging.get_logger("diffusers.hooks.faster_cache") - logger.setLevel(logging.INFO) - - # Check if warning is raise when no attention_weight_callback is provided - pipe = self.pipeline_class(**components) - with CaptureLogger(logger) as cap_logger: - config = FasterCacheConfig(spatial_attention_block_skip_range=2, attention_weight_callback=None) - apply_faster_cache(pipe.transformer, config) - self.assertTrue("No `attention_weight_callback` provided when enabling FasterCache" in cap_logger.out) - - # Check if error raised when unsupported tensor format used - pipe = self.pipeline_class(**components) - with self.assertRaises(ValueError): - config = FasterCacheConfig(spatial_attention_block_skip_range=2, tensor_format="BFHWC") - apply_faster_cache(pipe.transformer, config) - - def test_faster_cache_inference(self, expected_atol: float = 0.1): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - def create_pipe(): - torch.manual_seed(0) - num_layers = 2 - components = self.get_dummy_components(num_layers=num_layers) - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - return pipe - - def run_forward(pipe): - torch.manual_seed(0) - inputs = self.get_dummy_inputs(device) - inputs["num_inference_steps"] = 4 - return pipe(**inputs)[0] - - # Run inference without FasterCache - pipe = create_pipe() - output = run_forward(pipe).flatten() - original_image_slice = np.concatenate((output[:8], output[-8:])) - - # Run inference with FasterCache enabled - self.faster_cache_config.current_timestep_callback = lambda: pipe.current_timestep - pipe = create_pipe() - pipe.transformer.enable_cache(self.faster_cache_config) - output = run_forward(pipe).flatten() - image_slice_faster_cache_enabled = np.concatenate((output[:8], output[-8:])) - - # Run inference with FasterCache disabled - pipe.transformer.disable_cache() - output = run_forward(pipe).flatten() - image_slice_faster_cache_disabled = np.concatenate((output[:8], output[-8:])) - - assert np.allclose(original_image_slice, image_slice_faster_cache_enabled, atol=expected_atol), ( - "FasterCache outputs should not differ much in specified timestep range." - ) - assert np.allclose(original_image_slice, image_slice_faster_cache_disabled, atol=1e-4), ( - "Outputs from normal inference and after disabling cache should not differ." - ) - - def test_faster_cache_state(self): - from diffusers.hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK - - device = "cpu" # ensure determinism for the device-dependent torch.Generator - num_layers = 0 - num_single_layers = 0 - dummy_component_kwargs = {} - dummy_component_parameters = inspect.signature(self.get_dummy_components).parameters - if "num_layers" in dummy_component_parameters: - num_layers = 2 - dummy_component_kwargs["num_layers"] = num_layers - if "num_single_layers" in dummy_component_parameters: - num_single_layers = 2 - dummy_component_kwargs["num_single_layers"] = num_single_layers - - components = self.get_dummy_components(**dummy_component_kwargs) - pipe = self.pipeline_class(**components) - pipe.set_progress_bar_config(disable=None) - - self.faster_cache_config.current_timestep_callback = lambda: pipe.current_timestep - pipe.transformer.enable_cache(self.faster_cache_config) - - expected_hooks = 0 - if self.faster_cache_config.spatial_attention_block_skip_range is not None: - expected_hooks += num_layers + num_single_layers - if self.faster_cache_config.temporal_attention_block_skip_range is not None: - expected_hooks += num_layers + num_single_layers - - # Check if faster_cache denoiser hook is attached - denoiser = pipe.transformer if hasattr(pipe, "transformer") else pipe.unet - self.assertTrue( - hasattr(denoiser, "_diffusers_hook") - and isinstance(denoiser._diffusers_hook.get_hook(_FASTER_CACHE_DENOISER_HOOK), FasterCacheDenoiserHook), - "Hook should be of type FasterCacheDenoiserHook.", - ) - - # Check if all blocks have faster_cache block hook attached - count = 0 - for name, module in denoiser.named_modules(): - if hasattr(module, "_diffusers_hook"): - if name == "": - # Skip the root denoiser module - continue - count += 1 - self.assertTrue( - isinstance(module._diffusers_hook.get_hook(_FASTER_CACHE_BLOCK_HOOK), FasterCacheBlockHook), - "Hook should be of type FasterCacheBlockHook.", - ) - self.assertEqual(count, expected_hooks, "Number of hooks should match expected number.") - - # Perform inference to ensure that states are updated correctly - def faster_cache_state_check_callback(pipe, i, t, kwargs): - for name, module in denoiser.named_modules(): - if not hasattr(module, "_diffusers_hook"): - continue - if name == "": - # Root denoiser module - state = module._diffusers_hook.get_hook(_FASTER_CACHE_DENOISER_HOOK).state - if not self.faster_cache_config.is_guidance_distilled: - self.assertTrue(state.low_frequency_delta is not None, "Low frequency delta should be set.") - self.assertTrue(state.high_frequency_delta is not None, "High frequency delta should be set.") - else: - # Internal blocks - state = module._diffusers_hook.get_hook(_FASTER_CACHE_BLOCK_HOOK).state - self.assertTrue(state.cache is not None and len(state.cache) == 2, "Cache should be set.") - self.assertTrue(state.iteration == i + 1, "Hook iteration state should have updated during inference.") - return {} - - inputs = self.get_dummy_inputs(device) - inputs["num_inference_steps"] = 4 - inputs["callback_on_step_end"] = faster_cache_state_check_callback - _ = pipe(**inputs)[0] - - # After inference, reset_stateful_hooks is called within the pipeline, which should have reset the states - for name, module in denoiser.named_modules(): - if not hasattr(module, "_diffusers_hook"): - continue - - if name == "": - # Root denoiser module - state = module._diffusers_hook.get_hook(_FASTER_CACHE_DENOISER_HOOK).state - self.assertTrue(state.iteration == 0, "Iteration should be reset to 0.") - self.assertTrue(state.low_frequency_delta is None, "Low frequency delta should be reset to None.") - self.assertTrue(state.high_frequency_delta is None, "High frequency delta should be reset to None.") - else: - # Internal blocks - state = module._diffusers_hook.get_hook(_FASTER_CACHE_BLOCK_HOOK).state - self.assertTrue(state.iteration == 0, "Iteration should be reset to 0.") - self.assertTrue(state.batch_size is None, "Batch size should be reset to None.") - self.assertTrue(state.cache is None, "Cache should be reset to None.") - - -# TODO(aryan, dhruv): the cache tester mixins should probably be rewritten so that more models can be tested out -# of the box once there is better cache support/implementation -class FirstBlockCacheTesterMixin: - # threshold is intentionally set higher than usual values since we're testing with random unconverged models - # that will not satisfy the expected properties of the denoiser for caching to be effective - first_block_cache_config = FirstBlockCacheConfig(threshold=0.8) - - def test_first_block_cache_inference(self, expected_atol: float = 0.1): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - def create_pipe(): - torch.manual_seed(0) - num_layers = 2 - components = self.get_dummy_components(num_layers=num_layers) - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - return pipe - - def run_forward(pipe): - torch.manual_seed(0) - inputs = self.get_dummy_inputs(device) - inputs["num_inference_steps"] = 4 - return pipe(**inputs)[0] - - # Run inference without FirstBlockCache - pipe = create_pipe() - output = run_forward(pipe).flatten() - original_image_slice = np.concatenate((output[:8], output[-8:])) - - # Run inference with FirstBlockCache enabled - pipe = create_pipe() - pipe.transformer.enable_cache(self.first_block_cache_config) - output = run_forward(pipe).flatten() - image_slice_fbc_enabled = np.concatenate((output[:8], output[-8:])) - - # Run inference with FirstBlockCache disabled - pipe.transformer.disable_cache() - output = run_forward(pipe).flatten() - image_slice_fbc_disabled = np.concatenate((output[:8], output[-8:])) - - assert np.allclose(original_image_slice, image_slice_fbc_enabled, atol=expected_atol), ( - "FirstBlockCache outputs should not differ much." - ) - assert np.allclose(original_image_slice, image_slice_fbc_disabled, atol=1e-4), ( - "Outputs from normal inference and after disabling cache should not differ." - ) - - -class TaylorSeerCacheTesterMixin: - taylorseer_cache_config = TaylorSeerCacheConfig( - cache_interval=5, - disable_cache_before_step=10, - max_order=1, - taylor_factors_dtype=torch.bfloat16, - use_lite_mode=True, - ) - - def test_taylorseer_cache_inference(self, expected_atol: float = 0.1): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - def create_pipe(): - torch.manual_seed(0) - num_layers = 2 - components = self.get_dummy_components(num_layers=num_layers) - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - return pipe - - def run_forward(pipe): - torch.manual_seed(0) - inputs = self.get_dummy_inputs(device) - inputs["num_inference_steps"] = 50 - return pipe(**inputs)[0] - - # Run inference without TaylorSeerCache - pipe = create_pipe() - output = run_forward(pipe).flatten() - original_image_slice = np.concatenate((output[:8], output[-8:])) - - # Run inference with TaylorSeerCache enabled - pipe = create_pipe() - pipe.transformer.enable_cache(self.taylorseer_cache_config) - output = run_forward(pipe).flatten() - image_slice_fbc_enabled = np.concatenate((output[:8], output[-8:])) - - # Run inference with TaylorSeerCache disabled - pipe.transformer.disable_cache() - output = run_forward(pipe).flatten() - image_slice_fbc_disabled = np.concatenate((output[:8], output[-8:])) - - assert np.allclose(original_image_slice, image_slice_fbc_enabled, atol=expected_atol), ( - "TaylorSeerCache outputs should not differ much." - ) - assert np.allclose(original_image_slice, image_slice_fbc_disabled, atol=1e-4), ( - "Outputs from normal inference and after disabling cache should not differ." - ) - - -class MagCacheTesterMixin: - mag_cache_config = MagCacheConfig( - threshold=0.06, - max_skip_steps=3, - retention_ratio=0.2, - num_inference_steps=50, - mag_ratios=torch.ones(50), - ) - - def test_mag_cache_inference(self, expected_atol: float = 0.1): - device = "cpu" - - def create_pipe(): - torch.manual_seed(0) - num_layers = 2 - components = self.get_dummy_components(num_layers=num_layers) - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - return pipe - - def run_forward(pipe): - torch.manual_seed(0) - inputs = self.get_dummy_inputs(device) - # Match the config steps - inputs["num_inference_steps"] = 50 - return pipe(**inputs)[0] - - # 1. Run inference without MagCache (Baseline) - pipe = create_pipe() - output = run_forward(pipe).flatten() - original_image_slice = np.concatenate((output[:8], output[-8:])) - - # 2. Run inference with MagCache ENABLED - pipe = create_pipe() - pipe.transformer.enable_cache(self.mag_cache_config) - output = run_forward(pipe).flatten() - image_slice_enabled = np.concatenate((output[:8], output[-8:])) - - # 3. Run inference with MagCache DISABLED - pipe.transformer.disable_cache() - output = run_forward(pipe).flatten() - image_slice_disabled = np.concatenate((output[:8], output[-8:])) - - assert np.allclose(original_image_slice, image_slice_enabled, atol=expected_atol), ( - "MagCache outputs should not differ too much from baseline." - ) - - assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-4), ( - "Outputs after disabling cache should match original inference exactly." - ) diff --git a/tests/pipelines/testing_utils/from_pipe.py b/tests/pipelines/testing_utils/from_pipe.py index d76027d1cff9..569c5926d5b0 100644 --- a/tests/pipelines/testing_utils/from_pipe.py +++ b/tests/pipelines/testing_utils/from_pipe.py @@ -42,8 +42,7 @@ 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`. + `get_dummy_inputs()`. The pytest-style successor of the removed unittest-era `PipelineFromPipeTesterMixin`. """ # Set on the test class to pull the original pipeline from a repo other than the default for its class. diff --git a/utils/extract_tests_from_mixin.py b/utils/extract_tests_from_mixin.py index 04b157ff502c..29f6add213a4 100644 --- a/utils/extract_tests_from_mixin.py +++ b/utils/extract_tests_from_mixin.py @@ -43,9 +43,11 @@ def generate_pattern_for_mixins(mixin_classes: List[Type]) -> str: if __name__ == "__main__": mixin_classes = [] if args.type == "pipeline": - from tests.pipelines.test_pipelines_common import PipelineTesterMixin + # The pipeline tester suite is split across several mixins under `tests/pipelines/testing_utils`, + # so aggregate their test methods to reconstruct the full coverage. + from tests.pipelines.testing_utils import MemoryTesterMixin, PipelineTesterMixin - mixin_classes = [PipelineTesterMixin] + mixin_classes = [PipelineTesterMixin, MemoryTesterMixin] elif args.type == "models": # The model tester suite is split across several mixins under `tests/models/testing_utils`, From 7511684586b209ba5f31ebcdc2414502e8b91442 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 1 Sep 2026 07:32:07 +0000 Subject: [PATCH 5/6] simplify test policy wording --- .ai/references/testing.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.ai/references/testing.md b/.ai/references/testing.md index 437ed6a081a0..99f06c716ff8 100644 --- a/.ai/references/testing.md +++ b/.ai/references/testing.md @@ -35,7 +35,7 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers - `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. It replaces the unittest-era `PipelineFromPipeTesterMixin`, since removed. +- **`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. - **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. - **`encode_prompt` reading a component that isn't a text encoder or tokenizer?** `test_encode_prompt_works_in_isolation` rebuilds the pipeline with only the components whose names contain `text` or `tokenizer`. When `encode_prompt` also needs another one — a `processor` used for chat templating, say — list it in `text_stack_component_names` on the config class rather than re-implementing the test. From d75be5b0ed9005eed980129d1c7834375eac094c Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 1 Sep 2026 07:40:55 +0000 Subject: [PATCH 6/6] move around --- .github/workflows/pr_tests_gpu.yml | 1 + tests/pipelines/testing_utils/__init__.py | 3 +-- tests/pipelines/testing_utils/common.py | 11 +++++++++++ tests/pipelines/testing_utils/from_pipe.py | 2 +- tests/pipelines/testing_utils/utils.py | 13 ------------- 5 files changed, 14 insertions(+), 16 deletions(-) diff --git a/.github/workflows/pr_tests_gpu.yml b/.github/workflows/pr_tests_gpu.yml index 4527b96ebc74..0bf737708793 100644 --- a/.github/workflows/pr_tests_gpu.yml +++ b/.github/workflows/pr_tests_gpu.yml @@ -16,6 +16,7 @@ on: - "src/diffusers/loaders/peft.py" - "tests/pipelines/testing_utils/**.py" - "tests/models/test_modeling_common.py" + - "tests/models/testing_utils/**.py" - "examples/**/*.py" workflow_dispatch: diff --git a/tests/pipelines/testing_utils/__init__.py b/tests/pipelines/testing_utils/__init__.py index bfb93aaec03e..50cff4d3c06a 100644 --- a/tests/pipelines/testing_utils/__init__.py +++ b/tests/pipelines/testing_utils/__init__.py @@ -6,7 +6,7 @@ PyramidAttentionBroadcastTesterMixin, TaylorSeerCacheTesterMixin, ) -from .common import BasePipelineTesterConfig, PipelineTesterMixin +from .common import BasePipelineTesterConfig, PipelineTesterMixin, assert_mean_pixel_difference from .from_pipe import FromPipeTesterMixin from .ip_adapter import FluxIPAdapterTesterMixin, IPAdapterTesterMixin from .lora import LoraMemoryTesterMixin, LoraTesterMixin, UNetLoraTesterMixin @@ -17,7 +17,6 @@ PipelineOffloadTesterMixin, ) from .utils import ( - assert_mean_pixel_difference, check_qkv_fused_layers_exist, check_qkv_fusion_matches_attn_procs_length, check_qkv_fusion_processors_exist, diff --git a/tests/pipelines/testing_utils/common.py b/tests/pipelines/testing_utils/common.py index 771951341ded..76ac3dcb393d 100644 --- a/tests/pipelines/testing_utils/common.py +++ b/tests/pipelines/testing_utils/common.py @@ -19,6 +19,7 @@ import os from typing import Callable +import numpy as np import pytest import torch import torch.nn as nn @@ -79,6 +80,16 @@ def cast_pipeline_to_dtype(pipe, dtype): return pipe +# Some models (e.g. unCLIP) are extremely likely to significantly deviate depending on which hardware is used. +# This helper function is used to check that the image doesn't deviate on average more than 10 pixels from a +# reference image. +def assert_mean_pixel_difference(image, expected_image, expected_max_diff=10): + image = np.asarray(DiffusionPipeline.numpy_to_pil(image)[0], dtype=np.float32) + expected_image = np.asarray(DiffusionPipeline.numpy_to_pil(expected_image)[0], dtype=np.float32) + avg_diff = np.abs(image - expected_image).mean() + assert avg_diff < expected_max_diff, f"Error image deviates {avg_diff} pixels on average" + + class BasePipelineTesterConfig: """ Base class defining the configuration interface for pipeline testing. diff --git a/tests/pipelines/testing_utils/from_pipe.py b/tests/pipelines/testing_utils/from_pipe.py index 569c5926d5b0..3a4aab7e7155 100644 --- a/tests/pipelines/testing_utils/from_pipe.py +++ b/tests/pipelines/testing_utils/from_pipe.py @@ -42,7 +42,7 @@ 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 the removed unittest-era `PipelineFromPipeTesterMixin`. + `get_dummy_inputs()`. """ # Set on the test class to pull the original pipeline from a repo other than the default for its class. diff --git a/tests/pipelines/testing_utils/utils.py b/tests/pipelines/testing_utils/utils.py index 58bc15f2ab8b..a20c0ca9becb 100644 --- a/tests/pipelines/testing_utils/utils.py +++ b/tests/pipelines/testing_utils/utils.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import numpy as np - -from diffusers import DiffusionPipeline from diffusers.models.attention import AttentionModuleMixin @@ -55,13 +52,3 @@ def check_qkv_fused_layers_exist(model, layer_names): is_fused = is_fused_attribute_set and is_fused_layer is_fused_submodules.append(is_fused) return all(is_fused_submodules) - - -# Some models (e.g. unCLIP) are extremely likely to significantly deviate depending on which hardware is used. -# This helper function is used to check that the image doesn't deviate on average more than 10 pixels from a -# reference image. -def assert_mean_pixel_difference(image, expected_image, expected_max_diff=10): - image = np.asarray(DiffusionPipeline.numpy_to_pil(image)[0], dtype=np.float32) - expected_image = np.asarray(DiffusionPipeline.numpy_to_pil(expected_image)[0], dtype=np.float32) - avg_diff = np.abs(image - expected_image).mean() - assert avg_diff < expected_max_diff, f"Error image deviates {avg_diff} pixels on average"