[tests] refactor d series pipeline tests.

src/diffusers/pipelines/ddim/pipeline_ddim.py

@@ -90,7 +90,8 @@ def __call__(
                 If `True` or `False`, see documentation for [`DDIMScheduler.step`]. If `None`, nothing is passed
                 downstream to the scheduler (use `None` for schedulers which don't support this argument).
             output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+                The output format of the generated image. Choose between `"pil"` (`PIL.Image`), `"np"` (`np.array`) or
+                `"pt"` (`torch.Tensor`).
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
 
@@ -167,9 +168,13 @@ def __call__(
                 xm.mark_step()
 
         image = (image / 2 + 0.5).clamp(0, 1)
-        image = image.cpu().permute(0, 2, 3, 1).numpy()
-        if output_type == "pil":
-            image = self.numpy_to_pil(image)
+        if output_type != "pt":
+            image = image.cpu().permute(0, 2, 3, 1).numpy()
+            if output_type == "pil":
+                image = self.numpy_to_pil(image)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
 
         if not return_dict:
             return (image,)

src/diffusers/pipelines/ddpm/pipeline_ddpm.py

@@ -73,7 +73,8 @@ def __call__(
                 The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                 expense of slower inference.
             output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+                The output format of the generated image. Choose between `"pil"` (`PIL.Image`), `"np"` (`np.array`) or
+                `"pt"` (`torch.Tensor`).
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
 
@@ -108,12 +109,13 @@ def __call__(
         else:
             image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)
 
-        if self.device.type == "mps":
+        device = self._execution_device
+        if device.type == "mps":
             # randn does not work reproducibly on mps
             image = randn_tensor(image_shape, generator=generator, dtype=self.unet.dtype)
-            image = image.to(self.device)
+            image = image.to(device)
         else:
-            image = randn_tensor(image_shape, generator=generator, device=self.device, dtype=self.unet.dtype)
+            image = randn_tensor(image_shape, generator=generator, device=device, dtype=self.unet.dtype)
 
         # set step values
         self.scheduler.set_timesteps(num_inference_steps)
@@ -129,9 +131,13 @@ def __call__(
                 xm.mark_step()
 
         image = (image / 2 + 0.5).clamp(0, 1)
-        image = image.cpu().permute(0, 2, 3, 1).numpy()
-        if output_type == "pil":
-            image = self.numpy_to_pil(image)
+        if output_type != "pt":
+            image = image.cpu().permute(0, 2, 3, 1).numpy()
+            if output_type == "pil":
+                image = self.numpy_to_pil(image)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
 
         if not return_dict:
             return (image,)

src/diffusers/pipelines/dit/pipeline_dit.py

@@ -124,7 +124,8 @@ def __call__(
                 The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                 expense of slower inference.
             output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+                The output format of the generated image. Choose between `"pil"` (`PIL.Image`), `"np"` (`np.array`) or
+                `"pt"` (`torch.Tensor`).
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`ImagePipelineOutput`] instead of a plain tuple.
 
@@ -234,11 +235,12 @@ def __call__(
 
         samples = (samples / 2 + 0.5).clamp(0, 1)
 
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
-        samples = samples.cpu().permute(0, 2, 3, 1).float().numpy()
+        if output_type != "pt":
+            # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+            samples = samples.cpu().permute(0, 2, 3, 1).float().numpy()
 
-        if output_type == "pil":
-            samples = self.numpy_to_pil(samples)
+            if output_type == "pil":
+                samples = self.numpy_to_pil(samples)
 
         # Offload all models
         self.maybe_free_model_hooks()

tests/pipelines/ddim/test_ddim.py

@@ -13,31 +13,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-
 import numpy as np
 import torch
 
 from diffusers import DDIMPipeline, DDIMScheduler, UNet2DModel
 
-from ...testing_utils import enable_full_determinism, require_torch_accelerator, slow, torch_device
+from ...testing_utils import (
+    assert_tensors_close,
+    enable_full_determinism,
+    require_torch_accelerator,
+    slow,
+    torch_device,
+)
 from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
-from ..test_pipelines_common import PipelineTesterMixin
+from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
 
 
 enable_full_determinism()
 
 
-class DDIMPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+class DDIMPipelineTesterConfig(BasePipelineTesterConfig):
     pipeline_class = DDIMPipeline
-    params = UNCONDITIONAL_IMAGE_GENERATION_PARAMS
-    required_optional_params = PipelineTesterMixin.required_optional_params - {
-        "num_images_per_prompt",
-        "latents",
-        "callback",
-        "callback_steps",
-    }
-    batch_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
+    required_input_params_in_call_signature = UNCONDITIONAL_IMAGE_GENERATION_PARAMS
+    batch_input_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
+    # DDIM is unconditional and samples its own noise: there is no prompt to repeat
+    # (`num_images_per_prompt`) and no user-suppliable `latents`.
+    optional_input_params = BasePipelineTesterConfig.optional_input_params - {"num_images_per_prompt", "latents"}
+    output_shape = (3, 8, 8)
 
     def get_dummy_components(self):
         torch.manual_seed(0)
@@ -52,55 +54,53 @@ def get_dummy_components(self):
             up_block_types=("AttnUpBlock2D", "UpBlock2D"),
         )
         scheduler = DDIMScheduler()
-        components = {"unet": unet, "scheduler": scheduler}
-        return components
-
-    def get_dummy_inputs(self, device, seed=0):
-        if str(device).startswith("mps"):
-            generator = torch.manual_seed(seed)
-        else:
-            generator = torch.Generator(device=device).manual_seed(seed)
-        inputs = {
+        return {"unet": unet, "scheduler": scheduler}
+
+    def get_dummy_inputs(self):
+        return {
             "batch_size": 1,
-            "generator": generator,
+            "generator": self.get_generator(0),
             "num_inference_steps": 2,
-            "output_type": "np",
+            # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
+            "output_type": "pt",
         }
-        return inputs
 
+
+class TestDDIMPipeline(DDIMPipelineTesterConfig, PipelineTesterMixin):
     def test_inference(self):
-        device = "cpu"
+        # Run on CPU: the expected slice below is CPU-specific.
+        pipe = self.get_pipeline()
 
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe.to(device)
-        pipe.set_progress_bar_config(disable=None)
+        image = pipe(**self.get_dummy_inputs()).images
+        generated_image = image[0]
+        assert generated_image.shape == self.output_shape
 
-        inputs = self.get_dummy_inputs(device)
-        image = pipe(**inputs).images
-        image_slice = image[0, -3:, -3:, -1]
+        # fmt: off
+        expected_slice = torch.tensor([0.0, 9.979e-01, 0.0, 9.999e-01, 9.986e-01, 9.991e-01, 7.106e-04, 0.0, 0.0])
+        # fmt: on
 
-        self.assertEqual(image.shape, (1, 8, 8, 3))
-        expected_slice = np.array([0.0, 9.979e-01, 0.0, 9.999e-01, 9.986e-01, 9.991e-01, 7.106e-04, 0.0, 0.0])
-        max_diff = np.abs(image_slice.flatten() - expected_slice).max()
-        self.assertLessEqual(max_diff, 1e-3)
+        # `"pt"` images are `(channels, height, width)`, so the trailing-channel corner slice of the old
+        # `"np"` layout is the last channel's bottom-right 3x3 block here.
+        generated_slice = generated_image[-1, -3:, -3:].flatten()
+        assert_tensors_close(generated_slice, expected_slice, atol=1e-3)
 
     def test_dict_tuple_outputs_equivalent(self):
         super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3)
 
-    def test_save_load_local(self):
-        super().test_save_load_local(expected_max_difference=3e-3)
-
-    def test_save_load_optional_components(self):
-        super().test_save_load_optional_components(expected_max_difference=3e-3)
+    def test_save_load_local(self, tmp_path, base_pipe_output):
+        super().test_save_load_local(tmp_path, base_pipe_output, expected_max_difference=3e-3)
 
     def test_inference_batch_single_identical(self):
         super().test_inference_batch_single_identical(expected_max_diff=3e-3)
 
 
+class TestDDIMPipelineMemory(DDIMPipelineTesterConfig, MemoryTesterMixin):
+    """Memory optimization tests (CPU offload, group offload, layerwise casting) for the DDIM pipeline."""
+
+
 @slow
 @require_torch_accelerator
-class DDIMPipelineIntegrationTests(unittest.TestCase):
+class TestDDIMPipelineIntegration:
     def test_inference_cifar10(self):
         model_id = "google/ddpm-cifar10-32"
 

tests/pipelines/ddpm/test_ddpm.py

@@ -13,24 +13,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-
 import numpy as np
 import torch
 
 from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel
 
-from ...testing_utils import enable_full_determinism, require_torch_accelerator, slow, torch_device
+from ...testing_utils import (
+    assert_tensors_close,
+    enable_full_determinism,
+    require_torch_accelerator,
+    slow,
+    torch_device,
+)
+from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
+from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
 
 
 enable_full_determinism()
 
 
-class DDPMPipelineFastTests(unittest.TestCase):
-    @property
-    def dummy_uncond_unet(self):
+class DDPMPipelineTesterConfig(BasePipelineTesterConfig):
+    pipeline_class = DDPMPipeline
+    required_input_params_in_call_signature = UNCONDITIONAL_IMAGE_GENERATION_PARAMS
+    batch_input_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
+    # DDPM is unconditional and samples its own noise: there is no prompt to repeat
+    # (`num_images_per_prompt`) and no user-suppliable `latents`.
+    optional_input_params = BasePipelineTesterConfig.optional_input_params - {"num_images_per_prompt", "latents"}
+    output_shape = (3, 8, 8)
+
+    def get_dummy_components(self):
         torch.manual_seed(0)
-        model = UNet2DModel(
+        unet = UNet2DModel(
             block_out_channels=(4, 8),
             layers_per_block=1,
             norm_num_groups=4,
@@ -40,57 +53,61 @@ def dummy_uncond_unet(self):
             down_block_types=("DownBlock2D", "AttnDownBlock2D"),
             up_block_types=("AttnUpBlock2D", "UpBlock2D"),
         )
-        return model
-
-    def test_fast_inference(self):
-        device = "cpu"
-        unet = self.dummy_uncond_unet
         scheduler = DDPMScheduler()
+        return {"unet": unet, "scheduler": scheduler}
 
-        ddpm = DDPMPipeline(unet=unet, scheduler=scheduler)
-        ddpm.to(device)
-        ddpm.set_progress_bar_config(disable=None)
+    def get_dummy_inputs(self):
+        return {
+            "batch_size": 1,
+            "generator": self.get_generator(0),
+            "num_inference_steps": 2,
+            # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
+            "output_type": "pt",
+        }
 
-        generator = torch.Generator(device=device).manual_seed(0)
-        image = ddpm(generator=generator, num_inference_steps=2, output_type="np").images
 
-        generator = torch.Generator(device=device).manual_seed(0)
-        image_from_tuple = ddpm(generator=generator, num_inference_steps=2, output_type="np", return_dict=False)[0]
+class TestDDPMPipeline(DDPMPipelineTesterConfig, PipelineTesterMixin):
+    def test_inference(self):
+        # Run on CPU: the expected slice below is CPU-specific.
+        pipe = self.get_pipeline()
 
-        image_slice = image[0, -3:, -3:, -1]
-        image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1]
+        image = pipe(**self.get_dummy_inputs()).images
+        generated_image = image[0]
+        assert generated_image.shape == self.output_shape
 
-        assert image.shape == (1, 8, 8, 3)
-        expected_slice = np.array([0.0, 0.9996672, 0.00329116, 1.0, 0.9995991, 1.0, 0.0060907, 0.00115037, 0.0])
+        # fmt: off
+        expected_slice = torch.tensor([0.0, 0.9996672, 0.00329116, 1.0, 0.9995991, 1.0, 0.0060907, 0.00115037, 0.0])
+        # fmt: on
 
-        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
-        assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2
+        # `"pt"` images are `(channels, height, width)`, so the trailing-channel corner slice of the old
+        # `"np"` layout is the last channel's bottom-right 3x3 block here.
+        generated_slice = generated_image[-1, -3:, -3:].flatten()
+        assert_tensors_close(generated_slice, expected_slice, atol=1e-2)
 
     def test_inference_predict_sample(self):
-        unet = self.dummy_uncond_unet
-        scheduler = DDPMScheduler(prediction_type="sample")
+        # `prediction_type="sample"` makes the UNet output the denoised sample rather than the noise, so the
+        # scheduler consumes it differently and the pipeline must produce a different image than the default
+        # `epsilon` parameterization.
+        pipe = self.get_pipeline().to(torch_device)
+        output_epsilon = self.run_pipe(pipe)
 
-        ddpm = DDPMPipeline(unet=unet, scheduler=scheduler)
-        ddpm.to(torch_device)
-        ddpm.set_progress_bar_config(disable=None)
+        components = self.get_dummy_components()
+        components["scheduler"] = DDPMScheduler(prediction_type="sample")
+        pipe_sample = self.get_pipeline(**components).to(torch_device)
+        output_sample = self.run_pipe(pipe_sample)
 
-        generator = torch.manual_seed(0)
-        image = ddpm(generator=generator, num_inference_steps=2, output_type="np").images
+        assert output_sample.shape == output_epsilon.shape
+        assert not torch.isnan(output_sample).any()
+        assert not torch.allclose(output_sample, output_epsilon, atol=1e-3)
 
-        generator = torch.manual_seed(0)
-        image_eps = ddpm(generator=generator, num_inference_steps=2, output_type="np")[0]
-
-        image_slice = image[0, -3:, -3:, -1]
-        image_eps_slice = image_eps[0, -3:, -3:, -1]
 
-        assert image.shape == (1, 8, 8, 3)
-        tolerance = 1e-2 if torch_device != "mps" else 3e-2
-        assert np.abs(image_slice.flatten() - image_eps_slice.flatten()).max() < tolerance
+class TestDDPMPipelineMemory(DDPMPipelineTesterConfig, MemoryTesterMixin):
+    """Memory optimization tests (CPU offload, group offload, layerwise casting) for the DDPM pipeline."""
 
 
 @slow
 @require_torch_accelerator
-class DDPMPipelineIntegrationTests(unittest.TestCase):
+class TestDDPMPipelineIntegration:
     def test_inference_cifar10(self):
         model_id = "google/ddpm-cifar10-32"