refactor deepfloyd if pipeline tests to the new mixin structure

tests/pipelines/deepfloyd_if/test_if.py

@@ -14,15 +14,17 @@
 # limitations under the License.
 
 import gc
-import unittest
 
+import pytest
 import torch
+from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
 
-from diffusers import IFPipeline
+from diffusers import DDPMScheduler, IFPipeline, UNet2DConditionModel
 from diffusers.models.attention_processor import AttnAddedKVProcessor
-from diffusers.utils.import_utils import is_xformers_available
+from diffusers.pipelines.deepfloyd_if import IFWatermarker
 
 from ...testing_utils import (
+    assert_tensors_close,
     backend_empty_cache,
     backend_max_memory_allocated,
     backend_reset_max_memory_allocated,
@@ -34,77 +36,172 @@
     slow,
     torch_device,
 )
-from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
-from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
-from . import IFPipelineTesterMixin
+from ..test_pipelines_common import assert_mean_pixel_difference
+from ..testing_utils import (
+    BasePipelineTesterConfig,
+    PipelineOffloadTesterMixin,
+    PipelineTesterMixin,
+)
 
 
-@skip_mps
-class IFPipelineFastTests(PipelineTesterMixin, IFPipelineTesterMixin, unittest.TestCase):
+class IFPipelineTesterConfig(BasePipelineTesterConfig):
     pipeline_class = IFPipeline
-    params = TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
-    batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
-    required_optional_params = PipelineTesterMixin.required_optional_params - {"latents"}
+    required_input_params_in_call_signature = frozenset(
+        ["prompt", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"]
+    )
+    # IF pipelines take no `latents` argument (pixel-space UNet, no user-suppliable latents)
+    optional_input_params = BasePipelineTesterConfig.optional_input_params - {"latents"}
+    batch_input_params = frozenset(["prompt", "negative_prompt"])
+    output_shape = (3, 32, 32)
 
     def get_dummy_components(self):
-        return self._get_dummy_components()
+        torch.manual_seed(0)
+        config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
+        text_encoder = T5EncoderModel(config)
+
+        torch.manual_seed(0)
+        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
+
+        torch.manual_seed(0)
+        unet = UNet2DConditionModel(
+            sample_size=32,
+            layers_per_block=1,
+            block_out_channels=[32, 64],
+            down_block_types=[
+                "ResnetDownsampleBlock2D",
+                "SimpleCrossAttnDownBlock2D",
+            ],
+            mid_block_type="UNetMidBlock2DSimpleCrossAttn",
+            up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"],
+            in_channels=3,
+            out_channels=6,
+            cross_attention_dim=32,
+            encoder_hid_dim=32,
+            attention_head_dim=8,
+            addition_embed_type="text",
+            addition_embed_type_num_heads=2,
+            cross_attention_norm="group_norm",
+            resnet_time_scale_shift="scale_shift",
+            act_fn="gelu",
+        )
+        unet.set_attn_processor(AttnAddedKVProcessor())  # For reproducibility tests
+
+        torch.manual_seed(0)
+        scheduler = DDPMScheduler(
+            num_train_timesteps=1000,
+            beta_schedule="squaredcos_cap_v2",
+            beta_start=0.0001,
+            beta_end=0.02,
+            thresholding=True,
+            dynamic_thresholding_ratio=0.95,
+            sample_max_value=1.0,
+            prediction_type="epsilon",
+            variance_type="learned_range",
+        )
 
-    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)
+        torch.manual_seed(0)
+        watermarker = IFWatermarker()
+
+        return {
+            "text_encoder": text_encoder,
+            "tokenizer": tokenizer,
+            "unet": unet,
+            "scheduler": scheduler,
+            "watermarker": watermarker,
+            "safety_checker": None,
+            "feature_extractor": None,
+        }
 
