test_pipelines_common.py 95.9 KB
Newer Older
1
2
import gc
import inspect
3
4
import json
import os
5
6
import tempfile
import unittest
7
import uuid
Aryan's avatar
Aryan committed
8
from typing import Any, Callable, Dict, Union
9
10

import numpy as np
Anh71me's avatar
Anh71me committed
11
import PIL.Image
12
import torch
13
import torch.nn as nn
14
15
from huggingface_hub import ModelCard, delete_repo
from huggingface_hub.utils import is_jinja_available
16
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
17

18
import diffusers
19
20
21
22
23
24
25
from diffusers import (
    AsymmetricAutoencoderKL,
    AutoencoderKL,
    AutoencoderTiny,
    ConsistencyDecoderVAE,
    DDIMScheduler,
    DiffusionPipeline,
Álvaro Somoza's avatar
Álvaro Somoza committed
26
    KolorsPipeline,
27
    StableDiffusionPipeline,
28
    StableDiffusionXLPipeline,
29
30
    UNet2DConditionModel,
)
31
from diffusers.image_processor import VaeImageProcessor
Aryan's avatar
Aryan committed
32
from diffusers.loaders import IPAdapterMixin
33
from diffusers.models.attention_processor import AttnProcessor
34
from diffusers.models.controlnets.controlnet_xs import UNetControlNetXSModel
35
36
37
38
from diffusers.models.unets.unet_3d_condition import UNet3DConditionModel
from diffusers.models.unets.unet_i2vgen_xl import I2VGenXLUNet
from diffusers.models.unets.unet_motion_model import UNetMotionModel
from diffusers.pipelines.pipeline_utils import StableDiffusionMixin
39
from diffusers.schedulers import KarrasDiffusionSchedulers
40
from diffusers.utils import logging
41
from diffusers.utils.import_utils import is_xformers_available
42
43
from diffusers.utils.testing_utils import (
    CaptureLogger,
44
45
    require_accelerate_version_greater,
    require_accelerator,
46
47
48
49
    require_torch,
    skip_mps,
    torch_device,
)
50

51
from ..models.autoencoders.vae import (
52
53
54
55
56
    get_asym_autoencoder_kl_config,
    get_autoencoder_kl_config,
    get_autoencoder_tiny_config,
    get_consistency_vae_config,
)
57
58
59
60
from ..models.unets.test_models_unet_2d_condition import (
    create_ip_adapter_faceid_state_dict,
    create_ip_adapter_state_dict,
)
61
62
from ..others.test_utils import TOKEN, USER, is_staging_test

63

64
65
66
67
68
69
70
def to_np(tensor):
    if isinstance(tensor, torch.Tensor):
        tensor = tensor.detach().cpu().numpy()

    return tensor


71
72
73
74
75
def check_same_shape(tensor_list):
    shapes = [tensor.shape for tensor in tensor_list]
    return all(shape == shapes[0] for shape in shapes[1:])


76
77
78
79
80
81
82
83
84
85
86
def check_qkv_fusion_matches_attn_procs_length(model, original_attn_processors):
    current_attn_processors = model.attn_processors
    return len(current_attn_processors) == len(original_attn_processors)


def check_qkv_fusion_processors_exist(model):
    current_attn_processors = model.attn_processors
    proc_names = [v.__class__.__name__ for _, v in current_attn_processors.items()]
    return all(p.startswith("Fused") for p in proc_names)


87
88
89
90
91
92
class SDFunctionTesterMixin:
    """
    This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes.
    It provides a set of common tests for PyTorch pipeline that inherit from StableDiffusionMixin, e.g. vae_slicing, vae_tiling, freeu, etc.
    """

93
    def test_vae_slicing(self, image_count=4):
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
        device = "cpu"  # ensure determinism for the device-dependent torch.Generator
        components = self.get_dummy_components()
        # components["scheduler"] = LMSDiscreteScheduler.from_config(components["scheduler"].config)
        pipe = self.pipeline_class(**components)
        pipe = pipe.to(device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(device)
        inputs["prompt"] = [inputs["prompt"]] * image_count
        if "image" in inputs:  # fix batch size mismatch in I2V_Gen pipeline
            inputs["image"] = [inputs["image"]] * image_count
        output_1 = pipe(**inputs)

        # make sure sliced vae decode yields the same result
        pipe.enable_vae_slicing()
        inputs = self.get_dummy_inputs(device)
        inputs["prompt"] = [inputs["prompt"]] * image_count
        if "image" in inputs:
            inputs["image"] = [inputs["image"]] * image_count
        inputs["return_dict"] = False
        output_2 = pipe(**inputs)

        assert np.abs(output_2[0].flatten() - output_1[0].flatten()).max() < 1e-2

    def test_vae_tiling(self):
        components = self.get_dummy_components()

        # make sure here that pndm scheduler skips prk
        if "safety_checker" in components:
            components["safety_checker"] = None
        pipe = self.pipeline_class(**components)
125
        pipe = pipe.to(torch_device)
126
127
128
129
130
131
132
133
134
135
136
137
138
139
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(torch_device)
        inputs["return_dict"] = False

        # Test that tiled decode at 512x512 yields the same result as the non-tiled decode
        output_1 = pipe(**inputs)[0]

        # make sure tiled vae decode yields the same result
        pipe.enable_vae_tiling()
        inputs = self.get_dummy_inputs(torch_device)
        inputs["return_dict"] = False
        output_2 = pipe(**inputs)[0]

140
        assert np.abs(to_np(output_2) - to_np(output_1)).max() < 5e-1
141
142

        # test that tiled decode works with various shapes
143
        shapes = [(1, 4, 73, 97), (1, 4, 65, 49)]
YiYi Xu's avatar
YiYi Xu committed
144
145
146
147
        with torch.no_grad():
            for shape in shapes:
                zeros = torch.zeros(shape).to(torch_device)
                pipe.vae.decode(zeros)
148

149
    # MPS currently doesn't support ComplexFloats, which are required for FreeU - see https://github.com/huggingface/diffusers/issues/7569.
150
    @skip_mps
151
    def test_freeu(self):
152
153
154
155
156
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe = pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

157
        # Normal inference
158
159
        inputs = self.get_dummy_inputs(torch_device)
        inputs["return_dict"] = False
Dhruv Nair's avatar
Dhruv Nair committed
160
        inputs["output_type"] = "np"
161
162
        output = pipe(**inputs)[0]

163
        # FreeU-enabled inference
164
165
166
        pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4)
        inputs = self.get_dummy_inputs(torch_device)
        inputs["return_dict"] = False
Dhruv Nair's avatar
Dhruv Nair committed
167
        inputs["output_type"] = "np"
168
169
        output_freeu = pipe(**inputs)[0]

170
        # FreeU-disabled inference
171
172
173
174
175
176
177
178
        pipe.disable_freeu()
        freeu_keys = {"s1", "s2", "b1", "b2"}
        for upsample_block in pipe.unet.up_blocks:
            for key in freeu_keys:
                assert getattr(upsample_block, key) is None, f"Disabling of FreeU should have set {key} to None."

        inputs = self.get_dummy_inputs(torch_device)
        inputs["return_dict"] = False
Dhruv Nair's avatar
Dhruv Nair committed
179
        inputs["output_type"] = "np"
180
        output_no_freeu = pipe(**inputs)[0]
181
182
183
184

        assert not np.allclose(
            output[0, -3:, -3:, -1], output_freeu[0, -3:, -3:, -1]
        ), "Enabling of FreeU should lead to different results."
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
        assert np.allclose(
            output, output_no_freeu, atol=1e-2
        ), f"Disabling of FreeU should lead to results similar to the default pipeline results but Max Abs Error={np.abs(output_no_freeu - output).max()}."

    def test_fused_qkv_projections(self):
        device = "cpu"  # ensure determinism for the device-dependent torch.Generator
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe = pipe.to(device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(device)
        inputs["return_dict"] = False
        image = pipe(**inputs)[0]
        original_image_slice = image[0, -3:, -3:, -1]

        pipe.fuse_qkv_projections()
202
203
204
205
206
207
208
209
210
211
212
213
214
        for _, component in pipe.components.items():
            if (
                isinstance(component, nn.Module)
                and hasattr(component, "original_attn_processors")
                and component.original_attn_processors is not None
            ):
                assert check_qkv_fusion_processors_exist(
                    component
                ), "Something wrong with the fused attention processors. Expected all the attention processors to be fused."
                assert check_qkv_fusion_matches_attn_procs_length(
                    component, component.original_attn_processors
                ), "Something wrong with the attention processors concerning the fused QKV projections."

215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
        inputs = self.get_dummy_inputs(device)
        inputs["return_dict"] = False
        image_fused = pipe(**inputs)[0]
        image_slice_fused = image_fused[0, -3:, -3:, -1]

        pipe.unfuse_qkv_projections()
        inputs = self.get_dummy_inputs(device)
        inputs["return_dict"] = False
        image_disabled = pipe(**inputs)[0]
        image_slice_disabled = image_disabled[0, -3:, -3:, -1]

        assert np.allclose(
            original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2
        ), "Fusion of QKV projections shouldn't affect the outputs."
        assert np.allclose(
            image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2
        ), "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."
        assert np.allclose(
            original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2
        ), "Original outputs should match when fused QKV projections are disabled."


Aryan's avatar
Aryan committed
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
class IPAdapterTesterMixin:
    """
    This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes.
    It provides a set of common tests for pipelines that support IP Adapters.
    """

    def test_pipeline_signature(self):
        parameters = inspect.signature(self.pipeline_class.__call__).parameters

        assert issubclass(self.pipeline_class, IPAdapterMixin)
        self.assertIn(
            "ip_adapter_image",
            parameters,
            "`ip_adapter_image` argument must be supported by the `__call__` method",
        )
        self.assertIn(
            "ip_adapter_image_embeds",
            parameters,
            "`ip_adapter_image_embeds` argument must be supported by the `__call__` method",
        )

    def _get_dummy_image_embeds(self, cross_attention_dim: int = 32):
        return torch.randn((2, 1, cross_attention_dim), device=torch_device)

261
262
263
    def _get_dummy_faceid_image_embeds(self, cross_attention_dim: int = 32):
        return torch.randn((2, 1, 1, cross_attention_dim), device=torch_device)

264
265
266
267
268
    def _get_dummy_masks(self, input_size: int = 64):
        _masks = torch.zeros((1, 1, input_size, input_size), device=torch_device)
        _masks[0, :, :, : int(input_size / 2)] = 1
        return _masks

Aryan's avatar
Aryan committed
269
270
271
272
273
274
275
276
277
    def _modify_inputs_for_ip_adapter_test(self, inputs: Dict[str, Any]):
        parameters = inspect.signature(self.pipeline_class.__call__).parameters
        if "image" in parameters.keys() and "strength" in parameters.keys():
            inputs["num_inference_steps"] = 4

        inputs["output_type"] = "np"
        inputs["return_dict"] = False
        return inputs

278
279
280
281
282
283
284
285
286
    def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=None):
        r"""Tests for IP-Adapter.

        The following scenarios are tested:
          - Single IP-Adapter with scale=0 should produce same output as no IP-Adapter.
          - Multi IP-Adapter with scale=0 should produce same output as no IP-Adapter.
          - Single IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter.
          - Multi IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter.
        """
