test_pipelines.py 19.7 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# coding=utf-8
# Copyright 2022 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

16
import gc
17
import os
18
import random
19
20
21
22
23
24
25
26
import tempfile
import unittest

import numpy as np
import torch

import PIL
from diffusers import (
27
    AutoencoderKL,
28
29
30
31
32
    DDIMPipeline,
    DDIMScheduler,
    DDPMPipeline,
    DDPMScheduler,
    PNDMScheduler,
33
    StableDiffusionImg2ImgPipeline,
34
    StableDiffusionInpaintPipelineLegacy,
35
    StableDiffusionPipeline,
36
    UNet2DConditionModel,
37
    UNet2DModel,
38
    VQModel,
39
    logging,
40
41
)
from diffusers.pipeline_utils import DiffusionPipeline
42
from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
43
from diffusers.utils import CONFIG_NAME, WEIGHTS_NAME, floats_tensor, slow, torch_device
44
from diffusers.utils.testing_utils import CaptureLogger, get_tests_dir, require_torch_gpu
45
from parameterized import parameterized
46
from PIL import Image
Patrick von Platen's avatar
Patrick von Platen committed
47
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextConfig, CLIPTextModel, CLIPTokenizer
48
49
50
51
52


torch.backends.cuda.matmul.allow_tf32 = False


hysts's avatar
hysts committed
53
54
55
56
57
58
59
60
61
62
63
64
65
def test_progress_bar(capsys):
    model = UNet2DModel(
        block_out_channels=(32, 64),
        layers_per_block=2,
        sample_size=32,
        in_channels=3,
        out_channels=3,
        down_block_types=("DownBlock2D", "AttnDownBlock2D"),
        up_block_types=("AttnUpBlock2D", "UpBlock2D"),
    )
    scheduler = DDPMScheduler(num_train_timesteps=10)

    ddpm = DDPMPipeline(model, scheduler).to(torch_device)
66
    ddpm(output_type="numpy").images
hysts's avatar
hysts committed
67
68
69
70
    captured = capsys.readouterr()
    assert "10/10" in captured.err, "Progress bar has to be displayed"

    ddpm.set_progress_bar_config(disable=True)
71
    ddpm(output_type="numpy").images
hysts's avatar
hysts committed
72
73
74
75
    captured = capsys.readouterr()
    assert captured.err == "", "Progress bar should be disabled"


76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
class DownloadTests(unittest.TestCase):
    def test_download_only_pytorch(self):
        with tempfile.TemporaryDirectory() as tmpdirname:
            # pipeline has Flax weights
            _ = DiffusionPipeline.from_pretrained(
                "hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None, cache_dir=tmpdirname
            )

            all_root_files = [t[-1] for t in os.walk(os.path.join(tmpdirname, os.listdir(tmpdirname)[0], "snapshots"))]
            files = [item for sublist in all_root_files for item in sublist]

            # None of the downloaded files should be a flax file even if we have some here:
            # https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_flax_model.msgpack
            assert not any(f.endswith(".msgpack") for f in files)


Patrick von Platen's avatar
Patrick von Platen committed
92
93
94
95
96
class CustomPipelineTests(unittest.TestCase):
    def test_load_custom_pipeline(self):
        pipeline = DiffusionPipeline.from_pretrained(
            "google/ddpm-cifar10-32", custom_pipeline="hf-internal-testing/diffusers-dummy-pipeline"
        )
97
        pipeline = pipeline.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
98
99
100
101
102
103
104
105
        # NOTE that `"CustomPipeline"` is not a class that is defined in this library, but solely on the Hub
        # under https://huggingface.co/hf-internal-testing/diffusers-dummy-pipeline/blob/main/pipeline.py#L24
        assert pipeline.__class__.__name__ == "CustomPipeline"

    def test_run_custom_pipeline(self):
        pipeline = DiffusionPipeline.from_pretrained(
            "google/ddpm-cifar10-32", custom_pipeline="hf-internal-testing/diffusers-dummy-pipeline"
        )
106
        pipeline = pipeline.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
