test_modeling_utils.py 30.4 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.

patil-suraj's avatar
patil-suraj committed
16
import inspect
17
import math
18
19
20
import tempfile
import unittest

21
import numpy as np
22
23
import torch

24
import PIL
Patrick von Platen's avatar
Patrick von Platen committed
25
from diffusers import UNet2DConditionModel  # noqa: F401 TODO(Patrick) - need to write tests with it
Patrick von Platen's avatar
Patrick von Platen committed
26
from diffusers import (
patil-suraj's avatar
patil-suraj committed
27
    AutoencoderKL,
Patrick von Platen's avatar
Patrick von Platen committed
28
    DDIMPipeline,
29
    DDIMScheduler,
Patrick von Platen's avatar
Patrick von Platen committed
30
    DDPMPipeline,
31
    DDPMScheduler,
Patrick von Platen's avatar
Patrick von Platen committed
32
33
    LDMPipeline,
    LDMTextToImagePipeline,
Patrick von Platen's avatar
Patrick von Platen committed
34
    PNDMPipeline,
35
    PNDMScheduler,
Patrick von Platen's avatar
Patrick von Platen committed
36
37
    ScoreSdeVePipeline,
    ScoreSdeVeScheduler,
Patrick von Platen's avatar
Patrick von Platen committed
38
    UNet2DModel,
patil-suraj's avatar
patil-suraj committed
39
    VQModel,
40
)
41
from diffusers.configuration_utils import ConfigMixin, register_to_config
Patrick von Platen's avatar
Patrick von Platen committed
42
from diffusers.pipeline_utils import DiffusionPipeline
Patrick von Platen's avatar
Patrick von Platen committed
43
from diffusers.testing_utils import floats_tensor, slow, torch_device
44
from diffusers.training_utils import EMAModel
45
46


Patrick von Platen's avatar
Patrick von Platen committed
47
torch.backends.cuda.matmul.allow_tf32 = False
48
49


50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
class SampleObject(ConfigMixin):
    config_name = "config.json"

    @register_to_config
    def __init__(
        self,
        a=2,
        b=5,
        c=(2, 5),
        d="for diffusion",
        e=[1, 3],
    ):
        pass


65
66
67
68
69
class ConfigTester(unittest.TestCase):
    def test_load_not_from_mixin(self):
        with self.assertRaises(ValueError):
            ConfigMixin.from_config("dummy_path")

70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
    def test_register_to_config(self):
        obj = SampleObject()
        config = obj.config
        assert config["a"] == 2
        assert config["b"] == 5
        assert config["c"] == (2, 5)
        assert config["d"] == "for diffusion"
        assert config["e"] == [1, 3]

        # init ignore private arguments
        obj = SampleObject(_name_or_path="lalala")
        config = obj.config
        assert config["a"] == 2
        assert config["b"] == 5
        assert config["c"] == (2, 5)
        assert config["d"] == "for diffusion"
        assert config["e"] == [1, 3]
87

88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
        # can override default
        obj = SampleObject(c=6)
        config = obj.config
        assert config["a"] == 2
        assert config["b"] == 5
        assert config["c"] == 6
        assert config["d"] == "for diffusion"
        assert config["e"] == [1, 3]

        # can use positional arguments.
        obj = SampleObject(1, c=6)
        config = obj.config
        assert config["a"] == 1
        assert config["b"] == 5
        assert config["c"] == 6
        assert config["d"] == "for diffusion"
        assert config["e"] == [1, 3]

    def test_save_load(self):
107
108
109
110
111
112
113
114
115
116
117
118
119
120
        obj = SampleObject()
        config = obj.config

        assert config["a"] == 2
        assert config["b"] == 5
        assert config["c"] == (2, 5)
        assert config["d"] == "for diffusion"
        assert config["e"] == [1, 3]

        with tempfile.TemporaryDirectory() as tmpdirname:
            obj.save_config(tmpdirname)
            new_obj = SampleObject.from_config(tmpdirname)
            new_config = new_obj.config

Patrick von Platen's avatar
Patrick von Platen committed
121
122
123
124
        # unfreeze configs
        config = dict(config)
        new_config = dict(new_config)

125
126
127
128
129
        assert config.pop("c") == (2, 5)  # instantiated as tuple
        assert new_config.pop("c") == [2, 5]  # saved & loaded as list because of json
        assert config == new_config


patil-suraj's avatar
patil-suraj committed
130
class ModelTesterMixin:
131
    def test_from_pretrained_save_pretrained(self):
patil-suraj's avatar
patil-suraj committed
132
133
134
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

        model = self.model_class(**init_dict)
Patrick von Platen's avatar
Patrick von Platen committed
135
        model.to(torch_device)
patil-suraj's avatar
patil-suraj committed
136
        model.eval()
137
138
139

        with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_pretrained(tmpdirname)
patil-suraj's avatar
patil-suraj committed
140
            new_model = self.model_class.from_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
141
            new_model.to(torch_device)
142

patil-suraj's avatar
patil-suraj committed
143
144
        with torch.no_grad():
            image = model(**inputs_dict)
Patrick von Platen's avatar
Patrick von Platen committed
145
146
147
            if isinstance(image, dict):
                image = image["sample"]

patil-suraj's avatar
patil-suraj committed
148
            new_image = new_model(**inputs_dict)
149