287
288
289
290
        # Raising the tolerance for this test when it's run on a CPU because we
        # compare against static slices and that can be shaky (with a VVVV low probability).
        expected_max_diff = 9e-4 if torch_device == "cpu" else expected_max_diff

Aryan's avatar
Aryan committed
291
292
293
294
295
296
297
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components).to(torch_device)
        pipe.set_progress_bar_config(disable=None)
        cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32)

        # forward pass without ip adapter
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
298
299
300
301
        if expected_pipe_slice is None:
            output_without_adapter = pipe(**inputs)[0]
        else:
            output_without_adapter = expected_pipe_slice
Aryan's avatar
Aryan committed
302

303
        # 1. Single IP-Adapter test cases
Aryan's avatar
Aryan committed
304
305
306
307
308
309
310
311
        adapter_state_dict = create_ip_adapter_state_dict(pipe.unet)
        pipe.unet._load_ip_adapter_weights(adapter_state_dict)

        # forward pass with single ip adapter, but scale=0 which should have no effect
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
        pipe.set_ip_adapter_scale(0.0)
        output_without_adapter_scale = pipe(**inputs)[0]
312
313
        if expected_pipe_slice is not None:
            output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten()
Aryan's avatar
Aryan committed
314
315
316
317
318
319

        # forward pass with single ip adapter, but with scale of adapter weights
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
        pipe.set_ip_adapter_scale(42.0)
        output_with_adapter_scale = pipe(**inputs)[0]
320
321
        if expected_pipe_slice is not None:
            output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten()
Aryan's avatar
Aryan committed
322
323
324
325
326
327
328
329
330
331
332
333
334

        max_diff_without_adapter_scale = np.abs(output_without_adapter_scale - output_without_adapter).max()
        max_diff_with_adapter_scale = np.abs(output_with_adapter_scale - output_without_adapter).max()

        self.assertLess(
            max_diff_without_adapter_scale,
            expected_max_diff,
            "Output without ip-adapter must be same as normal inference",
        )
        self.assertGreater(
            max_diff_with_adapter_scale, 1e-2, "Output with ip-adapter must be different from normal inference"
        )

335
        # 2. Multi IP-Adapter test cases
Aryan's avatar
Aryan committed
336
337
338
339
340
341
342
343
344
        adapter_state_dict_1 = create_ip_adapter_state_dict(pipe.unet)
        adapter_state_dict_2 = create_ip_adapter_state_dict(pipe.unet)
        pipe.unet._load_ip_adapter_weights([adapter_state_dict_1, adapter_state_dict_2])

        # forward pass with multi ip adapter, but scale=0 which should have no effect
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] * 2
        pipe.set_ip_adapter_scale([0.0, 0.0])
        output_without_multi_adapter_scale = pipe(**inputs)[0]
345
346
        if expected_pipe_slice is not None:
            output_without_multi_adapter_scale = output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten()
Aryan's avatar
Aryan committed
347
348
349
350
351
352

        # forward pass with multi ip adapter, but with scale of adapter weights
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)] * 2
        pipe.set_ip_adapter_scale([42.0, 42.0])
        output_with_multi_adapter_scale = pipe(**inputs)[0]
353
354
        if expected_pipe_slice is not None:
            output_with_multi_adapter_scale = output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten()
Aryan's avatar
Aryan committed
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370

        max_diff_without_multi_adapter_scale = np.abs(
            output_without_multi_adapter_scale - output_without_adapter
        ).max()
        max_diff_with_multi_adapter_scale = np.abs(output_with_multi_adapter_scale - output_without_adapter).max()
        self.assertLess(
            max_diff_without_multi_adapter_scale,
            expected_max_diff,
            "Output without multi-ip-adapter must be same as normal inference",
        )
        self.assertGreater(
            max_diff_with_multi_adapter_scale,
            1e-2,
            "Output with multi-ip-adapter scale must be different from normal inference",
        )

371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
    def test_ip_adapter_cfg(self, expected_max_diff: float = 1e-4):
        parameters = inspect.signature(self.pipeline_class.__call__).parameters

        if "guidance_scale" not in parameters:
            return

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components).to(torch_device)
        pipe.set_progress_bar_config(disable=None)
        cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32)

        adapter_state_dict = create_ip_adapter_state_dict(pipe.unet)
        pipe.unet._load_ip_adapter_weights(adapter_state_dict)
        pipe.set_ip_adapter_scale(1.0)

        # forward pass with CFG not applied
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)[0].unsqueeze(0)]
        inputs["guidance_scale"] = 1.0
        out_no_cfg = pipe(**inputs)[0]

        # forward pass with CFG applied
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
        inputs["guidance_scale"] = 7.5
        out_cfg = pipe(**inputs)[0]

        assert out_cfg.shape == out_no_cfg.shape

400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
    def test_ip_adapter_masks(self, expected_max_diff: float = 1e-4):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components).to(torch_device)
        pipe.set_progress_bar_config(disable=None)
        cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32)
        sample_size = pipe.unet.config.get("sample_size", 32)
        block_out_channels = pipe.vae.config.get("block_out_channels", [128, 256, 512, 512])
        input_size = sample_size * (2 ** (len(block_out_channels) - 1))

        # forward pass without ip adapter
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        output_without_adapter = pipe(**inputs)[0]
        output_without_adapter = output_without_adapter[0, -3:, -3:, -1].flatten()

        adapter_state_dict = create_ip_adapter_state_dict(pipe.unet)
        pipe.unet._load_ip_adapter_weights(adapter_state_dict)

        # forward pass with single ip adapter and masks, but scale=0 which should have no effect
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
        inputs["cross_attention_kwargs"] = {"ip_adapter_masks": [self._get_dummy_masks(input_size)]}
        pipe.set_ip_adapter_scale(0.0)
        output_without_adapter_scale = pipe(**inputs)[0]
        output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten()

        # forward pass with single ip adapter and masks, but with scale of adapter weights
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(cross_attention_dim)]
        inputs["cross_attention_kwargs"] = {"ip_adapter_masks": [self._get_dummy_masks(input_size)]}
        pipe.set_ip_adapter_scale(42.0)
        output_with_adapter_scale = pipe(**inputs)[0]
        output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten()

        max_diff_without_adapter_scale = np.abs(output_without_adapter_scale - output_without_adapter).max()
        max_diff_with_adapter_scale = np.abs(output_with_adapter_scale - output_without_adapter).max()

        self.assertLess(
            max_diff_without_adapter_scale,
            expected_max_diff,
            "Output without ip-adapter must be same as normal inference",
        )
        self.assertGreater(
            max_diff_with_adapter_scale, 1e-3, "Output with ip-adapter must be different from normal inference"
        )

445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
    def test_ip_adapter_faceid(self, expected_max_diff: float = 1e-4):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components).to(torch_device)
        pipe.set_progress_bar_config(disable=None)
        cross_attention_dim = pipe.unet.config.get("cross_attention_dim", 32)

        # forward pass without ip adapter
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        output_without_adapter = pipe(**inputs)[0]
        output_without_adapter = output_without_adapter[0, -3:, -3:, -1].flatten()

        adapter_state_dict = create_ip_adapter_faceid_state_dict(pipe.unet)
        pipe.unet._load_ip_adapter_weights(adapter_state_dict)

        # forward pass with single ip adapter, but scale=0 which should have no effect
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_faceid_image_embeds(cross_attention_dim)]
        pipe.set_ip_adapter_scale(0.0)
        output_without_adapter_scale = pipe(**inputs)[0]
        output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten()

        # forward pass with single ip adapter, but with scale of adapter weights
        inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs(torch_device))
        inputs["ip_adapter_image_embeds"] = [self._get_dummy_faceid_image_embeds(cross_attention_dim)]
        pipe.set_ip_adapter_scale(42.0)
        output_with_adapter_scale = pipe(**inputs)[0]
        output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten()

        max_diff_without_adapter_scale = np.abs(output_without_adapter_scale - output_without_adapter).max()
        max_diff_with_adapter_scale = np.abs(output_with_adapter_scale - output_without_adapter).max()

        self.assertLess(
            max_diff_without_adapter_scale,
            expected_max_diff,
            "Output without ip-adapter must be same as normal inference",
        )
        self.assertGreater(
            max_diff_with_adapter_scale, 1e-3, "Output with ip-adapter must be different from normal inference"
        )

Aryan's avatar
Aryan committed
485

486
487
488
489
490
491
492
493
494
495
496
497
498
499
class PipelineLatentTesterMixin:
    """
    This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes.
    It provides a set of common tests for PyTorch pipeline that has vae, e.g.
    equivalence of different input and output types, etc.
    """

    @property
    def image_params(self) -> frozenset:
        raise NotImplementedError(
            "You need to set the attribute `image_params` in the child test class. "
            "`image_params` are tested for if all accepted input image types (i.e. `pt`,`pil`,`np`) are producing same results"
        )

500
501
502
503
504
505
506
    @property
    def image_latents_params(self) -> frozenset:
        raise NotImplementedError(
            "You need to set the attribute `image_latents_params` in the child test class. "
            "`image_latents_params` are tested for if passing latents directly are producing same results"
        )

507
508
509
    def get_dummy_inputs_by_type(self, device, seed=0, input_image_type="pt", output_type="np"):
        inputs = self.get_dummy_inputs(device, seed)

510
511
512
513
514
515
516
517
518
519
520
521
        def convert_to_pt(image):
            if isinstance(image, torch.Tensor):
                input_image = image
            elif isinstance(image, np.ndarray):
                input_image = VaeImageProcessor.numpy_to_pt(image)
            elif isinstance(image, PIL.Image.Image):
                input_image = VaeImageProcessor.pil_to_numpy(image)
                input_image = VaeImageProcessor.numpy_to_pt(input_image)
            else:
                raise ValueError(f"unsupported input_image_type {type(image)}")
            return input_image

