[tests] refactor k-series pipeline tests

tests/pipelines/kandinsky/test_kandinsky.py

@@ -14,34 +14,40 @@
 # limitations under the License.
 
 import gc
-import random
-import unittest
 
-import numpy as np
 import pytest
 import torch
 from transformers import XLMRobertaTokenizerFast
 
 from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNet2DConditionModel, VQModel
 from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
-from diffusers.utils import is_transformers_version
 
 from ...testing_utils import (
+    assert_tensors_close,
     backend_empty_cache,
     enable_full_determinism,
-    floats_tensor,
     load_numpy,
     require_torch_accelerator,
     slow,
     torch_device,
 )
-from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
+from ..test_pipelines_common import assert_mean_pixel_difference
+from ..testing_utils import (
+    BasePipelineTesterConfig,
+    MemoryTesterMixin,
+    PipelineTesterMixin,
+)
 
 
 enable_full_determinism()
 
 
-class Dummies:
+class KandinskyPipelineTesterConfig(BasePipelineTesterConfig):
+    pipeline_class = KandinskyPipeline
+    required_input_params_in_call_signature = frozenset(["prompt", "image_embeds", "negative_image_embeds"])
+    batch_input_params = frozenset(["prompt", "negative_prompt", "image_embeds", "negative_image_embeds"])
+    output_shape = (3, 64, 64)
+
     @property
     def text_embedder_hidden_size(self):
         return 32
@@ -162,136 +168,67 @@ def get_dummy_components(self):
         }
         return components
 
-    def get_dummy_inputs(self, device, seed=0):
-        image_embeds = floats_tensor((1, self.cross_attention_dim), rng=random.Random(seed)).to(device)
-        negative_image_embeds = floats_tensor((1, self.cross_attention_dim), rng=random.Random(seed + 1)).to(device)
-        if str(device).startswith("mps"):
-            generator = torch.manual_seed(seed)
-        else:
-            generator = torch.Generator(device=device).manual_seed(seed)
-        inputs = {
+    def get_dummy_inputs(self):
+        image_embeds = torch.randn((1, self.cross_attention_dim), generator=self.get_generator(0)).to(torch_device)
+        negative_image_embeds = torch.randn((1, self.cross_attention_dim), generator=self.get_generator(1)).to(
+            torch_device
+        )
+        return {
             "prompt": "horse",
             "image_embeds": image_embeds,
             "negative_image_embeds": negative_image_embeds,
-            "generator": generator,
+            "generator": self.get_generator(0),
             "height": 64,
             "width": 64,
             "guidance_scale": 4.0,
             "num_inference_steps": 2,
-            "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 KandinskyPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
-    pipeline_class = KandinskyPipeline
-    params = [
-        "prompt",
-        "image_embeds",
-        "negative_image_embeds",
-    ]
-    batch_params = ["prompt", "negative_prompt", "image_embeds", "negative_image_embeds"]
-    required_optional_params = [
-        "generator",
-        "height",
-        "width",
-        "latents",
-        "guidance_scale",
-        "negative_prompt",
-        "num_inference_steps",
-        "return_dict",
-        "guidance_scale",
-        "num_images_per_prompt",
-        "output_type",
-        "return_dict",
-    ]
-    test_xformers_attention = False
+class TestKandinskyPipeline(KandinskyPipelineTesterConfig, PipelineTesterMixin):
+    def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-1):
+        # Batched inference is only approximately equal to single inference here: the batch pads to the longest
+        # prompt and the tiny 2-step denoising loop amplifies the resulting attention differences. Tolerance set
+        # from the measured drift.
+        super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
 
-    def get_dummy_components(self):
-        dummy = Dummies()
-        return dummy.get_dummy_components()
-
-    def get_dummy_inputs(self, device, seed=0):
-        dummy = Dummies()
-        return dummy.get_dummy_inputs(device=device, seed=seed)
-
-    @pytest.mark.xfail(
-        condition=is_transformers_version(">=", "4.56.2"),
-        reason="Latest transformers changes the slices",
-        strict=False,
-    )
     def test_kandinsky(self):
-        device = "cpu"
-
-        components = self.get_dummy_components()
+        # Run on CPU: the expected slice below is CPU-specific.
+        pipe = self.get_pipeline()
 
-        pipe = self.pipeline_class(**components)
-        pipe = pipe.to(device)
+        image = pipe(**self.get_dummy_inputs()).images
+        image_from_tuple = pipe(**self.get_dummy_inputs(), return_dict=False)[0]
 
-        pipe.set_progress_bar_config(disable=None)
+        assert image.shape == (1, *self.output_shape)
 
-        output = pipe(**self.get_dummy_inputs(device))
-        image = output.images
+        # This pipeline only denormalizes the decoded image for `output_type` "np"/"pil", so `"pt"` hands back the
+        # raw decoder output. Map it into the [0, 1] range the expected slice below was recorded in.
+        image = (image * 0.5 + 0.5).clamp(0, 1)
+        image_from_tuple = (image_from_tuple * 0.5 + 0.5).clamp(0, 1)
 
-        image_from_tuple = pipe(
-            **self.get_dummy_inputs(device),
-            return_dict=False,
-        )[0]
+        # fmt: off
+        expected_slice = torch.tensor([0.4428, 0.7424, 0.3413, 1.0000, 0.7061, 0.3452, 0.5017, 0.3987, 0.5046])
+        # fmt: on
+        assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2)
+        assert_tensors_close(image_from_tuple[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2)
 
-        image_slice = image[0, -3:, -3:, -1]
-        image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1]
 