Patrick von Platen's avatar
Patrick von Platen committed
150
151
152
            if isinstance(new_image, dict):
                new_image = new_image["sample"]

patil-suraj's avatar
patil-suraj committed
153
        max_diff = (image - new_image).abs().sum().item()
Patrick von Platen's avatar
Patrick von Platen committed
154
        self.assertLessEqual(max_diff, 5e-5, "Models give different forward passes")
155

patil-suraj's avatar
patil-suraj committed
156
    def test_determinism(self):
patil-suraj's avatar
patil-suraj committed
157
158
159
160
161
162
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        model = self.model_class(**init_dict)
        model.to(torch_device)
        model.eval()
        with torch.no_grad():
            first = model(**inputs_dict)
Patrick von Platen's avatar
Patrick von Platen committed
163
164
165
            if isinstance(first, dict):
                first = first["sample"]

patil-suraj's avatar
patil-suraj committed
166
            second = model(**inputs_dict)
Patrick von Platen's avatar
Patrick von Platen committed
167
168
            if isinstance(second, dict):
                second = second["sample"]
patil-suraj's avatar
patil-suraj committed
169
170
171
172
173
174
175

        out_1 = first.cpu().numpy()
        out_2 = second.cpu().numpy()
        out_1 = out_1[~np.isnan(out_1)]
        out_2 = out_2[~np.isnan(out_2)]
        max_diff = np.amax(np.abs(out_1 - out_2))
        self.assertLessEqual(max_diff, 1e-5)
176

patil-suraj's avatar
patil-suraj committed
177
    def test_output(self):
patil-suraj's avatar
patil-suraj committed
178
179
180
181
182
183
184
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
        model = self.model_class(**init_dict)
        model.to(torch_device)
        model.eval()

        with torch.no_grad():
            output = model(**inputs_dict)
185

Patrick von Platen's avatar
Patrick von Platen committed
186
187
188
            if isinstance(output, dict):
                output = output["sample"]

patil-suraj's avatar
patil-suraj committed
189
        self.assertIsNotNone(output)
190
        expected_shape = inputs_dict["sample"].shape
patil-suraj's avatar
patil-suraj committed
191
        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
192

patil-suraj's avatar
patil-suraj committed
193
    def test_forward_signature(self):
patil-suraj's avatar
patil-suraj committed
194
195
196
197
198
199
200
        init_dict, _ = self.prepare_init_args_and_inputs_for_common()

        model = self.model_class(**init_dict)
        signature = inspect.signature(model.forward)
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        arg_names = [*signature.parameters.keys()]

201
        expected_arg_names = ["sample", "timestep"]
patil-suraj's avatar
patil-suraj committed
202
        self.assertListEqual(arg_names[:2], expected_arg_names)
203

patil-suraj's avatar
patil-suraj committed
204
    def test_model_from_config(self):
patil-suraj's avatar
patil-suraj committed
205
206
207
208
209
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

        model = self.model_class(**init_dict)
        model.to(torch_device)
        model.eval()
210

patil-suraj's avatar
patil-suraj committed
211
212
213
214
215
216
217
        # test if the model can be loaded from the config
        # and has all the expected shape
        with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_config(tmpdirname)
            new_model = self.model_class.from_config(tmpdirname)
            new_model.to(torch_device)
            new_model.eval()
218

patil-suraj's avatar
patil-suraj committed
219
220
221
222
223
        # check if all paramters shape are the same
        for param_name in model.state_dict().keys():
            param_1 = model.state_dict()[param_name]
            param_2 = new_model.state_dict()[param_name]
            self.assertEqual(param_1.shape, param_2.shape)
224

patil-suraj's avatar
patil-suraj committed
225
226
        with torch.no_grad():
            output_1 = model(**inputs_dict)
Patrick von Platen's avatar
Patrick von Platen committed
227
228
229
230

            if isinstance(output_1, dict):
                output_1 = output_1["sample"]

patil-suraj's avatar
patil-suraj committed
231
            output_2 = new_model(**inputs_dict)
232

Patrick von Platen's avatar
Patrick von Platen committed
233
234
235
            if isinstance(output_2, dict):
                output_2 = output_2["sample"]

patil-suraj's avatar
patil-suraj committed
236
        self.assertEqual(output_1.shape, output_2.shape)
patil-suraj's avatar
patil-suraj committed
237
238

    def test_training(self):
patil-suraj's avatar
patil-suraj committed
239
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
240

patil-suraj's avatar
patil-suraj committed
241
242
243
244
        model = self.model_class(**init_dict)
        model.to(torch_device)
        model.train()
        output = model(**inputs_dict)
Patrick von Platen's avatar
Patrick von Platen committed
245
246
247
248

        if isinstance(output, dict):
            output = output["sample"]

249
        noise = torch.randn((inputs_dict["sample"].shape[0],) + self.output_shape).to(torch_device)
patil-suraj's avatar
patil-suraj committed
250
251
        loss = torch.nn.functional.mse_loss(output, noise)
        loss.backward()
252

253
254
255
256
257
258
259
260
261
    def test_ema_training(self):
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

        model = self.model_class(**init_dict)
        model.to(torch_device)
        model.train()
        ema_model = EMAModel(model, device=torch_device)

        output = model(**inputs_dict)
Patrick von Platen's avatar
Patrick von Platen committed
262
263
264
265

        if isinstance(output, dict):
            output = output["sample"]