522
523
524
525
526
527
528
529
530
531
532
533
534
535
        def convert_pt_to_type(image, input_image_type):
            if input_image_type == "pt":
                input_image = image
            elif input_image_type == "np":
                input_image = VaeImageProcessor.pt_to_numpy(image)
            elif input_image_type == "pil":
                input_image = VaeImageProcessor.pt_to_numpy(image)
                input_image = VaeImageProcessor.numpy_to_pil(input_image)
            else:
                raise ValueError(f"unsupported input_image_type {input_image_type}.")
            return input_image

        for image_param in self.image_params:
            if image_param in inputs.keys():
536
537
538
                inputs[image_param] = convert_pt_to_type(
                    convert_to_pt(inputs[image_param]).to(device), input_image_type
                )
539
540
541
542
543

        inputs["output_type"] = output_type

        return inputs

544
    def test_pt_np_pil_outputs_equivalent(self, expected_max_diff=1e-4):
545
546
547
        self._test_pt_np_pil_outputs_equivalent(expected_max_diff=expected_max_diff)

    def _test_pt_np_pil_outputs_equivalent(self, expected_max_diff=1e-4, input_image_type="pt"):
548
549
550
551
552
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe = pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

553
554
555
556
557
558
559
560
561
        output_pt = pipe(
            **self.get_dummy_inputs_by_type(torch_device, input_image_type=input_image_type, output_type="pt")
        )[0]
        output_np = pipe(
            **self.get_dummy_inputs_by_type(torch_device, input_image_type=input_image_type, output_type="np")
        )[0]
        output_pil = pipe(
            **self.get_dummy_inputs_by_type(torch_device, input_image_type=input_image_type, output_type="pil")
        )[0]
562
563

        max_diff = np.abs(output_pt.cpu().numpy().transpose(0, 2, 3, 1) - output_np).max()
564
565
566
        self.assertLess(
            max_diff, expected_max_diff, "`output_type=='pt'` generate different results from `output_type=='np'`"
        )
567
568

        max_diff = np.abs(np.array(output_pil[0]) - (output_np * 255).round()).max()
569
        self.assertLess(max_diff, 2.0, "`output_type=='pil'` generate different results from `output_type=='np'`")
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588

    def test_pt_np_pil_inputs_equivalent(self):
        if len(self.image_params) == 0:
            return

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe = pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        out_input_pt = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="pt"))[0]
        out_input_np = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="np"))[0]
        out_input_pil = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="pil"))[0]

        max_diff = np.abs(out_input_pt - out_input_np).max()
        self.assertLess(max_diff, 1e-4, "`input_type=='pt'` generate different result from `input_type=='np'`")
        max_diff = np.abs(out_input_pil - out_input_np).max()
        self.assertLess(max_diff, 1e-2, "`input_type=='pt'` generate different result from `input_type=='np'`")

589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
    def test_latents_input(self):
        if len(self.image_latents_params) == 0:
            return

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.image_processor = VaeImageProcessor(do_resize=False, do_normalize=False)
        pipe = pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        out = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="pt"))[0]

        vae = components["vae"]
        inputs = self.get_dummy_inputs_by_type(torch_device, input_image_type="pt")
        generator = inputs["generator"]
        for image_param in self.image_latents_params:
            if image_param in inputs.keys():
                inputs[image_param] = (
                    vae.encode(inputs[image_param]).latent_dist.sample(generator) * vae.config.scaling_factor
                )
        out_latents_inputs = pipe(**inputs)[0]

        max_diff = np.abs(out - out_latents_inputs).max()
        self.assertLess(max_diff, 1e-4, "passing latents as image input generate different result from passing image")

614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
    def test_multi_vae(self):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe = pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        block_out_channels = pipe.vae.config.block_out_channels
        norm_num_groups = pipe.vae.config.norm_num_groups

        vae_classes = [AutoencoderKL, AsymmetricAutoencoderKL, ConsistencyDecoderVAE, AutoencoderTiny]
        configs = [
            get_autoencoder_kl_config(block_out_channels, norm_num_groups),
            get_asym_autoencoder_kl_config(block_out_channels, norm_num_groups),
            get_consistency_vae_config(block_out_channels, norm_num_groups),
            get_autoencoder_tiny_config(block_out_channels),
        ]

        out_np = pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="np"))[0]

        for vae_cls, config in zip(vae_classes, configs):
            vae = vae_cls(**config)
            vae = vae.to(torch_device)
            components["vae"] = vae
            vae_pipe = self.pipeline_class(**components)
            out_vae_np = vae_pipe(**self.get_dummy_inputs_by_type(torch_device, input_image_type="np"))[0]

            assert out_vae_np.shape == out_np.shape

642

643
644
645
646
647
648
@require_torch
class PipelineFromPipeTesterMixin:
    @property
    def original_pipeline_class(self):
        if "xl" in self.pipeline_class.__name__.lower():
            original_pipeline_class = StableDiffusionXLPipeline
Álvaro Somoza's avatar
Álvaro Somoza committed
649
650
        elif "kolors" in self.pipeline_class.__name__.lower():
            original_pipeline_class = KolorsPipeline
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
        else:
            original_pipeline_class = StableDiffusionPipeline

        return original_pipeline_class

    def get_dummy_inputs_pipe(self, device, seed=0):
        inputs = self.get_dummy_inputs(device, seed=seed)
        inputs["output_type"] = "np"
        inputs["return_dict"] = False
        return inputs

    def get_dummy_inputs_for_pipe_original(self, device, seed=0):
        inputs = {}
        for k, v in self.get_dummy_inputs_pipe(device, seed=seed).items():
            if k in set(inspect.signature(self.original_pipeline_class.__call__).parameters.keys()):
                inputs[k] = v
        return inputs

    def test_from_pipe_consistent_config(self):
        if self.original_pipeline_class == StableDiffusionPipeline:
            original_repo = "hf-internal-testing/tiny-stable-diffusion-pipe"
            original_kwargs = {"requires_safety_checker": False}
        elif self.original_pipeline_class == StableDiffusionXLPipeline:
            original_repo = "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
            original_kwargs = {"requires_aesthetics_score": True, "force_zeros_for_empty_prompt": False}
Álvaro Somoza's avatar
Álvaro Somoza committed
676
677
678
        elif self.original_pipeline_class == KolorsPipeline:
            original_repo = "hf-internal-testing/tiny-kolors-pipe"
            original_kwargs = {"force_zeros_for_empty_prompt": False}
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
        else:
            raise ValueError(
                "original_pipeline_class must be either StableDiffusionPipeline or StableDiffusionXLPipeline"
            )

        # create original_pipeline_class(sd/sdxl)
        pipe_original = self.original_pipeline_class.from_pretrained(original_repo, **original_kwargs)

        # original_pipeline_class(sd/sdxl) -> pipeline_class
        pipe_components = self.get_dummy_components()
        pipe_additional_components = {}
        for name, component in pipe_components.items():
            if name not in pipe_original.components:
                pipe_additional_components[name] = component

        pipe = self.pipeline_class.from_pipe(pipe_original, **pipe_additional_components)

        # pipeline_class -> original_pipeline_class(sd/sdxl)
        original_pipe_additional_components = {}
        for name, component in pipe_original.components.items():
            if name not in pipe.components or not isinstance(component, pipe.components[name].__class__):
                original_pipe_additional_components[name] = component

        pipe_original_2 = self.original_pipeline_class.from_pipe(pipe, **original_pipe_additional_components)

        # compare the config
        original_config = {k: v for k, v in pipe_original.config.items() if not k.startswith("_")}
        original_config_2 = {k: v for k, v in pipe_original_2.config.items() if not k.startswith("_")}
        assert original_config_2 == original_config

    def test_from_pipe_consistent_forward_pass(self, expected_max_diff=1e-3):
        components = self.get_dummy_components()
        original_expected_modules, _ = self.original_pipeline_class._get_signature_keys(self.original_pipeline_class)

        # pipeline components that are also expected to be in the original pipeline
        original_pipe_components = {}
        # additional components that are not in the pipeline, but expected in the original pipeline
        original_pipe_additional_components = {}
        # additional components that are in the pipeline, but not expected in the original pipeline
        current_pipe_additional_components = {}

        for name, component in components.items():
            if name in original_expected_modules:
                original_pipe_components[name] = component
            else:
                current_pipe_additional_components[name] = component
        for name in original_expected_modules:
            if name not in original_pipe_components:
                if name in self.original_pipeline_class._optional_components:
                    original_pipe_additional_components[name] = None
                else:
                    raise ValueError(f"missing required module for {self.original_pipeline_class.__class__}: {name}")

        pipe_original = self.original_pipeline_class(**original_pipe_components, **original_pipe_additional_components)
        for component in pipe_original.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
        pipe_original.to(torch_device)
        pipe_original.set_progress_bar_config(disable=None)
        inputs = self.get_dummy_inputs_for_pipe_original(torch_device)
        output_original = pipe_original(**inputs)[0]

        pipe = self.pipeline_class(**components)
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)
        inputs = self.get_dummy_inputs_pipe(torch_device)
        output = pipe(**inputs)[0]

        pipe_from_original = self.pipeline_class.from_pipe(pipe_original, **current_pipe_additional_components)
        pipe_from_original.to(torch_device)
        pipe_from_original.set_progress_bar_config(disable=None)
        inputs = self.get_dummy_inputs_pipe(torch_device)
        output_from_original = pipe_from_original(**inputs)[0]

        max_diff = np.abs(to_np(output) - to_np(output_from_original)).max()
        self.assertLess(
            max_diff,
            expected_max_diff,
            "The outputs of the pipelines created with `from_pipe` and `__init__` are different.",
        )

        inputs = self.get_dummy_inputs_for_pipe_original(torch_device)
        output_original_2 = pipe_original(**inputs)[0]

        max_diff = np.abs(to_np(output_original) - to_np(output_original_2)).max()
        self.assertLess(max_diff, expected_max_diff, "`from_pipe` should not change the output of original pipeline.")

        for component in pipe_original.components.values():
            if hasattr(component, "attn_processors"):
                assert all(
                    type(proc) == AttnProcessor for proc in component.attn_processors.values()
                ), "`from_pipe` changed the attention processor in original pipeline."