107
108
109
        images, output_str = pipeline(num_inference_steps=2, output_type="np")

        assert images[0].shape == (1, 32, 32, 3)
110

Patrick von Platen's avatar
Patrick von Platen committed
111
112
113
114
115
116
117
118
        # compare output to https://huggingface.co/hf-internal-testing/diffusers-dummy-pipeline/blob/main/pipeline.py#L102
        assert output_str == "This is a test"

    def test_local_custom_pipeline(self):
        local_custom_pipeline_path = get_tests_dir("fixtures/custom_pipeline")
        pipeline = DiffusionPipeline.from_pretrained(
            "google/ddpm-cifar10-32", custom_pipeline=local_custom_pipeline_path
        )
119
        pipeline = pipeline.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
120
121
122
123
124
125
126
127
        images, output_str = pipeline(num_inference_steps=2, output_type="np")

        assert pipeline.__class__.__name__ == "CustomLocalPipeline"
        assert images[0].shape == (1, 32, 32, 3)
        # compare to https://github.com/huggingface/diffusers/blob/main/tests/fixtures/custom_pipeline/pipeline.py#L102
        assert output_str == "This is a local test"

    @slow
128
    @require_torch_gpu
Patrick von Platen's avatar
Patrick von Platen committed
129
130
131
    def test_load_pipeline_from_git(self):
        clip_model_id = "laion/CLIP-ViT-B-32-laion2B-s34B-b79K"

132
        feature_extractor = CLIPFeatureExtractor.from_pretrained(clip_model_id)
133
        clip_model = CLIPModel.from_pretrained(clip_model_id, torch_dtype=torch.float16)
Patrick von Platen's avatar
Patrick von Platen committed
134
135
136
137
138
139

        pipeline = DiffusionPipeline.from_pretrained(
            "CompVis/stable-diffusion-v1-4",
            custom_pipeline="clip_guided_stable_diffusion",
            clip_model=clip_model,
            feature_extractor=feature_extractor,
140
141
            torch_dtype=torch.float16,
            revision="fp16",
Patrick von Platen's avatar
Patrick von Platen committed
142
        )
143
        pipeline.enable_attention_slicing()
Patrick von Platen's avatar
Patrick von Platen committed
144
145
146
147
148
149
150
151
152
153
        pipeline = pipeline.to(torch_device)

        # NOTE that `"CLIPGuidedStableDiffusion"` is not a class that is defined in the pypi package of th e library, but solely on the community examples folder of GitHub under:
        # https://github.com/huggingface/diffusers/blob/main/examples/community/clip_guided_stable_diffusion.py
        assert pipeline.__class__.__name__ == "CLIPGuidedStableDiffusion"

        image = pipeline("a prompt", num_inference_steps=2, output_type="np").images[0]
        assert image.shape == (512, 512, 3)


154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
class PipelineFastTests(unittest.TestCase):
    @property
    def dummy_image(self):
        batch_size = 1
        num_channels = 3
        sizes = (32, 32)

        image = floats_tensor((batch_size, num_channels) + sizes, rng=random.Random(0)).to(torch_device)
        return image

    @property
    def dummy_uncond_unet(self):
        torch.manual_seed(0)
        model = UNet2DModel(
            block_out_channels=(32, 64),
            layers_per_block=2,
            sample_size=32,
            in_channels=3,
            out_channels=3,
            down_block_types=("DownBlock2D", "AttnDownBlock2D"),
            up_block_types=("AttnUpBlock2D", "UpBlock2D"),
        )
        return model

    @property
    def dummy_cond_unet(self):
        torch.manual_seed(0)
        model = 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,
        )
        return model

193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
    @property
    def dummy_cond_unet_inpaint(self):
        torch.manual_seed(0)
        model = UNet2DConditionModel(
            block_out_channels=(32, 64),
            layers_per_block=2,
            sample_size=32,
            in_channels=9,
            out_channels=4,
            down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
            up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
            cross_attention_dim=32,
        )
        return model