266
        noise = torch.randn((inputs_dict["sample"].shape[0],) + self.output_shape).to(torch_device)
267
268
269
270
        loss = torch.nn.functional.mse_loss(output, noise)
        loss.backward()
        ema_model.step(model)

patil-suraj's avatar
patil-suraj committed
271
272

class UnetModelTests(ModelTesterMixin, unittest.TestCase):
Patrick von Platen's avatar
Patrick von Platen committed
273
    model_class = UNet2DModel
patil-suraj's avatar
patil-suraj committed
274
275
276
277
278
279
280
281
282
283

    @property
    def dummy_input(self):
        batch_size = 4
        num_channels = 3
        sizes = (32, 32)

        noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
        time_step = torch.tensor([10]).to(torch_device)

284
        return {"sample": noise, "timestep": time_step}
285

patil-suraj's avatar
patil-suraj committed
286
    @property
Patrick von Platen's avatar
Patrick von Platen committed
287
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
288
        return (3, 32, 32)
289

patil-suraj's avatar
patil-suraj committed
290
    @property
Patrick von Platen's avatar
Patrick von Platen committed
291
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
292
        return (3, 32, 32)
patil-suraj's avatar
patil-suraj committed
293
294
295

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
Patrick von Platen's avatar
Patrick von Platen committed
296
297
298
299
            "block_out_channels": (32, 64),
            "down_block_types": ("DownBlock2D", "AttnDownBlock2D"),
            "up_block_types": ("AttnUpBlock2D", "UpBlock2D"),
            "attention_head_dim": None,
300
301
            "out_channels": 3,
            "in_channels": 3,
Patrick von Platen's avatar
Patrick von Platen committed
302
303
            "layers_per_block": 2,
            "sample_size": 32,
patil-suraj's avatar
patil-suraj committed
304
305
306
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict
307

patil-suraj's avatar
patil-suraj committed
308

Patrick von Platen's avatar
upload  
Patrick von Platen committed
309
310
#    TODO(Patrick) - Re-add this test after having correctly added the final VE checkpoints
#    def test_output_pretrained(self):
Patrick von Platen's avatar
Patrick von Platen committed
311
#        model = UNet2DModel.from_pretrained("fusing/ddpm_dummy_update", subfolder="unet")
Patrick von Platen's avatar
upload  
Patrick von Platen committed
312
313
314
315
316
317
#        model.eval()
#
#        torch.manual_seed(0)
#        if torch.cuda.is_available():
#            torch.cuda.manual_seed_all(0)
#
Patrick von Platen's avatar
Patrick von Platen committed
318
#        noise = torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
Patrick von Platen's avatar
upload  
Patrick von Platen committed
319
320
321
322
323
324
325
326
327
328
#        time_step = torch.tensor([10])
#
#        with torch.no_grad():
#            output = model(noise, time_step)["sample"]
#
#        output_slice = output[0, -1, -3:, -3:].flatten()
# fmt: off
#        expected_output_slice = torch.tensor([0.2891, -0.1899, 0.2595, -0.6214, 0.0968, -0.2622, 0.4688, 0.1311, 0.0053])
# fmt: on
#        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
329
330


patil-suraj's avatar
patil-suraj committed
331
class UNetLDMModelTests(ModelTesterMixin, unittest.TestCase):
Patrick von Platen's avatar
Patrick von Platen committed
332
    model_class = UNet2DModel
patil-suraj's avatar
patil-suraj committed
333
334
335
336
337
338
339
340
341
342

    @property
    def dummy_input(self):
        batch_size = 4
        num_channels = 4
        sizes = (32, 32)

        noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
        time_step = torch.tensor([10]).to(torch_device)

343
        return {"sample": noise, "timestep": time_step}
patil-suraj's avatar
patil-suraj committed
344
345

    @property
Patrick von Platen's avatar
Patrick von Platen committed
346
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
347
348
349
        return (4, 32, 32)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
350
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
351
352
353
354
        return (4, 32, 32)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
Patrick von Platen's avatar
Patrick von Platen committed
355
            "sample_size": 32,
patil-suraj's avatar
patil-suraj committed
356
357
            "in_channels": 4,
            "out_channels": 4,
Patrick von Platen's avatar
Patrick von Platen committed
358
359
360
361
362
            "layers_per_block": 2,
            "block_out_channels": (32, 64),
            "attention_head_dim": 32,
            "down_block_types": ("DownBlock2D", "DownBlock2D"),
            "up_block_types": ("UpBlock2D", "UpBlock2D"),
patil-suraj's avatar
patil-suraj committed
363
364
365
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict
anton-l's avatar
anton-l committed
366

patil-suraj's avatar
patil-suraj committed
367
    def test_from_pretrained_hub(self):
Patrick von Platen's avatar
Patrick von Platen committed
368
        model, loading_info = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update", output_loading_info=True)
Patrick von Platen's avatar
Patrick von Platen committed
369

patil-suraj's avatar
patil-suraj committed
370
        self.assertIsNotNone(model)
Patrick von Platen's avatar
upload  
Patrick von Platen committed
371
        self.assertEqual(len(loading_info["missing_keys"]), 0)
patil-suraj's avatar
patil-suraj committed
372
373

        model.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
374
        image = model(**self.dummy_input)["sample"]