775
776
    @require_accelerator
    @require_accelerate_version_greater("0.14.0")
777
778
779
780
781
782
    def test_from_pipe_consistent_forward_pass_cpu_offload(self, expected_max_diff=1e-3):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
783
        pipe.enable_model_cpu_offload(device=torch_device)
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
        pipe.set_progress_bar_config(disable=None)
        inputs = self.get_dummy_inputs_pipe(torch_device)
        output = pipe(**inputs)[0]

        original_expected_modules, _ = self.original_pipeline_class._get_signature_keys(self.original_pipeline_class)
        # pipeline components that are also expected to be in the original pipeline
        original_pipe_components = {}
        # additional components that are not in the pipeline, but expected in the original pipeline
        original_pipe_additional_components = {}
        # additional components that are in the pipeline, but not expected in the original pipeline
        current_pipe_additional_components = {}
        for name, component in components.items():
            if name in original_expected_modules:
                original_pipe_components[name] = component
            else:
                current_pipe_additional_components[name] = component
        for name in original_expected_modules:
            if name not in original_pipe_components:
                if name in self.original_pipeline_class._optional_components:
                    original_pipe_additional_components[name] = None
                else:
                    raise ValueError(f"missing required module for {self.original_pipeline_class.__class__}: {name}")

        pipe_original = self.original_pipeline_class(**original_pipe_components, **original_pipe_additional_components)
        for component in pipe_original.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
        pipe_original.set_progress_bar_config(disable=None)
812

813
        pipe_from_original = self.pipeline_class.from_pipe(pipe_original, **current_pipe_additional_components)
814
815
816
817
        for component in pipe_from_original.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()

818
        pipe_from_original.enable_model_cpu_offload(device=torch_device)
819
820
821
822
823
824
825
826
827
828
829
830
        pipe_from_original.set_progress_bar_config(disable=None)
        inputs = self.get_dummy_inputs_pipe(torch_device)
        output_from_original = pipe_from_original(**inputs)[0]

        max_diff = np.abs(to_np(output) - to_np(output_from_original)).max()
        self.assertLess(
            max_diff,
            expected_max_diff,
            "The outputs of the pipelines created with `from_pipe` and `__init__` are different.",
        )


831
832
833
834
835
836
837
838
@require_torch
class PipelineKarrasSchedulerTesterMixin:
    """
    This mixin is designed to be used with unittest.TestCase classes.
    It provides a set of common tests for each PyTorch pipeline that makes use of KarrasDiffusionSchedulers
    equivalence of dict and tuple outputs, etc.
    """

839
840
841
    def test_karras_schedulers_shape(
        self, num_inference_steps_for_strength=4, num_inference_steps_for_strength_for_iterations=5
    ):
842
843
844
845
846
847
848
849
850
851
852
853
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)

        # make sure that PNDM does not need warm-up
        pipe.scheduler.register_to_config(skip_prk_steps=True)

        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)
        inputs = self.get_dummy_inputs(torch_device)
        inputs["num_inference_steps"] = 2

        if "strength" in inputs:
854
            inputs["num_inference_steps"] = num_inference_steps_for_strength
855
856
857
858
859
            inputs["strength"] = 0.5

        outputs = []
        for scheduler_enum in KarrasDiffusionSchedulers:
            if "KDPM2" in scheduler_enum.name:
860
                inputs["num_inference_steps"] = num_inference_steps_for_strength_for_iterations
861
862
863
864
865
866
867
868
869
870
871
872

            scheduler_cls = getattr(diffusers, scheduler_enum.name)
            pipe.scheduler = scheduler_cls.from_config(pipe.scheduler.config)
            output = pipe(**inputs)[0]
            outputs.append(output)

            if "KDPM2" in scheduler_enum.name:
                inputs["num_inference_steps"] = 2

        assert check_same_shape(outputs)


873
874
875
876
877
878
879
880
@require_torch
class PipelineTesterMixin:
    """
    This mixin is designed to be used with unittest.TestCase classes.
    It provides a set of common tests for each PyTorch pipeline, e.g. saving and loading the pipeline,
    equivalence of dict and tuple outputs, etc.
    """

881
882
883
884
885
886
887
888
889
890
891
892
893
    # Canonical parameters that are passed to `__call__` regardless
    # of the type of pipeline. They are always optional and have common
    # sense default values.
    required_optional_params = frozenset(
        [
            "num_inference_steps",
            "num_images_per_prompt",
            "generator",
            "latents",
            "output_type",
            "return_dict",
        ]
    )
894

895
896
    # set these parameters to False in the child class if the pipeline does not support the corresponding functionality
    test_attention_slicing = True
897

898
899
    test_xformers_attention = True

900
901
902
903
904
    def get_generator(self, seed):
        device = torch_device if torch_device != "mps" else "cpu"
        generator = torch.Generator(device).manual_seed(seed)
        return generator

905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
    @property
    def pipeline_class(self) -> Union[Callable, DiffusionPipeline]:
        raise NotImplementedError(
            "You need to set the attribute `pipeline_class = ClassNameOfPipeline` in the child test class. "
            "See existing pipeline tests for reference."
        )

    def get_dummy_components(self):
        raise NotImplementedError(
            "You need to implement `get_dummy_components(self)` in the child test class. "
            "See existing pipeline tests for reference."
        )

    def get_dummy_inputs(self, device, seed=0):
        raise NotImplementedError(
            "You need to implement `get_dummy_inputs(self, device, seed)` in the child test class. "
            "See existing pipeline tests for reference."
        )

924
925
926
927
928
    @property
    def params(self) -> frozenset:
        raise NotImplementedError(
            "You need to set the attribute `params` in the child test class. "
            "`params` are checked for if all values are present in `__call__`'s signature."
929
            " You can set `params` using one of the common set of parameters defined in `pipeline_params.py`"
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
            " e.g., `TEXT_TO_IMAGE_PARAMS` defines the common parameters used in text to  "
            "image pipelines, including prompts and prompt embedding overrides."
            "If your pipeline's set of arguments has minor changes from one of the common sets of arguments, "
            "do not make modifications to the existing common sets of arguments. I.e. a text to image pipeline "
            "with non-configurable height and width arguments should set the attribute as "
            "`params = TEXT_TO_IMAGE_PARAMS - {'height', 'width'}`. "
            "See existing pipeline tests for reference."
        )

    @property
    def batch_params(self) -> frozenset:
        raise NotImplementedError(
            "You need to set the attribute `batch_params` in the child test class. "
            "`batch_params` are the parameters required to be batched when passed to the pipeline's "
            "`__call__` method. `pipeline_params.py` provides some common sets of parameters such as "
            "`TEXT_TO_IMAGE_BATCH_PARAMS`, `IMAGE_VARIATION_BATCH_PARAMS`, etc... If your pipeline's "
            "set of batch arguments has minor changes from one of the common sets of batch arguments, "
            "do not make modifications to the existing common sets of batch arguments. I.e. a text to "
            "image pipeline `negative_prompt` is not batched should set the attribute as "
            "`batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - {'negative_prompt'}`. "
            "See existing pipeline tests for reference."
        )

953
954
955
956
957
958
959
960
961
    @property
    def callback_cfg_params(self) -> frozenset:
        raise NotImplementedError(
            "You need to set the attribute `callback_cfg_params` in the child test class that requires to run test_callback_cfg. "
            "`callback_cfg_params` are the parameters that needs to be passed to the pipeline's callback "
            "function when dynamically adjusting `guidance_scale`. They are variables that require special"
            "treatment when `do_classifier_free_guidance` is `True`. `pipeline_params.py` provides some common"
            " sets of parameters such as `TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS`. If your pipeline's "
            "set of cfg arguments has minor changes from one of the common sets of cfg arguments, "
M. Tolga Cangöz's avatar
M. Tolga Cangöz committed
962
            "do not make modifications to the existing common sets of cfg arguments. I.e. for inpaint pipeline, you "
963
964
965
966
            " need to adjust batch size of `mask` and `masked_image_latents` so should set the attribute as"
            "`callback_cfg_params = TEXT_TO_IMAGE_CFG_PARAMS.union({'mask', 'masked_image_latents'})`"
        )

967
968
969
970
971
972
    def setUp(self):
        # clean up the VRAM before each test
        super().setUp()
        gc.collect()
        torch.cuda.empty_cache()

973
974
975
976
977
978
    def tearDown(self):
        # clean up the VRAM after each test in case of CUDA runtime errors
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

979
    def test_save_load_local(self, expected_max_difference=5e-4):
980
981
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
982
983
984
985
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()

986
987
988
989
990
991
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(torch_device)
        output = pipe(**inputs)[0]

992
993
994
        logger = logging.get_logger("diffusers.pipelines.pipeline_utils")
        logger.setLevel(diffusers.logging.INFO)

995
        with tempfile.TemporaryDirectory() as tmpdir:
996
            pipe.save_pretrained(tmpdir, safe_serialization=False)
997
998
999
1000

            with CaptureLogger(logger) as cap_logger:
                pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)

1001
1002
1003
1004
            for component in pipe_loaded.components.values():
                if hasattr(component, "set_default_attn_processor"):
                    component.set_default_attn_processor()

1005
1006
1007
1008
            for name in pipe_loaded.components.keys():
                if name not in pipe_loaded._optional_components:
                    assert name in str(cap_logger)

1009
1010
1011
1012
1013
1014
            pipe_loaded.to(torch_device)
            pipe_loaded.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(torch_device)
        output_loaded = pipe_loaded(**inputs)[0]

1015
        max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
1016
        self.assertLess(max_diff, expected_max_difference)
1017

1018
1019
1020
1021
1022
    def test_pipeline_call_signature(self):
        self.assertTrue(
            hasattr(self.pipeline_class, "__call__"), f"{self.pipeline_class} should have a `__call__` method"
        )

1023
1024
        parameters = inspect.signature(self.pipeline_class.__call__).parameters

1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
        optional_parameters = set()

        for k, v in parameters.items():
            if v.default != inspect._empty:
                optional_parameters.add(k)

        parameters = set(parameters.keys())
        parameters.remove("self")
        parameters.discard("kwargs")  # kwargs can be added if arguments of pipeline call function are deprecated

        remaining_required_parameters = set()

        for param in self.params:
            if param not in parameters:
                remaining_required_parameters.add(param)
1040

1041
1042
1043
1044
1045
1046
        self.assertTrue(
            len(remaining_required_parameters) == 0,
            f"Required parameters not present: {remaining_required_parameters}",
        )

        remaining_required_optional_parameters = set()