208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
    @property
    def dummy_vq_model(self):
        torch.manual_seed(0)
        model = VQModel(
            block_out_channels=[32, 64],
            in_channels=3,
            out_channels=3,
            down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
            up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
            latent_channels=3,
        )
        return model

    @property
    def dummy_vae(self):
        torch.manual_seed(0)
        model = 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,
        )
        return model

    @property
    def dummy_text_encoder(self):
        torch.manual_seed(0)
        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,
        )
        return CLIPTextModel(config)

    @property
    def dummy_extractor(self):
        def extract(*args, **kwargs):
            class Out:
                def __init__(self):
                    self.pixel_values = torch.ones([0])

                def to(self, device):
                    self.pixel_values.to(device)
                    return self

            return Out()

        return extract

265
266
    def test_components(self):
        """Test that components property works correctly"""
267
        unet = self.dummy_cond_unet
268
        scheduler = PNDMScheduler(skip_prk_steps=True)
269
270
271
272
        vae = self.dummy_vae
        bert = self.dummy_text_encoder
        tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

273
274
275
        image = self.dummy_image.cpu().permute(0, 2, 3, 1)[0]
        init_image = Image.fromarray(np.uint8(image)).convert("RGB")
        mask_image = Image.fromarray(np.uint8(image + 4)).convert("RGB").resize((128, 128))
276
277

        # make sure here that pndm scheduler skips prk
278
        inpaint = StableDiffusionInpaintPipelineLegacy(
279
280
281
282
283
            unet=unet,
            scheduler=scheduler,
            vae=vae,
            text_encoder=bert,
            tokenizer=tokenizer,
284
            safety_checker=None,
285
            feature_extractor=self.dummy_extractor,
286
287
288
        ).to(torch_device)
        img2img = StableDiffusionImg2ImgPipeline(**inpaint.components).to(torch_device)
        text2img = StableDiffusionPipeline(**inpaint.components).to(torch_device)
289
290

        prompt = "A painting of a squirrel eating a burger"
291
292
293
294
295
296
297

        # Device type MPS is not supported for torch.Generator() api.
        if torch_device == "mps":
            generator = torch.manual_seed(0)
        else:
            generator = torch.Generator(device=torch_device).manual_seed(0)

298
        image_inpaint = inpaint(
299
300
301
302
303
            [prompt],
            generator=generator,
            num_inference_steps=2,
            output_type="np",
            init_image=init_image,
304
305
306
            mask_image=mask_image,
        ).images
        image_img2img = img2img(
307
308
309
310
311
            [prompt],
            generator=generator,
            num_inference_steps=2,
            output_type="np",
            init_image=init_image,
312
313
314
        ).images
        image_text2img = text2img(
            [prompt],
315
316
317
            generator=generator,
            num_inference_steps=2,
            output_type="np",
318
        ).images
319

320
321
322
        assert image_inpaint.shape == (1, 32, 32, 3)
        assert image_img2img.shape == (1, 32, 32, 3)
        assert image_text2img.shape == (1, 128, 128, 3)
323

324

325
326
@slow
class PipelineSlowTests(unittest.TestCase):
327
328
329
330
331
332
    def tearDown(self):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

333
334
335
    def test_smart_download(self):
        model_id = "hf-internal-testing/unet-pipeline-dummy"
        with tempfile.TemporaryDirectory() as tmpdirname:
336
            _ = DiffusionPipeline.from_pretrained(model_id, cache_dir=tmpdirname, force_download=True)
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
            local_repo_name = "--".join(["models"] + model_id.split("/"))
            snapshot_dir = os.path.join(tmpdirname, local_repo_name, "snapshots")
            snapshot_dir = os.path.join(snapshot_dir, os.listdir(snapshot_dir)[0])

            # inspect all downloaded files to make sure that everything is included
            assert os.path.isfile(os.path.join(snapshot_dir, DiffusionPipeline.config_name))
            assert os.path.isfile(os.path.join(snapshot_dir, CONFIG_NAME))
            assert os.path.isfile(os.path.join(snapshot_dir, SCHEDULER_CONFIG_NAME))
            assert os.path.isfile(os.path.join(snapshot_dir, WEIGHTS_NAME))
            assert os.path.isfile(os.path.join(snapshot_dir, "scheduler", SCHEDULER_CONFIG_NAME))
            assert os.path.isfile(os.path.join(snapshot_dir, "unet", WEIGHTS_NAME))
            assert os.path.isfile(os.path.join(snapshot_dir, "unet", WEIGHTS_NAME))
            # let's make sure the super large numpy file:
            # https://huggingface.co/hf-internal-testing/unet-pipeline-dummy/blob/main/big_array.npy
            # is not downloaded, but all the expected ones
            assert not os.path.isfile(os.path.join(snapshot_dir, "big_array.npy"))