patil-suraj's avatar
patil-suraj committed
375
376
377
378

        assert image is not None, "Make sure output is not None"

    def test_output_pretrained(self):
Patrick von Platen's avatar
Patrick von Platen committed
379
        model = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update")
patil-suraj's avatar
patil-suraj committed
380
381
382
383
384
385
        model.eval()

        torch.manual_seed(0)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(0)

Patrick von Platen's avatar
Patrick von Platen committed
386
        noise = torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
patil-suraj's avatar
patil-suraj committed
387
388
389
        time_step = torch.tensor([10] * noise.shape[0])

        with torch.no_grad():
Patrick von Platen's avatar
Patrick von Platen committed
390
            output = model(noise, time_step)["sample"]
patil-suraj's avatar
patil-suraj committed
391
392
393
394
395
396
397
398

        output_slice = output[0, -1, -3:, -3:].flatten()
        # fmt: off
        expected_output_slice = torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800])
        # fmt: on

        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))

Patrick von Platen's avatar
Patrick von Platen committed
399

Patrick von Platen's avatar
upload  
Patrick von Platen committed
400
401
402
403
404
405
406
407
408
#    TODO(Patrick) - Re-add this test after having cleaned up LDM
#    def test_output_pretrained_spatial_transformer(self):
#        model = UNetLDMModel.from_pretrained("fusing/unet-ldm-dummy-spatial")
#        model.eval()
#
#        torch.manual_seed(0)
#        if torch.cuda.is_available():
#            torch.cuda.manual_seed_all(0)
#
Patrick von Platen's avatar
Patrick von Platen committed
409
#        noise = torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
Patrick von Platen's avatar
upload  
Patrick von Platen committed
410
411
412
413
414
415
416
417
418
419
420
421
422
#        context = torch.ones((1, 16, 64), dtype=torch.float32)
#        time_step = torch.tensor([10] * noise.shape[0])
#
#        with torch.no_grad():
#            output = model(noise, time_step, context=context)
#
#        output_slice = output[0, -1, -3:, -3:].flatten()
# fmt: off
#        expected_output_slice = torch.tensor([61.3445, 56.9005, 29.4339, 59.5497, 60.7375, 34.1719, 48.1951, 42.6569, 25.0890])
# fmt: on
#
#        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
#
Patrick von Platen's avatar
Patrick von Platen committed
423

patil-suraj's avatar
patil-suraj committed
424

425
class NCSNppModelTests(ModelTesterMixin, unittest.TestCase):
Patrick von Platen's avatar
Patrick von Platen committed
426
    model_class = UNet2DModel
427
428

    @property
Patrick von Platen's avatar
Patrick von Platen committed
429
    def dummy_input(self, sizes=(32, 32)):
430
431
432
433
434
435
        batch_size = 4
        num_channels = 3

        noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
        time_step = torch.tensor(batch_size * [10]).to(torch_device)

436
        return {"sample": noise, "timestep": time_step}
437
438

    @property
Patrick von Platen's avatar
Patrick von Platen committed
439
    def input_shape(self):
440
441
442
        return (3, 32, 32)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
443
    def output_shape(self):
444
445
446
447
        return (3, 32, 32)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
Patrick von Platen's avatar
Patrick von Platen committed
448
            "block_out_channels": [32, 64, 64, 64],
449
            "in_channels": 3,
Patrick von Platen's avatar
Patrick von Platen committed
450
            "layers_per_block": 1,
451
452
            "out_channels": 3,
            "time_embedding_type": "fourier",
Patrick von Platen's avatar
Patrick von Platen committed
453
            "norm_eps": 1e-6,
454
            "mid_block_scale_factor": math.sqrt(2.0),
Patrick von Platen's avatar
Patrick von Platen committed
455
456
457
458
459
460
            "norm_num_groups": None,
            "down_block_types": [
                "SkipDownBlock2D",
                "AttnSkipDownBlock2D",
                "SkipDownBlock2D",
                "SkipDownBlock2D",
461
            ],
Patrick von Platen's avatar
Patrick von Platen committed
462
463
464
465
466
            "up_block_types": [
                "SkipUpBlock2D",
                "SkipUpBlock2D",
                "AttnSkipUpBlock2D",
                "SkipUpBlock2D",
467
            ],
468
469
470
471
472
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict

    def test_from_pretrained_hub(self):
Patrick von Platen's avatar
Patrick von Platen committed
473
        model, loading_info = UNet2DModel.from_pretrained("google/ncsnpp-celebahq-256", output_loading_info=True)
474
        self.assertIsNotNone(model)
Patrick von Platen's avatar
upload  
Patrick von Platen committed
475
        self.assertEqual(len(loading_info["missing_keys"]), 0)
476
477

        model.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
478
479
480
481
        inputs = self.dummy_input
        noise = floats_tensor((4, 3) + (256, 256)).to(torch_device)
        inputs["sample"] = noise
        image = model(**inputs)
482
483
484

        assert image is not None, "Make sure output is not None"

485
    def test_output_pretrained_ve_mid(self):