-        inputs = {
+    def get_dummy_inputs(self):
+        return {
             "prompt": "A painting of a squirrel eating a burger",
-            "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
 
-    @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU")
-    @require_accelerator
-    def test_save_load_float16(self):
-        # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
-        super().test_save_load_float16(expected_max_diff=1e-1)
+@skip_mps
+class TestIFPipeline(IFPipelineTesterConfig, PipelineTesterMixin):
+    def test_inference(self):
+        # Run on CPU: the expected slice below is CPU-specific.
+        pipe = self.get_pipeline()
+
+        inputs = self.get_dummy_inputs()
+        image = pipe(**inputs).images
+        generated_image = image[0]
+        assert generated_image.shape == self.output_shape
 
-    def test_attention_slicing_forward_pass(self):
-        self._test_attention_slicing_forward_pass(expected_max_diff=1e-2)
+        # fmt: off
+        expected_slice = torch.tensor([-1.0000, -1.0000, -0.2177, -0.8947, 0.8860, 1.0000, -0.5022, -1.0000, 0.9550, -0.9627, -1.0000, 0.9768, -1.0000, -0.9924, -1.0000, -1.0000])
+        # fmt: on
 
-    def test_save_load_local(self):
-        self._test_save_load_local()
+        generated_slice = generated_image.flatten()
+        generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]])
+        assert_tensors_close(generated_slice, expected_slice, atol=1e-3)
+
+    @pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU")
+    @require_accelerator
+    def test_save_load_float16(self, tmp_path):
+        # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
+        super().test_save_load_float16(tmp_path, expected_max_diff=1e-1)
 
     def test_inference_batch_single_identical(self):
-        self._test_inference_batch_single_identical(
-            expected_max_diff=1e-2,
+        super().test_inference_batch_single_identical(expected_max_diff=1e-2)
+
+    def test_save_load_optional_components(self, tmp_path):
+        # The text encoder is optional so a pre-encoded prompt can be passed directly; the base test would
+        # pass the raw prompt with `text_encoder=None`, so encode it first (the intended usage).
+        pipe = self.get_pipeline().to(torch_device)
+
+        inputs = self.get_dummy_inputs()
+        prompt = inputs.pop("prompt")
+        prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(prompt)
+
+        for optional_component in pipe._optional_components:
+            setattr(pipe, optional_component, None)
+
+        inputs["prompt_embeds"] = prompt_embeds
+        inputs["negative_prompt_embeds"] = negative_prompt_embeds
+        torch.manual_seed(0)
+        output = pipe(**inputs)[0]
+
+        pipe.save_pretrained(tmp_path, safe_serialization=False)
+        pipe_loaded = self.pipeline_class.from_pretrained(tmp_path)
+        pipe_loaded.to(torch_device)
+        pipe_loaded.set_progress_bar_config(disable=None)
+
+        for optional_component in pipe._optional_components:
+            assert getattr(pipe_loaded, optional_component) is None, (
+                f"`{optional_component}` did not stay set to None after loading."
+            )
+
+        inputs = self.get_dummy_inputs()
+        inputs.pop("prompt")
+        inputs["prompt_embeds"] = prompt_embeds
+        inputs["negative_prompt_embeds"] = negative_prompt_embeds
+        torch.manual_seed(0)
+        output_loaded = pipe_loaded(**inputs)[0]
+
+        assert_tensors_close(
+            output_loaded, output, atol=1e-4, msg="Output changed after dropping optional components."
         )
 
-    @unittest.skipIf(
-        torch_device != "cuda" or not is_xformers_available(),
-        reason="XFormers attention is only available with CUDA and `xformers` installed",
-    )
-    def test_xformers_attention_forwardGenerator_pass(self):
-        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3)
 
-    @unittest.skip("Functionality is tested elsewhere.")
-    def test_save_load_optional_components(self):
-        pass
+@skip_mps
+class TestIFPipelineMemory(IFPipelineTesterConfig, PipelineOffloadTesterMixin):
+    pass
 
 
 @slow
 @require_torch_accelerator
-class IFPipelineSlowTests(unittest.TestCase):
-    def setUp(self):
-        # clean up the VRAM before each test
-        super().setUp()
+class TestIFPipelineSlow:
+    @pytest.fixture(autouse=True)
+    def cleanup(self):
         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)