354
355
356
357
358
    def test_warning_unused_kwargs(self):
        model_id = "hf-internal-testing/unet-pipeline-dummy"
        logger = logging.get_logger("diffusers.pipeline_utils")
        with tempfile.TemporaryDirectory() as tmpdirname:
            with CaptureLogger(logger) as cap_logger:
359
                DiffusionPipeline.from_pretrained(
360
361
362
363
                    model_id,
                    not_used=True,
                    cache_dir=tmpdirname,
                    force_download=True,
364
                )
365
366
367

        assert cap_logger.out == "Keyword arguments {'not_used': True} not recognized.\n"

368
369
370
371
372
373
374
375
376
377
378
379
380
381
    def test_from_pretrained_save_pretrained(self):
        # 1. Load models
        model = UNet2DModel(
            block_out_channels=(32, 64),
            layers_per_block=2,
            sample_size=32,
            in_channels=3,
            out_channels=3,
            down_block_types=("DownBlock2D", "AttnDownBlock2D"),
            up_block_types=("AttnUpBlock2D", "UpBlock2D"),
        )
        schedular = DDPMScheduler(num_train_timesteps=10)

        ddpm = DDPMPipeline(model, schedular)
382
        ddpm.to(torch_device)
383
        ddpm.set_progress_bar_config(disable=None)
384
385
386

        with tempfile.TemporaryDirectory() as tmpdirname:
            ddpm.save_pretrained(tmpdirname)
387
            new_ddpm = DDPMPipeline.from_pretrained(tmpdirname)
388
            new_ddpm.to(torch_device)
389
390

        generator = torch.manual_seed(0)
391
        image = ddpm(generator=generator, output_type="numpy").images
392

393
        generator = generator.manual_seed(0)
394
        new_image = new_ddpm(generator=generator, output_type="numpy").images
395
396
397
398
399
400

        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"

    def test_from_pretrained_hub(self):
        model_path = "google/ddpm-cifar10-32"

401
        scheduler = DDPMScheduler(num_train_timesteps=10)
402

403
        ddpm = DDPMPipeline.from_pretrained(model_path, scheduler=scheduler)
404
        ddpm = ddpm.to(torch_device)
405
        ddpm.set_progress_bar_config(disable=None)
406

407
        ddpm_from_hub = DiffusionPipeline.from_pretrained(model_path, scheduler=scheduler)
408
        ddpm_from_hub = ddpm_from_hub.to(torch_device)
409
        ddpm_from_hub.set_progress_bar_config(disable=None)
410
411

        generator = torch.manual_seed(0)
412
        image = ddpm(generator=generator, output_type="numpy").images
413

414
        generator = generator.manual_seed(0)
415
        new_image = ddpm_from_hub(generator=generator, output_type="numpy").images
416
417
418
419
420
421

        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"

    def test_from_pretrained_hub_pass_model(self):
        model_path = "google/ddpm-cifar10-32"

422
423
        scheduler = DDPMScheduler(num_train_timesteps=10)

424
        # pass unet into DiffusionPipeline
425
426
        unet = UNet2DModel.from_pretrained(model_path)
        ddpm_from_hub_custom_model = DiffusionPipeline.from_pretrained(model_path, unet=unet, scheduler=scheduler)
427
        ddpm_from_hub_custom_model = ddpm_from_hub_custom_model.to(torch_device)
428
        ddpm_from_hub_custom_model.set_progress_bar_config(disable=None)
429

430
        ddpm_from_hub = DiffusionPipeline.from_pretrained(model_path, scheduler=scheduler)
431
        ddpm_from_hub = ddpm_from_hub.to(torch_device)