1047

1048
        for param in self.required_optional_params:
1049
1050
1051
1052
1053
1054
1055
            if param not in optional_parameters:
                remaining_required_optional_parameters.add(param)

        self.assertTrue(
            len(remaining_required_optional_parameters) == 0,
            f"Required optional parameters not present: {remaining_required_optional_parameters}",
        )
1056

1057
    def test_inference_batch_consistent(self, batch_sizes=[2]):
1058
        self._test_inference_batch_consistent(batch_sizes=batch_sizes)
1059

1060
    def _test_inference_batch_consistent(
Will Berman's avatar
Will Berman committed
1061
        self, batch_sizes=[2], additional_params_copy_to_batched_inputs=["num_inference_steps"], batch_generator=True
1062
    ):
1063
1064
1065
1066
1067
1068
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(torch_device)
1069
        inputs["generator"] = self.get_generator(0)
1070
1071
1072
1073

        logger = logging.get_logger(pipe.__module__)
        logger.setLevel(level=diffusers.logging.FATAL)

1074
1075
        # prepare batched inputs
        batched_inputs = []
1076
        for batch_size in batch_sizes:
1077
1078
            batched_input = {}
            batched_input.update(inputs)
1079

1080
1081
1082
            for name in self.batch_params:
                if name not in inputs:
                    continue
1083

1084
1085
1086
1087
1088
                value = inputs[name]
                if name == "prompt":
                    len_prompt = len(value)
                    # make unequal batch sizes
                    batched_input[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)]
1089

1090
1091
                    # make last batch super long
                    batched_input[name][-1] = 100 * "very long"
1092

1093
1094
                else:
                    batched_input[name] = batch_size * [value]
1095

Will Berman's avatar
Will Berman committed
1096
            if batch_generator and "generator" in inputs:
1097
                batched_input["generator"] = [self.get_generator(i) for i in range(batch_size)]
1098

1099
1100
            if "batch_size" in inputs:
                batched_input["batch_size"] = batch_size
1101

1102
            batched_inputs.append(batched_input)
1103
1104

        logger.setLevel(level=diffusers.logging.WARNING)
1105
        for batch_size, batched_input in zip(batch_sizes, batched_inputs):
Aryan's avatar
Aryan committed
1106
            print(batch_size, batched_input)
1107
1108
            output = pipe(**batched_input)
            assert len(output[0]) == batch_size
1109

1110
1111
    def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-4):
        self._test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
1112
1113

    def _test_inference_batch_single_identical(
1114
        self,
1115
        batch_size=2,
1116
        expected_max_diff=1e-4,
1117
        additional_params_copy_to_batched_inputs=["num_inference_steps"],
1118
    ):
1119
1120
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
1121
1122
1123
1124
        for components in pipe.components.values():
            if hasattr(components, "set_default_attn_processor"):
                components.set_default_attn_processor()

1125
1126
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)
1127
1128
1129
        inputs = self.get_dummy_inputs(torch_device)
        # Reset generator in case it is has been used in self.get_dummy_inputs
        inputs["generator"] = self.get_generator(0)
1130
1131
1132
1133
1134
1135

        logger = logging.get_logger(pipe.__module__)
        logger.setLevel(level=diffusers.logging.FATAL)

        # batchify inputs
        batched_inputs = {}
1136
        batched_inputs.update(inputs)
1137

1138
1139
1140
        for name in self.batch_params:
            if name not in inputs:
                continue
1141

1142
1143
1144
1145
1146
            value = inputs[name]
            if name == "prompt":
                len_prompt = len(value)
                batched_inputs[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)]
                batched_inputs[name][-1] = 100 * "very long"
1147

1148
1149
            else:
                batched_inputs[name] = batch_size * [value]
1150

1151
1152
        if "generator" in inputs:
            batched_inputs["generator"] = [self.get_generator(i) for i in range(batch_size)]
1153

1154
1155
1156
1157
1158
        if "batch_size" in inputs:
            batched_inputs["batch_size"] = batch_size

        for arg in additional_params_copy_to_batched_inputs:
            batched_inputs[arg] = inputs[arg]
1159
1160

        output = pipe(**inputs)
1161
        output_batch = pipe(**batched_inputs)
1162

1163
        assert output_batch[0].shape[0] == batch_size
1164

YiYi Xu's avatar
YiYi Xu committed
1165
        max_diff = np.abs(to_np(output_batch[0][0]) - to_np(output[0][0])).max()
1166
        assert max_diff < expected_max_diff
1167

1168
    def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4):
1169
1170
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
1171
1172
1173
1174
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()

1175
1176
1177
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

1178
        generator_device = "cpu"
1179
1180
1181
1182
1183
        if expected_slice is None:
            output = pipe(**self.get_dummy_inputs(generator_device))[0]
        else:
            output = expected_slice

1184
        output_tuple = pipe(**self.get_dummy_inputs(generator_device), return_dict=False)[0]
1185

1186
1187
1188
1189
1190
1191
1192
1193
        if expected_slice is None:
            max_diff = np.abs(to_np(output) - to_np(output_tuple)).max()
        else:
            if output_tuple.ndim != 5:
                max_diff = np.abs(to_np(output) - to_np(output_tuple)[0, -3:, -3:, -1].flatten()).max()
            else:
                max_diff = np.abs(to_np(output) - to_np(output_tuple)[0, -3:, -3:, -1, -1].flatten()).max()

1194
        self.assertLess(max_diff, expected_max_difference)
1195
1196
1197

    def test_components_function(self):
        init_components = self.get_dummy_components()
Kashif Rasul's avatar
Kashif Rasul committed
1198
1199
        init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float))}

1200
1201
1202
1203
1204
        pipe = self.pipeline_class(**init_components)

        self.assertTrue(hasattr(pipe, "components"))
        self.assertTrue(set(pipe.components.keys()) == set(init_components.keys()))

1205
1206
    @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU")
    @require_accelerator
1207
    def test_float16_inference(self, expected_max_diff=5e-2):
1208
1209
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
1210
1211
1212
1213
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()

1214
1215
1216
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

YiYi Xu's avatar
YiYi Xu committed
1217
        components = self.get_dummy_components()
1218
        pipe_fp16 = self.pipeline_class(**components)
1219
1220
1221
1222
        for component in pipe_fp16.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()

1223
        pipe_fp16.to(torch_device, torch.float16)
1224
1225
        pipe_fp16.set_progress_bar_config(disable=None)

1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
        inputs = self.get_dummy_inputs(torch_device)
        # Reset generator in case it is used inside dummy inputs
        if "generator" in inputs:
            inputs["generator"] = self.get_generator(0)

        output = pipe(**inputs)[0]

        fp16_inputs = self.get_dummy_inputs(torch_device)
        # Reset generator in case it is used inside dummy inputs
        if "generator" in fp16_inputs:
            fp16_inputs["generator"] = self.get_generator(0)

        output_fp16 = pipe_fp16(**fp16_inputs)[0]
1239

1240
        max_diff = np.abs(to_np(output) - to_np(output_fp16)).max()
1241
        self.assertLess(max_diff, expected_max_diff, "The outputs of the fp16 and fp32 pipelines are too different.")
1242

1243
1244
    @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU")
    @require_accelerator
1245
    def test_save_load_float16(self, expected_max_diff=1e-2):
1246
1247
1248
1249
        components = self.get_dummy_components()
        for name, module in components.items():
            if hasattr(module, "half"):
                components[name] = module.to(torch_device).half()
1250

1251
        pipe = self.pipeline_class(**components)
1252
1253
1254
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
1255
1256
1257
1258
1259
1260
1261
1262
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(torch_device)
        output = pipe(**inputs)[0]

        with tempfile.TemporaryDirectory() as tmpdir:
            pipe.save_pretrained(tmpdir)
1263
            pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, torch_dtype=torch.float16)
1264
1265
1266
            for component in pipe_loaded.components.values():
                if hasattr(component, "set_default_attn_processor"):
                    component.set_default_attn_processor()
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
            pipe_loaded.to(torch_device)
            pipe_loaded.set_progress_bar_config(disable=None)

        for name, component in pipe_loaded.components.items():
            if hasattr(component, "dtype"):
                self.assertTrue(
                    component.dtype == torch.float16,
                    f"`{name}.dtype` switched from `float16` to {component.dtype} after loading.",
                )

        inputs = self.get_dummy_inputs(torch_device)
        output_loaded = pipe_loaded(**inputs)[0]
1279
        max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
1280
1281
1282
        self.assertLess(
            max_diff, expected_max_diff, "The output of the fp16 pipeline changed after saving and loading."
        )
1283

1284
    def test_save_load_optional_components(self, expected_max_difference=1e-4):
1285
1286
1287
1288
1289
        if not hasattr(self.pipeline_class, "_optional_components"):
            return

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
1290
1291
1292
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
1293
1294
1295
1296
1297
1298
1299
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        # set all optional components to None
        for optional_component in pipe._optional_components:
            setattr(pipe, optional_component, None)

Dhruv Nair's avatar
Dhruv Nair committed
1300
1301
        generator_device = "cpu"
        inputs = self.get_dummy_inputs(generator_device)
1302
1303
1304
        output = pipe(**inputs)[0]

        with tempfile.TemporaryDirectory() as tmpdir:
1305
            pipe.save_pretrained(tmpdir, safe_serialization=False)
1306
            pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
1307
1308
1309
            for component in pipe_loaded.components.values():
                if hasattr(component, "set_default_attn_processor"):
                    component.set_default_attn_processor()
1310
1311
1312
1313
1314
1315
1316
1317
1318
            pipe_loaded.to(torch_device)
            pipe_loaded.set_progress_bar_config(disable=None)

        for optional_component in pipe._optional_components:
            self.assertTrue(
                getattr(pipe_loaded, optional_component) is None,
                f"`{optional_component}` did not stay set to None after loading.",
            )

Dhruv Nair's avatar
Dhruv Nair committed
1319
        inputs = self.get_dummy_inputs(generator_device)
1320
1321
        output_loaded = pipe_loaded(**inputs)[0]

1322
        max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
1323
        self.assertLess(max_diff, expected_max_difference)
1324

1325
    @require_accelerator
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
    def test_to_device(self):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.set_progress_bar_config(disable=None)

        pipe.to("cpu")
        model_devices = [component.device.type for component in components.values() if hasattr(component, "device")]
        self.assertTrue(all(device == "cpu" for device in model_devices))

        output_cpu = pipe(**self.get_dummy_inputs("cpu"))[0]
        self.assertTrue(np.isnan(output_cpu).sum() == 0)