Patrick von Platen's avatar
Patrick von Platen committed
486
        model = UNet2DModel.from_pretrained("google/ncsnpp-celebahq-256")
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
        model.to(torch_device)

        torch.manual_seed(0)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(0)

        batch_size = 4
        num_channels = 3
        sizes = (256, 256)

        noise = torch.ones((batch_size, num_channels) + sizes).to(torch_device)
        time_step = torch.tensor(batch_size * [1e-4]).to(torch_device)

        with torch.no_grad():
            output = model(noise, time_step)["sample"]

        output_slice = output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        expected_output_slice = torch.tensor([-4836.2231, -6487.1387, -3816.7969, -7964.9253, -10966.2842, -20043.6016, 8137.0571, 2340.3499, 544.6114])
        # fmt: on

        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))

510
    def test_output_pretrained_ve_large(self):
Patrick von Platen's avatar
Patrick von Platen committed
511
        model = UNet2DModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update")
512
513
514
515
516
517
518
519
520
521
        model.to(torch_device)

        torch.manual_seed(0)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(0)

        batch_size = 4
        num_channels = 3
        sizes = (32, 32)

Patrick von Platen's avatar
Patrick von Platen committed
522
523
        noise = torch.ones((batch_size, num_channels) + sizes).to(torch_device)
        time_step = torch.tensor(batch_size * [1e-4]).to(torch_device)
524
525

        with torch.no_grad():
526
            output = model(noise, time_step)["sample"]
527
528
529

        output_slice = output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
Patrick von Platen's avatar
Patrick von Platen committed
530
        expected_output_slice = torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256])
531
532
        # fmt: on

Patrick von Platen's avatar
Patrick von Platen committed
533
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
534
535


patil-suraj's avatar
patil-suraj committed
536
537
538
539
class VQModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = VQModel

    @property
Patrick von Platen's avatar
Patrick von Platen committed
540
    def dummy_input(self, sizes=(32, 32)):
patil-suraj's avatar
patil-suraj committed
541
542
543
544
545
        batch_size = 4
        num_channels = 3

        image = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)

546
        return {"sample": image}
patil-suraj's avatar
patil-suraj committed
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561

    @property
    def input_shape(self):
        return (3, 32, 32)

    @property
    def output_shape(self):
        return (3, 32, 32)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "ch": 64,
            "out_ch": 3,
            "num_res_blocks": 1,
            "in_channels": 3,
562
            "attn_resolutions": [],
patil-suraj's avatar
patil-suraj committed
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
            "resolution": 32,
            "z_channels": 3,
            "n_embed": 256,
            "embed_dim": 3,
            "sane_index_shape": False,
            "ch_mult": (1,),
            "double_z": False,
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict

    def test_forward_signature(self):
        pass

    def test_training(self):
        pass

    def test_from_pretrained_hub(self):
Patrick von Platen's avatar
Patrick von Platen committed
581
        model, loading_info = VQModel.from_pretrained("fusing/vqgan-dummy", output_loading_info=True)
patil-suraj's avatar
patil-suraj committed
582
        self.assertIsNotNone(model)
Patrick von Platen's avatar
upload  
Patrick von Platen committed
583
        self.assertEqual(len(loading_info["missing_keys"]), 0)
patil-suraj's avatar
patil-suraj committed
584
585
586
587
588
589
590

        model.to(torch_device)
        image = model(**self.dummy_input)

        assert image is not None, "Make sure output is not None"

    def test_output_pretrained(self):
Patrick von Platen's avatar
Patrick von Platen committed
591
        model = VQModel.from_pretrained("fusing/vqgan-dummy")
patil-suraj's avatar
patil-suraj committed
592
593
594
595
596
597
598
599
600
601
602
603
        model.eval()

        torch.manual_seed(0)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(0)

        image = torch.randn(1, model.config.in_channels, model.config.resolution, model.config.resolution)
        with torch.no_grad():
            output = model(image)

        output_slice = output[0, -1, -3:, -3:].flatten()
        # fmt: off
Patrick von Platen's avatar
up  
Patrick von Platen committed
604
        expected_output_slice = torch.tensor([-1.1321, 0.1056, 0.3505, -0.6461, -0.2014, 0.0419, -0.5763, -0.8462, -0.4218])
patil-suraj's avatar
patil-suraj committed
605
        # fmt: on
Patrick von Platen's avatar
up  
Patrick von Platen committed
606
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
patil-suraj's avatar
patil-suraj committed
607
608


Patrick von Platen's avatar
Patrick von Platen committed
609
class AutoencoderKLTests(ModelTesterMixin, unittest.TestCase):
patil-suraj's avatar
patil-suraj committed
610
611
612
613
614
615
616
617
618
619
    model_class = AutoencoderKL

    @property
    def dummy_input(self):
        batch_size = 4
        num_channels = 3
        sizes = (32, 32)

        image = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)

620
        return {"sample": image}
patil-suraj's avatar
patil-suraj committed
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635

    @property
    def input_shape(self):
        return (3, 32, 32)

    @property
    def output_shape(self):
        return (3, 32, 32)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "ch": 64,
            "ch_mult": (1,),
            "embed_dim": 4,
            "in_channels": 3,
636
            "attn_resolutions": [],
patil-suraj's avatar
patil-suraj committed
637
638
639
640
641
642
643
644
645
646
647
648
649
            "num_res_blocks": 1,
            "out_ch": 3,
            "resolution": 32,
            "z_channels": 4,
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict

    def test_forward_signature(self):
        pass

    def test_training(self):
        pass
