From 82f4c61d5fd4eb61c6454564fa11dd25be2bc065 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 30 Jul 2026 12:21:39 +0000 Subject: [PATCH] refactor stable diffusion pipeline tests --- .../stable_diffusion/ip_adapter_tester.py | 241 ++++++++ .../stable_diffusion/test_stable_diffusion.py | 534 +++++++----------- .../test_stable_diffusion_img2img.py | 258 ++++----- .../test_stable_diffusion_inpaint.py | 440 ++++++--------- ...st_stable_diffusion_instruction_pix2pix.py | 190 ++----- 5 files changed, 769 insertions(+), 894 deletions(-) create mode 100644 tests/pipelines/stable_diffusion/ip_adapter_tester.py diff --git a/tests/pipelines/stable_diffusion/ip_adapter_tester.py b/tests/pipelines/stable_diffusion/ip_adapter_tester.py new file mode 100644 index 000000000000..a7e614879e7e --- /dev/null +++ b/tests/pipelines/stable_diffusion/ip_adapter_tester.py @@ -0,0 +1,241 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect + +import pytest +import torch + +from diffusers.loaders import IPAdapterMixin + +from ...testing_utils import assert_tensors_close, is_ip_adapter, torch_device +from ..testing_utils.common import BasePipelineOutputMixin + + +@is_ip_adapter +class IPAdapterTesterMixin(BasePipelineOutputMixin): + """IP-Adapter tests shared by the Stable Diffusion pipelines in this directory. + + 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. + """ + + 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, :, :, : input_size // 2] = 1 + return masks + + def _get_ip_adapter_inputs(self): + inputs = self.get_dummy_inputs() + parameters = inspect.signature(self.pipeline_class.__call__).parameters + if "image" in parameters and "strength" in parameters: + inputs["num_inference_steps"] = 4 + inputs["return_dict"] = False + return inputs + + def _load_ip_adapters(self, pipe, num_adapters=1, faceid=False): + # The state dict builders are imported here rather than at module scope: they live 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.unets.test_models_unet_2d_condition import ( + create_ip_adapter_faceid_state_dict, + create_ip_adapter_state_dict, + ) + + create_state_dict = create_ip_adapter_faceid_state_dict if faceid else create_ip_adapter_state_dict + state_dicts = [create_state_dict(pipe.unet) for _ in range(num_adapters)] + + # Load through the pipeline's public IP-Adapter API. `image_encoder_folder=None` skips fetching a CLIP image + # encoder since these tests feed pre-computed `ip_adapter_image_embeds` directly. + pipe.load_ip_adapter( + state_dicts, + subfolder=[""] * num_adapters, + weight_name=[""] * num_adapters, + image_encoder_folder=None, + ) + + def test_pipeline_signature(self): + parameters = inspect.signature(self.pipeline_class.__call__).parameters + + assert issubclass(self.pipeline_class, IPAdapterMixin) + 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 test_ip_adapter(self, expected_max_diff: float = 1e-4): + 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. + """ + pipe = self.get_pipeline().to(torch_device) + cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32) + + # forward pass without ip adapter + output_without_adapter = pipe(**self._get_ip_adapter_inputs())[0] + + # 1. Single IP-Adapter test cases + self._load_ip_adapters(pipe) + + # forward pass with single ip adapter, but scale=0 which should have no effect + inputs = self._get_ip_adapter_inputs() + 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] + + # forward pass with single ip adapter, but with scale of adapter weights + inputs = self._get_ip_adapter_inputs() + 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] + + 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 + self._load_ip_adapters(pipe, num_adapters=2) + + # forward pass with multi ip adapter, but scale=0 which should have no effect + inputs = self._get_ip_adapter_inputs() + 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] + + # forward pass with multi ip adapter, but with scale of adapter weights + inputs = self._get_ip_adapter_inputs() + 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] + + 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" + ) + + def test_ip_adapter_cfg(self): + if "guidance_scale" not in inspect.signature(self.pipeline_class.__call__).parameters: + pytest.skip( + f"Skipping test because `guidance_scale` wasn't found in the args accepted in {self.pipeline_class}'s call." + ) + + pipe = self.get_pipeline().to(torch_device) + cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32) + + self._load_ip_adapters(pipe) + pipe.set_ip_adapter_scale(1.0) + + # forward pass with CFG not applied + inputs = self._get_ip_adapter_inputs() + 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._get_ip_adapter_inputs() + 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): + pipe = self.get_pipeline().to(torch_device) + 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 + output_without_adapter = pipe(**self._get_ip_adapter_inputs())[0] + + self._load_ip_adapters(pipe) + + # forward pass with single ip adapter and masks, but scale=0 which should have no effect + inputs = self._get_ip_adapter_inputs() + 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] + + # forward pass with single ip adapter and masks, but with scale of adapter weights + inputs = self._get_ip_adapter_inputs() + 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] + + 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-3, "Output with ip-adapter must be different from normal inference" + + def test_ip_adapter_faceid(self, expected_max_diff: float = 1e-4): + pipe = self.get_pipeline().to(torch_device) + cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32) + + # forward pass without ip adapter + output_without_adapter = pipe(**self._get_ip_adapter_inputs())[0] + + self._load_ip_adapters(pipe, faceid=True) + + # forward pass with single ip adapter, but scale=0 which should have no effect + inputs = self._get_ip_adapter_inputs() + 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] + + # forward pass with single ip adapter, but with scale of adapter weights + inputs = self._get_ip_adapter_inputs() + 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] + + 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-3, "Output with ip-adapter must be different from normal inference" diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py index 23607a39ecd0..31a8aec3c58b 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -15,11 +15,10 @@ import gc -import tempfile import time -import unittest import numpy as np +import pytest import torch from huggingface_hub import hf_hub_download from transformers import ( @@ -44,11 +43,11 @@ from ...testing_utils import ( CaptureLogger, + assert_tensors_close, backend_empty_cache, backend_max_memory_allocated, backend_reset_max_memory_allocated, backend_reset_peak_memory_stats, - enable_full_determinism, load_numpy, nightly, numpy_cosine_similarity_distance, @@ -61,36 +60,20 @@ ) from ..pipeline_params import ( TEXT_TO_IMAGE_BATCH_PARAMS, - TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS, ) -from ..test_pipelines_common import ( - IPAdapterTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, PipelineTesterMixin, ) +from .ip_adapter_tester import IPAdapterTesterMixin -enable_full_determinism() - - -class StableDiffusionPipelineFastTests( - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineTesterMixin, - unittest.TestCase, -): +class StableDiffusionPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionPipeline - params = TEXT_TO_IMAGE_PARAMS - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS def get_dummy_components(self, time_cond_proj_dim=None): cross_attention_dim = 8 @@ -152,80 +135,61 @@ def get_dummy_components(self, time_cond_proj_dim=None): } return components - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) + def get_dummy_inputs(self): inputs = { "prompt": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs - def test_stable_diffusion_ddim(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - 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] +class TestStableDiffusionPipeline(StableDiffusionPipelineTesterConfig, PipelineTesterMixin): + def test_stable_diffusion_ddim(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.1641, 0.4640, 0.4864, 0.2683, 0.3652, 0.4507, 0.5290, 0.3307, 0.3977]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 64, 64) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # fmt: off + expected_slice = torch.tensor([0.1641, 0.4640, 0.4864, 0.2683, 0.3652, 0.4507, 0.5290, 0.3307, 0.3977]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) def test_stable_diffusion_lcm(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionPipeline(**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, 3, 64, 64) - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.2368, 0.4900, 0.5019, 0.2723, 0.4473, 0.4578, 0.4551, 0.3532, 0.4133]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # fmt: off + expected_slice = torch.tensor([0.2368, 0.4900, 0.5019, 0.2723, 0.4473, 0.4578, 0.4551, 0.3532, 0.4133]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) def test_stable_diffusion_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 = StableDiffusionPipeline(**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] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.2368, 0.4900, 0.5019, 0.2723, 0.4473, 0.4578, 0.4551, 0.3532, 0.4133]) + image = sd_pipe(**inputs).images + assert image.shape == (1, 3, 64, 64) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # Custom timesteps matching the default 2-step schedule reproduce `test_stable_diffusion_lcm`'s output. + # fmt: off + expected_slice = torch.tensor([0.2368, 0.4900, 0.5019, 0.2723, 0.4473, 0.4578, 0.4551, 0.3532, 0.4133]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) def test_stable_diffusion_ays(self): from diffusers.schedulers import AysSchedules @@ -233,53 +197,38 @@ def test_stable_diffusion_ays(self): timestep_schedule = AysSchedules["StableDiffusionTimesteps"] sigma_schedule = AysSchedules["StableDiffusionSigmas"] - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionPipeline(**components) + sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)).to(torch_device) sd_pipe.scheduler = EulerDiscreteScheduler.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() inputs["num_inference_steps"] = 10 output = sd_pipe(**inputs).images - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["num_inference_steps"] = None inputs["timesteps"] = timestep_schedule output_ts = sd_pipe(**inputs).images - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["num_inference_steps"] = None inputs["sigmas"] = sigma_schedule output_sigmas = sd_pipe(**inputs).images - assert np.abs(output_sigmas.flatten() - output_ts.flatten()).max() < 1e-3, ( + assert (output_sigmas - output_ts).abs().max() < 1e-3, ( "ays timesteps and ays sigmas should have the same outputs" ) - assert np.abs(output.flatten() - output_ts.flatten()).max() > 1e-3, ( - "use ays timesteps should have different outputs" - ) - assert np.abs(output.flatten() - output_sigmas.flatten()).max() > 1e-3, ( - "use ays sigmas should have different outputs" - ) + assert (output - output_ts).abs().max() > 1e-3, "use ays timesteps should have different outputs" + assert (output - output_sigmas).abs().max() > 1e-3, "use ays sigmas should have different outputs" def test_stable_diffusion_prompt_embeds(self): - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt"] = 3 * [inputs["prompt"]] - - # forward output = sd_pipe(**inputs) - image_slice_1 = output.images[0, -3:, -3:, -1] + image_slice_1 = output.images[0, -1, -3:, -3:] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() prompt = 3 * [inputs.pop("prompt")] text_inputs = sd_pipe.tokenizer( @@ -290,34 +239,26 @@ def test_stable_diffusion_prompt_embeds(self): return_tensors="pt", ) text_inputs = text_inputs["input_ids"].to(torch_device) + inputs["prompt_embeds"] = sd_pipe.text_encoder(text_inputs)[0] - prompt_embeds = sd_pipe.text_encoder(text_inputs)[0] - - inputs["prompt_embeds"] = prompt_embeds - - # forward output = sd_pipe(**inputs) - image_slice_2 = output.images[0, -3:, -3:, -1] + image_slice_2 = output.images[0, -1, -3:, -3:] - assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 + assert_tensors_close( + image_slice_2, image_slice_1, atol=1e-4, msg="Passing `prompt_embeds` changed the output." + ) def test_stable_diffusion_negative_prompt_embeds(self): - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() negative_prompt = 3 * ["this is a negative prompt"] inputs["negative_prompt"] = negative_prompt inputs["prompt"] = 3 * [inputs["prompt"]] - - # forward output = sd_pipe(**inputs) - image_slice_1 = output.images[0, -3:, -3:, -1] + image_slice_1 = output.images[0, -1, -3:, -3:] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() prompt = 3 * [inputs.pop("prompt")] embeds = [] @@ -330,55 +271,96 @@ def test_stable_diffusion_negative_prompt_embeds(self): return_tensors="pt", ) text_inputs = text_inputs["input_ids"].to(torch_device) - embeds.append(sd_pipe.text_encoder(text_inputs)[0]) inputs["prompt_embeds"], inputs["negative_prompt_embeds"] = embeds - # forward output = sd_pipe(**inputs) - image_slice_2 = output.images[0, -3:, -3:, -1] + image_slice_2 = output.images[0, -1, -3:, -3:] - assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 + assert_tensors_close( + image_slice_2, image_slice_1, atol=1e-4, msg="Passing `negative_prompt_embeds` changed the output." + ) def test_stable_diffusion_ddim_factor_8(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + image = sd_pipe(**self.get_dummy_inputs(), height=136, width=136).images + assert image.shape == (1, 3, 136, 136) - inputs = self.get_dummy_inputs(device) - output = sd_pipe(**inputs, height=136, width=136) - image = output.images + # fmt: off + expected_slice = torch.tensor([0.4587, 0.5238, 0.5006, 0.3869, 0.4538, 0.4228, 0.5666, 0.5814, 0.5468]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) - image_slice = image[0, -3:, -3:, -1] + def test_stable_diffusion_pndm(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() + sd_pipe.scheduler = PNDMScheduler(skip_prk_steps=True) - assert image.shape == (1, 136, 136, 3) - expected_slice = np.array([0.4587, 0.5238, 0.5006, 0.3869, 0.4538, 0.4228, 0.5666, 0.5814, 0.5468]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 64, 64) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # fmt: off + expected_slice = torch.tensor([0.1791, 0.4611, 0.4741, 0.2336, 0.4175, 0.4484, 0.5582, 0.3556, 0.3951]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) - def test_stable_diffusion_pndm(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe.scheduler = PNDMScheduler(skip_prk_steps=True) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + def test_stable_diffusion_k_lms(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() + sd_pipe.scheduler = LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config) - 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, 3, 64, 64) + + # fmt: off + expected_slice = torch.tensor([0.2185, 0.4538, 0.4660, 0.2454, 0.4408, 0.4377, 0.5629, 0.3754, 0.3927]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) + + def test_stable_diffusion_k_euler_ancestral(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() + sd_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(sd_pipe.scheduler.config) + + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 64, 64) + + # fmt: off + expected_slice = torch.tensor([0.2185, 0.4534, 0.4657, 0.2452, 0.4406, 0.4375, 0.5631, 0.3755, 0.3927]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) + + def test_stable_diffusion_k_euler(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() + sd_pipe.scheduler = EulerDiscreteScheduler.from_config(sd_pipe.scheduler.config) + + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 64, 64) + + # fmt: off + expected_slice = torch.tensor([0.2185, 0.4538, 0.4660, 0.2454, 0.4408, 0.4377, 0.5629, 0.3754, 0.3927]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) + + def test_stable_diffusion_negative_prompt(self): + # Run on CPU: the expected slice below is CPU-specific. + components = self.get_dummy_components() + components["scheduler"] = PNDMScheduler(skip_prk_steps=True) + sd_pipe = self.get_pipeline(**components) - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.1791, 0.4611, 0.4741, 0.2336, 0.4175, 0.4484, 0.5582, 0.3556, 0.3951]) + image = sd_pipe(**self.get_dummy_inputs(), negative_prompt="french fries").images + assert image.shape == (1, 3, 64, 64) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # fmt: off + expected_slice = torch.tensor([0.1772, 0.4578, 0.4700, 0.2363, 0.4173, 0.4462, 0.5595, 0.3588, 0.3959]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) - def test_stable_diffusion_no_safety_checker(self): + def test_stable_diffusion_no_safety_checker(self, tmp_path): pipe = StableDiffusionPipeline.from_pretrained( "hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None ) @@ -390,150 +372,56 @@ def test_stable_diffusion_no_safety_checker(self): assert image is not None # check that there's no error when saving a pipeline with one of the models being None - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) - pipe = StableDiffusionPipeline.from_pretrained(tmpdirname) + pipe.save_pretrained(tmp_path) + pipe = StableDiffusionPipeline.from_pretrained(tmp_path) # sanity check that the pipeline still works assert pipe.safety_checker is None image = pipe("example prompt", num_inference_steps=2).images[0] assert image is not None - def test_stable_diffusion_k_lms(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe.scheduler = LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(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] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.2185, 0.4538, 0.4660, 0.2454, 0.4408, 0.4377, 0.5629, 0.3754, 0.3927]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - - def test_stable_diffusion_k_euler_ancestral(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(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] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.2185, 0.4534, 0.4657, 0.2452, 0.4406, 0.4375, 0.5631, 0.3755, 0.3927]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - - def test_stable_diffusion_k_euler(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe.scheduler = EulerDiscreteScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(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] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.2185, 0.4538, 0.4660, 0.2454, 0.4408, 0.4377, 0.5629, 0.3754, 0.3927]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - def test_stable_diffusion_vae_slicing(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator + # Run on CPU: sliced VAE decoding is compared against a full-batch decode of the same run. components = self.get_dummy_components() components["scheduler"] = LMSDiscreteScheduler.from_config(components["scheduler"].config) - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline(**components) image_count = 4 - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["prompt"] = [inputs["prompt"]] * image_count output_1 = sd_pipe(**inputs) # make sure sliced vae decode yields the same result sd_pipe.vae.enable_slicing() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["prompt"] = [inputs["prompt"]] * image_count output_2 = sd_pipe(**inputs) # there is a small discrepancy at image borders vs. full batch decode - assert np.abs(output_2.images.flatten() - output_1.images.flatten()).max() < 3e-3 + assert (output_2.images - output_1.images).abs().max() < 3e-3 def test_stable_diffusion_vae_tiling(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - - # make sure here that pndm scheduler skips prk - components["safety_checker"] = None - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) - - prompt = "A painting of a squirrel eating a burger" + # Run on CPU: tiled VAE decoding is compared against a non-tiled decode of the same run. + sd_pipe = self.get_pipeline() # Test that tiled decode at 512x512 yields the same result as the non-tiled decode - generator = torch.Generator(device=device).manual_seed(0) - output_1 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np") + output_1 = sd_pipe(**self.get_dummy_inputs()) # make sure tiled vae decode yields the same result sd_pipe.vae.enable_tiling() - generator = torch.Generator(device=device).manual_seed(0) - output_2 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np") + output_2 = sd_pipe(**self.get_dummy_inputs()) - assert np.abs(output_2.images.flatten() - output_1.images.flatten()).max() < 5e-1 + assert (output_2.images - output_1.images).abs().max() < 5e-1 # test that tiled decode works with various shapes - shapes = [(1, 4, 73, 97), (1, 4, 97, 73), (1, 4, 49, 65), (1, 4, 65, 49)] - for shape in shapes: - zeros = torch.zeros(shape).to(device) - sd_pipe.vae.decode(zeros) - - def test_stable_diffusion_negative_prompt(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - components["scheduler"] = PNDMScheduler(skip_prk_steps=True) - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - negative_prompt = "french fries" - output = sd_pipe(**inputs, negative_prompt=negative_prompt) - - image = output.images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.1772, 0.4578, 0.4700, 0.2363, 0.4173, 0.4462, 0.5595, 0.3588, 0.3959]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + for shape in [(1, 4, 73, 97), (1, 4, 97, 73), (1, 4, 49, 65), (1, 4, 65, 49)]: + sd_pipe.vae.decode(torch.zeros(shape)) def test_stable_diffusion_long_prompt(self): components = self.get_dummy_components() components["scheduler"] = LMSDiscreteScheduler.from_config(components["scheduler"].config) - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline(**components).to(torch_device) do_classifier_free_guidance = True negative_prompt = None @@ -577,29 +465,21 @@ def test_stable_diffusion_long_prompt(self): def test_stable_diffusion_height_width_opt(self): components = self.get_dummy_components() components["scheduler"] = LMSDiscreteScheduler.from_config(components["scheduler"].config) - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline(**components).to(torch_device) prompt = "hey" - output = sd_pipe(prompt, num_inference_steps=1, output_type="np") - image_shape = output.images[0].shape[:2] - assert image_shape == (64, 64) + output = sd_pipe(prompt, num_inference_steps=1, output_type="pt") + assert output.images[0].shape[-2:] == (64, 64) - output = sd_pipe(prompt, num_inference_steps=1, height=96, width=96, output_type="np") - image_shape = output.images[0].shape[:2] - assert image_shape == (96, 96) + output = sd_pipe(prompt, num_inference_steps=1, height=96, width=96, output_type="pt") + assert output.images[0].shape[-2:] == (96, 96) config = dict(sd_pipe.unet.config) config["sample_size"] = 96 sd_pipe.unet = UNet2DConditionModel.from_config(config).to(torch_device) - output = sd_pipe(prompt, num_inference_steps=1, output_type="np") - image_shape = output.images[0].shape[:2] - assert image_shape == (192, 192) - - def test_attention_slicing_forward_pass(self): - super().test_attention_slicing_forward_pass(expected_max_diff=3e-3) + output = sd_pipe(prompt, num_inference_steps=1, output_type="pt") + assert output.images[0].shape[-2:] == (192, 192) def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-3) @@ -607,82 +487,63 @@ def test_inference_batch_single_identical(self): # MPS currently doesn't support ComplexFloats, which are required for freeU - see https://github.com/huggingface/diffusers/issues/7569. @skip_mps def test_freeu_enabled(self): - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline().to(torch_device) prompt = "hey" - output = sd_pipe(prompt, num_inference_steps=1, output_type="np", generator=torch.manual_seed(0)).images + output = sd_pipe(prompt, num_inference_steps=1, output_type="pt", generator=torch.manual_seed(0)).images sd_pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4) - output_freeu = sd_pipe(prompt, num_inference_steps=1, output_type="np", generator=torch.manual_seed(0)).images + output_freeu = sd_pipe(prompt, num_inference_steps=1, output_type="pt", generator=torch.manual_seed(0)).images - assert not np.allclose(output[0, -3:, -3:, -1], output_freeu[0, -3:, -3:, -1]), ( + assert not torch.allclose(output[0, -1, -3:, -3:], output_freeu[0, -1, -3:, -3:]), ( "Enabling of FreeU should lead to different results." ) def test_freeu_disabled(self): - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline().to(torch_device) prompt = "hey" - output = sd_pipe(prompt, num_inference_steps=1, output_type="np", generator=torch.manual_seed(0)).images + output = sd_pipe(prompt, num_inference_steps=1, output_type="pt", generator=torch.manual_seed(0)).images sd_pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4) sd_pipe.disable_freeu() - freeu_keys = {"s1", "s2", "b1", "b2"} for upsample_block in sd_pipe.unet.up_blocks: - for key in freeu_keys: + for key in {"s1", "s2", "b1", "b2"}: assert getattr(upsample_block, key) is None, f"Disabling of FreeU should have set {key} to None." output_no_freeu = sd_pipe( - prompt, num_inference_steps=1, output_type="np", generator=torch.manual_seed(0) + prompt, num_inference_steps=1, output_type="pt", generator=torch.manual_seed(0) ).images - assert np.allclose(output[0, -3:, -3:, -1], output_no_freeu[0, -3:, -3:, -1]), ( + assert torch.allclose(output[0, -1, -3:, -3:], output_no_freeu[0, -1, -3:, -3:]), ( "Disabling of FreeU should lead to results similar to the default pipeline results." ) def test_fused_qkv_projections(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + # Run on CPU: fused and unfused attention are compared against each other. + sd_pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = sd_pipe(**inputs).images - original_image_slice = image[0, -3:, -3:, -1] + original_image_slice = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] sd_pipe.fuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - image = sd_pipe(**inputs).images - image_slice_fused = image[0, -3:, -3:, -1] + image_slice_fused = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] sd_pipe.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - image = sd_pipe(**inputs).images - image_slice_disabled = image[0, -3:, -3:, -1] + image_slice_disabled = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( + assert torch.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), ( + assert torch.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), ( + assert torch.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( "Original outputs should match when fused QKV projections are disabled." ) def test_pipeline_interrupt(self): - components = self.get_dummy_components() - sd_pipe = StableDiffusionPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline().to(torch_device) prompt = "hey" num_inference_steps = 3 @@ -700,7 +561,7 @@ def apply(self, pipe, i, t, callback_kwargs): sd_pipe( prompt, num_inference_steps=num_inference_steps, - output_type="np", + output_type="pt", generator=torch.Generator("cpu").manual_seed(0), callback_on_step_end=pipe_state.apply, ).images @@ -741,15 +602,27 @@ def test_pipeline_accept_tuple_type_unet_sample_size(self): def test_encode_prompt_works_in_isolation(self): extra_required_param_value_dict = { "device": torch.device(torch_device).type, - "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) +class TestStableDiffusionPipelineMemory(StableDiffusionPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Stable Diffusion pipeline.""" + + +class TestStableDiffusionPipelineIPAdapter(StableDiffusionPipelineTesterConfig, IPAdapterTesterMixin): + """IP-Adapter tests for the Stable Diffusion pipeline.""" + + @slow @require_torch_accelerator -class StableDiffusionPipelineSlowTests(unittest.TestCase): - def setUp(self): +class TestStableDiffusionPipelineSlow: + @pytest.fixture(autouse=True) + def cleanup(self): + gc.collect() + backend_empty_cache(torch_device) + yield gc.collect() backend_empty_cache(torch_device) @@ -1189,14 +1062,12 @@ def test_stable_diffusion_textual_inversion_with_sequential_cpu_offload(self): @slow @require_torch_accelerator -class StableDiffusionPipelineCkptTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestStableDiffusionPipelineCkpt: + @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) @@ -1234,14 +1105,12 @@ def test_download_local(self): @nightly @require_torch_accelerator -class StableDiffusionPipelineNightlyTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestStableDiffusionPipelineNightly: + @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) @@ -1339,9 +1208,10 @@ def test_stable_diffusion_euler(self): @slow @require_torch_multi_accelerator @require_accelerate_version_greater("0.27.0") -class StableDiffusionPipelineDeviceMapTests(unittest.TestCase): - def tearDown(self): - super().tearDown() +class TestStableDiffusionPipelineDeviceMap: + @pytest.fixture(autouse=True) + def cleanup(self): + yield gc.collect() backend_empty_cache(torch_device) diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py index fcaf21e1b04a..dcf104c1cc7d 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py @@ -15,9 +15,9 @@ import gc import random -import unittest import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -35,51 +35,39 @@ ) from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, backend_max_memory_allocated, backend_reset_max_memory_allocated, backend_reset_peak_memory_stats, - enable_full_determinism, floats_tensor, load_image, load_numpy, nightly, require_torch_accelerator, - skip_mps, slow, torch_device, ) from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS, - TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( - IPAdapterTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, PipelineTesterMixin, ) +from .ip_adapter_tester import IPAdapterTesterMixin -enable_full_determinism() - - -class StableDiffusionImg2ImgPipelineFastTests( - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineTesterMixin, - unittest.TestCase, -): +class StableDiffusionImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionImg2ImgPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"} - required_optional_params = PipelineTesterMixin.required_optional_params - {"latents"} - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"} + batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS + # img2img derives its latents from the input image, so `__call__` takes no `latents` argument. + optional_input_params = frozenset( + ["num_inference_steps", "num_images_per_prompt", "generator", "output_type", "return_dict"] + ) def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) @@ -134,182 +122,128 @@ def get_dummy_components(self, time_cond_proj_dim=None): def get_dummy_tiny_autoencoder(self): return AutoencoderTiny(in_channels=3, out_channels=3, latent_channels=4) - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) image = image / 2 + 0.5 - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) inputs = { "prompt": "A painting of a squirrel eating a burger", "image": image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs - def test_stable_diffusion_img2img_default_case(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionImg2ImgPipeline(**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 TestStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipelineTesterConfig, PipelineTesterMixin): + def test_stable_diffusion_img2img_default_case(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.4517, 0.3640, 0.4036, 0.4014, 0.4381, 0.3990, 0.3824, 0.4628, 0.4415]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 32, 32) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + # fmt: off + expected_slice = torch.tensor([0.4517, 0.3640, 0.4036, 0.4014, 0.4381, 0.3990, 0.3824, 0.4628, 0.4415]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_stable_diffusion_img2img_default_case_lcm(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionImg2ImgPipeline(**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(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] - - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.5915, 0.4625, 0.4547, 0.6083, 0.5414, 0.6927, 0.6253, 0.5223, 0.5449]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 32, 32) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + # fmt: off + expected_slice = torch.tensor([0.5915, 0.4625, 0.4547, 0.6083, 0.5414, 0.6927, 0.6253, 0.5223, 0.5449]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_stable_diffusion_img2img_default_case_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 = StableDiffusionImg2ImgPipeline(**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(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] image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] + assert image.shape == (1, 3, 32, 32) - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.5915, 0.4625, 0.4547, 0.6083, 0.5414, 0.6927, 0.6253, 0.5223, 0.5449]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + # Custom timesteps matching the default schedule reproduce `..._default_case_lcm`'s output. + # fmt: off + expected_slice = torch.tensor([0.5915, 0.4625, 0.4547, 0.6083, 0.5414, 0.6927, 0.6253, 0.5223, 0.5449]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_stable_diffusion_img2img_negative_prompt(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionImg2ImgPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - negative_prompt = "french fries" - output = sd_pipe(**inputs, negative_prompt=negative_prompt) - image = output.images - image_slice = image[0, -3:, -3:, -1] + image = sd_pipe(**self.get_dummy_inputs(), negative_prompt="french fries").images + assert image.shape == (1, 3, 32, 32) - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.4642, 0.3576, 0.4112, 0.4289, 0.4396, 0.4068, 0.3845, 0.4603, 0.4427]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - - def test_ip_adapter(self): - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array([0.4923, 0.5012, 0.5041, 0.5462, 0.5665, 0.6621, 0.6451, 0.5046, 0.5431]) - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) + # fmt: off + expected_slice = torch.tensor([0.4642, 0.3576, 0.4112, 0.4289, 0.4396, 0.4068, 0.3845, 0.4603, 0.4427]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_stable_diffusion_img2img_multiple_init_images(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionImg2ImgPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["prompt"] = [inputs["prompt"]] * 2 inputs["image"] = inputs["image"].repeat(2, 1, 1, 1) image = sd_pipe(**inputs).images - image_slice = image[-1, -3:, -3:, -1] - - assert image.shape == (2, 32, 32, 3) - expected_slice = np.array([0.4385, 0.3589, 0.4905, 0.4607, 0.4040, 0.5551, 0.5114, 0.5269, 0.5267]) + assert image.shape == (2, 3, 32, 32) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + # fmt: off + expected_slice = torch.tensor([0.4385, 0.3589, 0.4905, 0.4607, 0.4040, 0.5551, 0.5114, 0.5269, 0.5267]) + # fmt: on + assert_tensors_close(image[-1, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_stable_diffusion_img2img_k_lms(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator + # Run on CPU: the expected slice below is CPU-specific. components = self.get_dummy_components() components["scheduler"] = LMSDiscreteScheduler( beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear" ) - sd_pipe = StableDiffusionImg2ImgPipeline(**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] + sd_pipe = self.get_pipeline(**components) - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.4409, 0.4952, 0.4357, 0.6574, 0.5583, 0.4375, 0.5777, 0.5960, 0.5168]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 32, 32) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + # fmt: off + expected_slice = torch.tensor([0.4409, 0.4952, 0.4357, 0.6574, 0.5583, 0.4375, 0.5777, 0.5960, 0.5168]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_stable_diffusion_img2img_tiny_autoencoder(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionImg2ImgPipeline(**components) + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() sd_pipe.vae = self.get_dummy_tiny_autoencoder() - 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] - - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.00551, 0.01295, 0.01731, 0.00725, 0.01140, 0.01334, 0.00152, 0.00997, 0.01265]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - - @skip_mps - def test_save_load_local(self): - return super().test_save_load_local() - - @skip_mps - def test_dict_tuple_outputs_equivalent(self): - return super().test_dict_tuple_outputs_equivalent() - @skip_mps - def test_save_load_optional_components(self): - return super().test_save_load_optional_components() + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 32, 32) - @skip_mps - def test_attention_slicing_forward_pass(self): - return super().test_attention_slicing_forward_pass(expected_max_diff=5e-3) + # fmt: off + expected_slice = torch.tensor([0.00551, 0.01295, 0.01731, 0.00725, 0.01140, 0.01334, 0.00152, 0.00997, 0.01265]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-3) - def test_float16_inference(self): - super().test_float16_inference(expected_max_diff=5e-1) - def test_pipeline_interrupt(self): - components = self.get_dummy_components() - sd_pipe = StableDiffusionImg2ImgPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() prompt = "hey" num_inference_steps = 3 @@ -328,7 +262,7 @@ def apply(self, pipe, i, t, callback_kwargs): prompt, image=inputs["image"], num_inference_steps=num_inference_steps, - output_type="np", + output_type="pt", generator=torch.Generator("cpu").manual_seed(0), callback_on_step_end=pipe_state.apply, ).images @@ -362,21 +296,27 @@ def callback_on_step_end(pipe, i, t, callback_kwargs): def test_encode_prompt_works_in_isolation(self): extra_required_param_value_dict = { "device": torch.device(torch_device).type, - "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) +class TestStableDiffusionImg2ImgPipelineMemory(StableDiffusionImg2ImgPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the img2img pipeline.""" + + +class TestStableDiffusionImg2ImgPipelineIPAdapter(StableDiffusionImg2ImgPipelineTesterConfig, IPAdapterTesterMixin): + """IP-Adapter tests for the img2img pipeline.""" + + @slow @require_torch_accelerator -class StableDiffusionImg2ImgPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestStableDiffusionImg2ImgPipelineSlow: + @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) @@ -621,14 +561,12 @@ def test_img2img_safety_checker_works(self): @nightly @require_torch_accelerator -class StableDiffusionImg2ImgPipelineNightlyTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestStableDiffusionImg2ImgPipelineNightly: + @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/stable_diffusion/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py index a69f71c37530..706d7dc2c923 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py @@ -15,9 +15,9 @@ import gc import random -import unittest import numpy as np +import pytest import torch from huggingface_hub import hf_hub_download from PIL import Image @@ -38,11 +38,11 @@ from ...testing_utils import ( Expectations, + assert_tensors_close, backend_empty_cache, backend_max_memory_allocated, backend_reset_max_memory_allocated, backend_reset_peak_memory_stats, - enable_full_determinism, floats_tensor, load_image, load_numpy, @@ -54,33 +54,19 @@ from ..pipeline_params import ( TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS, - TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( - IPAdapterTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, PipelineTesterMixin, ) +from .ip_adapter_tester import IPAdapterTesterMixin -enable_full_determinism() - - -class StableDiffusionInpaintPipelineFastTests( - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineTesterMixin, - unittest.TestCase, -): +class StableDiffusionInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionInpaintPipeline - params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - image_params = frozenset([]) - # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess - image_latents_params = frozenset([]) - callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"mask", "masked_image_latents"}) + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS + batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) @@ -132,150 +118,128 @@ def get_dummy_components(self, time_cond_proj_dim=None): } return components - def get_dummy_inputs(self, device, seed=0, img_res=64, output_pil=True): + def get_dummy_inputs(self, img_res=64): # TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched - if output_pil: - # Get random floats in [0, 1] as image - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - image = image.cpu().permute(0, 2, 3, 1)[0] - mask_image = torch.ones_like(image) - # Convert image and mask_image to [0, 255] - image = 255 * image - mask_image = 255 * mask_image - # Convert to PIL image - init_image = Image.fromarray(np.uint8(image)).convert("RGB").resize((img_res, img_res)) - mask_image = Image.fromarray(np.uint8(mask_image)).convert("RGB").resize((img_res, img_res)) - else: - # Get random floats in [0, 1] as image with spatial size (img_res, img_res) - image = floats_tensor((1, 3, img_res, img_res), rng=random.Random(seed)).to(device) - # Convert image to [-1, 1] - init_image = 2.0 * image - 1.0 - mask_image = torch.ones((1, 1, img_res, img_res), device=device) - - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) + # Get random floats in [0, 1] as image + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)) + image = image.permute(0, 2, 3, 1)[0] + mask_image = torch.ones_like(image) + # Convert image and mask_image to [0, 255] and then to PIL + init_image = Image.fromarray(np.uint8(255 * image)).convert("RGB").resize((img_res, img_res)) + mask_image = Image.fromarray(np.uint8(255 * mask_image)).convert("RGB").resize((img_res, img_res)) inputs = { "prompt": "A painting of a squirrel eating a burger", "image": init_image, "mask_image": mask_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs - def test_stable_diffusion_inpaint(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionInpaintPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + def get_dummy_tensor_inputs(self, img_res=64): + """Same as `get_dummy_inputs`, with the image and the mask passed as tensors instead of PIL images.""" + # Get random floats in [0, 1] as image with spatial size (img_res, img_res) and convert them to [-1, 1] + image = floats_tensor((1, 3, img_res, img_res), rng=random.Random(0)) + + inputs = self.get_dummy_inputs() + inputs["image"] = 2.0 * image - 1.0 + inputs["mask_image"] = torch.ones((1, 1, img_res, img_res)) + return inputs - 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.4816, 0.5766, 0.3897, 0.5448, 0.5982, 0.4359, 0.5142, 0.4836, 0.4536]) +class TestStableDiffusionInpaintPipeline(StableDiffusionInpaintPipelineTesterConfig, PipelineTesterMixin): + def test_stable_diffusion_inpaint(self): + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 64, 64) + + # fmt: off + expected_slice = torch.tensor([0.4816, 0.5766, 0.3897, 0.5448, 0.5982, 0.4359, 0.5142, 0.4836, 0.4536]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) def test_stable_diffusion_inpaint_lcm(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionInpaintPipeline(**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(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] - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4979, 0.5907, 0.4608, 0.5595, 0.6673, 0.5135, 0.6035, 0.5464, 0.5325]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 64, 64) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # fmt: off + expected_slice = torch.tensor([0.4979, 0.5907, 0.4608, 0.5595, 0.6673, 0.5135, 0.6035, 0.5464, 0.5325]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) def test_stable_diffusion_inpaint_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 = StableDiffusionInpaintPipeline(**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(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] image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4979, 0.5907, 0.4608, 0.5595, 0.6673, 0.5135, 0.6035, 0.5464, 0.5325]) + assert image.shape == (1, 3, 64, 64) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # Custom timesteps matching the default schedule reproduce `test_stable_diffusion_inpaint_lcm`'s output. + # fmt: off + expected_slice = torch.tensor([0.4979, 0.5907, 0.4608, 0.5595, 0.6673, 0.5135, 0.6035, 0.5464, 0.5325]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) def test_stable_diffusion_inpaint_image_tensor(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionInpaintPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + # Run on CPU: PIL and tensor inputs are compared against each other. + sd_pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - output = sd_pipe(**inputs) - out_pil = output.images + inputs = self.get_dummy_inputs() + out_pil = sd_pipe(**inputs).images - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["image"] = torch.tensor(np.array(inputs["image"]) / 127.5 - 1).permute(2, 0, 1).unsqueeze(0) inputs["mask_image"] = torch.tensor(np.array(inputs["mask_image"]) / 255).permute(2, 0, 1)[:1].unsqueeze(0) - output = sd_pipe(**inputs) - out_tensor = output.images + out_tensor = sd_pipe(**inputs).images - assert out_pil.shape == (1, 64, 64, 3) - assert np.abs(out_pil.flatten() - out_tensor.flatten()).max() < 5e-2 + assert out_pil.shape == (1, 3, 64, 64) + assert (out_pil - out_tensor).abs().max() < 5e-2 def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-3) def test_stable_diffusion_inpaint_strength_zero_test(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionInpaintPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() # check that the pipeline raises value error when num_inference_steps is < 1 inputs["strength"] = 0.01 - with self.assertRaises(ValueError): - sd_pipe(**inputs).images + with pytest.raises(ValueError): + sd_pipe(**inputs) def test_stable_diffusion_inpaint_mask_latents(self): - device = "cpu" - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components).to(device) - sd_pipe.set_progress_bar_config(disable=None) + # Run on CPU: the two runs below seed their VAE sampling with a CPU generator. + sd_pipe = self.get_pipeline() # normal mask + normal image ## `image`: pil, `mask_image``: pil, `masked_image_latents``: None - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["strength"] = 0.9 out_0 = sd_pipe(**inputs).images # image latents + mask latents - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = sd_pipe.image_processor.preprocess(inputs["image"]).to(sd_pipe.device) mask = sd_pipe.mask_processor.preprocess(inputs["mask_image"]).to(sd_pipe.device) masked_image = image * (mask < 0.5) - generator = torch.Generator(device=device).manual_seed(0) + generator = self.get_generator(0) image_latents = ( sd_pipe.vae.encode(image).latent_dist.sample(generator=generator) * sd_pipe.vae.config.scaling_factor ) @@ -288,19 +252,17 @@ def test_stable_diffusion_inpaint_mask_latents(self): inputs["masked_image_latents"] = mask_latents inputs["mask_image"] = mask inputs["strength"] = 0.9 - generator = torch.Generator(device=device).manual_seed(0) + generator = self.get_generator(0) torch.randn((1, 4, 32, 32), generator=generator) inputs["generator"] = generator out_1 = sd_pipe(**inputs).images - assert np.abs(out_0 - out_1).max() < 1e-2 + + assert (out_0 - out_1).abs().max() < 1e-2 def test_pipeline_interrupt(self): - components = self.get_dummy_components() - sd_pipe = StableDiffusionInpaintPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() prompt = "hey" num_inference_steps = 3 @@ -320,7 +282,7 @@ def apply(self, pipe, i, t, callback_kwargs): image=inputs["image"], mask_image=inputs["mask_image"], num_inference_steps=num_inference_steps, - output_type="np", + output_type="pt", generator=torch.Generator("cpu").manual_seed(0), callback_on_step_end=pipe_state.apply, ).images @@ -355,21 +317,19 @@ def callback_on_step_end(pipe, i, t, callback_kwargs): def test_encode_prompt_works_in_isolation(self): extra_required_param_value_dict = { "device": torch.device(torch_device).type, - "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict, atol=1e-3, rtol=1e-3) -class StableDiffusionSimpleInpaintPipelineFastTests(StableDiffusionInpaintPipelineFastTests): - pipeline_class = StableDiffusionInpaintPipeline - params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - image_params = frozenset([]) - # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess +class StableDiffusionSimpleInpaintPipelineTesterConfig(StableDiffusionInpaintPipelineTesterConfig): + """Same pipeline driven by a regular (non-inpaint-specific) UNet, i.e. one taking 4 input channels.""" def get_dummy_components(self, time_cond_proj_dim=None): + components = super().get_dummy_components(time_cond_proj_dim=time_cond_proj_dim) + torch.manual_seed(0) - unet = UNet2DConditionModel( + components["unet"] = UNet2DConditionModel( block_out_channels=(32, 64), layers_per_block=2, time_cond_proj_dim=time_cond_proj_dim, @@ -380,187 +340,127 @@ def get_dummy_components(self, time_cond_proj_dim=None): up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), cross_attention_dim=32, ) - scheduler = PNDMScheduler(skip_prk_steps=True) - torch.manual_seed(0) - vae = AutoencoderKL( - block_out_channels=[32, 64], - in_channels=3, - out_channels=3, - down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], - 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") - - components = { - "unet": unet, - "scheduler": scheduler, - "vae": vae, - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "safety_checker": None, - "feature_extractor": None, - "image_encoder": None, - } return components - def get_dummy_inputs_2images(self, device, seed=0, img_res=64): - # Get random floats in [0, 1] as image with spatial size (img_res, img_res) - image1 = floats_tensor((1, 3, img_res, img_res), rng=random.Random(seed)).to(device) - image2 = floats_tensor((1, 3, img_res, img_res), rng=random.Random(seed + 22)).to(device) - # Convert images to [-1, 1] - init_image1 = 2.0 * image1 - 1.0 - init_image2 = 2.0 * image2 - 1.0 + def get_dummy_inputs_2images(self, img_res=64): + """Two different images (and empty masks) batched into a single call.""" + # Get random floats in [0, 1] as image with spatial size (img_res, img_res) and convert them to [-1, 1] + image1 = floats_tensor((1, 3, img_res, img_res), rng=random.Random(0)) + image2 = floats_tensor((1, 3, img_res, img_res), rng=random.Random(22)) + inputs = self.get_dummy_inputs() + inputs["prompt"] = [inputs["prompt"]] * 2 + inputs["image"] = [2.0 * image1 - 1.0, 2.0 * image2 - 1.0] # empty mask - mask_image = torch.zeros((1, 1, img_res, img_res), device=device) + inputs["mask_image"] = [torch.zeros((1, 1, img_res, img_res))] * 2 + inputs["generator"] = [self.get_generator(0), self.get_generator(0)] + return inputs - if str(device).startswith("mps"): - generator1 = torch.manual_seed(seed) - generator2 = torch.manual_seed(seed) - else: - generator1 = torch.Generator(device=device).manual_seed(seed) - generator2 = torch.Generator(device=device).manual_seed(seed) - inputs = { - "prompt": ["A painting of a squirrel eating a burger"] * 2, - "image": [init_image1, init_image2], - "mask_image": [mask_image] * 2, - "generator": [generator1, generator2], - "num_inference_steps": 2, - "guidance_scale": 6.0, - "output_type": "np", - } - return inputs +class TestStableDiffusionSimpleInpaintPipeline( + StableDiffusionSimpleInpaintPipelineTesterConfig, TestStableDiffusionInpaintPipeline +): + """Reruns the inpaint tests against the 4-channel UNet, with its own expected slices.""" def test_stable_diffusion_inpaint(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionInpaintPipeline(**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] + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.6584, 0.5424, 0.5649, 0.5449, 0.5897, 0.6111, 0.5404, 0.5463, 0.5214]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 64, 64) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # fmt: off + expected_slice = torch.tensor([0.6584, 0.5424, 0.5649, 0.5449, 0.5897, 0.6111, 0.5404, 0.5463, 0.5214]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) def test_stable_diffusion_inpaint_lcm(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionInpaintPipeline(**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(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] + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 64, 64) - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.6240, 0.5355, 0.5649, 0.5378, 0.5374, 0.6242, 0.5132, 0.5347, 0.5396]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # fmt: off + expected_slice = torch.tensor([0.6240, 0.5355, 0.5649, 0.5378, 0.5374, 0.6242, 0.5132, 0.5347, 0.5396]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) def test_stable_diffusion_inpaint_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 = StableDiffusionInpaintPipeline(**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(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] image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.6240, 0.5355, 0.5649, 0.5378, 0.5374, 0.6242, 0.5132, 0.5347, 0.5396]) + assert image.shape == (1, 3, 64, 64) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # Custom timesteps matching the default schedule reproduce `test_stable_diffusion_inpaint_lcm`'s output. + # fmt: off + expected_slice = torch.tensor([0.6240, 0.5355, 0.5649, 0.5378, 0.5374, 0.6242, 0.5132, 0.5347, 0.5396]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2) def test_stable_diffusion_inpaint_2_images(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + # Run on CPU: the two runs below are compared against each other. + sd_pipe = self.get_pipeline() # test to confirm if we pass two same image, we will get same output - inputs = self.get_dummy_inputs(device) - gen1 = torch.Generator(device=device).manual_seed(0) - gen2 = torch.Generator(device=device).manual_seed(0) + inputs = self.get_dummy_inputs() for name in ["prompt", "image", "mask_image"]: inputs[name] = [inputs[name]] * 2 - inputs["generator"] = [gen1, gen2] + inputs["generator"] = [self.get_generator(0), self.get_generator(0)] images = sd_pipe(**inputs).images + assert images.shape == (2, 3, 64, 64) - assert images.shape == (2, 64, 64, 3) - - image_slice1 = images[0, -3:, -3:, -1] - image_slice2 = images[1, -3:, -3:, -1] - assert np.abs(image_slice1.flatten() - image_slice2.flatten()).max() < 1e-4 + image_slice1 = images[0, -1, -3:, -3:] + image_slice2 = images[1, -1, -3:, -3:] + assert (image_slice1 - image_slice2).abs().max() < 1e-4 # test to confirm that if we pass two different images, we will get different output - inputs = self.get_dummy_inputs_2images(device) - images = sd_pipe(**inputs).images - assert images.shape == (2, 64, 64, 3) + images = sd_pipe(**self.get_dummy_inputs_2images()).images + assert images.shape == (2, 3, 64, 64) - image_slice1 = images[0, -3:, -3:, -1] - image_slice2 = images[1, -3:, -3:, -1] - assert np.abs(image_slice1.flatten() - image_slice2.flatten()).max() > 1e-2 + image_slice1 = images[0, -1, -3:, -3:] + image_slice2 = images[1, -1, -3:, -3:] + assert (image_slice1 - image_slice2).abs().max() > 1e-2 def test_stable_diffusion_inpaint_euler(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionInpaintPipeline(**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 = EulerAncestralDiscreteScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(device, output_pil=False) + inputs = self.get_dummy_tensor_inputs() half_dim = inputs["image"].shape[2] // 2 inputs["mask_image"][0, 0, :half_dim, :half_dim] = 0 - inputs["num_inference_steps"] = 4 + image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] + assert image.shape == (1, 3, 64, 64) - assert image.shape == (1, 64, 64, 3) + # fmt: off + expected_slice = torch.tensor([0.6443541, 0.55819386, 0.58656645, 0.5574256, 0.55064464, 0.64331865, 0.5232371, 0.5466772, 0.5478098]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-4) - expected_slice = np.array( - [[0.6443541, 0.55819386, 0.58656645, 0.5574256, 0.55064464, 0.64331865, 0.5232371, 0.5466772, 0.5478098]] - ) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-4 + +class TestStableDiffusionInpaintPipelineMemory(StableDiffusionInpaintPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the inpaint pipeline.""" + + +class TestStableDiffusionInpaintPipelineIPAdapter(StableDiffusionInpaintPipelineTesterConfig, IPAdapterTesterMixin): + """IP-Adapter tests for the inpaint pipeline.""" @slow @require_torch_accelerator -class StableDiffusionInpaintPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() - - def tearDown(self): - super().tearDown() +class TestStableDiffusionInpaintPipelineSlow: + @pytest.fixture(autouse=True) + def cleanup(self): + yield gc.collect() backend_empty_cache(torch_device) @@ -735,12 +635,10 @@ def test_stable_diffusion_simple_inpaint_ddim(self): @slow @require_torch_accelerator -class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() - - def tearDown(self): - super().tearDown() +class TestStableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlow: + @pytest.fixture(autouse=True) + def cleanup(self): + yield gc.collect() backend_empty_cache(torch_device) @@ -986,14 +884,12 @@ def test_download_local(self): @nightly @require_torch_accelerator -class StableDiffusionInpaintPipelineNightlyTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestStableDiffusionInpaintPipelineNightly: + @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/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py index 0852965c2fe8..9045e92b8be9 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py @@ -15,9 +15,9 @@ import gc import random -import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -31,14 +31,13 @@ StableDiffusionInstructPix2PixPipeline, UNet2DConditionModel, ) -from diffusers.image_processor import VaeImageProcessor from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, backend_max_memory_allocated, backend_reset_max_memory_allocated, backend_reset_peak_memory_stats, - enable_full_determinism, floats_tensor, load_image, require_torch_accelerator, @@ -46,30 +45,20 @@ torch_device, ) from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS, - TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, PipelineTesterMixin, ) -enable_full_determinism() - - -class StableDiffusionInstructPix2PixPipelineFastTests( - PipelineLatentTesterMixin, PipelineKarrasSchedulerTesterMixin, PipelineTesterMixin, unittest.TestCase -): +class StableDiffusionInstructPix2PixPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionInstructPix2PixPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "cross_attention_kwargs"} - batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union({"image_latents"}) - {"negative_prompt_embeds"} + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"} + batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS def get_dummy_components(self): torch.manual_seed(0) @@ -120,165 +109,106 @@ def get_dummy_components(self): } return components - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - image = image.cpu().permute(0, 2, 3, 1)[0] + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)) + image = image.permute(0, 2, 3, 1)[0] image = Image.fromarray(np.uint8(image)).convert("RGB") - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) inputs = { "prompt": "A painting of a squirrel eating a burger", "image": image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "image_guidance_scale": 1, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs + +class TestStableDiffusionInstructPix2PixPipeline( + StableDiffusionInstructPix2PixPipelineTesterConfig, PipelineTesterMixin +): def test_stable_diffusion_pix2pix_default_case(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionInstructPix2PixPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.7024, 0.3769, 0.4392, 0.5858, 0.5620, 0.4827, 0.4174, 0.5506, 0.4806]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 32, 32) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + # fmt: off + expected_slice = torch.tensor([0.7024, 0.3769, 0.4392, 0.5858, 0.5620, 0.4827, 0.4174, 0.5506, 0.4806]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_stable_diffusion_pix2pix_negative_prompt(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionInstructPix2PixPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - negative_prompt = "french fries" - output = sd_pipe(**inputs, negative_prompt=negative_prompt) - image = output.images - image_slice = image[0, -3:, -3:, -1] + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.7059, 0.3731, 0.4428, 0.5922, 0.5649, 0.4871, 0.4203, 0.5507, 0.4809]) + image = sd_pipe(**self.get_dummy_inputs(), negative_prompt="french fries").images + assert image.shape == (1, 3, 32, 32) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + # fmt: off + expected_slice = torch.tensor([0.7059, 0.3731, 0.4428, 0.5922, 0.5649, 0.4871, 0.4203, 0.5507, 0.4809]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_stable_diffusion_pix2pix_multiple_init_images(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionInstructPix2PixPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + sd_pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["prompt"] = [inputs["prompt"]] * 2 image = np.array(inputs["image"]).astype(np.float32) / 255.0 - image = torch.from_numpy(image).unsqueeze(0).to(device) + image = torch.from_numpy(image).unsqueeze(0) image = image / 2 + 0.5 image = image.permute(0, 3, 1, 2) inputs["image"] = image.repeat(2, 1, 1, 1) image = sd_pipe(**inputs).images - image_slice = image[-1, -3:, -3:, -1] - - assert image.shape == (2, 32, 32, 3) - expected_slice = np.array([0.5710, 0.5834, 0.5128, 0.5892, 0.5773, 0.7047, 0.6797, 0.5411, 0.5536]) + assert image.shape == (2, 3, 32, 32) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + # fmt: off + expected_slice = torch.tensor([0.5710, 0.5834, 0.5128, 0.5892, 0.5773, 0.7047, 0.6797, 0.5411, 0.5536]) + # fmt: on + assert_tensors_close(image[-1, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_stable_diffusion_pix2pix_euler(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator + # Run on CPU: the expected slice below is CPU-specific. components = self.get_dummy_components() components["scheduler"] = EulerAncestralDiscreteScheduler( beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear" ) - sd_pipe = StableDiffusionInstructPix2PixPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) + sd_pipe = self.get_pipeline(**components) - inputs = self.get_dummy_inputs(device) - image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.6674, 0.3879, 0.4447, 0.5375, 0.5464, 0.4881, 0.3896, 0.5467, 0.4923]) + image = sd_pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, 3, 32, 32) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + # fmt: off + expected_slice = torch.tensor([0.6674, 0.3879, 0.4447, 0.5375, 0.5464, 0.4881, 0.3896, 0.5467, 0.4923]) + # fmt: on + assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3) def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-3) - # Overwrite the default test_latents_inputs because pix2pix encode the image differently - def test_latents_input(self): - components = self.get_dummy_components() - pipe = StableDiffusionInstructPix2PixPipeline(**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") - - for image_param in self.image_latents_params: - if image_param in inputs.keys(): - inputs[image_param] = vae.encode(inputs[image_param]).latent_dist.mode() - - 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") - - # Override the default test_callback_cfg because pix2pix create inputs for cfg differently - def test_callback_cfg(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - def callback_no_cfg(pipe, i, t, callback_kwargs): - if i == 1: - for k, w in callback_kwargs.items(): - if k in self.callback_cfg_params: - callback_kwargs[k] = callback_kwargs[k].chunk(3)[0] - pipe._guidance_scale = 1.0 - - return callback_kwargs - inputs = self.get_dummy_inputs(torch_device) - inputs["guidance_scale"] = 1.0 - inputs["num_inference_steps"] = 2 - out_no_cfg = pipe(**inputs)[0] - - inputs["guidance_scale"] = 7.5 - inputs["callback_on_step_end"] = callback_no_cfg - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - out_callback_no_cfg = pipe(**inputs)[0] - - assert out_no_cfg.shape == out_callback_no_cfg.shape +class TestStableDiffusionInstructPix2PixPipelineMemory( + StableDiffusionInstructPix2PixPipelineTesterConfig, MemoryTesterMixin +): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the InstructPix2Pix pipeline.""" @slow @require_torch_accelerator -class StableDiffusionInstructPix2PixPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestStableDiffusionInstructPix2PixPipelineSlow: + @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)