1338
        pipe.to(torch_device)
1339
        model_devices = [component.device.type for component in components.values() if hasattr(component, "device")]
1340
        self.assertTrue(all(device == torch_device for device in model_devices))
1341

1342
        output_cuda = pipe(**self.get_dummy_inputs(torch_device))[0]
1343
        self.assertTrue(np.isnan(to_np(output_cuda)).sum() == 0)
1344

1345
1346
1347
1348
1349
1350
1351
1352
    def test_to_dtype(self):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.set_progress_bar_config(disable=None)

        model_dtypes = [component.dtype for component in components.values() if hasattr(component, "dtype")]
        self.assertTrue(all(dtype == torch.float32 for dtype in model_dtypes))

1353
        pipe.to(dtype=torch.float16)
1354
1355
1356
        model_dtypes = [component.dtype for component in components.values() if hasattr(component, "dtype")]
        self.assertTrue(all(dtype == torch.float16 for dtype in model_dtypes))

1357
1358
    def test_attention_slicing_forward_pass(self, expected_max_diff=1e-3):
        self._test_attention_slicing_forward_pass(expected_max_diff=expected_max_diff)
1359

1360
1361
1362
    def _test_attention_slicing_forward_pass(
        self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
    ):
1363
1364
1365
1366
1367
        if not self.test_attention_slicing:
            return

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
1368
1369
1370
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
1371
1372
1373
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

1374
1375
        generator_device = "cpu"
        inputs = self.get_dummy_inputs(generator_device)
1376
1377
1378
        output_without_slicing = pipe(**inputs)[0]

        pipe.enable_attention_slicing(slice_size=1)
1379
        inputs = self.get_dummy_inputs(generator_device)
1380
1381
1382
1383
1384
        output_with_slicing1 = pipe(**inputs)[0]

        pipe.enable_attention_slicing(slice_size=2)
        inputs = self.get_dummy_inputs(generator_device)
        output_with_slicing2 = pipe(**inputs)[0]
1385

1386
        if test_max_difference:
1387
1388
1389
1390
1391
1392
1393
            max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
            max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
            self.assertLess(
                max(max_diff1, max_diff2),
                expected_max_diff,
                "Attention slicing should not affect the inference results",
            )
1394

1395
        if test_mean_pixel_difference:
1396
1397
            assert_mean_pixel_difference(to_np(output_with_slicing1[0]), to_np(output_without_slicing[0]))
            assert_mean_pixel_difference(to_np(output_with_slicing2[0]), to_np(output_without_slicing[0]))
1398

1399
1400
    @require_accelerator
    @require_accelerate_version_greater("0.14.0")
1401
    def test_sequential_cpu_offload_forward_pass(self, expected_max_diff=1e-4):
1402
1403
        import accelerate

1404
1405
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
1406
1407
1408
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
1409
1410
1411
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

1412
1413
        generator_device = "cpu"
        inputs = self.get_dummy_inputs(generator_device)
1414
1415
        output_without_offload = pipe(**inputs)[0]

1416
1417
        pipe.enable_sequential_cpu_offload(device=torch_device)
        assert pipe._execution_device.type == torch_device
1418
1419
1420
1421
1422
1423
1424

        inputs = self.get_dummy_inputs(generator_device)
        output_with_offload = pipe(**inputs)[0]

        max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max()
        self.assertLess(max_diff, expected_max_diff, "CPU offloading should not affect the inference results")

1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
        # make sure all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are offloaded correctly
        offloaded_modules = {
            k: v
            for k, v in pipe.components.items()
            if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload
        }
        # 1. all offloaded modules should be saved to cpu and moved to meta device
        self.assertTrue(
            all(v.device.type == "meta" for v in offloaded_modules.values()),
            f"Not offloaded: {[k for k, v in offloaded_modules.items() if v.device.type != 'meta']}",
        )
        # 2. all offloaded modules should have hook installed
        self.assertTrue(
            all(hasattr(v, "_hf_hook") for k, v in offloaded_modules.items()),
            f"No hook attached: {[k for k, v in offloaded_modules.items() if not hasattr(v, '_hf_hook')]}",
        )
        # 3. all offloaded modules should have correct hooks installed, should be either one of these two
        #    - `AlignDevicesHook`
        #    - a SequentialHook` that contains `AlignDevicesHook`
        offloaded_modules_with_incorrect_hooks = {}
        for k, v in offloaded_modules.items():
            if hasattr(v, "_hf_hook"):
                if isinstance(v._hf_hook, accelerate.hooks.SequentialHook):
                    # if it is a `SequentialHook`, we loop through its `hooks` attribute to check if it only contains `AlignDevicesHook`
                    for hook in v._hf_hook.hooks:
                        if not isinstance(hook, accelerate.hooks.AlignDevicesHook):
                            offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook.hooks[0])
                elif not isinstance(v._hf_hook, accelerate.hooks.AlignDevicesHook):
                    offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook)

        self.assertTrue(
            len(offloaded_modules_with_incorrect_hooks) == 0,
            f"Not installed correct hook: {offloaded_modules_with_incorrect_hooks}",
        )

1460
1461
    @require_accelerator
    @require_accelerate_version_greater("0.17.0")
1462
    def test_model_cpu_offload_forward_pass(self, expected_max_diff=2e-4):
1463
1464
        import accelerate

1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
        generator_device = "cpu"
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)

        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()

        pipe = pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(generator_device)
        output_without_offload = pipe(**inputs)[0]

1479
1480
        pipe.enable_model_cpu_offload(device=torch_device)
        assert pipe._execution_device.type == torch_device
1481

1482
        inputs = self.get_dummy_inputs(generator_device)
1483
1484
        output_with_offload = pipe(**inputs)[0]

1485
        max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max()
1486
        self.assertLess(max_diff, expected_max_diff, "CPU offloading should not affect the inference results")
1487
1488
1489
1490

        # make sure all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are offloaded correctly
        offloaded_modules = {
            k: v
1491
1492
            for k, v in pipe.components.items()
            if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
        }
        # 1. check if all offloaded modules are saved to cpu
        self.assertTrue(
            all(v.device.type == "cpu" for v in offloaded_modules.values()),
            f"Not offloaded: {[k for k, v in offloaded_modules.items() if v.device.type != 'cpu']}",
        )
        # 2. check if all offloaded modules have hooks installed
        self.assertTrue(
            all(hasattr(v, "_hf_hook") for k, v in offloaded_modules.items()),
            f"No hook attached: {[k for k, v in offloaded_modules.items() if not hasattr(v, '_hf_hook')]}",
        )
        # 3. check if all offloaded modules have correct type of hooks installed, should be `CpuOffload`
        offloaded_modules_with_incorrect_hooks = {}
        for k, v in offloaded_modules.items():
            if hasattr(v, "_hf_hook") and not isinstance(v._hf_hook, accelerate.hooks.CpuOffload):
                offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook)

        self.assertTrue(
            len(offloaded_modules_with_incorrect_hooks) == 0,
            f"Not installed correct hook: {offloaded_modules_with_incorrect_hooks}",
1513
        )
1514

1515
1516
    @require_accelerator
    @require_accelerate_version_greater("0.17.0")
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
    def test_cpu_offload_forward_pass_twice(self, expected_max_diff=2e-4):
        import accelerate

        generator_device = "cpu"
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)

        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()

        pipe.set_progress_bar_config(disable=None)

1530
        pipe.enable_model_cpu_offload(device=torch_device)
1531
1532
1533
        inputs = self.get_dummy_inputs(generator_device)
        output_with_offload = pipe(**inputs)[0]

1534
        pipe.enable_model_cpu_offload(device=torch_device)
1535
1536
1537
1538
1539
1540
1541
        inputs = self.get_dummy_inputs(generator_device)
        output_with_offload_twice = pipe(**inputs)[0]

        max_diff = np.abs(to_np(output_with_offload) - to_np(output_with_offload_twice)).max()
        self.assertLess(
            max_diff, expected_max_diff, "running CPU offloading 2nd time should not affect the inference results"
        )
1542
1543

        # make sure all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are offloaded correctly
YiYi Xu's avatar
YiYi Xu committed
1544
1545
        offloaded_modules = {
            k: v
1546
1547
            for k, v in pipe.components.items()
            if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload
YiYi Xu's avatar
YiYi Xu committed
1548
        }
1549
        # 1. check if all offloaded modules are saved to cpu
YiYi Xu's avatar
YiYi Xu committed
1550
1551
1552
        self.assertTrue(
            all(v.device.type == "cpu" for v in offloaded_modules.values()),
            f"Not offloaded: {[k for k, v in offloaded_modules.items() if v.device.type != 'cpu']}",
1553
        )
1554
1555
1556
1557
1558
1559
        # 2. check if all offloaded modules have hooks installed
        self.assertTrue(
            all(hasattr(v, "_hf_hook") for k, v in offloaded_modules.items()),
            f"No hook attached: {[k for k, v in offloaded_modules.items() if not hasattr(v, '_hf_hook')]}",
        )
        # 3. check if all offloaded modules have correct type of hooks installed, should be `CpuOffload`
YiYi Xu's avatar
YiYi Xu committed
1560
1561
1562
1563
1564
1565
1566
1567
        offloaded_modules_with_incorrect_hooks = {}
        for k, v in offloaded_modules.items():
            if hasattr(v, "_hf_hook") and not isinstance(v._hf_hook, accelerate.hooks.CpuOffload):
                offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook)

        self.assertTrue(
            len(offloaded_modules_with_incorrect_hooks) == 0,
            f"Not installed correct hook: {offloaded_modules_with_incorrect_hooks}",
1568
1569
        )

1570
1571
    @require_accelerator
    @require_accelerate_version_greater("0.14.0")
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
    def test_sequential_offload_forward_pass_twice(self, expected_max_diff=2e-4):
        import accelerate

        generator_device = "cpu"
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)

        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()

        pipe.set_progress_bar_config(disable=None)

1585
        pipe.enable_sequential_cpu_offload(device=torch_device)
1586
1587
1588
        inputs = self.get_dummy_inputs(generator_device)
        output_with_offload = pipe(**inputs)[0]

1589
        pipe.enable_sequential_cpu_offload(device=torch_device)