-        assert image.shape == (1, 64, 64, 3)
-
-        expected_slice = np.array([1.0000, 1.0000, 0.2766, 1.0000, 0.5447, 0.1737, 1.0000, 0.4316, 0.9024])
-
-        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, (
-            f" expected_slice {expected_slice}, but got {image_slice.flatten()}"
-        )
-        assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2, (
-            f" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"
-        )
-
-    @require_torch_accelerator
-    def test_offloads(self):
-        pipes = []
-        components = self.get_dummy_components()
-        sd_pipe = self.pipeline_class(**components).to(torch_device)
-        pipes.append(sd_pipe)
-
-        components = self.get_dummy_components()
-        sd_pipe = self.pipeline_class(**components)
-        sd_pipe.enable_model_cpu_offload(device=torch_device)
-        pipes.append(sd_pipe)
-
-        components = self.get_dummy_components()
-        sd_pipe = self.pipeline_class(**components)
-        sd_pipe.enable_sequential_cpu_offload(device=torch_device)
-        pipes.append(sd_pipe)
-
-        image_slices = []
-        for pipe in pipes:
-            inputs = self.get_dummy_inputs(torch_device)
-            image = pipe(**inputs).images
-
-            image_slices.append(image[0, -3:, -3:, -1].flatten())
-
-        assert np.abs(image_slices[0] - image_slices[1]).max() < 1e-3
-        assert np.abs(image_slices[0] - image_slices[2]).max() < 1e-3
+class TestKandinskyPipelineMemory(KandinskyPipelineTesterConfig, MemoryTesterMixin):
+    """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Kandinsky pipeline."""
 
 
 @slow
 @require_torch_accelerator
-class KandinskyPipelineIntegrationTests(unittest.TestCase):
-    def setUp(self):
-        # clean up the VRAM before each test
-        super().setUp()
+class TestKandinskyPipelineIntegration:
+    @pytest.fixture(autouse=True)
+    def cleanup(self):
+        # clean up the VRAM before and after each test
         gc.collect()
         backend_empty_cache(torch_device)
-
-    def tearDown(self):
-        # clean up the VRAM after each test
-        super().tearDown()
+        yield
         gc.collect()
         backend_empty_cache(torch_device)