feat(diffusers/pipelines): add fast tests for pag pipes x6 (v0.30.3)

tests/diffusers/pipelines/pag/test_pag_controlnet_sd.py

@@ -0,0 +1,204 @@
+import random
+import unittest
+
+import numpy as np
+import torch
+from ddt import data, ddt, unpack
+from transformers import CLIPTextConfig
+
+import mindspore as ms
+
+from ..pipeline_test_utils import (
+    THRESHOLD_FP16,
+    THRESHOLD_FP32,
+    PipelineTesterMixin,
+    floats_tensor,
+    get_module,
+    get_pipeline_components,
+)
+
+test_cases = [
+    {"mode": ms.PYNATIVE_MODE, "dtype": "float32"},
+    {"mode": ms.PYNATIVE_MODE, "dtype": "float16"},
+    {"mode": ms.GRAPH_MODE, "dtype": "float32"},
+    {"mode": ms.GRAPH_MODE, "dtype": "float16"},
+]
+
+
+@ddt
+class StableDiffusionControlNetPAGPipelineFastTests(
+    PipelineTesterMixin,
+    unittest.TestCase,
+):
+    time_cond_proj_dim = None
+    pipeline_config = [
+        [
+            "unet",
+            "diffusers.models.unets.unet_2d_condition.UNet2DConditionModel",
+            "mindone.diffusers.models.unets.unet_2d_condition.UNet2DConditionModel",
+            dict(
+                block_out_channels=(4, 8),
+                layers_per_block=2,
+                sample_size=32,
+                in_channels=4,
+                out_channels=4,
+                down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
+                up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
+                cross_attention_dim=8,
+                time_cond_proj_dim=time_cond_proj_dim,
+                norm_num_groups=2,
+            ),
+        ],
+        [
+            "controlnet",
+            "diffusers.models.controlnet.ControlNetModel",
+            "mindone.diffusers.models.controlnet.ControlNetModel",
+            dict(
+                block_out_channels=(4, 8),
+                layers_per_block=2,
+                in_channels=4,
+                down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
+                conditioning_embedding_out_channels=(2, 4),
+                cross_attention_dim=8,
+                norm_num_groups=2,
+            ),
+        ],
+        [
+            "scheduler",
+            "diffusers.schedulers.scheduling_ddim.DDIMScheduler",
+            "mindone.diffusers.schedulers.scheduling_ddim.DDIMScheduler",
+            dict(
+                beta_start=0.00085,
+                beta_end=0.012,
+                beta_schedule="scaled_linear",
+                clip_sample=False,
+                set_alpha_to_one=False,
+            ),
+        ],
+        [
+            "vae",
+            "diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL",
+            "mindone.diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL",
+            dict(
+                block_out_channels=[4, 8],
+                in_channels=3,
+                out_channels=3,
+                down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
+                up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
+                latent_channels=4,
+                norm_num_groups=2,
+            ),
+        ],
+        [
+            "text_encoder",
+            "transformers.models.clip.modeling_clip.CLIPTextModel",
+            "mindone.transformers.models.clip.modeling_clip.CLIPTextModel",
+            dict(
+                config=CLIPTextConfig(
+                    bos_token_id=0,
+                    eos_token_id=2,
+                    hidden_size=8,
+                    intermediate_size=16,
+                    layer_norm_eps=1e-05,
+                    num_attention_heads=2,
+                    num_hidden_layers=2,
+                    pad_token_id=1,
+                    vocab_size=1000,
+                ),
+            ),
+        ],
+        [
+            "tokenizer",
+            "transformers.models.clip.tokenization_clip.CLIPTokenizer",
+            "transformers.models.clip.tokenization_clip.CLIPTokenizer",
+            dict(
+                pretrained_model_name_or_path="hf-internal-testing/tiny-random-clip",
+            ),
+        ],
+    ]
+
+    def get_dummy_components(self):
+        components = {
+            key: None
+            for key in [
+                "unet",
+                "controlnet",
+                "scheduler",
+                "vae",
+                "text_encoder",
+                "tokenizer",
+                "feature_extractor",
+                "image_encoder",
+            ]
+        }
+
+        return get_pipeline_components(components, self.pipeline_config)
+
+    def get_dummy_inputs(self, seed=0):
+        generator = torch.manual_seed(seed)
+
+        controlnet_embedder_scale_factor = 2
+        pt_image = floats_tensor(
+            (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor),
+            rng=random.Random(seed),
+        )
+        ms_image = ms.Tensor(pt_image.numpy())
+
+        pt_inputs = {
+            "prompt": "A painting of a squirrel eating a burger",
+            "generator": generator,
+            "num_inference_steps": 2,
+            "guidance_scale": 6.0,
+            "pag_scale": 3.0,
+            "output_type": "np",
+            "image": pt_image,
+        }
+
+        ms_inputs = {
+            "prompt": "A painting of a squirrel eating a burger",
+            "generator": generator,
+            "num_inference_steps": 2,
+            "guidance_scale": 6.0,
+            "pag_scale": 3.0,
+            "output_type": "np",
+            "image": ms_image,
+        }
+
+        return pt_inputs, ms_inputs
+
+    @data(*test_cases)
+    @unpack
+    def test_inference(self, mode, dtype):
+        ms.set_context(mode=mode)
+
+        pt_components, ms_components = self.get_dummy_components()
+        pt_pipe_cls = get_module(
+            "diffusers.pipelines.pag.pipeline_pag_controlnet_sd.StableDiffusionControlNetPAGPipeline"
+        )
+        ms_pipe_cls = get_module(
+            "mindone.diffusers.pipelines.pag.pipeline_pag_controlnet_sd.StableDiffusionControlNetPAGPipeline"
+        )
+
+        pt_pipe = pt_pipe_cls(**pt_components, pag_applied_layers=["mid", "up", "down"])
+        ms_pipe = ms_pipe_cls(**ms_components, pag_applied_layers=["mid", "up", "down"])
+
+        pt_pipe.set_progress_bar_config(disable=None)
+        ms_pipe.set_progress_bar_config(disable=None)
+
+        ms_dtype, pt_dtype = getattr(ms, dtype), getattr(torch, dtype)
+        pt_pipe = pt_pipe.to(pt_dtype)
+        ms_pipe = ms_pipe.to(ms_dtype)
+
+        pt_inputs, ms_inputs = self.get_dummy_inputs()
+
+        torch.manual_seed(0)
+        pt_image = pt_pipe(**pt_inputs)
+        torch.manual_seed(0)
+        ms_image = ms_pipe(**ms_inputs)
+
+        pt_image_slice = pt_image.images[0, -3:, -3:, -1]
+        ms_image_slice = ms_image[0][0, -3:, -3:, -1]
+
+        threshold = THRESHOLD_FP32 if dtype == "float32" else THRESHOLD_FP16
+
+        assert np.max(np.linalg.norm(pt_image_slice - ms_image_slice) / np.linalg.norm(pt_image_slice)) < threshold