1590
1591
1592
1593
1594
1595
1596
        inputs = self.get_dummy_inputs(generator_device)
        output_with_offload_twice = pipe(**inputs)[0]

        max_diff = np.abs(to_np(output_with_offload) - to_np(output_with_offload_twice)).max()
        self.assertLess(
            max_diff, expected_max_diff, "running sequential offloading second time should have the inference results"
        )
1597
1598

        # make sure all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are offloaded correctly
YiYi Xu's avatar
YiYi Xu committed
1599
1600
        offloaded_modules = {
            k: v
1601
1602
            for k, v in pipe.components.items()
            if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload
YiYi Xu's avatar
YiYi Xu committed
1603
        }
1604
        # 1. check if all offloaded modules are moved to meta device
YiYi Xu's avatar
YiYi Xu committed
1605
1606
1607
        self.assertTrue(
            all(v.device.type == "meta" for v in offloaded_modules.values()),
            f"Not offloaded: {[k for k, v in offloaded_modules.items() if v.device.type != 'meta']}",
1608
        )
1609
1610
1611
1612
1613
1614
1615
1616
        # 2. check if all offloaded modules have hook installed
        self.assertTrue(
            all(hasattr(v, "_hf_hook") for k, v in offloaded_modules.items()),
            f"No hook attached: {[k for k, v in offloaded_modules.items() if not hasattr(v, '_hf_hook')]}",
        )
        # 3. check if all offloaded modules have correct hooks installed, should be either one of these two
        #    - `AlignDevicesHook`
        #    - a SequentialHook` that contains `AlignDevicesHook`
YiYi Xu's avatar
YiYi Xu committed
1617
1618
        offloaded_modules_with_incorrect_hooks = {}
        for k, v in offloaded_modules.items():
1619
1620
1621
1622
1623
1624
1625
1626
            if hasattr(v, "_hf_hook"):
                if isinstance(v._hf_hook, accelerate.hooks.SequentialHook):
                    # if it is a `SequentialHook`, we loop through its `hooks` attribute to check if it only contains `AlignDevicesHook`
                    for hook in v._hf_hook.hooks:
                        if not isinstance(hook, accelerate.hooks.AlignDevicesHook):
                            offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook.hooks[0])
                elif not isinstance(v._hf_hook, accelerate.hooks.AlignDevicesHook):
                    offloaded_modules_with_incorrect_hooks[k] = type(v._hf_hook)
1627

YiYi Xu's avatar
YiYi Xu committed
1628
1629
1630
        self.assertTrue(
            len(offloaded_modules_with_incorrect_hooks) == 0,
            f"Not installed correct hook: {offloaded_modules_with_incorrect_hooks}",
1631
1632
        )

1633
1634
1635
1636
    @unittest.skipIf(
        torch_device != "cuda" or not is_xformers_available(),
        reason="XFormers attention is only available with CUDA and `xformers` installed",
    )
Kashif Rasul's avatar
Kashif Rasul committed
1637
    def test_xformers_attention_forwardGenerator_pass(self):
Will Berman's avatar
Will Berman committed
1638
1639
        self._test_xformers_attention_forwardGenerator_pass()

1640
1641
1642
    def _test_xformers_attention_forwardGenerator_pass(
        self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-4
    ):
1643
1644
1645
1646
1647
        if not self.test_xformers_attention:
            return

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
1648
1649
1650
        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
1651
1652
1653
1654
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(torch_device)
1655
        output_without_offload = pipe(**inputs)[0]
1656
1657
1658
        output_without_offload = (
            output_without_offload.cpu() if torch.is_tensor(output_without_offload) else output_without_offload
        )
1659
1660
1661

        pipe.enable_xformers_memory_efficient_attention()
        inputs = self.get_dummy_inputs(torch_device)
1662
        output_with_offload = pipe(**inputs)[0]
1663
1664
1665
        output_with_offload = (
            output_with_offload.cpu() if torch.is_tensor(output_with_offload) else output_without_offload
        )
1666

Will Berman's avatar
Will Berman committed
1667
        if test_max_difference:
1668
            max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max()
1669
            self.assertLess(max_diff, expected_max_diff, "XFormers attention should not affect the inference results")
Will Berman's avatar
Will Berman committed
1670

1671
1672
        if test_mean_pixel_difference:
            assert_mean_pixel_difference(output_with_offload[0], output_without_offload[0])
1673

1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
    def test_num_images_per_prompt(self):
        sig = inspect.signature(self.pipeline_class.__call__)

        if "num_images_per_prompt" not in sig.parameters:
            return

        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 inputs.keys():
                    if key in self.batch_params:
                        inputs[key] = batch_size * [inputs[key]]

1696
                images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt)[0]
1697
1698
1699

                assert images.shape[0] == batch_size * num_images_per_prompt

1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
    def test_cfg(self):
        sig = inspect.signature(self.pipeline_class.__call__)

        if "guidance_scale" not in sig.parameters:
            return

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe = pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        inputs = self.get_dummy_inputs(torch_device)

        inputs["guidance_scale"] = 1.0
        out_no_cfg = pipe(**inputs)[0]

        inputs["guidance_scale"] = 7.5
        out_cfg = pipe(**inputs)[0]

        assert out_cfg.shape == out_no_cfg.shape

1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
    def test_callback_inputs(self):
        sig = inspect.signature(self.pipeline_class.__call__)
        has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
        has_callback_step_end = "callback_on_step_end" in sig.parameters

        if not (has_callback_tensor_inputs and has_callback_step_end):
            return

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe = pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)
        self.assertTrue(
            hasattr(pipe, "_callback_tensor_inputs"),
            f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
        )

        def callback_inputs_subset(pipe, i, t, callback_kwargs):
M. Tolga Cangöz's avatar
M. Tolga Cangöz committed
1739
            # iterate over callback args
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
            for tensor_name, tensor_value in callback_kwargs.items():
                # check that we're only passing in allowed tensor inputs
                assert tensor_name in pipe._callback_tensor_inputs

            return callback_kwargs

        def callback_inputs_all(pipe, i, t, callback_kwargs):
            for tensor_name in pipe._callback_tensor_inputs:
                assert tensor_name in callback_kwargs

M. Tolga Cangöz's avatar
M. Tolga Cangöz committed
1750
            # iterate over callback args
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
            for tensor_name, tensor_value in callback_kwargs.items():
                # check that we're only passing in allowed tensor inputs
                assert tensor_name in pipe._callback_tensor_inputs

            return callback_kwargs

        inputs = self.get_dummy_inputs(torch_device)

        # Test passing in a subset
        inputs["callback_on_step_end"] = callback_inputs_subset
        inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
        inputs["output_type"] = "latent"
        output = pipe(**inputs)[0]

        # Test passing in a everything
        inputs["callback_on_step_end"] = callback_inputs_all
        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
        inputs["output_type"] = "latent"
        output = pipe(**inputs)[0]

        def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
            is_last = i == (pipe.num_timesteps - 1)
            if is_last:
                callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
            return callback_kwargs

        inputs["callback_on_step_end"] = callback_inputs_change_tensor
        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
        inputs["output_type"] = "latent"
        output = pipe(**inputs)[0]
        assert output.abs().sum() == 0

    def test_callback_cfg(self):
        sig = inspect.signature(self.pipeline_class.__call__)
        has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
        has_callback_step_end = "callback_on_step_end" in sig.parameters

        if not (has_callback_tensor_inputs and has_callback_step_end):
            return

        if "guidance_scale" not in sig.parameters:
            return

        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)
        self.assertTrue(
            hasattr(pipe, "_callback_tensor_inputs"),
            f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
        )

        def callback_increase_guidance(pipe, i, t, callback_kwargs):
            pipe._guidance_scale += 1.0

            return callback_kwargs

        inputs = self.get_dummy_inputs(torch_device)

        # use cfg guidance because some pipelines modify the shape of the latents
        # outside of the denoising loop
        inputs["guidance_scale"] = 2.0
        inputs["callback_on_step_end"] = callback_increase_guidance
        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
        _ = pipe(**inputs)[0]

        # we increase the guidance scale by 1.0 at every step
        # check that the guidance scale is increased by the number of scheduler timesteps
        # accounts for models that modify the number of inference steps based on strength
        assert pipe.guidance_scale == (inputs["guidance_scale"] + pipe.num_timesteps)

1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
    def test_serialization_with_variants(self):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        model_components = [
            component_name for component_name, component in pipe.components.items() if isinstance(component, nn.Module)
        ]
        variant = "fp16"

        with tempfile.TemporaryDirectory() as tmpdir:
            pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False)

            with open(f"{tmpdir}/model_index.json", "r") as f:
                config = json.load(f)

            for subfolder in os.listdir(tmpdir):
                if not os.path.isfile(subfolder) and subfolder in model_components:
                    folder_path = os.path.join(tmpdir, subfolder)
                    is_folder = os.path.isdir(folder_path) and subfolder in config
                    assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path))

    def test_loading_with_variants(self):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        variant = "fp16"

        def is_nan(tensor):
            if tensor.ndimension() == 0:
                has_nan = torch.isnan(tensor).item()
            else:
                has_nan = torch.isnan(tensor).any()
            return has_nan

        with tempfile.TemporaryDirectory() as tmpdir:
            pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False)
            pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, variant=variant)

            model_components_pipe = {
                component_name: component
                for component_name, component in pipe.components.items()
                if isinstance(component, nn.Module)
            }
            model_components_pipe_loaded = {
                component_name: component
                for component_name, component in pipe_loaded.components.items()
                if isinstance(component, nn.Module)
            }
            for component_name in model_components_pipe:
                pipe_component = model_components_pipe[component_name]
                pipe_loaded_component = model_components_pipe_loaded[component_name]
                for p1, p2 in zip(pipe_component.parameters(), pipe_loaded_component.parameters()):
                    # nan check for luminanext (mps).
                    if not (is_nan(p1) and is_nan(p2)):
                        self.assertTrue(torch.equal(p1, p2))

    def test_loading_with_incorrect_variants_raises_error(self):
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        variant = "fp16"

        with tempfile.TemporaryDirectory() as tmpdir:
            # Don't save with variants.
            pipe.save_pretrained(tmpdir, safe_serialization=False)

            with self.assertRaises(ValueError) as error:
                _ = self.pipeline_class.from_pretrained(tmpdir, variant=variant)

            assert f"You are trying to load the model files of the `variant={variant}`" in str(error.exception)

