From 8e07788fc68d09af2de9cacfda51da9f580767c8 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 30 Jul 2026 12:21:39 +0000 Subject: [PATCH 1/3] 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) From 95e9ef2766f5967996f419a5c7913f5270eb5c53 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 8 Aug 2026 06:35:44 +0000 Subject: [PATCH 2/3] fix minimax h3 failing tests on GPU. --- .../autoencoders/autoencoder_kl_minimax_h3.py | 6 +++- .../autoencoder_kl_minimax_h3_audio.py | 4 +++ .../transformers/transformer_minimax_h3.py | 4 ++- .../test_models_autoencoder_kl_kvae_video.py | 28 ++++++++--------- .../test_models_autoencoder_kl_minimax_h3.py | 26 ++++++++++++++++ ..._models_autoencoder_kl_minimax_h3_audio.py | 28 +++++++++++++---- .../test_models_autoencoder_vidtok.py | 30 +++++++++---------- tests/models/testing_utils/__init__.py | 2 ++ tests/models/testing_utils/utils.py | 16 ++++++++++ 9 files changed, 106 insertions(+), 38 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index acc04f2c179c..23ad2c725c00 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py @@ -24,7 +24,7 @@ from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..modeling_outputs import AutoencoderKLOutput -from ..modeling_utils import ModelMixin +from ..modeling_utils import ModelMixin, get_parameter_dtype from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution @@ -857,6 +857,9 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutp The latent distribution of the encoded videos. Note that MiniMax-H3 normalizes the sampled latents with `latents_mean` / `latents_std` afterwards. """ + # Every module is pinned to float32 by `_keep_in_fp32_modules`, so a pipeline running in a lower `torch_dtype` + # hands over lower-precision pixels; align them with the weights, like the audio autoencoder does. + x = x.to(get_parameter_dtype(self.encoder)) if self.use_slicing and x.shape[0] > 1: moments = torch.cat([self._encode(x_slice) for x_slice in x.split(1)]) else: @@ -881,6 +884,7 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: The decoded videos, shape `(batch_size, out_channels, num_frames, height, width)`. """ + z = z.to(get_parameter_dtype(self.decoder)) if self.use_slicing and z.shape[0] > 1: decoded = torch.cat([self._decode(z_slice) for z_slice in z.split(1)]) else: diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py index 957774f791bc..a70c947712ae 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py @@ -524,6 +524,10 @@ class AutoencoderKLMiniMaxH3Audio(ModelMixin, ConfigMixin, AttentionMixin): """ _supports_gradient_checkpointing = False + # `weight_norm` recomputes `weight` from `weight_g` / `weight_v` in a forward pre-hook, which runs before the + # leaf-level group offloading hook has onloaded them, so the convolution would see a CPU weight. Same reason the + # other weight-normalized audio autoencoders (`AutoencoderOobleck`, `Cosmos3AVAEAudioTokenizer`) opt out. + _supports_group_offloading = False # The released checkpoint is float32 and the DAC/BigVGAN stack (weight-normalized convolutions, Snake # activations) degrades audibly under bfloat16 (roughly 20 dB quieter decodes), so a pipeline-level # `torch_dtype=torch.bfloat16` must not downcast the weights. diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 5b2be83bdbbd..bbabe7c6c7fe 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -51,7 +51,9 @@ class MiniMaxH3TransformerOutput(BaseOutput): """ sample: torch.Tensor - audio_sample: torch.Tensor + # `forward` always populates `audio_sample`; the default is what lets the output be rebuilt from a plain dict of + # its fields, which is how the accelerate offload hooks move a `BaseOutput` back to the input device. + audio_sample: torch.Tensor | None = None def _apply_rotary_emb(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py index a7097e128a13..e97dd8fd81ba 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py @@ -20,23 +20,19 @@ from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device -from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TrainingTesterMixin, + run_nondeterministic, +) from .testing_utils import NewAutoencoderTesterMixin enable_full_determinism() -def _run_nondeterministic(fn): - # reflection_pad3d_backward_out_cuda has no deterministic CUDA implementation; - # temporarily relax the requirement for tests that do backward passes. - torch.use_deterministic_algorithms(False) - try: - fn() - finally: - torch.use_deterministic_algorithms(True) - - class AutoencoderKLKVAEVideoTesterConfig(BaseModelTesterConfig): @property def model_class(self): @@ -91,14 +87,16 @@ def test_gradient_checkpointing_is_applied(self): expected_set = {"KVAECachedEncoder3D", "KVAECachedDecoder3D"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + # reflection_pad3d_backward_out_cuda has no deterministic implementation, so every test below that runs a + # backward pass has to relax determinism. def test_training(self): - _run_nondeterministic(super().test_training) + run_nondeterministic(super().test_training) def test_training_with_ema(self): - _run_nondeterministic(super().test_training_with_ema) + run_nondeterministic(super().test_training_with_ema) def test_mixed_precision_training(self): - _run_nondeterministic(super().test_mixed_precision_training) + run_nondeterministic(super().test_mixed_precision_training) @pytest.mark.skip( "Gradient checkpointing recomputes the forward pass, but the model uses a stateful cache_dict " @@ -113,7 +111,7 @@ class TestAutoencoderKLKVAEVideoMemory(AutoencoderKLKVAEVideoTesterConfig, Memor """Memory optimization tests for AutoencoderKLKVAEVideo.""" def test_layerwise_casting_training(self): - _run_nondeterministic(super().test_layerwise_casting_training) + run_nondeterministic(super().test_layerwise_casting_training) class TestAutoencoderKLKVAEVideoSlicingTiling(AutoencoderKLKVAEVideoTesterConfig, NewAutoencoderTesterMixin): diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py index cb6c4ebb12df..497614b38141 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py @@ -27,6 +27,7 @@ ModelTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, + run_nondeterministic, ) from .testing_utils import NewAutoencoderTesterMixin @@ -134,10 +135,35 @@ def test_encode_decode_temporal_geometry(self): class TestAutoencoderKLMiniMaxH3Memory(AutoencoderKLMiniMaxH3TesterConfig, MemoryTesterMixin): """Memory optimization tests for the MiniMax-H3 video autoencoder.""" + @pytest.mark.skip( + "`_keep_in_fp32_modules` pins every module of this autoencoder, so layerwise casting has nothing left to " + "cast and the memory footprint cannot go down." + ) + def test_layerwise_casting_memory(self): + pass + + # The encoder pads spatially with `mode="reflect"`, whose reflection_pad3d_backward_out_cuda has no deterministic + # implementation, so every test below that runs a backward pass has to relax determinism. + def test_layerwise_casting_training(self): + run_nondeterministic(super().test_layerwise_casting_training) + class TestAutoencoderKLMiniMaxH3Training(AutoencoderKLMiniMaxH3TesterConfig, TrainingTesterMixin): """Training tests for the MiniMax-H3 video autoencoder.""" + # See `TestAutoencoderKLMiniMaxH3Memory` for why these relax determinism. + def test_training(self): + run_nondeterministic(super().test_training) + + def test_training_with_ema(self): + run_nondeterministic(super().test_training_with_ema) + + def test_mixed_precision_training(self): + run_nondeterministic(super().test_mixed_precision_training) + + def test_gradient_checkpointing_equivalence(self): + run_nondeterministic(super().test_gradient_checkpointing_equivalence) + def test_gradient_checkpointing_is_applied(self): super().test_gradient_checkpointing_is_applied( expected_set={"MiniMaxH3VideoDownBlock3d", "MiniMaxH3VideoViTDecoder3d"} diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py index 676cae58d335..ab8956096664 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3_audio.py @@ -25,8 +25,8 @@ BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, - TorchCompileTesterMixin, TrainingTesterMixin, + run_nondeterministic, ) @@ -130,10 +130,32 @@ def test_encode_pads_to_the_hop_length(self): class TestAutoencoderKLMiniMaxH3AudioMemory(AutoencoderKLMiniMaxH3AudioTesterConfig, MemoryTesterMixin): """Memory optimization tests for the MiniMax-H3 audio autoencoder.""" + @pytest.mark.skip( + "`_keep_in_fp32_modules` pins every module of this autoencoder, so layerwise casting has nothing left to " + "cast and the memory footprint cannot go down." + ) + def test_layerwise_casting_memory(self): + pass + + # The latent projection pools with `F.adaptive_avg_pool1d`, whose adaptive_avg_pool2d_backward_cuda has no + # deterministic implementation, so every test below that runs a backward pass has to relax determinism. + def test_layerwise_casting_training(self): + run_nondeterministic(super().test_layerwise_casting_training) + class TestAutoencoderKLMiniMaxH3AudioTraining(AutoencoderKLMiniMaxH3AudioTesterConfig, TrainingTesterMixin): """Training tests for the MiniMax-H3 audio autoencoder.""" + # See `TestAutoencoderKLMiniMaxH3AudioMemory` for why these relax determinism. + def test_training(self): + run_nondeterministic(super().test_training) + + def test_training_with_ema(self): + run_nondeterministic(super().test_training_with_ema) + + def test_mixed_precision_training(self): + run_nondeterministic(super().test_mixed_precision_training) + class TestAutoencoderKLMiniMaxH3AudioAttention(AutoencoderKLMiniMaxH3AudioTesterConfig, AttentionTesterMixin): """Attention processor tests for the MiniMax-H3 audio autoencoder.""" @@ -144,7 +166,3 @@ class TestAutoencoderKLMiniMaxH3AudioAttention(AutoencoderKLMiniMaxH3AudioTester ) def test_attention_processor_count_mismatch_raises_error(self): pass - - -class TestAutoencoderKLMiniMaxH3AudioTorchCompile(AutoencoderKLMiniMaxH3AudioTesterConfig, TorchCompileTesterMixin): - """Torch compile tests for the MiniMax-H3 audio autoencoder.""" diff --git a/tests/models/autoencoders/test_models_autoencoder_vidtok.py b/tests/models/autoencoders/test_models_autoencoder_vidtok.py index 9810296a07d9..61695b377f44 100644 --- a/tests/models/autoencoders/test_models_autoencoder_vidtok.py +++ b/tests/models/autoencoders/test_models_autoencoder_vidtok.py @@ -19,23 +19,19 @@ from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device -from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TrainingTesterMixin, + run_nondeterministic, +) from .testing_utils import NewAutoencoderTesterMixin enable_full_determinism() -def _run_nondeterministic(fn): - # avg_pool3d_backward_cuda has no deterministic CUDA implementation; - # temporarily relax the requirement for tests that do backward passes. - torch.use_deterministic_algorithms(False) - try: - fn() - finally: - torch.use_deterministic_algorithms(True) - - class AutoencoderVidTokTesterConfig(BaseModelTesterConfig): @property def model_class(self): @@ -90,24 +86,26 @@ def test_gradient_checkpointing_is_applied(self): expected_set = {"VidTokEncoder3D", "VidTokDecoder3D"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + # avg_pool3d_backward_cuda has no deterministic implementation, so every test below that runs a backward pass + # has to relax determinism. def test_training(self): - _run_nondeterministic(super().test_training) + run_nondeterministic(super().test_training) def test_training_with_ema(self): - _run_nondeterministic(super().test_training_with_ema) + run_nondeterministic(super().test_training_with_ema) def test_mixed_precision_training(self): - _run_nondeterministic(super().test_mixed_precision_training) + run_nondeterministic(super().test_mixed_precision_training) def test_gradient_checkpointing_equivalence(self): - _run_nondeterministic(super().test_gradient_checkpointing_equivalence) + run_nondeterministic(super().test_gradient_checkpointing_equivalence) class TestAutoencoderVidTokMemory(AutoencoderVidTokTesterConfig, MemoryTesterMixin): """Memory optimization tests for AutoencoderVidTok.""" def test_layerwise_casting_training(self): - _run_nondeterministic(super().test_layerwise_casting_training) + run_nondeterministic(super().test_layerwise_casting_training) class TestAutoencoderVidTokSlicingTiling(AutoencoderVidTokTesterConfig, NewAutoencoderTesterMixin): diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 23aa871b80a2..67161159732a 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -48,6 +48,7 @@ ) from .single_file import SingleFileTesterMixin from .training import TrainingTesterMixin +from .utils import run_nondeterministic __all__ = [ @@ -94,6 +95,7 @@ "QuantoCompileTesterMixin", "QuantoConfigMixin", "QuantoTesterMixin", + "run_nondeterministic", "SDNQCompileTesterMixin", "SDNQConfigMixin", "SDNQTesterMixin", diff --git a/tests/models/testing_utils/utils.py b/tests/models/testing_utils/utils.py index 3beb59ed1a66..6c11552ae011 100644 --- a/tests/models/testing_utils/utils.py +++ b/tests/models/testing_utils/utils.py @@ -24,3 +24,19 @@ def _maybe_cast_to_bf16(backend, model, inputs_dict): for k, v in inputs_dict.items() } return model, inputs_dict + + +def run_nondeterministic(fn): + """ + Run `fn` with `enable_full_determinism`'s deterministic-algorithm requirement lifted. + + Several models reach a backward kernel that has no deterministic CUDA implementation (reflection/replication + padding, average pooling), which makes every test doing a backward pass raise under + `torch.use_deterministic_algorithms(True)`. Wrap those tests instead of relaxing determinism for the whole module, + and name the offending op at the call site. + """ + torch.use_deterministic_algorithms(False) + try: + fn() + finally: + torch.use_deterministic_algorithms(True) From d85dafba102d47c68c75ded3c6c134aaf088ed10 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 8 Aug 2026 06:36:27 +0000 Subject: [PATCH 3/3] Revert "refactor stable diffusion pipeline tests" This reverts commit 8e07788fc68d09af2de9cacfda51da9f580767c8. --- .../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, 894 insertions(+), 769 deletions(-) delete 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 deleted file mode 100644 index a7e614879e7e..000000000000 --- a/tests/pipelines/stable_diffusion/ip_adapter_tester.py +++ /dev/null @@ -1,241 +0,0 @@ -# coding=utf-8 -# Copyright 2026 HuggingFace Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import inspect - -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 31a8aec3c58b..23607a39ecd0 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -15,10 +15,11 @@ 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 ( @@ -43,11 +44,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, @@ -60,20 +61,36 @@ ) 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 ..testing_utils import ( - BasePipelineTesterConfig, - MemoryTesterMixin, +from ..test_pipelines_common import ( + IPAdapterTesterMixin, + PipelineKarrasSchedulerTesterMixin, + PipelineLatentTesterMixin, PipelineTesterMixin, ) -from .ip_adapter_tester import IPAdapterTesterMixin -class StableDiffusionPipelineTesterConfig(BasePipelineTesterConfig): +enable_full_determinism() + + +class StableDiffusionPipelineFastTests( + IPAdapterTesterMixin, + PipelineLatentTesterMixin, + PipelineKarrasSchedulerTesterMixin, + PipelineTesterMixin, + unittest.TestCase, +): pipeline_class = StableDiffusionPipeline - required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + 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 def get_dummy_components(self, time_cond_proj_dim=None): cross_attention_dim = 8 @@ -135,61 +152,80 @@ def get_dummy_components(self, time_cond_proj_dim=None): } return components - def get_dummy_inputs(self): + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) inputs = { "prompt": "A painting of a squirrel eating a burger", - "generator": self.get_generator(0), + "generator": generator, "num_inference_steps": 2, "guidance_scale": 6.0, - # 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", + "output_type": "np", } return inputs - -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() + 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) - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 64, 64) + inputs = self.get_dummy_inputs(device) + output = sd_pipe(**inputs) + image = output.images + + image_slice = image[0, -3:, -3:, -1] - # 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) + 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]) + + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 def test_stable_diffusion_lcm(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + 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.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 = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 64, 64) + image_slice = image[0, -3:, -3:, -1] - # 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) + 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 def test_stable_diffusion_lcm_custom_timesteps(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + 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.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() + inputs = self.get_dummy_inputs(device) del inputs["num_inference_steps"] inputs["timesteps"] = [999, 499] - image = sd_pipe(**inputs).images - assert image.shape == (1, 3, 64, 64) + output = sd_pipe(**inputs) + image = output.images + + image_slice = image[0, -3:, -3:, -1] - # 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) + 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 def test_stable_diffusion_ays(self): from diffusers.schedulers import AysSchedules @@ -197,38 +233,53 @@ def test_stable_diffusion_ays(self): timestep_schedule = AysSchedules["StableDiffusionTimesteps"] sigma_schedule = AysSchedules["StableDiffusionSigmas"] - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)).to(torch_device) + 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.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() + inputs = self.get_dummy_inputs(device) inputs["num_inference_steps"] = 10 output = sd_pipe(**inputs).images - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(device) inputs["num_inference_steps"] = None inputs["timesteps"] = timestep_schedule output_ts = sd_pipe(**inputs).images - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(device) inputs["num_inference_steps"] = None inputs["sigmas"] = sigma_schedule output_sigmas = sd_pipe(**inputs).images - assert (output_sigmas - output_ts).abs().max() < 1e-3, ( + assert np.abs(output_sigmas.flatten() - output_ts.flatten()).max() < 1e-3, ( "ays timesteps and ays sigmas should have the same 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" + 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" + ) def test_stable_diffusion_prompt_embeds(self): - sd_pipe = self.get_pipeline().to(torch_device) + 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) - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(torch_device) inputs["prompt"] = 3 * [inputs["prompt"]] + + # forward output = sd_pipe(**inputs) - image_slice_1 = output.images[0, -1, -3:, -3:] + image_slice_1 = output.images[0, -3:, -3:, -1] - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(torch_device) prompt = 3 * [inputs.pop("prompt")] text_inputs = sd_pipe.tokenizer( @@ -239,26 +290,34 @@ 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, -1, -3:, -3:] + image_slice_2 = output.images[0, -3:, -3:, -1] - assert_tensors_close( - image_slice_2, image_slice_1, atol=1e-4, msg="Passing `prompt_embeds` changed the output." - ) + assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 def test_stable_diffusion_negative_prompt_embeds(self): - sd_pipe = self.get_pipeline().to(torch_device) + 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) - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(torch_device) 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, -1, -3:, -3:] + image_slice_1 = output.images[0, -3:, -3:, -1] - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(torch_device) prompt = 3 * [inputs.pop("prompt")] embeds = [] @@ -271,96 +330,55 @@ 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, -1, -3:, -3:] + image_slice_2 = output.images[0, -3:, -3:, -1] - assert_tensors_close( - image_slice_2, image_slice_1, atol=1e-4, msg="Passing `negative_prompt_embeds` changed the output." - ) + assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 def test_stable_diffusion_ddim_factor_8(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline() - - image = sd_pipe(**self.get_dummy_inputs(), height=136, width=136).images - assert image.shape == (1, 3, 136, 136) - - # 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) - - 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) - - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 64, 64) - - # 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_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) - - 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) + device = "cpu" # ensure determinism for the device-dependent torch.Generator - 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) + 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()).images - assert image.shape == (1, 3, 64, 64) + inputs = self.get_dummy_inputs(device) + output = sd_pipe(**inputs, height=136, width=136) + image = output.images - # 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) + image_slice = image[0, -3:, -3:, -1] - 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) + 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.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. + def test_stable_diffusion_pndm(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 = self.get_pipeline(**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) + + 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(), negative_prompt="french fries").images - assert image.shape == (1, 3, 64, 64) + 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]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - def test_stable_diffusion_no_safety_checker(self, tmp_path): + def test_stable_diffusion_no_safety_checker(self): pipe = StableDiffusionPipeline.from_pretrained( "hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None ) @@ -372,56 +390,150 @@ def test_stable_diffusion_no_safety_checker(self, tmp_path): assert image is not None # check that there's no error when saving a pipeline with one of the models being None - pipe.save_pretrained(tmp_path) - pipe = StableDiffusionPipeline.from_pretrained(tmp_path) + with tempfile.TemporaryDirectory() as tmpdirname: + pipe.save_pretrained(tmpdirname) + pipe = StableDiffusionPipeline.from_pretrained(tmpdirname) # 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): - # Run on CPU: sliced VAE decoding is compared against a full-batch decode of the same run. + device = "cpu" # ensure determinism for the device-dependent torch.Generator components = self.get_dummy_components() components["scheduler"] = LMSDiscreteScheduler.from_config(components["scheduler"].config) - sd_pipe = self.get_pipeline(**components) + sd_pipe = StableDiffusionPipeline(**components) + sd_pipe = sd_pipe.to(device) + sd_pipe.set_progress_bar_config(disable=None) image_count = 4 - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(device) 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() + inputs = self.get_dummy_inputs(device) inputs["prompt"] = [inputs["prompt"]] * image_count output_2 = sd_pipe(**inputs) # there is a small discrepancy at image borders vs. full batch decode - assert (output_2.images - output_1.images).abs().max() < 3e-3 + assert np.abs(output_2.images.flatten() - output_1.images.flatten()).max() < 3e-3 def test_stable_diffusion_vae_tiling(self): - # Run on CPU: tiled VAE decoding is compared against a non-tiled decode of the same run. - sd_pipe = self.get_pipeline() + 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" # Test that tiled decode at 512x512 yields the same result as the non-tiled decode - output_1 = sd_pipe(**self.get_dummy_inputs()) + 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") # make sure tiled vae decode yields the same result sd_pipe.vae.enable_tiling() - output_2 = sd_pipe(**self.get_dummy_inputs()) + 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") - assert (output_2.images - output_1.images).abs().max() < 5e-1 + assert np.abs(output_2.images.flatten() - output_1.images.flatten()).max() < 5e-1 # test that tiled decode works with various shapes - 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)) + 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 def test_stable_diffusion_long_prompt(self): components = self.get_dummy_components() components["scheduler"] = LMSDiscreteScheduler.from_config(components["scheduler"].config) - sd_pipe = self.get_pipeline(**components).to(torch_device) + sd_pipe = StableDiffusionPipeline(**components) + sd_pipe = sd_pipe.to(torch_device) + sd_pipe.set_progress_bar_config(disable=None) do_classifier_free_guidance = True negative_prompt = None @@ -465,21 +577,29 @@ 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 = self.get_pipeline(**components).to(torch_device) + sd_pipe = StableDiffusionPipeline(**components) + sd_pipe = sd_pipe.to(torch_device) + sd_pipe.set_progress_bar_config(disable=None) prompt = "hey" - 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, output_type="np") + image_shape = output.images[0].shape[:2] + assert image_shape == (64, 64) - output = sd_pipe(prompt, num_inference_steps=1, height=96, width=96, output_type="pt") - assert output.images[0].shape[-2:] == (96, 96) + 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) 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="pt") - assert output.images[0].shape[-2:] == (192, 192) + 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) def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=3e-3) @@ -487,63 +607,82 @@ 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): - sd_pipe = self.get_pipeline().to(torch_device) + components = self.get_dummy_components() + sd_pipe = StableDiffusionPipeline(**components) + sd_pipe = sd_pipe.to(torch_device) + sd_pipe.set_progress_bar_config(disable=None) prompt = "hey" - output = sd_pipe(prompt, num_inference_steps=1, output_type="pt", generator=torch.manual_seed(0)).images + output = sd_pipe(prompt, num_inference_steps=1, output_type="np", 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="pt", generator=torch.manual_seed(0)).images + output_freeu = sd_pipe(prompt, num_inference_steps=1, output_type="np", generator=torch.manual_seed(0)).images - assert not torch.allclose(output[0, -1, -3:, -3:], output_freeu[0, -1, -3:, -3:]), ( + assert not np.allclose(output[0, -3:, -3:, -1], output_freeu[0, -3:, -3:, -1]), ( "Enabling of FreeU should lead to different results." ) def test_freeu_disabled(self): - sd_pipe = self.get_pipeline().to(torch_device) + components = self.get_dummy_components() + sd_pipe = StableDiffusionPipeline(**components) + sd_pipe = sd_pipe.to(torch_device) + sd_pipe.set_progress_bar_config(disable=None) prompt = "hey" - output = sd_pipe(prompt, num_inference_steps=1, output_type="pt", generator=torch.manual_seed(0)).images + output = sd_pipe(prompt, num_inference_steps=1, output_type="np", 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 {"s1", "s2", "b1", "b2"}: + for key in freeu_keys: 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="pt", generator=torch.manual_seed(0) + prompt, num_inference_steps=1, output_type="np", generator=torch.manual_seed(0) ).images - assert torch.allclose(output[0, -1, -3:, -3:], output_no_freeu[0, -1, -3:, -3:]), ( + assert np.allclose(output[0, -3:, -3:, -1], output_no_freeu[0, -3:, -3:, -1]), ( "Disabling of FreeU should lead to results similar to the default pipeline results." ) def test_fused_qkv_projections(self): - # Run on CPU: fused and unfused attention are compared against each other. - sd_pipe = self.get_pipeline() + 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) - original_image_slice = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] + inputs = self.get_dummy_inputs(device) + image = sd_pipe(**inputs).images + original_image_slice = image[0, -3:, -3:, -1] sd_pipe.fuse_qkv_projections() - image_slice_fused = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] + inputs = self.get_dummy_inputs(device) + image = sd_pipe(**inputs).images + image_slice_fused = image[0, -3:, -3:, -1] sd_pipe.unfuse_qkv_projections() - image_slice_disabled = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] + inputs = self.get_dummy_inputs(device) + image = sd_pipe(**inputs).images + image_slice_disabled = image[0, -3:, -3:, -1] - assert torch.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( + assert np.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( "Fusion of QKV projections shouldn't affect the outputs." ) - assert torch.allclose(image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2), ( + assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2), ( "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." ) - assert torch.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( + assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( "Original outputs should match when fused QKV projections are disabled." ) def test_pipeline_interrupt(self): - sd_pipe = self.get_pipeline().to(torch_device) + components = self.get_dummy_components() + sd_pipe = StableDiffusionPipeline(**components) + sd_pipe = sd_pipe.to(torch_device) + sd_pipe.set_progress_bar_config(disable=None) prompt = "hey" num_inference_steps = 3 @@ -561,7 +700,7 @@ def apply(self, pipe, i, t, callback_kwargs): sd_pipe( prompt, num_inference_steps=num_inference_steps, - output_type="pt", + output_type="np", generator=torch.Generator("cpu").manual_seed(0), callback_on_step_end=pipe_state.apply, ).images @@ -602,27 +741,15 @@ 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().get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) -class 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 TestStableDiffusionPipelineSlow: - @pytest.fixture(autouse=True) - def cleanup(self): - gc.collect() - backend_empty_cache(torch_device) - yield +class StableDiffusionPipelineSlowTests(unittest.TestCase): + def setUp(self): gc.collect() backend_empty_cache(torch_device) @@ -1062,12 +1189,14 @@ def test_stable_diffusion_textual_inversion_with_sequential_cpu_offload(self): @slow @require_torch_accelerator -class TestStableDiffusionPipelineCkpt: - @pytest.fixture(autouse=True) - def cleanup(self): +class StableDiffusionPipelineCkptTests(unittest.TestCase): + def setUp(self): + super().setUp() gc.collect() backend_empty_cache(torch_device) - yield + + def tearDown(self): + super().tearDown() gc.collect() backend_empty_cache(torch_device) @@ -1105,12 +1234,14 @@ def test_download_local(self): @nightly @require_torch_accelerator -class TestStableDiffusionPipelineNightly: - @pytest.fixture(autouse=True) - def cleanup(self): +class StableDiffusionPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() gc.collect() backend_empty_cache(torch_device) - yield + + def tearDown(self): + super().tearDown() gc.collect() backend_empty_cache(torch_device) @@ -1208,10 +1339,9 @@ def test_stable_diffusion_euler(self): @slow @require_torch_multi_accelerator @require_accelerate_version_greater("0.27.0") -class TestStableDiffusionPipelineDeviceMap: - @pytest.fixture(autouse=True) - def cleanup(self): - yield +class StableDiffusionPipelineDeviceMapTests(unittest.TestCase): + def tearDown(self): + super().tearDown() 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 dcf104c1cc7d..fcaf21e1b04a 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,39 +35,51 @@ ) 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 ..testing_utils import ( - BasePipelineTesterConfig, - MemoryTesterMixin, +from ..test_pipelines_common import ( + IPAdapterTesterMixin, + PipelineKarrasSchedulerTesterMixin, + PipelineLatentTesterMixin, PipelineTesterMixin, ) -from .ip_adapter_tester import IPAdapterTesterMixin -class StableDiffusionImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): +enable_full_determinism() + + +class StableDiffusionImg2ImgPipelineFastTests( + IPAdapterTesterMixin, + PipelineLatentTesterMixin, + PipelineKarrasSchedulerTesterMixin, + PipelineTesterMixin, + unittest.TestCase, +): pipeline_class = StableDiffusionImg2ImgPipeline - 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"] - ) + 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 def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) @@ -122,128 +134,182 @@ 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): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) + def get_dummy_inputs(self, device, seed=0): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(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": self.get_generator(0), + "generator": generator, "num_inference_steps": 2, "guidance_scale": 6.0, - # 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", + "output_type": "np", } return inputs - -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() + 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) - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 32, 32) + inputs = self.get_dummy_inputs(device) + image = sd_pipe(**inputs).images + image_slice = image[0, -3:, -3:, -1] - # 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) + 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]) + + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_stable_diffusion_img2img_default_case_lcm(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components(time_cond_proj_dim=256) + sd_pipe = StableDiffusionImg2ImgPipeline(**components) sd_pipe.scheduler = LCMScheduler.from_config(sd_pipe.scheduler.config) + sd_pipe = sd_pipe.to(device) + sd_pipe.set_progress_bar_config(disable=None) - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 32, 32) + 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]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_stable_diffusion_img2img_default_case_lcm_custom_timesteps(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components(time_cond_proj_dim=256) + sd_pipe = StableDiffusionImg2ImgPipeline(**components) 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() + inputs = self.get_dummy_inputs(device) del inputs["num_inference_steps"] inputs["timesteps"] = [999, 499] image = sd_pipe(**inputs).images - assert image.shape == (1, 3, 32, 32) + image_slice = image[0, -3:, -3:, -1] - # 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) + 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 def test_stable_diffusion_img2img_negative_prompt(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline() + 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) - image = sd_pipe(**self.get_dummy_inputs(), negative_prompt="french fries").images - assert image.shape == (1, 3, 32, 32) + 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] - # 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) + 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) def test_stable_diffusion_img2img_multiple_init_images(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline() + 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() + inputs = self.get_dummy_inputs(device) inputs["prompt"] = [inputs["prompt"]] * 2 inputs["image"] = inputs["image"].repeat(2, 1, 1, 1) image = sd_pipe(**inputs).images - assert image.shape == (2, 3, 32, 32) + 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]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_stable_diffusion_img2img_k_lms(self): - # Run on CPU: the expected slice below is CPU-specific. + device = "cpu" # ensure determinism for the device-dependent torch.Generator components = self.get_dummy_components() components["scheduler"] = LMSDiscreteScheduler( beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear" ) - sd_pipe = self.get_pipeline(**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] - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 32, 32) + 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]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_stable_diffusion_img2img_tiny_autoencoder(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline() + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components() + sd_pipe = StableDiffusionImg2ImgPipeline(**components) 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() - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 32, 32) + @skip_mps + def test_save_load_optional_components(self): + return super().test_save_load_optional_components() - # 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) + @skip_mps + def test_attention_slicing_forward_pass(self): + return super().test_attention_slicing_forward_pass(expected_max_diff=5e-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): - sd_pipe = self.get_pipeline().to(torch_device) + components = self.get_dummy_components() + sd_pipe = StableDiffusionImg2ImgPipeline(**components) + sd_pipe = sd_pipe.to(torch_device) + sd_pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(torch_device) prompt = "hey" num_inference_steps = 3 @@ -262,7 +328,7 @@ def apply(self, pipe, i, t, callback_kwargs): prompt, image=inputs["image"], num_inference_steps=num_inference_steps, - output_type="pt", + output_type="np", generator=torch.Generator("cpu").manual_seed(0), callback_on_step_end=pipe_state.apply, ).images @@ -296,27 +362,21 @@ 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().get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict) -class 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 TestStableDiffusionImg2ImgPipelineSlow: - @pytest.fixture(autouse=True) - def cleanup(self): +class StableDiffusionImg2ImgPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() gc.collect() backend_empty_cache(torch_device) - yield + + def tearDown(self): + super().tearDown() gc.collect() backend_empty_cache(torch_device) @@ -561,12 +621,14 @@ def test_img2img_safety_checker_works(self): @nightly @require_torch_accelerator -class TestStableDiffusionImg2ImgPipelineNightly: - @pytest.fixture(autouse=True) - def cleanup(self): +class StableDiffusionImg2ImgPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() gc.collect() backend_empty_cache(torch_device) - yield + + def tearDown(self): + super().tearDown() 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 706d7dc2c923..a69f71c37530 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,19 +54,33 @@ from ..pipeline_params import ( TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS, + TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..testing_utils import ( - BasePipelineTesterConfig, - MemoryTesterMixin, +from ..test_pipelines_common import ( + IPAdapterTesterMixin, + PipelineKarrasSchedulerTesterMixin, + PipelineLatentTesterMixin, PipelineTesterMixin, ) -from .ip_adapter_tester import IPAdapterTesterMixin -class StableDiffusionInpaintPipelineTesterConfig(BasePipelineTesterConfig): +enable_full_determinism() + + +class StableDiffusionInpaintPipelineFastTests( + IPAdapterTesterMixin, + PipelineLatentTesterMixin, + PipelineKarrasSchedulerTesterMixin, + PipelineTesterMixin, + unittest.TestCase, +): pipeline_class = StableDiffusionInpaintPipeline - required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS + 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"}) def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) @@ -118,128 +132,150 @@ def get_dummy_components(self, time_cond_proj_dim=None): } return components - def get_dummy_inputs(self, img_res=64): + def get_dummy_inputs(self, device, seed=0, img_res=64, output_pil=True): # TODO: use tensor inputs instead of PIL, this is here just to leave the old expected_slices untouched - # 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)) + 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) inputs = { "prompt": "A painting of a squirrel eating a burger", "image": init_image, "mask_image": mask_image, - "generator": self.get_generator(0), + "generator": generator, "num_inference_steps": 2, "guidance_scale": 6.0, - # 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", + "output_type": "np", } return inputs - 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 - - -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() + 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) - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 64, 64) + inputs = self.get_dummy_inputs(device) + image = sd_pipe(**inputs).images + image_slice = image[0, -3:, -3:, -1] - # 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) + 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]) + + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 def test_stable_diffusion_inpaint_lcm(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components(time_cond_proj_dim=256) + sd_pipe = StableDiffusionInpaintPipeline(**components) 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.4979, 0.5907, 0.4608, 0.5595, 0.6673, 0.5135, 0.6035, 0.5464, 0.5325]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 def test_stable_diffusion_inpaint_lcm_custom_timesteps(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components(time_cond_proj_dim=256) + sd_pipe = StableDiffusionInpaintPipeline(**components) 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() + inputs = self.get_dummy_inputs(device) del inputs["num_inference_steps"] inputs["timesteps"] = [999, 499] image = sd_pipe(**inputs).images - assert image.shape == (1, 3, 64, 64) + 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]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 def test_stable_diffusion_inpaint_image_tensor(self): - # Run on CPU: PIL and tensor inputs are compared against each other. - sd_pipe = self.get_pipeline() + 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() - out_pil = sd_pipe(**inputs).images + inputs = self.get_dummy_inputs(device) + output = sd_pipe(**inputs) + out_pil = output.images - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(device) 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) - out_tensor = sd_pipe(**inputs).images + output = sd_pipe(**inputs) + out_tensor = output.images - assert out_pil.shape == (1, 3, 64, 64) - assert (out_pil - out_tensor).abs().max() < 5e-2 + assert out_pil.shape == (1, 64, 64, 3) + assert np.abs(out_pil.flatten() - out_tensor.flatten()).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): - sd_pipe = self.get_pipeline().to(torch_device) + 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() + inputs = self.get_dummy_inputs(device) # check that the pipeline raises value error when num_inference_steps is < 1 inputs["strength"] = 0.01 - with pytest.raises(ValueError): - sd_pipe(**inputs) + with self.assertRaises(ValueError): + sd_pipe(**inputs).images def test_stable_diffusion_inpaint_mask_latents(self): - # Run on CPU: the two runs below seed their VAE sampling with a CPU generator. - sd_pipe = self.get_pipeline() + device = "cpu" + components = self.get_dummy_components() + sd_pipe = self.pipeline_class(**components).to(device) + sd_pipe.set_progress_bar_config(disable=None) # normal mask + normal image ## `image`: pil, `mask_image``: pil, `masked_image_latents``: None - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(device) inputs["strength"] = 0.9 out_0 = sd_pipe(**inputs).images # image latents + mask latents - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(device) 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 = self.get_generator(0) + generator = torch.Generator(device=device).manual_seed(0) image_latents = ( sd_pipe.vae.encode(image).latent_dist.sample(generator=generator) * sd_pipe.vae.config.scaling_factor ) @@ -252,17 +288,19 @@ def test_stable_diffusion_inpaint_mask_latents(self): inputs["masked_image_latents"] = mask_latents inputs["mask_image"] = mask inputs["strength"] = 0.9 - generator = self.get_generator(0) + generator = torch.Generator(device=device).manual_seed(0) torch.randn((1, 4, 32, 32), generator=generator) inputs["generator"] = generator out_1 = sd_pipe(**inputs).images - - assert (out_0 - out_1).abs().max() < 1e-2 + assert np.abs(out_0 - out_1).max() < 1e-2 def test_pipeline_interrupt(self): - sd_pipe = self.get_pipeline().to(torch_device) + components = self.get_dummy_components() + sd_pipe = StableDiffusionInpaintPipeline(**components) + sd_pipe = sd_pipe.to(torch_device) + sd_pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(torch_device) prompt = "hey" num_inference_steps = 3 @@ -282,7 +320,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="pt", + output_type="np", generator=torch.Generator("cpu").manual_seed(0), callback_on_step_end=pipe_state.apply, ).images @@ -317,19 +355,21 @@ 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().get("guidance_scale", 1.0) > 1.0, + "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0, } return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict, atol=1e-3, rtol=1e-3) -class StableDiffusionSimpleInpaintPipelineTesterConfig(StableDiffusionInpaintPipelineTesterConfig): - """Same pipeline driven by a regular (non-inpaint-specific) UNet, i.e. one taking 4 input channels.""" +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 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) - components["unet"] = UNet2DConditionModel( + unet = UNet2DConditionModel( block_out_channels=(32, 64), layers_per_block=2, time_cond_proj_dim=time_cond_proj_dim, @@ -340,127 +380,187 @@ 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, 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)) + 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 - inputs = self.get_dummy_inputs() - inputs["prompt"] = [inputs["prompt"]] * 2 - inputs["image"] = [2.0 * image1 - 1.0, 2.0 * image2 - 1.0] # empty mask - inputs["mask_image"] = [torch.zeros((1, 1, img_res, img_res))] * 2 - inputs["generator"] = [self.get_generator(0), self.get_generator(0)] - return inputs + mask_image = torch.zeros((1, 1, img_res, img_res), device=device) + 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) -class TestStableDiffusionSimpleInpaintPipeline( - StableDiffusionSimpleInpaintPipelineTesterConfig, TestStableDiffusionInpaintPipeline -): - """Reruns the inpaint tests against the 4-channel UNet, with its own expected slices.""" + 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 def test_stable_diffusion_inpaint(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline() + 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] - 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.6584, 0.5424, 0.5649, 0.5449, 0.5897, 0.6111, 0.5404, 0.5463, 0.5214]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 def test_stable_diffusion_inpaint_lcm(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components(time_cond_proj_dim=256) + sd_pipe = StableDiffusionInpaintPipeline(**components) sd_pipe.scheduler = LCMScheduler.from_config(sd_pipe.scheduler.config) + sd_pipe = sd_pipe.to(device) + sd_pipe.set_progress_bar_config(disable=None) - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 64, 64) + inputs = self.get_dummy_inputs(device) + image = sd_pipe(**inputs).images + image_slice = image[0, -3:, -3:, -1] - # 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) + 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 def test_stable_diffusion_inpaint_lcm_custom_timesteps(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components(time_cond_proj_dim=256) + sd_pipe = StableDiffusionInpaintPipeline(**components) 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() + inputs = self.get_dummy_inputs(device) del inputs["num_inference_steps"] inputs["timesteps"] = [999, 499] image = sd_pipe(**inputs).images - assert image.shape == (1, 3, 64, 64) + 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]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 def test_stable_diffusion_inpaint_2_images(self): - # Run on CPU: the two runs below are compared against each other. - sd_pipe = self.get_pipeline() + 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) # test to confirm if we pass two same image, we will get same output - inputs = self.get_dummy_inputs() + inputs = self.get_dummy_inputs(device) + gen1 = torch.Generator(device=device).manual_seed(0) + gen2 = torch.Generator(device=device).manual_seed(0) for name in ["prompt", "image", "mask_image"]: inputs[name] = [inputs[name]] * 2 - inputs["generator"] = [self.get_generator(0), self.get_generator(0)] + inputs["generator"] = [gen1, gen2] images = sd_pipe(**inputs).images - assert images.shape == (2, 3, 64, 64) - image_slice1 = images[0, -1, -3:, -3:] - image_slice2 = images[1, -1, -3:, -3:] - assert (image_slice1 - image_slice2).abs().max() < 1e-4 + 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 # test to confirm that if we pass two different images, we will get different output - images = sd_pipe(**self.get_dummy_inputs_2images()).images - assert images.shape == (2, 3, 64, 64) + inputs = self.get_dummy_inputs_2images(device) + images = sd_pipe(**inputs).images + assert images.shape == (2, 64, 64, 3) - image_slice1 = images[0, -1, -3:, -3:] - image_slice2 = images[1, -1, -3:, -3:] - assert (image_slice1 - image_slice2).abs().max() > 1e-2 + 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 def test_stable_diffusion_inpaint_euler(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components(time_cond_proj_dim=256) + sd_pipe = StableDiffusionInpaintPipeline(**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_tensor_inputs() + inputs = self.get_dummy_inputs(device, output_pil=False) half_dim = inputs["image"].shape[2] // 2 inputs["mask_image"][0, 0, :half_dim, :half_dim] = 0 - inputs["num_inference_steps"] = 4 + inputs["num_inference_steps"] = 4 image = sd_pipe(**inputs).images - assert image.shape == (1, 3, 64, 64) - - # 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) - + image_slice = image[0, -3:, -3:, -1] -class TestStableDiffusionInpaintPipelineMemory(StableDiffusionInpaintPipelineTesterConfig, MemoryTesterMixin): - """Memory optimization tests (CPU offload, group offload, layerwise casting) for the inpaint pipeline.""" + assert image.shape == (1, 64, 64, 3) - -class TestStableDiffusionInpaintPipelineIPAdapter(StableDiffusionInpaintPipelineTesterConfig, IPAdapterTesterMixin): - """IP-Adapter tests for the inpaint pipeline.""" + 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 @slow @require_torch_accelerator -class TestStableDiffusionInpaintPipelineSlow: - @pytest.fixture(autouse=True) - def cleanup(self): - yield +class StableDiffusionInpaintPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + + def tearDown(self): + super().tearDown() gc.collect() backend_empty_cache(torch_device) @@ -635,10 +735,12 @@ def test_stable_diffusion_simple_inpaint_ddim(self): @slow @require_torch_accelerator -class TestStableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlow: - @pytest.fixture(autouse=True) - def cleanup(self): - yield +class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() + + def tearDown(self): + super().tearDown() gc.collect() backend_empty_cache(torch_device) @@ -884,12 +986,14 @@ def test_download_local(self): @nightly @require_torch_accelerator -class TestStableDiffusionInpaintPipelineNightly: - @pytest.fixture(autouse=True) - def cleanup(self): +class StableDiffusionInpaintPipelineNightlyTests(unittest.TestCase): + def setUp(self): + super().setUp() gc.collect() backend_empty_cache(torch_device) - yield + + def tearDown(self): + super().tearDown() 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 9045e92b8be9..0852965c2fe8 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,13 +31,14 @@ 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, @@ -45,20 +46,30 @@ 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 ..testing_utils import ( - BasePipelineTesterConfig, - MemoryTesterMixin, +from ..test_pipelines_common import ( + PipelineKarrasSchedulerTesterMixin, + PipelineLatentTesterMixin, PipelineTesterMixin, ) -class StableDiffusionInstructPix2PixPipelineTesterConfig(BasePipelineTesterConfig): +enable_full_determinism() + + +class StableDiffusionInstructPix2PixPipelineFastTests( + PipelineLatentTesterMixin, PipelineKarrasSchedulerTesterMixin, PipelineTesterMixin, unittest.TestCase +): pipeline_class = StableDiffusionInstructPix2PixPipeline - required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"} - batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS + 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"} def get_dummy_components(self): torch.manual_seed(0) @@ -109,106 +120,165 @@ def get_dummy_components(self): } return components - def get_dummy_inputs(self): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)) - image = image.permute(0, 2, 3, 1)[0] + 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] 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": self.get_generator(0), + "generator": generator, "num_inference_steps": 2, "guidance_scale": 6.0, "image_guidance_scale": 1, - # 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", + "output_type": "np", } return inputs - -class TestStableDiffusionInstructPix2PixPipeline( - StableDiffusionInstructPix2PixPipelineTesterConfig, PipelineTesterMixin -): def test_stable_diffusion_pix2pix_default_case(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline() + 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) - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 32, 32) + 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]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_stable_diffusion_pix2pix_negative_prompt(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline() + 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] - 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.7059, 0.3731, 0.4428, 0.5922, 0.5649, 0.4871, 0.4203, 0.5507, 0.4809]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_stable_diffusion_pix2pix_multiple_init_images(self): - # Run on CPU: the expected slice below is CPU-specific. - sd_pipe = self.get_pipeline() + 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() + inputs = self.get_dummy_inputs(device) inputs["prompt"] = [inputs["prompt"]] * 2 image = np.array(inputs["image"]).astype(np.float32) / 255.0 - image = torch.from_numpy(image).unsqueeze(0) + image = torch.from_numpy(image).unsqueeze(0).to(device) 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 - assert image.shape == (2, 3, 32, 32) + 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]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_stable_diffusion_pix2pix_euler(self): - # Run on CPU: the expected slice below is CPU-specific. + device = "cpu" # ensure determinism for the device-dependent torch.Generator components = self.get_dummy_components() components["scheduler"] = EulerAncestralDiscreteScheduler( beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear" ) - sd_pipe = self.get_pipeline(**components) + sd_pipe = StableDiffusionInstructPix2PixPipeline(**components) + sd_pipe = sd_pipe.to(device) + sd_pipe.set_progress_bar_config(disable=None) - image = sd_pipe(**self.get_dummy_inputs()).images - assert image.shape == (1, 3, 32, 32) + 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]) - # 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) + assert np.abs(image_slice.flatten() - expected_slice).max() < 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) -class TestStableDiffusionInstructPix2PixPipelineMemory( - StableDiffusionInstructPix2PixPipelineTesterConfig, MemoryTesterMixin -): - """Memory optimization tests (CPU offload, group offload, layerwise casting) for the InstructPix2Pix pipeline.""" + 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 @slow @require_torch_accelerator -class TestStableDiffusionInstructPix2PixPipelineSlow: - @pytest.fixture(autouse=True) - def cleanup(self): +class StableDiffusionInstructPix2PixPipelineSlowTests(unittest.TestCase): + def setUp(self): + super().setUp() gc.collect() backend_empty_cache(torch_device) - yield + + def tearDown(self): + super().tearDown() gc.collect() backend_empty_cache(torch_device)