432
        ddpm_from_hub_custom_model.set_progress_bar_config(disable=None)
433
434

        generator = torch.manual_seed(0)
435
        image = ddpm_from_hub_custom_model(generator=generator, output_type="numpy").images
436

437
        generator = generator.manual_seed(0)
438
        new_image = ddpm_from_hub(generator=generator, output_type="numpy").images
439
440
441
442
443
444

        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"

    def test_output_format(self):
        model_path = "google/ddpm-cifar10-32"

445
        pipe = DDIMPipeline.from_pretrained(model_path)
446
        pipe.to(torch_device)
447
        pipe.set_progress_bar_config(disable=None)
448
449

        generator = torch.manual_seed(0)
450
        images = pipe(generator=generator, output_type="numpy").images
451
452
453
        assert images.shape == (1, 32, 32, 3)
        assert isinstance(images, np.ndarray)

454
        images = pipe(generator=generator, output_type="pil").images
455
456
457
458
459
        assert isinstance(images, list)
        assert len(images) == 1
        assert isinstance(images[0], PIL.Image.Image)

        # use PIL by default
460
        images = pipe(generator=generator).images
461
462
463
        assert isinstance(images, list)
        assert isinstance(images[0], PIL.Image.Image)

464
465
466
    # Make sure the test passes for different values of random seed
    @parameterized.expand([(0,), (4,)])
    def test_ddpm_ddim_equality(self, seed):
467
        model_id = "google/ddpm-cifar10-32"
468

469
        unet = UNet2DModel.from_pretrained(model_id)
470
471
        ddpm_scheduler = DDPMScheduler()
        ddim_scheduler = DDIMScheduler()
472

473
474
475
476
477
478
        ddpm = DDPMPipeline(unet=unet, scheduler=ddpm_scheduler)
        ddpm.to(torch_device)
        ddpm.set_progress_bar_config(disable=None)
        ddim = DDIMPipeline(unet=unet, scheduler=ddim_scheduler)
        ddim.to(torch_device)
        ddim.set_progress_bar_config(disable=None)
479

480
        generator = torch.manual_seed(seed)
481
        ddpm_image = ddpm(generator=generator, output_type="numpy").images
482

483
484
485
486
487
488
489
490
        generator = torch.manual_seed(seed)
        ddim_image = ddim(
            generator=generator,
            num_inference_steps=1000,
            eta=1.0,
            output_type="numpy",
            use_clipped_model_output=True,  # Need this to make DDIM match DDPM
        ).images
491

492
493
        # the values aren't exactly equal, but the images look the same visually
        assert np.abs(ddpm_image - ddim_image).max() < 1e-1
494

495
496
497
    # Make sure the test passes for different values of random seed
    @parameterized.expand([(0,), (4,)])
    def test_ddpm_ddim_equality_batched(self, seed):
498
        model_id = "google/ddpm-cifar10-32"
499

500
        unet = UNet2DModel.from_pretrained(model_id)
501
502
        ddpm_scheduler = DDPMScheduler()
        ddim_scheduler = DDIMScheduler()
503

504
505
506
        ddpm = DDPMPipeline(unet=unet, scheduler=ddpm_scheduler)
        ddpm.to(torch_device)
        ddpm.set_progress_bar_config(disable=None)
507

508
509
510
        ddim = DDIMPipeline(unet=unet, scheduler=ddim_scheduler)
        ddim.to(torch_device)
        ddim.set_progress_bar_config(disable=None)
511

512
        generator = torch.manual_seed(seed)
513
514
        ddpm_images = ddpm(batch_size=4, generator=generator, output_type="numpy").images

515
        generator = torch.manual_seed(seed)
516
        ddim_images = ddim(
517
518
519
520
521
522
            batch_size=4,
            generator=generator,
            num_inference_steps=1000,
            eta=1.0,
            output_type="numpy",
            use_clipped_model_output=True,  # Need this to make DDIM match DDPM
523
        ).images
524

525
526
        # the values aren't exactly equal, but the images look the same visually
        assert np.abs(ddpm_images - ddim_images).max() < 1e-1