1890
1891
1892
1893
1894
1895
1896
1897
1898
    def test_StableDiffusionMixin_component(self):
        """Any pipeline that have LDMFuncMixin should have vae and unet components."""
        if not issubclass(self.pipeline_class, StableDiffusionMixin):
            return
        components = self.get_dummy_components()
        pipe = self.pipeline_class(**components)
        self.assertTrue(hasattr(pipe, "vae") and isinstance(pipe.vae, (AutoencoderKL, AutoencoderTiny)))
        self.assertTrue(
            hasattr(pipe, "unet")
1899
1900
1901
1902
            and isinstance(
                pipe.unet,
                (UNet2DConditionModel, UNet3DConditionModel, I2VGenXLUNet, UNetMotionModel, UNetControlNetXSModel),
            )
1903
1904
        )

1905

1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
@is_staging_test
class PipelinePushToHubTester(unittest.TestCase):
    identifier = uuid.uuid4()
    repo_id = f"test-pipeline-{identifier}"
    org_repo_id = f"valid_org/{repo_id}-org"

    def get_pipeline_components(self):
        unet = UNet2DConditionModel(
            block_out_channels=(32, 64),
            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=32,
        )

        scheduler = DDIMScheduler(
            beta_start=0.00085,
            beta_end=0.012,
            beta_schedule="scaled_linear",
            clip_sample=False,
            set_alpha_to_one=False,
        )

        vae = AutoencoderKL(
            block_out_channels=[32, 64],
            in_channels=3,
            out_channels=3,
            down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
            up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
            latent_channels=4,
        )

        text_encoder_config = CLIPTextConfig(
            bos_token_id=0,
            eos_token_id=2,
            hidden_size=32,
            intermediate_size=37,
            layer_norm_eps=1e-05,
            num_attention_heads=4,
            num_hidden_layers=5,
            pad_token_id=1,
            vocab_size=1000,
        )
        text_encoder = CLIPTextModel(text_encoder_config)

        with tempfile.TemporaryDirectory() as tmpdir:
            dummy_vocab = {"<|startoftext|>": 0, "<|endoftext|>": 1, "!": 2}
            vocab_path = os.path.join(tmpdir, "vocab.json")
            with open(vocab_path, "w") as f:
                json.dump(dummy_vocab, f)

            merges = "Ġ t\nĠt h"
            merges_path = os.path.join(tmpdir, "merges.txt")
            with open(merges_path, "w") as f:
                f.writelines(merges)
            tokenizer = CLIPTokenizer(vocab_file=vocab_path, merges_file=merges_path)

        components = {
            "unet": unet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        return components

    def test_push_to_hub(self):
        components = self.get_pipeline_components()
        pipeline = StableDiffusionPipeline(**components)
        pipeline.push_to_hub(self.repo_id, token=TOKEN)

        new_model = UNet2DConditionModel.from_pretrained(f"{USER}/{self.repo_id}", subfolder="unet")
        unet = components["unet"]
        for p1, p2 in zip(unet.parameters(), new_model.parameters()):
            self.assertTrue(torch.equal(p1, p2))

        # Reset repo
        delete_repo(token=TOKEN, repo_id=self.repo_id)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
            pipeline.save_pretrained(tmp_dir, repo_id=self.repo_id, push_to_hub=True, token=TOKEN)

        new_model = UNet2DConditionModel.from_pretrained(f"{USER}/{self.repo_id}", subfolder="unet")
        for p1, p2 in zip(unet.parameters(), new_model.parameters()):
            self.assertTrue(torch.equal(p1, p2))

        # Reset repo
        delete_repo(self.repo_id, token=TOKEN)

    def test_push_to_hub_in_organization(self):
        components = self.get_pipeline_components()
        pipeline = StableDiffusionPipeline(**components)
        pipeline.push_to_hub(self.org_repo_id, token=TOKEN)

        new_model = UNet2DConditionModel.from_pretrained(self.org_repo_id, subfolder="unet")
        unet = components["unet"]
        for p1, p2 in zip(unet.parameters(), new_model.parameters()):
            self.assertTrue(torch.equal(p1, p2))

        # Reset repo
        delete_repo(token=TOKEN, repo_id=self.org_repo_id)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
            pipeline.save_pretrained(tmp_dir, push_to_hub=True, token=TOKEN, repo_id=self.org_repo_id)

        new_model = UNet2DConditionModel.from_pretrained(self.org_repo_id, subfolder="unet")
        for p1, p2 in zip(unet.parameters(), new_model.parameters()):
            self.assertTrue(torch.equal(p1, p2))

        # Reset repo
        delete_repo(self.org_repo_id, token=TOKEN)
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038

    @unittest.skipIf(
        not is_jinja_available(),
        reason="Model card tests cannot be performed without Jinja installed.",
    )
    def test_push_to_hub_library_name(self):
        components = self.get_pipeline_components()
        pipeline = StableDiffusionPipeline(**components)
        pipeline.push_to_hub(self.repo_id, token=TOKEN)

        model_card = ModelCard.load(f"{USER}/{self.repo_id}", token=TOKEN).data
        assert model_card.library_name == "diffusers"

        # Reset repo
        delete_repo(self.repo_id, token=TOKEN)
2039
2040


2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
# For SDXL and its derivative pipelines (such as ControlNet), we have the text encoders
# and the tokenizers as optional components. So, we need to override the `test_save_load_optional_components()`
# test for all such pipelines. This requires us to use a custom `encode_prompt()` function.
class SDXLOptionalComponentsTesterMixin:
    def encode_prompt(
        self, tokenizers, text_encoders, prompt: str, num_images_per_prompt: int = 1, negative_prompt: str = None
    ):
        device = text_encoders[0].device

        if isinstance(prompt, str):
            prompt = [prompt]
        batch_size = len(prompt)

        prompt_embeds_list = []
        for tokenizer, text_encoder in zip(tokenizers, text_encoders):
            text_inputs = tokenizer(
                prompt,
                padding="max_length",
                max_length=tokenizer.model_max_length,
                truncation=True,
                return_tensors="pt",
            )

            text_input_ids = text_inputs.input_ids

            prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
            pooled_prompt_embeds = prompt_embeds[0]
            prompt_embeds = prompt_embeds.hidden_states[-2]
            prompt_embeds_list.append(prompt_embeds)

        prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)

        if negative_prompt is None:
            negative_prompt_embeds = torch.zeros_like(prompt_embeds)
            negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
        else:
            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt

            negative_prompt_embeds_list = []
            for tokenizer, text_encoder in zip(tokenizers, text_encoders):
                uncond_input = tokenizer(
                    negative_prompt,
                    padding="max_length",
                    max_length=tokenizer.model_max_length,
                    truncation=True,
                    return_tensors="pt",
                )

                negative_prompt_embeds = text_encoder(uncond_input.input_ids.to(device), output_hidden_states=True)
                negative_pooled_prompt_embeds = negative_prompt_embeds[0]
                negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
                negative_prompt_embeds_list.append(negative_prompt_embeds)

            negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)

        bs_embed, seq_len, _ = prompt_embeds.shape

        # duplicate text embeddings for each generation per prompt, using mps friendly method
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)

        # for classifier-free guidance
        # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
        seq_len = negative_prompt_embeds.shape[1]

        negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
        negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)

        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
            bs_embed * num_images_per_prompt, -1
        )

        # for classifier-free guidance
        negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
            bs_embed * num_images_per_prompt, -1
        )

        return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds

    def _test_save_load_optional_components(self, expected_max_difference=1e-4):
        components = self.get_dummy_components()

        pipe = self.pipeline_class(**components)
        for optional_component in pipe._optional_components:
            setattr(pipe, optional_component, None)

        for component in pipe.components.values():
            if hasattr(component, "set_default_attn_processor"):
                component.set_default_attn_processor()
        pipe.to(torch_device)
        pipe.set_progress_bar_config(disable=None)

        generator_device = "cpu"
        inputs = self.get_dummy_inputs(generator_device)

        tokenizer = components.pop("tokenizer")
        tokenizer_2 = components.pop("tokenizer_2")
        text_encoder = components.pop("text_encoder")
        text_encoder_2 = components.pop("text_encoder_2")

        tokenizers = [tokenizer, tokenizer_2] if tokenizer is not None else [tokenizer_2]
        text_encoders = [text_encoder, text_encoder_2] if text_encoder is not None else [text_encoder_2]
        prompt = inputs.pop("prompt")
        (
            prompt_embeds,
            negative_prompt_embeds,
            pooled_prompt_embeds,
            negative_pooled_prompt_embeds,
        ) = self.encode_prompt(tokenizers, text_encoders, prompt)
        inputs["prompt_embeds"] = prompt_embeds
        inputs["negative_prompt_embeds"] = negative_prompt_embeds
        inputs["pooled_prompt_embeds"] = pooled_prompt_embeds
        inputs["negative_pooled_prompt_embeds"] = negative_pooled_prompt_embeds

        output = pipe(**inputs)[0]

        with tempfile.TemporaryDirectory() as tmpdir:
            pipe.save_pretrained(tmpdir)
            pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
            for component in pipe_loaded.components.values():
                if hasattr(component, "set_default_attn_processor"):
                    component.set_default_attn_processor()
            pipe_loaded.to(torch_device)
            pipe_loaded.set_progress_bar_config(disable=None)

        for optional_component in pipe._optional_components:
            self.assertTrue(
                getattr(pipe_loaded, optional_component) is None,
                f"`{optional_component}` did not stay set to None after loading.",
            )

        inputs = self.get_dummy_inputs(generator_device)
        _ = inputs.pop("prompt")
        inputs["prompt_embeds"] = prompt_embeds
        inputs["negative_prompt_embeds"] = negative_prompt_embeds
        inputs["pooled_prompt_embeds"] = pooled_prompt_embeds
        inputs["negative_pooled_prompt_embeds"] = negative_pooled_prompt_embeds

        output_loaded = pipe_loaded(**inputs)[0]

        max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
        self.assertLess(max_diff, expected_max_difference)


2185
2186
2187
# 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
# reference image.
2188
def assert_mean_pixel_difference(image, expected_image, expected_max_diff=10):
2189
2190
2191
    image = np.asarray(DiffusionPipeline.numpy_to_pil(image)[0], dtype=np.float32)
    expected_image = np.asarray(DiffusionPipeline.numpy_to_pil(expected_image)[0], dtype=np.float32)
    avg_diff = np.abs(image - expected_image).mean()
2192
    assert avg_diff < expected_max_diff, f"Error image deviates {avg_diff} pixels on average"