patil-suraj's avatar
patil-suraj committed
650

patil-suraj's avatar
patil-suraj committed
651
    def test_from_pretrained_hub(self):
Patrick von Platen's avatar
Patrick von Platen committed
652
        model, loading_info = AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy", output_loading_info=True)
patil-suraj's avatar
patil-suraj committed
653
        self.assertIsNotNone(model)
Patrick von Platen's avatar
upload  
Patrick von Platen committed
654
        self.assertEqual(len(loading_info["missing_keys"]), 0)
patil-suraj's avatar
patil-suraj committed
655
656
657
658
659
660
661

        model.to(torch_device)
        image = model(**self.dummy_input)

        assert image is not None, "Make sure output is not None"

    def test_output_pretrained(self):
Patrick von Platen's avatar
Patrick von Platen committed
662
        model = AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy")
patil-suraj's avatar
patil-suraj committed
663
664
665
666
667
668
669
670
671
672
673
674
        model.eval()

        torch.manual_seed(0)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(0)

        image = torch.randn(1, model.config.in_channels, model.config.resolution, model.config.resolution)
        with torch.no_grad():
            output = model(image, sample_posterior=True)

        output_slice = output[0, -1, -3:, -3:].flatten()
        # fmt: off
Patrick von Platen's avatar
up  
Patrick von Platen committed
675
        expected_output_slice = torch.tensor([-0.0814, -0.0229, -0.1320, -0.4123, -0.0366, -0.3473, 0.0438, -0.1662, 0.1750])
patil-suraj's avatar
patil-suraj committed
676
        # fmt: on
Patrick von Platen's avatar
up  
Patrick von Platen committed
677
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
patil-suraj's avatar
patil-suraj committed
678
679


680
681
682
class PipelineTesterMixin(unittest.TestCase):
    def test_from_pretrained_save_pretrained(self):
        # 1. Load models
Patrick von Platen's avatar
Patrick von Platen committed
683
684
685
686
        model = UNet2DModel(
            block_out_channels=(32, 64),
            layers_per_block=2,
            sample_size=32,
Patrick von Platen's avatar
Patrick von Platen committed
687
688
            in_channels=3,
            out_channels=3,
Patrick von Platen's avatar
Patrick von Platen committed
689
690
            down_block_types=("DownBlock2D", "AttnDownBlock2D"),
            up_block_types=("AttnUpBlock2D", "UpBlock2D"),
691
        )
Patrick von Platen's avatar
Patrick von Platen committed
692
        schedular = DDPMScheduler(num_train_timesteps=10)
693

Patrick von Platen's avatar
Patrick von Platen committed
694
        ddpm = DDPMPipeline(model, schedular)
695
696
697

        with tempfile.TemporaryDirectory() as tmpdirname:
            ddpm.save_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
698
            new_ddpm = DDPMPipeline.from_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
699
700

        generator = torch.manual_seed(0)
701

anton-l's avatar
anton-l committed
702
        image = ddpm(generator=generator, output_type="numpy")["sample"]
Patrick von Platen's avatar
Patrick von Platen committed
703
        generator = generator.manual_seed(0)
anton-l's avatar
anton-l committed
704
        new_image = new_ddpm(generator=generator, output_type="numpy")["sample"]
705

anton-l's avatar
anton-l committed
706
        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"
707
708
709

    @slow
    def test_from_pretrained_hub(self):
Patrick von Platen's avatar
Patrick von Platen committed
710
        model_path = "google/ddpm-cifar10-32"
711

Patrick von Platen's avatar
Patrick von Platen committed
712
        ddpm = DDPMPipeline.from_pretrained(model_path)
713
714
        ddpm_from_hub = DiffusionPipeline.from_pretrained(model_path)

715
716
        ddpm.scheduler.num_timesteps = 10
        ddpm_from_hub.scheduler.num_timesteps = 10
717

Patrick von Platen's avatar
Patrick von Platen committed
718
        generator = torch.manual_seed(0)
719

anton-l's avatar
anton-l committed
720
        image = ddpm(generator=generator, output_type="numpy")["sample"]
Patrick von Platen's avatar
Patrick von Platen committed
721
        generator = generator.manual_seed(0)
anton-l's avatar
anton-l committed
722
        new_image = ddpm_from_hub(generator=generator, output_type="numpy")["sample"]
723

anton-l's avatar
anton-l committed
724
        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"
Patrick von Platen's avatar
Patrick von Platen committed
725

726
727
    @slow
    def test_output_format(self):
Patrick von Platen's avatar
Patrick von Platen committed
728
        model_path = "google/ddpm-cifar10-32"
729
730
731
732
733
734
735
736
737
738
739
740
741

        pipe = DDIMPipeline.from_pretrained(model_path)

        generator = torch.manual_seed(0)
        images = pipe(generator=generator, output_type="numpy")["sample"]
        assert images.shape == (1, 32, 32, 3)
        assert isinstance(images, np.ndarray)

        images = pipe(generator=generator, output_type="pil")["sample"]
        assert isinstance(images, list)
        assert len(images) == 1
        assert isinstance(images[0], PIL.Image.Image)

anton-l's avatar
anton-l committed
742
743
744
745
746
        # use PIL by default
        images = pipe(generator=generator)["sample"]
        assert isinstance(images, list)
        assert isinstance(images[0], PIL.Image.Image)

