From dde70958434fd1e3468c1f1a5921330173672032 Mon Sep 17 00:00:00 2001 From: William Berman Date: Sat, 25 Feb 2023 19:14:27 -0800 Subject: [PATCH 1/4] attend and excite batch test causing timeouts --- .../test_stable_diffusion_attend_and_excite.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py index 1025b52636a6..ca300dc0ce11 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py @@ -135,8 +135,9 @@ def test_inference(self): max_diff = np.abs(image_slice.flatten() - expected_slice).max() self.assertLessEqual(max_diff, 1e-3) - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(relax_max_difference=False) + def test_inference_batch_consistent(self): + # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches + self._test_inference_batch_consistent(batch_sizes=[2, 4]) @require_torch_gpu From 1ae7ac367417281ed41cb49cbfb034fff9a13fe2 Mon Sep 17 00:00:00 2001 From: William Berman Date: Fri, 24 Feb 2023 15:55:58 -0800 Subject: [PATCH 2/4] move test num_images_per_prompt to pipeline mixin --- ...line_stable_diffusion_attend_and_excite.py | 30 +++++++++-- .../paint_by_example/test_paint_by_example.py | 14 +---- .../stable_diffusion/test_cycle_diffusion.py | 3 ++ .../stable_diffusion/test_stable_diffusion.py | 37 ------------- .../test_stable_diffusion_image_variation.py | 37 +------------ .../test_stable_diffusion_img2img.py | 36 ------------- .../test_stable_diffusion_inpaint.py | 13 ----- ...st_stable_diffusion_instruction_pix2pix.py | 36 ------------- .../test_stable_diffusion_panorama.py | 36 ------------- .../test_stable_diffusion_pix2pix_zero.py | 28 ---------- ...test_stable_diffusion_attend_and_excite.py | 3 ++ .../test_stable_diffusion_depth.py | 36 ------------- .../unclip/test_unclip_image_variation.py | 54 +------------------ tests/test_pipelines_common.py | 33 +++++++++++- 14 files changed, 71 insertions(+), 325 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_attend_and_excite.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_attend_and_excite.py index 8e6330c9a983..90e70ab5843b 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_attend_and_excite.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_attend_and_excite.py @@ -513,8 +513,30 @@ def check_inputs( f" {negative_prompt_embeds.shape}." ) - if (indices is None) or (indices is not None and not isinstance(indices, List)): - raise ValueError(f"`indices` has to be a list but is {type(indices)}") + indices_is_list_ints = isinstance(indices, list) and isinstance(indices[0], int) + indices_is_list_list_ints = ( + isinstance(indices, list) and isinstance(indices[0], list) and isinstance(indices[0][0], int) + ) + + if not indices_is_list_ints and not indices_is_list_list_ints: + raise TypeError("`indices` must be a list of ints or a list of a list of ints") + + if indices_is_list_ints: + indices_batch_size = 1 + elif indices_is_list_list_ints: + indices_batch_size = len(indices) + + if prompt is not None and isinstance(prompt, str): + prompt_batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + prompt_batch_size = len(prompt) + elif prompt_embeds is not None: + prompt_batch_size = prompt_embeds.shape[0] + + if indices_batch_size != prompt_batch_size: + raise ValueError( + f"indices batch size must be same as prompt batch size. indices batch size: {indices_batch_size}, prompt batch size: {prompt_batch_size}" + ) # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): @@ -671,7 +693,7 @@ def get_indices(self, prompt: str) -> Dict[str, int]: def __call__( self, prompt: Union[str, List[str]], - token_indices: List[int], + token_indices: Union[List[int], List[List[int]]], height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, @@ -847,7 +869,9 @@ def __call__( if isinstance(token_indices[0], int): token_indices = [token_indices] + indices = [] + for ind in token_indices: indices = indices + [ind] * num_images_per_prompt diff --git a/tests/pipelines/paint_by_example/test_paint_by_example.py b/tests/pipelines/paint_by_example/test_paint_by_example.py index a2e04d20a067..137c0e99aeac 100644 --- a/tests/pipelines/paint_by_example/test_paint_by_example.py +++ b/tests/pipelines/paint_by_example/test_paint_by_example.py @@ -163,18 +163,8 @@ def test_paint_by_example_image_tensor(self): assert out_1.shape == (1, 64, 64, 3) assert np.abs(out_1.flatten() - out_2.flatten()).max() < 5e-2 - def test_paint_by_example_inpaint_with_num_images_per_prompt(self): - device = "cpu" - pipe = PaintByExamplePipeline(**self.get_dummy_components()) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs() - - images = pipe(**inputs, num_images_per_prompt=2).images - - # check if the output is a list of 2 images - assert len(images) == 2 + def test_num_images_per_prompt(self): + self._test_num_images_per_prompt(prompt_key=["image", "example_image", "mask_image"]) @slow diff --git a/tests/pipelines/stable_diffusion/test_cycle_diffusion.py b/tests/pipelines/stable_diffusion/test_cycle_diffusion.py index 948a39786dcb..828f06e8d6b6 100644 --- a/tests/pipelines/stable_diffusion/test_cycle_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_cycle_diffusion.py @@ -149,6 +149,9 @@ def test_stable_diffusion_cycle_fp16(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + def test_num_images_per_prompt(self): + self._test_num_images_per_prompt(prompt_key=["prompt", "source_prompt", "image"]) + @slow @require_torch_gpu diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py index 75e2a44f018f..29f0425bfdde 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -451,43 +451,6 @@ def test_stable_diffusion_negative_prompt(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - def test_stable_diffusion_num_images_per_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) - - prompt = "A painting of a squirrel eating a burger" - - # test num_images_per_prompt=1 (default) - images = sd_pipe(prompt, num_inference_steps=2, output_type="np").images - - assert images.shape == (1, 64, 64, 3) - - # test num_images_per_prompt=1 (default) for batch of prompts - batch_size = 2 - images = sd_pipe([prompt] * batch_size, num_inference_steps=2, output_type="np").images - - assert images.shape == (batch_size, 64, 64, 3) - - # test num_images_per_prompt for single prompt - num_images_per_prompt = 2 - images = sd_pipe( - prompt, num_inference_steps=2, output_type="np", num_images_per_prompt=num_images_per_prompt - ).images - - assert images.shape == (num_images_per_prompt, 64, 64, 3) - - # test num_images_per_prompt for batch of prompts - batch_size = 2 - images = sd_pipe( - [prompt] * batch_size, num_inference_steps=2, output_type="np", num_images_per_prompt=num_images_per_prompt - ).images - - assert images.shape == (batch_size * num_images_per_prompt, 64, 64, 3) - def test_stable_diffusion_long_prompt(self): components = self.get_dummy_components() components["scheduler"] = LMSDiscreteScheduler.from_config(components["scheduler"].config) diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_image_variation.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_image_variation.py index a7aa4051774d..06e680d964d3 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_image_variation.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_image_variation.py @@ -140,41 +140,8 @@ def test_stable_diffusion_img_variation_multiple_images(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - def test_stable_diffusion_img_variation_num_images_per_prompt(self): - device = "cpu" - components = self.get_dummy_components() - sd_pipe = StableDiffusionImageVariationPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) - - # test num_images_per_prompt=1 (default) - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs).images - - assert images.shape == (1, 64, 64, 3) - - # test num_images_per_prompt=1 (default) for batch of images - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["image"] = batch_size * [inputs["image"]] - images = sd_pipe(**inputs).images - - assert images.shape == (batch_size, 64, 64, 3) - - # test num_images_per_prompt for single prompt - num_images_per_prompt = 2 - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (num_images_per_prompt, 64, 64, 3) - - # test num_images_per_prompt for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["image"] = batch_size * [inputs["image"]] - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (batch_size * num_images_per_prompt, 64, 64, 3) + def test_num_images_per_prompt(self): + self._test_num_images_per_prompt(prompt_key="image") @slow diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py index e10b608734b6..58db9c11561a 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py @@ -177,42 +177,6 @@ def test_stable_diffusion_img2img_k_lms(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - def test_stable_diffusion_img2img_num_images_per_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) - - # test num_images_per_prompt=1 (default) - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs).images - - assert images.shape == (1, 32, 32, 3) - - # test num_images_per_prompt=1 (default) for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * batch_size - images = sd_pipe(**inputs).images - - assert images.shape == (batch_size, 32, 32, 3) - - # test num_images_per_prompt for single prompt - num_images_per_prompt = 2 - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (num_images_per_prompt, 32, 32, 3) - - # test num_images_per_prompt for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * batch_size - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (batch_size * num_images_per_prompt, 32, 32, 3) - @slow @require_torch_gpu diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py index c44101d13c5a..c585be1e13f7 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py @@ -148,19 +148,6 @@ def test_stable_diffusion_inpaint_image_tensor(self): assert out_pil.shape == (1, 64, 64, 3) assert np.abs(out_pil.flatten() - out_tensor.flatten()).max() < 5e-2 - def test_stable_diffusion_inpaint_with_num_images_per_prompt(self): - device = "cpu" - 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) - images = sd_pipe(**inputs, num_images_per_prompt=2).images - - # check if the output is a list of 2 images - assert len(images) == 2 - @slow @require_torch_gpu 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 4c232b573b4f..c453ae92ac14 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py @@ -188,42 +188,6 @@ def test_stable_diffusion_pix2pix_euler(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - def test_stable_diffusion_pix2pix_num_images_per_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) - - # test num_images_per_prompt=1 (default) - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs).images - - assert images.shape == (1, 32, 32, 3) - - # test num_images_per_prompt=1 (default) for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * batch_size - images = sd_pipe(**inputs).images - - assert images.shape == (batch_size, 32, 32, 3) - - # test num_images_per_prompt for single prompt - num_images_per_prompt = 2 - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (num_images_per_prompt, 32, 32, 3) - - # test num_images_per_prompt for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * batch_size - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (batch_size * num_images_per_prompt, 32, 32, 3) - @slow @require_torch_gpu diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_panorama.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_panorama.py index 366cd39da53a..556d115da438 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_panorama.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_panorama.py @@ -174,42 +174,6 @@ def test_stable_diffusion_panorama_pndm(self): with self.assertRaises(ValueError): _ = sd_pipe(**inputs).images - def test_stable_diffusion_panorama_num_images_per_prompt(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionPanoramaPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) - - # test num_images_per_prompt=1 (default) - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs).images - - assert images.shape == (1, 64, 64, 3) - - # test num_images_per_prompt=1 (default) for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * batch_size - images = sd_pipe(**inputs).images - - assert images.shape == (batch_size, 64, 64, 3) - - # test num_images_per_prompt for single prompt - num_images_per_prompt = 2 - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (num_images_per_prompt, 64, 64, 3) - - # test num_images_per_prompt for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * batch_size - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (batch_size * num_images_per_prompt, 64, 64, 3) - @slow @require_torch_gpu diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_pix2pix_zero.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_pix2pix_zero.py index 103430aefb91..08e954f15ca9 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_pix2pix_zero.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_pix2pix_zero.py @@ -195,34 +195,6 @@ def test_stable_diffusion_pix2pix_zero_ddpm(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - def test_stable_diffusion_pix2pix_zero_num_images_per_prompt(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = StableDiffusionPix2PixZeroPipeline(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) - - # test num_images_per_prompt=1 (default) - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs).images - - assert images.shape == (1, 64, 64, 3) - - # test num_images_per_prompt=2 for a single prompt - num_images_per_prompt = 2 - inputs = self.get_dummy_inputs(device) - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (num_images_per_prompt, 64, 64, 3) - - # test num_images_per_prompt for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * batch_size - images = sd_pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (batch_size * num_images_per_prompt, 64, 64, 3) - @slow @require_torch_gpu diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py index ca300dc0ce11..ffb6a8ca8576 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py @@ -139,6 +139,9 @@ def test_inference_batch_consistent(self): # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches self._test_inference_batch_consistent(batch_sizes=[2, 4]) + def test_num_images_per_prompt(self): + self._test_num_images_per_prompt(prompt_key=["prompt", "token_indices"]) + @require_torch_gpu @slow diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py index 0bcd2de62982..a5042a2f0050 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py @@ -327,42 +327,6 @@ def test_stable_diffusion_depth2img_multiple_init_images(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - def test_stable_diffusion_depth2img_num_images_per_prompt(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = StableDiffusionDepth2ImgPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) - - # test num_images_per_prompt=1 (default) - inputs = self.get_dummy_inputs(device) - images = pipe(**inputs).images - - assert images.shape == (1, 32, 32, 3) - - # test num_images_per_prompt=1 (default) for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * batch_size - images = pipe(**inputs).images - - assert images.shape == (batch_size, 32, 32, 3) - - # test num_images_per_prompt for single prompt - num_images_per_prompt = 2 - inputs = self.get_dummy_inputs(device) - images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (num_images_per_prompt, 32, 32, 3) - - # test num_images_per_prompt for batch of prompts - batch_size = 2 - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"]] * batch_size - images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images - - assert images.shape == (batch_size * num_images_per_prompt, 32, 32, 3) - def test_stable_diffusion_depth2img_pil(self): device = "cpu" # ensure determinism for the device-dependent torch.Generator components = self.get_dummy_components() diff --git a/tests/pipelines/unclip/test_unclip_image_variation.py b/tests/pipelines/unclip/test_unclip_image_variation.py index 865cf11dd825..5111acf34f2a 100644 --- a/tests/pipelines/unclip/test_unclip_image_variation.py +++ b/tests/pipelines/unclip/test_unclip_image_variation.py @@ -362,58 +362,8 @@ def test_unclip_image_variation_input_list_images(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 - def test_unclip_image_variation_input_num_images_per_prompt(self): - device = "cpu" - - components = self.get_dummy_components() - - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - - pipe.set_progress_bar_config(disable=None) - - pipeline_inputs = self.get_dummy_inputs(device, pil_image=True) - pipeline_inputs["image"] = [ - pipeline_inputs["image"], - pipeline_inputs["image"], - ] - - output = pipe(**pipeline_inputs, num_images_per_prompt=2) - image = output.images - - tuple_pipeline_inputs = self.get_dummy_inputs(device, pil_image=True) - tuple_pipeline_inputs["image"] = [ - tuple_pipeline_inputs["image"], - tuple_pipeline_inputs["image"], - ] - - image_from_tuple = pipe( - **tuple_pipeline_inputs, - num_images_per_prompt=2, - return_dict=False, - )[0] - - image_slice = image[0, -3:, -3:, -1] - image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1] - - assert image.shape == (4, 64, 64, 3) - - expected_slice = np.array( - [ - 0.9980, - 0.9997, - 0.0023, - 0.0029, - 0.9997, - 0.9985, - 0.9997, - 0.0010, - 0.9995, - ] - ) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 + def test_num_images_per_prompt(self): + self._test_num_images_per_prompt(prompt_key="image") def test_unclip_passed_image_embed(self): device = torch.device("cpu") diff --git a/tests/test_pipelines_common.py b/tests/test_pipelines_common.py index 7c145be21b34..4a47cc2ccbbc 100644 --- a/tests/test_pipelines_common.py +++ b/tests/test_pipelines_common.py @@ -5,7 +5,7 @@ import re import tempfile import unittest -from typing import Callable, Union +from typing import Callable, List, Union import numpy as np import torch @@ -546,6 +546,37 @@ def test_progress_bar(self): _ = pipe(**inputs) self.assertTrue(stderr.getvalue() == "", "Progress bar should be disabled") + def test_num_images_per_prompt(self): + self._test_num_images_per_prompt() + + def _test_num_images_per_prompt(self, prompt_key: Union[str, List[str]] = "prompt"): + sig = inspect.signature(self.pipeline_class.__call__) + + if "num_images_per_prompt" not in sig.parameters: + return + + if not isinstance(prompt_key, list): + prompt_key = [prompt_key] + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + + batch_sizes = [1, 2] + num_images_per_prompts = [1, 2] + + for batch_size in batch_sizes: + for num_images_per_prompt in num_images_per_prompts: + inputs = self.get_dummy_inputs(torch_device) + + for key in prompt_key: + inputs[key] = batch_size * [inputs[key]] + + images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images + + assert images.shape[0] == batch_size * num_images_per_prompt + # Some models (e.g. unCLIP) are extremely likely to significantly deviate depending on which hardware is used. # This helper function is used to check that the image doesn't deviate on average more than 10 pixels from a From 6647d556737751212251c1a4f6ad76672ff796e5 Mon Sep 17 00:00:00 2001 From: William Berman Date: Fri, 3 Mar 2023 11:01:55 -0800 Subject: [PATCH 3/4] style --- .../pipelines/stable_diffusion/test_stable_diffusion_img2img.py | 1 + .../stable_diffusion/test_stable_diffusion_pix2pix_zero.py | 1 + 2 files changed, 2 insertions(+) diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py index c69a3730203f..77dfa9be1d1e 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py @@ -197,6 +197,7 @@ def test_save_load_optional_components(self): def test_attention_slicing_forward_pass(self): return super().test_attention_slicing_forward_pass() + @slow @require_torch_gpu class StableDiffusionImg2ImgPipelineSlowTests(unittest.TestCase): diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_pix2pix_zero.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_pix2pix_zero.py index f511f63ff7d2..9e80ef7452a8 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_pix2pix_zero.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_pix2pix_zero.py @@ -196,6 +196,7 @@ def test_stable_diffusion_pix2pix_zero_ddpm(self): def test_inference_batch_single_identical(self): return super().test_inference_batch_single_identical() + @slow @require_torch_gpu class StableDiffusionPix2PixZeroPipelineSlowTests(unittest.TestCase): From 4ce1cd66b6cbc2179bbe33ca9a7078f772b5d9e7 Mon Sep 17 00:00:00 2001 From: William Berman Date: Fri, 3 Mar 2023 11:15:00 -0800 Subject: [PATCH 4/4] prompt_key -> self.batch_params --- .../paint_by_example/test_paint_by_example.py | 3 --- .../stable_diffusion/test_cycle_diffusion.py | 5 +---- .../test_stable_diffusion_image_variation.py | 3 --- .../test_stable_diffusion_attend_and_excite.py | 3 --- .../test_stable_diffusion_depth.py | 2 +- .../pipelines/unclip/test_unclip_image_variation.py | 3 --- tests/test_pipelines_common.py | 13 ++++--------- 7 files changed, 6 insertions(+), 26 deletions(-) diff --git a/tests/pipelines/paint_by_example/test_paint_by_example.py b/tests/pipelines/paint_by_example/test_paint_by_example.py index 9c8e31b2d22c..81d1989200ac 100644 --- a/tests/pipelines/paint_by_example/test_paint_by_example.py +++ b/tests/pipelines/paint_by_example/test_paint_by_example.py @@ -160,9 +160,6 @@ def test_paint_by_example_image_tensor(self): assert out_1.shape == (1, 64, 64, 3) assert np.abs(out_1.flatten() - out_2.flatten()).max() < 5e-2 - def test_num_images_per_prompt(self): - self._test_num_images_per_prompt(prompt_key=["image", "example_image", "mask_image"]) - @slow @require_torch_gpu diff --git a/tests/pipelines/stable_diffusion/test_cycle_diffusion.py b/tests/pipelines/stable_diffusion/test_cycle_diffusion.py index 9ff5fd435eb2..5282cfd8dd24 100644 --- a/tests/pipelines/stable_diffusion/test_cycle_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_cycle_diffusion.py @@ -41,7 +41,7 @@ class CycleDiffusionPipelineFastTests(PipelineTesterMixin, unittest.TestCase): "negative_prompt_embeds", } required_optional_params = PipelineTesterMixin.required_optional_params - {"latents"} - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS + batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"source_prompt"}) def get_dummy_components(self): torch.manual_seed(0) @@ -158,9 +158,6 @@ def test_stable_diffusion_cycle_fp16(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - def test_num_images_per_prompt(self): - self._test_num_images_per_prompt(prompt_key=["prompt", "source_prompt", "image"]) - @skip_mps def test_save_load_local(self): return super().test_save_load_local() diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_image_variation.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_image_variation.py index ca88e87be2ae..01c2e22e4816 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_image_variation.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_image_variation.py @@ -143,9 +143,6 @@ def test_stable_diffusion_img_variation_multiple_images(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - def test_num_images_per_prompt(self): - self._test_num_images_per_prompt(prompt_key="image") - @slow @require_torch_gpu diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py index fa5f357b6d47..780abf304a46 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_attend_and_excite.py @@ -142,9 +142,6 @@ def test_inference_batch_consistent(self): # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches self._test_inference_batch_consistent(batch_sizes=[2, 4]) - def test_num_images_per_prompt(self): - self._test_num_images_per_prompt(prompt_key=["prompt", "token_indices"]) - @require_torch_gpu @slow diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py index 12e7113399a8..110dbbd7f80c 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py @@ -64,7 +64,7 @@ class StableDiffusionDepth2ImgPipelineFastTests(PipelineTesterMixin, unittest.Te test_save_load_optional_components = False params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"} required_optional_params = PipelineTesterMixin.required_optional_params - {"latents"} - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS + batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - {"image"} def get_dummy_components(self): torch.manual_seed(0) diff --git a/tests/pipelines/unclip/test_unclip_image_variation.py b/tests/pipelines/unclip/test_unclip_image_variation.py index dec509a3e675..ff32ac5f9aaf 100644 --- a/tests/pipelines/unclip/test_unclip_image_variation.py +++ b/tests/pipelines/unclip/test_unclip_image_variation.py @@ -361,9 +361,6 @@ def test_unclip_image_variation_input_list_images(self): assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 - def test_num_images_per_prompt(self): - self._test_num_images_per_prompt(prompt_key="image") - def test_unclip_passed_image_embed(self): device = torch.device("cpu") diff --git a/tests/test_pipelines_common.py b/tests/test_pipelines_common.py index 5b8187e38725..8025bdd564b8 100644 --- a/tests/test_pipelines_common.py +++ b/tests/test_pipelines_common.py @@ -5,7 +5,7 @@ import re import tempfile import unittest -from typing import Callable, List, Union +from typing import Callable, Union import numpy as np import torch @@ -551,17 +551,11 @@ def test_progress_bar(self): self.assertTrue(stderr.getvalue() == "", "Progress bar should be disabled") def test_num_images_per_prompt(self): - self._test_num_images_per_prompt() - - def _test_num_images_per_prompt(self, prompt_key: Union[str, List[str]] = "prompt"): sig = inspect.signature(self.pipeline_class.__call__) if "num_images_per_prompt" not in sig.parameters: return - if not isinstance(prompt_key, list): - prompt_key = [prompt_key] - components = self.get_dummy_components() pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) @@ -574,8 +568,9 @@ def _test_num_images_per_prompt(self, prompt_key: Union[str, List[str]] = "promp for num_images_per_prompt in num_images_per_prompts: inputs = self.get_dummy_inputs(torch_device) - for key in prompt_key: - inputs[key] = batch_size * [inputs[key]] + for key in inputs.keys(): + if key in self.batch_params: + inputs[key] = batch_size * [inputs[key]] images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt).images