Patrick von Platen's avatar
Patrick von Platen committed
747
748
    @slow
    def test_ddpm_cifar10(self):
Patrick von Platen's avatar
Patrick von Platen committed
749
        model_id = "google/ddpm-cifar10-32"
Patrick von Platen's avatar
Patrick von Platen committed
750

Patrick von Platen's avatar
Patrick von Platen committed
751
        unet = UNet2DModel.from_pretrained(model_id)
752
753
        scheduler = DDPMScheduler.from_config(model_id)
        scheduler = scheduler.set_format("pt")
Patrick von Platen's avatar
Patrick von Platen committed
754

755
        ddpm = DDPMPipeline(unet=unet, scheduler=scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
756
757

        generator = torch.manual_seed(0)
anton-l's avatar
anton-l committed
758
        image = ddpm(generator=generator, output_type="numpy")["sample"]
Patrick von Platen's avatar
Patrick von Platen committed
759

760
        image_slice = image[0, -3:, -3:, -1]
Patrick von Platen's avatar
Patrick von Platen committed
761

762
763
764
        assert image.shape == (1, 32, 32, 3)
        expected_slice = np.array([0.41995, 0.35885, 0.19385, 0.38475, 0.3382, 0.2647, 0.41545, 0.3582, 0.33845])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
765
766
767

    @slow
    def test_ddim_lsun(self):
Patrick von Platen's avatar
Patrick von Platen committed
768
        model_id = "google/ddpm-ema-bedroom-256"
769

Patrick von Platen's avatar
Patrick von Platen committed
770
        unet = UNet2DModel.from_pretrained(model_id)
771
        scheduler = DDIMScheduler.from_config(model_id)
772

773
        ddpm = DDIMPipeline(unet=unet, scheduler=scheduler)
774
775

        generator = torch.manual_seed(0)
anton-l's avatar
anton-l committed
776
        image = ddpm(generator=generator, output_type="numpy")["sample"]
777

778
        image_slice = image[0, -3:, -3:, -1]
779

780
781
782
        assert image.shape == (1, 256, 256, 3)
        expected_slice = np.array([0.00605, 0.0201, 0.0344, 0.00235, 0.00185, 0.00025, 0.00215, 0.0, 0.00685])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
Patrick von Platen's avatar
Patrick von Platen committed
783
784
785

    @slow
    def test_ddim_cifar10(self):
Patrick von Platen's avatar
Patrick von Platen committed
786
        model_id = "google/ddpm-cifar10-32"
Patrick von Platen's avatar
Patrick von Platen committed
787

Patrick von Platen's avatar
Patrick von Platen committed
788
        unet = UNet2DModel.from_pretrained(model_id)
789
        scheduler = DDIMScheduler(tensor_format="pt")
Patrick von Platen's avatar
Patrick von Platen committed
790

791
        ddim = DDIMPipeline(unet=unet, scheduler=scheduler)
792
793

        generator = torch.manual_seed(0)
anton-l's avatar
anton-l committed
794
        image = ddim(generator=generator, eta=0.0, output_type="numpy")["sample"]
Patrick von Platen's avatar
Patrick von Platen committed
795

796
        image_slice = image[0, -3:, -3:, -1]
Patrick von Platen's avatar
Patrick von Platen committed
797

798
799
800
        assert image.shape == (1, 32, 32, 3)
        expected_slice = np.array([0.17235, 0.16175, 0.16005, 0.16255, 0.1497, 0.1513, 0.15045, 0.1442, 0.1453])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
patil-suraj's avatar
patil-suraj committed
801

Patrick von Platen's avatar
Patrick von Platen committed
802
803
    @slow
    def test_pndm_cifar10(self):
Patrick von Platen's avatar
Patrick von Platen committed
804
        model_id = "google/ddpm-cifar10-32"
Patrick von Platen's avatar
Patrick von Platen committed
805

Patrick von Platen's avatar
Patrick von Platen committed
806
        unet = UNet2DModel.from_pretrained(model_id)
807
        scheduler = PNDMScheduler(tensor_format="pt")
Patrick von Platen's avatar
Patrick von Platen committed
808

809
        pndm = PNDMPipeline(unet=unet, scheduler=scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
810
        generator = torch.manual_seed(0)
anton-l's avatar
anton-l committed
811
        image = pndm(generator=generator, output_type="numpy")["sample"]
Patrick von Platen's avatar
Patrick von Platen committed
812

813
        image_slice = image[0, -3:, -3:, -1]
Patrick von Platen's avatar
Patrick von Platen committed
814

815
816
817
        assert image.shape == (1, 32, 32, 3)
        expected_slice = np.array([0.1564, 0.14645, 0.1406, 0.14715, 0.12425, 0.14045, 0.13115, 0.12175, 0.125])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
Patrick von Platen's avatar
Patrick von Platen committed
818

patil-suraj's avatar
patil-suraj committed
819
820
    @slow
    def test_ldm_text2img(self):
Patrick von Platen's avatar
Patrick von Platen committed
821
        ldm = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")
patil-suraj's avatar
patil-suraj committed
822
823
824

        prompt = "A painting of a squirrel eating a burger"
        generator = torch.manual_seed(0)
anton-l's avatar
anton-l committed
825
826
827
        image = ldm([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=20, output_type="numpy")[
            "sample"
        ]
patil-suraj's avatar
patil-suraj committed
828

829
        image_slice = image[0, -3:, -3:, -1]
patil-suraj's avatar
patil-suraj committed
830

831
832
833
        assert image.shape == (1, 256, 256, 3)
        expected_slice = np.array([0.9256, 0.9340, 0.8933, 0.9361, 0.9113, 0.8727, 0.9122, 0.8745, 0.8099])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
834

patil-suraj's avatar
patil-suraj committed
835
836
    @slow
    def test_ldm_text2img_fast(self):
Patrick von Platen's avatar
Patrick von Platen committed
837
        ldm = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")
patil-suraj's avatar
patil-suraj committed
838
839
840

        prompt = "A painting of a squirrel eating a burger"
        generator = torch.manual_seed(0)
anton-l's avatar
anton-l committed
841
        image = ldm([prompt], generator=generator, num_inference_steps=1, output_type="numpy")["sample"]
patil-suraj's avatar
patil-suraj committed
842

843
        image_slice = image[0, -3:, -3:, -1]
patil-suraj's avatar
patil-suraj committed
844

845
846
847
        assert image.shape == (1, 256, 256, 3)
        expected_slice = np.array([0.3163, 0.8670, 0.6465, 0.1865, 0.6291, 0.5139, 0.2824, 0.3723, 0.4344])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
patil-suraj's avatar
patil-suraj committed
848

Patrick von Platen's avatar
Patrick von Platen committed
849
850
    @slow
    def test_score_sde_ve_pipeline(self):
Patrick von Platen's avatar
Patrick von Platen committed
851
852
        model_id = "google/ncsnpp-church-256"
        model = UNet2DModel.from_pretrained(model_id)
853

Patrick von Platen's avatar
Patrick von Platen committed
854
        scheduler = ScoreSdeVeScheduler.from_config(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
855

Patrick von Platen's avatar
Patrick von Platen committed
856
        sde_ve = ScoreSdeVePipeline(unet=model, scheduler=scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
857

858
        torch.manual_seed(0)
anton-l's avatar
anton-l committed
859
        image = sde_ve(num_inference_steps=300, output_type="numpy")["sample"]
Nathan Lambert's avatar
Nathan Lambert committed
860

861
        image_slice = image[0, -3:, -3:, -1]
Patrick von Platen's avatar
Patrick von Platen committed
862

863
864
865
        assert image.shape == (1, 256, 256, 3)
        expected_slice = np.array([0.64363, 0.5868, 0.3031, 0.2284, 0.7409, 0.3216, 0.25643, 0.6557, 0.2633])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
Patrick von Platen's avatar
Patrick von Platen committed
866

patil-suraj's avatar
patil-suraj committed
867
868
    @slow
    def test_ldm_uncond(self):
Patrick von Platen's avatar
Patrick von Platen committed
869
        ldm = LDMPipeline.from_pretrained("CompVis/ldm-celebahq-256")
patil-suraj's avatar
patil-suraj committed
870
871

        generator = torch.manual_seed(0)
anton-l's avatar
anton-l committed
872
        image = ldm(generator=generator, num_inference_steps=5, output_type="numpy")["sample"]
patil-suraj's avatar
patil-suraj committed
873

874
        image_slice = image[0, -3:, -3:, -1]
patil-suraj's avatar
patil-suraj committed
875

876
877
878
        assert image.shape == (1, 256, 256, 3)
        expected_slice = np.array([0.4399, 0.44975, 0.46825, 0.474, 0.4359, 0.4581, 0.45095, 0.4341, 0.4447])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920

    @slow
    def test_ddpm_ddim_equality(self):
        model_id = "google/ddpm-cifar10-32"

        unet = UNet2DModel.from_pretrained(model_id)
        ddpm_scheduler = DDPMScheduler(tensor_format="pt")
        ddim_scheduler = DDIMScheduler(tensor_format="pt")

        ddpm = DDPMPipeline(unet=unet, scheduler=ddpm_scheduler)
        ddim = DDIMPipeline(unet=unet, scheduler=ddim_scheduler)

        generator = torch.manual_seed(0)
        ddpm_image = ddpm(generator=generator, output_type="numpy")["sample"]

        generator = torch.manual_seed(0)
        ddim_image = ddim(generator=generator, num_inference_steps=1000, eta=1.0, output_type="numpy")["sample"]

        # the values aren't exactly equal, but the images look the same upon visual inspection
        assert np.abs(ddpm_image - ddim_image).max() < 1e-1

    @slow
    def test_ddpm_ddim_equality_batched(self):
        model_id = "google/ddpm-cifar10-32"

        unet = UNet2DModel.from_pretrained(model_id)
        ddpm_scheduler = DDPMScheduler(tensor_format="pt")
        ddim_scheduler = DDIMScheduler(tensor_format="pt")

        ddpm = DDPMPipeline(unet=unet, scheduler=ddpm_scheduler)
        ddim = DDIMPipeline(unet=unet, scheduler=ddim_scheduler)

        generator = torch.manual_seed(0)
        ddpm_images = ddpm(batch_size=2, generator=generator, output_type="numpy")["sample"]

        generator = torch.manual_seed(0)
        ddim_images = ddim(batch_size=2, generator=generator, num_inference_steps=1000, eta=1.0, output_type="numpy")[
            "sample"
        ]

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