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

Patrick von Platen's avatar
Patrick von Platen committed
16

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

21
import numpy as np
22
23
import torch

Patrick von Platen's avatar
Patrick von Platen committed
24
from diffusers import (
patil-suraj's avatar
patil-suraj committed
25
    AutoencoderKL,
Patrick von Platen's avatar
Patrick von Platen committed
26
27
    BDDMPipeline,
    DDIMPipeline,
28
    DDIMScheduler,
Patrick von Platen's avatar
Patrick von Platen committed
29
    DDPMPipeline,
30
    DDPMScheduler,
Patrick von Platen's avatar
Patrick von Platen committed
31
    GlidePipeline,
Patrick von Platen's avatar
Patrick von Platen committed
32
33
    GlideSuperResUNetModel,
    GlideTextToImageUNetModel,
Patrick von Platen's avatar
Patrick von Platen committed
34
    GradTTSPipeline,
35
    GradTTSScheduler,
Patrick von Platen's avatar
Patrick von Platen committed
36
    LatentDiffusionPipeline,
patil-suraj's avatar
patil-suraj committed
37
    LatentDiffusionUncondPipeline,
Patrick von Platen's avatar
Patrick von Platen committed
38
    NCSNpp,
Patrick von Platen's avatar
Patrick von Platen committed
39
    PNDMPipeline,
40
    PNDMScheduler,
Patrick von Platen's avatar
Patrick von Platen committed
41
42
    ScoreSdeVePipeline,
    ScoreSdeVeScheduler,
Patrick von Platen's avatar
Patrick von Platen committed
43
44
    ScoreSdeVpPipeline,
    ScoreSdeVpScheduler,
45
    TemporalUNet,
patil-suraj's avatar
patil-suraj committed
46
    UNetGradTTSModel,
anton-l's avatar
anton-l committed
47
48
    UNetLDMModel,
    UNetModel,
49
    UNetUnconditionalModel,
patil-suraj's avatar
patil-suraj committed
50
    VQModel,
51
)
52
from diffusers.configuration_utils import ConfigMixin
Patrick von Platen's avatar
Patrick von Platen committed
53
from diffusers.pipeline_utils import DiffusionPipeline
54
from diffusers.pipelines.bddm.pipeline_bddm import DiffWave
Patrick von Platen's avatar
Patrick von Platen committed
55
from diffusers.testing_utils import floats_tensor, slow, torch_device
56
from diffusers.training_utils import EMAModel
57
58


Patrick von Platen's avatar
Patrick von Platen committed
59
torch.backends.cuda.matmul.allow_tf32 = False
60
61


62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
class ConfigTester(unittest.TestCase):
    def test_load_not_from_mixin(self):
        with self.assertRaises(ValueError):
            ConfigMixin.from_config("dummy_path")

    def test_save_load(self):
        class SampleObject(ConfigMixin):
            config_name = "config.json"

            def __init__(
                self,
                a=2,
                b=5,
                c=(2, 5),
                d="for diffusion",
                e=[1, 3],
            ):
79
                self.register_to_config(a=a, b=b, c=c, d=d, e=e)
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

        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
95
96
97
98
        # unfreeze configs
        config = dict(config)
        new_config = dict(new_config)

99
100
101
102
103
        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
104
class ModelTesterMixin:
105
    def test_from_pretrained_save_pretrained(self):
patil-suraj's avatar
patil-suraj committed
106
107
108
        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
109
        model.to(torch_device)
patil-suraj's avatar
patil-suraj committed
110
        model.eval()
111
112
113

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

patil-suraj's avatar
patil-suraj committed
117
118
119
        with torch.no_grad():
            image = model(**inputs_dict)
            new_image = new_model(**inputs_dict)
120

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

patil-suraj's avatar
patil-suraj committed
124
    def test_determinism(self):
patil-suraj's avatar
patil-suraj committed
125
126
127
128
129
130
131
132
133
134
135
136
137
138
        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)
            second = model(**inputs_dict)

        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)
139

patil-suraj's avatar
patil-suraj committed
140
    def test_output(self):
patil-suraj's avatar
patil-suraj committed
141
142
143
144
145
146
147
        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)
148

patil-suraj's avatar
patil-suraj committed
149
        self.assertIsNotNone(output)
150
        expected_shape = inputs_dict["sample"].shape
patil-suraj's avatar
patil-suraj committed
151
        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
152

patil-suraj's avatar
patil-suraj committed
153
    def test_forward_signature(self):
patil-suraj's avatar
patil-suraj committed
154
155
156
157
158
159
160
        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()]

161
        expected_arg_names = ["sample", "timesteps"]
patil-suraj's avatar
patil-suraj committed
162
        self.assertListEqual(arg_names[:2], expected_arg_names)
163

patil-suraj's avatar
patil-suraj committed
164
    def test_model_from_config(self):
patil-suraj's avatar
patil-suraj committed
165
166
167
168
169
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

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

patil-suraj's avatar
patil-suraj committed
171
172
173
174
175
176
177
        # 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()
178

patil-suraj's avatar
patil-suraj committed
179
180
181
182
183
        # 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)
184

patil-suraj's avatar
patil-suraj committed
185
186
187
        with torch.no_grad():
            output_1 = model(**inputs_dict)
            output_2 = new_model(**inputs_dict)
188

patil-suraj's avatar
patil-suraj committed
189
        self.assertEqual(output_1.shape, output_2.shape)
patil-suraj's avatar
patil-suraj committed
190
191

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

patil-suraj's avatar
patil-suraj committed
194
195
196
197
        model = self.model_class(**init_dict)
        model.to(torch_device)
        model.train()
        output = model(**inputs_dict)
198
        noise = torch.randn((inputs_dict["sample"].shape[0],) + self.output_shape).to(torch_device)
patil-suraj's avatar
patil-suraj committed
199
200
        loss = torch.nn.functional.mse_loss(output, noise)
        loss.backward()
201

202
203
204
205
206
207
208
209
210
    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)
211
        noise = torch.randn((inputs_dict["sample"].shape[0],) + self.output_shape).to(torch_device)
212
213
214
215
        loss = torch.nn.functional.mse_loss(output, noise)
        loss.backward()
        ema_model.step(model)

patil-suraj's avatar
patil-suraj committed
216
217
218
219
220
221
222
223
224
225
226
227
228

class UnetModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = UNetModel

    @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)

229
        return {"sample": noise, "timesteps": time_step}
230

patil-suraj's avatar
patil-suraj committed
231
    @property
Patrick von Platen's avatar
Patrick von Platen committed
232
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
233
        return (3, 32, 32)
234

patil-suraj's avatar
patil-suraj committed
235
    @property
Patrick von Platen's avatar
Patrick von Platen committed
236
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
237
        return (3, 32, 32)
patil-suraj's avatar
patil-suraj committed
238
239
240
241
242
243
244
245
246
247
248

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "ch": 32,
            "ch_mult": (1, 2),
            "num_res_blocks": 2,
            "attn_resolutions": (16,),
            "resolution": 32,
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict
249

patil-suraj's avatar
patil-suraj committed
250
    def test_from_pretrained_hub(self):
patil-suraj's avatar
patil-suraj committed
251
252
253
        model, loading_info = UNetModel.from_pretrained("fusing/ddpm_dummy", output_loading_info=True)
        self.assertIsNotNone(model)
        self.assertEqual(len(loading_info["missing_keys"]), 0)
patil-suraj's avatar
patil-suraj committed
254

patil-suraj's avatar
patil-suraj committed
255
        model.to(torch_device)
patil-suraj's avatar
patil-suraj committed
256
257
258
        image = model(**self.dummy_input)

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

patil-suraj's avatar
patil-suraj committed
260
261
262
263
264
265
266
    def test_output_pretrained(self):
        model = UNetModel.from_pretrained("fusing/ddpm_dummy")
        model.eval()

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

patil-suraj's avatar
patil-suraj committed
268
269
        noise = torch.randn(1, model.config.in_channels, model.config.resolution, model.config.resolution)
        time_step = torch.tensor([10])
270

patil-suraj's avatar
patil-suraj committed
271
272
        with torch.no_grad():
            output = model(noise, time_step)
273

patil-suraj's avatar
patil-suraj committed
274
275
        output_slice = output[0, -1, -3:, -3:].flatten()
        # fmt: off
Patrick von Platen's avatar
Patrick von Platen committed
276
        expected_output_slice = torch.tensor([0.2891, -0.1899, 0.2595, -0.6214, 0.0968, -0.2622, 0.4688, 0.1311, 0.0053])
patil-suraj's avatar
patil-suraj committed
277
        # fmt: on
Patrick von Platen's avatar
Patrick von Platen committed
278
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
patil-suraj's avatar
patil-suraj committed
279

280

Patrick von Platen's avatar
Patrick von Platen committed
281
282
class GlideSuperResUNetTests(ModelTesterMixin, unittest.TestCase):
    model_class = GlideSuperResUNetModel
patil-suraj's avatar
patil-suraj committed
283
284
285
286
287
288
289
290
291
292
293
294

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

        noise = torch.randn((batch_size, num_channels // 2) + sizes).to(torch_device)
        low_res = torch.randn((batch_size, 3) + low_res_size).to(torch_device)
        time_step = torch.tensor([10] * noise.shape[0], device=torch_device)

295
        return {"sample": noise, "timesteps": time_step, "low_res": low_res}
296

patil-suraj's avatar
patil-suraj committed
297
    @property
Patrick von Platen's avatar
Patrick von Platen committed
298
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
299
        return (3, 32, 32)
300

patil-suraj's avatar
patil-suraj committed
301
    @property
Patrick von Platen's avatar
Patrick von Platen committed
302
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
303
        return (6, 32, 32)
304

patil-suraj's avatar
patil-suraj committed
305
306
307
    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "attention_resolutions": (2,),
308
            "channel_mult": (1, 2),
patil-suraj's avatar
patil-suraj committed
309
310
311
312
313
314
315
316
            "in_channels": 6,
            "out_channels": 6,
            "model_channels": 32,
            "num_head_channels": 8,
            "num_heads_upsample": 1,
            "num_res_blocks": 2,
            "resblock_updown": True,
            "resolution": 32,
317
            "use_scale_shift_norm": True,
patil-suraj's avatar
patil-suraj committed
318
319
320
321
322
323
324
325
326
327
328
329
330
331
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict

    def test_output(self):
        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)

        output, _ = torch.split(output, 3, dim=1)
332

patil-suraj's avatar
patil-suraj committed
333
        self.assertIsNotNone(output)
334
        expected_shape = inputs_dict["sample"].shape
patil-suraj's avatar
patil-suraj committed
335
        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
336

patil-suraj's avatar
patil-suraj committed
337
    def test_from_pretrained_hub(self):
Patrick von Platen's avatar
Patrick von Platen committed
338
        model, loading_info = GlideSuperResUNetModel.from_pretrained(
339
340
            "fusing/glide-super-res-dummy", output_loading_info=True
        )
patil-suraj's avatar
patil-suraj committed
341
342
343
344
345
346
347
        self.assertIsNotNone(model)
        self.assertEqual(len(loading_info["missing_keys"]), 0)

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

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

patil-suraj's avatar
patil-suraj committed
349
    def test_output_pretrained(self):
Patrick von Platen's avatar
Patrick von Platen committed
350
        model = GlideSuperResUNetModel.from_pretrained("fusing/glide-super-res-dummy")
patil-suraj's avatar
patil-suraj committed
351
352
353
354

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

356
        noise = torch.randn(1, 3, 64, 64)
patil-suraj's avatar
patil-suraj committed
357
358
        low_res = torch.randn(1, 3, 4, 4)
        time_step = torch.tensor([42] * noise.shape[0])
359

patil-suraj's avatar
patil-suraj committed
360
361
        with torch.no_grad():
            output = model(noise, time_step, low_res)
362

patil-suraj's avatar
patil-suraj committed
363
364
365
        output, _ = torch.split(output, 3, dim=1)
        output_slice = output[0, -1, -3:, -3:].flatten()
        # fmt: off
366
        expected_output_slice = torch.tensor([-22.8782, -23.2652, -15.3966, -22.8034, -23.3159, -15.5640, -15.3970, -15.4614, - 10.4370])
patil-suraj's avatar
patil-suraj committed
367
368
        # fmt: on
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
patil-suraj's avatar
patil-suraj committed
369

anton-l's avatar
anton-l committed
370

Patrick von Platen's avatar
Patrick von Platen committed
371
372
class GlideTextToImageUNetModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = GlideTextToImageUNetModel
373
374
375
376
377
378
379
380
381
382
383
384
385

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

        noise = torch.randn((batch_size, num_channels) + sizes).to(torch_device)
        emb = torch.randn((batch_size, seq_len, transformer_dim)).to(torch_device)
        time_step = torch.tensor([10] * noise.shape[0], device=torch_device)

386
        return {"sample": noise, "timesteps": time_step, "transformer_out": emb}
387
388

    @property
Patrick von Platen's avatar
Patrick von Platen committed
389
    def input_shape(self):
390
391
392
        return (3, 32, 32)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
393
    def output_shape(self):
394
395
396
397
398
399
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
        return (6, 32, 32)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "attention_resolutions": (2,),
            "channel_mult": (1, 2),
            "in_channels": 3,
            "out_channels": 6,
            "model_channels": 32,
            "num_head_channels": 8,
            "num_heads_upsample": 1,
            "num_res_blocks": 2,
            "resblock_updown": True,
            "resolution": 32,
            "use_scale_shift_norm": True,
            "transformer_dim": 32,
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict

    def test_output(self):
        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)

        output, _ = torch.split(output, 3, dim=1)

        self.assertIsNotNone(output)
426
        expected_shape = inputs_dict["sample"].shape
427
428
429
        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")

    def test_from_pretrained_hub(self):
Patrick von Platen's avatar
Patrick von Platen committed
430
        model, loading_info = GlideTextToImageUNetModel.from_pretrained(
431
432
433
434
435
436
437
438
439
440
441
            "fusing/unet-glide-text2im-dummy", output_loading_info=True
        )
        self.assertIsNotNone(model)
        self.assertEqual(len(loading_info["missing_keys"]), 0)

        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
442
        model = GlideTextToImageUNetModel.from_pretrained("fusing/unet-glide-text2im-dummy")
443
444
445
446
447
448
449
450
451
452
453

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

        noise = torch.randn((1, model.config.in_channels, model.config.resolution, model.config.resolution)).to(
            torch_device
        )
        emb = torch.randn((1, 16, model.config.transformer_dim)).to(torch_device)
        time_step = torch.tensor([10] * noise.shape[0], device=torch_device)

Patrick von Platen's avatar
Patrick von Platen committed
454
        model.to(torch_device)
455
456
457
458
        with torch.no_grad():
            output = model(noise, time_step, emb)

        output, _ = torch.split(output, 3, dim=1)
Patrick von Platen's avatar
Patrick von Platen committed
459
        output_slice = output[0, -1, -3:, -3:].cpu().flatten()
460
        # fmt: off
Patrick von Platen's avatar
Patrick von Platen committed
461
        expected_output_slice = torch.tensor([2.7766, -10.3558, -14.9149, -0.9376, -14.9175, -17.7679, -5.5565, -12.9521, -12.9845])
462
463
464
465
        # fmt: on
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))


patil-suraj's avatar
patil-suraj committed
466
class UNetLDMModelTests(ModelTesterMixin, unittest.TestCase):
467
    model_class = UNetUnconditionalModel
patil-suraj's avatar
patil-suraj committed
468
469
470
471
472
473
474
475
476
477

    @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)

478
        return {"sample": noise, "timesteps": time_step}
patil-suraj's avatar
patil-suraj committed
479
480

    @property
Patrick von Platen's avatar
Patrick von Platen committed
481
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
482
483
484
        return (4, 32, 32)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
485
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
486
487
488
489
490
491
492
493
494
        return (4, 32, 32)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "image_size": 32,
            "in_channels": 4,
            "out_channels": 4,
            "num_res_blocks": 2,
            "attention_resolutions": (16,),
Patrick von Platen's avatar
Patrick von Platen committed
495
496
            "resnet_input_channels": [32, 32],
            "resnet_output_channels": [32, 64],
497
            "num_head_channels": 32,
patil-suraj's avatar
patil-suraj committed
498
499
500
501
            "conv_resample": True,
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict
anton-l's avatar
anton-l committed
502

patil-suraj's avatar
patil-suraj committed
503
    def test_from_pretrained_hub(self):
504
        model, loading_info = UNetUnconditionalModel.from_pretrained("fusing/unet-ldm-dummy", output_loading_info=True)
patil-suraj's avatar
patil-suraj committed
505
506
507
508
509
510
511
512
513
        self.assertIsNotNone(model)
        self.assertEqual(len(loading_info["missing_keys"]), 0)

        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):
514
        model = UNetUnconditionalModel.from_pretrained("fusing/unet-ldm-dummy")
patil-suraj's avatar
patil-suraj committed
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
        model.eval()

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

        noise = torch.randn(1, model.config.in_channels, model.config.image_size, model.config.image_size)
        time_step = torch.tensor([10] * noise.shape[0])

        with torch.no_grad():
            output = model(noise, time_step)

        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
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
    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)

        noise = torch.randn(1, model.config.in_channels, model.config.image_size, model.config.image_size)
        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))

patil-suraj's avatar
patil-suraj committed
556

patil-suraj's avatar
patil-suraj committed
557
558
559
560
561
562
563
564
565
566
567
568
569
570
class UNetGradTTSModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = UNetGradTTSModel

    @property
    def dummy_input(self):
        batch_size = 4
        num_features = 32
        seq_len = 16

        noise = floats_tensor((batch_size, num_features, seq_len)).to(torch_device)
        condition = floats_tensor((batch_size, num_features, seq_len)).to(torch_device)
        mask = floats_tensor((batch_size, 1, seq_len)).to(torch_device)
        time_step = torch.tensor([10] * batch_size).to(torch_device)

571
        return {"sample": noise, "timesteps": time_step, "mu": condition, "mask": mask}
patil-suraj's avatar
patil-suraj committed
572
573

    @property
Patrick von Platen's avatar
Patrick von Platen committed
574
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
575
576
577
        return (4, 32, 16)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
578
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
579
580
581
582
583
584
585
586
587
588
589
590
591
        return (4, 32, 16)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "dim": 64,
            "groups": 4,
            "dim_mults": (1, 2),
            "n_feats": 32,
            "pe_scale": 1000,
            "n_spks": 1,
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict
anton-l's avatar
anton-l committed
592

patil-suraj's avatar
patil-suraj committed
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
    def test_from_pretrained_hub(self):
        model, loading_info = UNetGradTTSModel.from_pretrained("fusing/unet-grad-tts-dummy", output_loading_info=True)
        self.assertIsNotNone(model)
        self.assertEqual(len(loading_info["missing_keys"]), 0)

        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):
        model = UNetGradTTSModel.from_pretrained("fusing/unet-grad-tts-dummy")
        model.eval()

        torch.manual_seed(0)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(0)
anton-l's avatar
anton-l committed
610

patil-suraj's avatar
patil-suraj committed
611
612
613
614
615
616
617
618
619
620
621
622
        num_features = model.config.n_feats
        seq_len = 16
        noise = torch.randn((1, num_features, seq_len))
        condition = torch.randn((1, num_features, seq_len))
        mask = torch.randn((1, 1, seq_len))
        time_step = torch.tensor([10])

        with torch.no_grad():
            output = model(noise, time_step, condition, mask)

        output_slice = output[0, -3:, -3:].flatten()
        # fmt: off
Patrick von Platen's avatar
Patrick von Platen committed
623
        expected_output_slice = torch.tensor([-0.0690, -0.0531, 0.0633, -0.0660, -0.0541, 0.0650, -0.0656, -0.0555, 0.0617])
patil-suraj's avatar
patil-suraj committed
624
625
        # fmt: on

Patrick von Platen's avatar
up  
Patrick von Platen committed
626
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-3))
627
628
629
630
631


class TemporalUNetModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = TemporalUNet

Patrick von Platen's avatar
Patrick von Platen committed
632
633
634
635
636
637
638
639
640
    @property
    def dummy_input(self):
        batch_size = 4
        num_features = 14
        seq_len = 16

        noise = floats_tensor((batch_size, seq_len, num_features)).to(torch_device)
        time_step = torch.tensor([10] * batch_size).to(torch_device)

641
        return {"sample": noise, "timesteps": time_step}
Patrick von Platen's avatar
Patrick von Platen committed
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663

    @property
    def input_shape(self):
        return (4, 16, 14)

    @property
    def output_shape(self):
        return (4, 16, 14)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "training_horizon": 128,
            "dim": 32,
            "dim_mults": [1, 4, 8],
            "predict_epsilon": False,
            "clip_denoised": True,
            "transition_dim": 14,
            "cond_dim": 3,
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict

664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
    def test_from_pretrained_hub(self):
        model, loading_info = TemporalUNet.from_pretrained(
            "fusing/ddpm-unet-rl-hopper-hor128", output_loading_info=True
        )
        self.assertIsNotNone(model)
        self.assertEqual(len(loading_info["missing_keys"]), 0)

        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):
        model = TemporalUNet.from_pretrained("fusing/ddpm-unet-rl-hopper-hor128")
        model.eval()

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

        num_features = model.transition_dim
        seq_len = 16
        noise = torch.randn((1, seq_len, num_features))
        time_step = torch.full((num_features,), 0)

        with torch.no_grad():
            output = model(noise, time_step)

        output_slice = output[0, -3:, -3:].flatten()
        # fmt: off
Patrick von Platen's avatar
Patrick von Platen committed
694
        expected_output_slice = torch.tensor([-0.2714, 0.1042, -0.0794, -0.2820, 0.0803, -0.0811, -0.2345, 0.0580, -0.0584])
695
696
        # fmt: on

Patrick von Platen's avatar
up  
Patrick von Platen committed
697
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-3))
patil-suraj's avatar
patil-suraj committed
698
699


700
701
702
703
704
705
706
707
708
709
710
711
class NCSNppModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = NCSNpp

    @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(batch_size * [10]).to(torch_device)

712
        return {"sample": noise, "timesteps": time_step}
713
714

    @property
Patrick von Platen's avatar
Patrick von Platen committed
715
    def input_shape(self):
716
717
718
        return (3, 32, 32)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
719
    def output_shape(self):
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
        return (3, 32, 32)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "image_size": 32,
            "ch_mult": [1, 2, 2, 2],
            "nf": 32,
            "fir": True,
            "progressive": "output_skip",
            "progressive_combine": "sum",
            "progressive_input": "input_skip",
            "scale_by_sigma": True,
            "skip_rescale": True,
            "embedding_type": "fourier",
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict

    def test_from_pretrained_hub(self):
        model, loading_info = NCSNpp.from_pretrained("fusing/cifar10-ncsnpp-ve", output_loading_info=True)
        self.assertIsNotNone(model)
        self.assertEqual(len(loading_info["missing_keys"]), 0)

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

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

    def test_output_pretrained_ve_small(self):
        model = NCSNpp.from_pretrained("fusing/ncsnpp-cifar10-ve-dummy")
        model.eval()
        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
761
762
        noise = torch.ones((batch_size, num_channels) + sizes).to(torch_device)
        time_step = torch.tensor(batch_size * [1e-4]).to(torch_device)
763
764
765
766
767
768

        with torch.no_grad():
            output = model(noise, time_step)

        output_slice = output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
Patrick von Platen's avatar
Patrick von Platen committed
769
        expected_output_slice = torch.tensor([0.1315, 0.0741, 0.0393, 0.0455, 0.0556, 0.0180, -0.0832, -0.0644, -0.0856])
770
771
        # fmt: on

Patrick von Platen's avatar
Patrick von Platen committed
772
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
773
774
775
776
777
778
779
780
781
782
783
784
785
786

    def test_output_pretrained_ve_large(self):
        model = NCSNpp.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy")
        model.eval()
        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
787
788
        noise = torch.ones((batch_size, num_channels) + sizes).to(torch_device)
        time_step = torch.tensor(batch_size * [1e-4]).to(torch_device)
789
790
791
792
793
794

        with torch.no_grad():
            output = model(noise, time_step)

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

Patrick von Platen's avatar
Patrick von Platen committed
798
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
799
800

    def test_output_pretrained_vp(self):
Patrick von Platen's avatar
Patrick von Platen committed
801
        model = NCSNpp.from_pretrained("fusing/cifar10-ddpmpp-vp")
802
803
804
805
806
807
808
809
810
811
812
        model.eval()
        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
813
        noise = torch.randn((batch_size, num_channels) + sizes).to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
814
        time_step = torch.tensor(batch_size * [9.0]).to(torch_device)
815
816
817
818
819
820

        with torch.no_grad():
            output = model(noise, time_step)

        output_slice = output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
Patrick von Platen's avatar
Patrick von Platen committed
821
        expected_output_slice = torch.tensor([0.3303, -0.2275, -2.8872, -0.1309, -1.2861, 3.4567, -1.0083, 2.5325, -1.3866])
822
823
        # fmt: on

Patrick von Platen's avatar
Patrick von Platen committed
824
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
825
826


patil-suraj's avatar
patil-suraj committed
827
828
829
830
831
832
833
834
835
836
837
class VQModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = VQModel

    @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)

838
        return {"sample": image}
patil-suraj's avatar
patil-suraj committed
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896

    @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,
            "attn_resolutions": [],
            "in_channels": 3,
            "resolution": 32,
            "z_channels": 3,
            "n_embed": 256,
            "embed_dim": 3,
            "sane_index_shape": False,
            "ch_mult": (1,),
            "dropout": 0.0,
            "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):
        model, loading_info = VQModel.from_pretrained("fusing/vqgan-dummy", output_loading_info=True)
        self.assertIsNotNone(model)
        self.assertEqual(len(loading_info["missing_keys"]), 0)

        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):
        model = VQModel.from_pretrained("fusing/vqgan-dummy")
        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
897
        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
898
        # fmt: on
Patrick von Platen's avatar
up  
Patrick von Platen committed
899
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
patil-suraj's avatar
patil-suraj committed
900
901


patil-suraj's avatar
patil-suraj committed
902
903
904
905
906
907
908
909
910
911
912
class AutoEncoderKLTests(ModelTesterMixin, unittest.TestCase):
    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)

913
        return {"sample": image}
patil-suraj's avatar
patil-suraj committed
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932

    @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,
            "num_res_blocks": 1,
            "out_ch": 3,
            "resolution": 32,
            "z_channels": 4,
patil-suraj's avatar
patil-suraj committed
933
            "attn_resolutions": [],
patil-suraj's avatar
patil-suraj committed
934
935
936
937
938
939
940
941
942
        }
        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
943

patil-suraj's avatar
patil-suraj committed
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
    def test_from_pretrained_hub(self):
        model, loading_info = AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy", output_loading_info=True)
        self.assertIsNotNone(model)
        self.assertEqual(len(loading_info["missing_keys"]), 0)

        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):
        model = AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy")
        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
968
        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
969
        # fmt: on
Patrick von Platen's avatar
up  
Patrick von Platen committed
970
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
patil-suraj's avatar
patil-suraj committed
971
972


973
974
975
976
class PipelineTesterMixin(unittest.TestCase):
    def test_from_pretrained_save_pretrained(self):
        # 1. Load models
        model = UNetModel(ch=32, ch_mult=(1, 2), num_res_blocks=2, attn_resolutions=(16,), resolution=32)
Patrick von Platen's avatar
Patrick von Platen committed
977
        schedular = DDPMScheduler(timesteps=10)
978

Patrick von Platen's avatar
Patrick von Platen committed
979
        ddpm = DDPMPipeline(model, schedular)
980
981
982

        with tempfile.TemporaryDirectory() as tmpdirname:
            ddpm.save_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
983
            new_ddpm = DDPMPipeline.from_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
984
985

        generator = torch.manual_seed(0)
986

patil-suraj's avatar
patil-suraj committed
987
        image = ddpm(generator=generator)
Patrick von Platen's avatar
Patrick von Platen committed
988
        generator = generator.manual_seed(0)
patil-suraj's avatar
patil-suraj committed
989
        new_image = new_ddpm(generator=generator)
990
991
992
993
994
995
996

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

    @slow
    def test_from_pretrained_hub(self):
        model_path = "fusing/ddpm-cifar10"

Patrick von Platen's avatar
Patrick von Platen committed
997
        ddpm = DDPMPipeline.from_pretrained(model_path)
998
999
1000
1001
1002
        ddpm_from_hub = DiffusionPipeline.from_pretrained(model_path)

        ddpm.noise_scheduler.num_timesteps = 10
        ddpm_from_hub.noise_scheduler.num_timesteps = 10

Patrick von Platen's avatar
Patrick von Platen committed
1003
        generator = torch.manual_seed(0)
1004

patil-suraj's avatar
patil-suraj committed
1005
        image = ddpm(generator=generator)
Patrick von Platen's avatar
Patrick von Platen committed
1006
        generator = generator.manual_seed(0)
patil-suraj's avatar
patil-suraj committed
1007
        new_image = ddpm_from_hub(generator=generator)
1008
1009

        assert (image - new_image).abs().sum() < 1e-5, "Models don't give the same forward pass"
Patrick von Platen's avatar
Patrick von Platen committed
1010
1011
1012
1013
1014

    @slow
    def test_ddpm_cifar10(self):
        model_id = "fusing/ddpm-cifar10"

Patrick von Platen's avatar
Patrick von Platen committed
1015
        unet = UNetModel.from_pretrained(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
1016
        noise_scheduler = DDPMScheduler.from_config(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
1017
        noise_scheduler = noise_scheduler.set_format("pt")
Patrick von Platen's avatar
Patrick von Platen committed
1018

Patrick von Platen's avatar
Patrick von Platen committed
1019
        ddpm = DDPMPipeline(unet=unet, noise_scheduler=noise_scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
1020
1021

        generator = torch.manual_seed(0)
Patrick von Platen's avatar
Patrick von Platen committed
1022
1023
1024
1025
1026
        image = ddpm(generator=generator)

        image_slice = image[0, -1, -3:, -3:].cpu()

        assert image.shape == (1, 3, 32, 32)
Patrick von Platen's avatar
Patrick von Platen committed
1027
1028
1029
        expected_slice = torch.tensor(
            [-0.5712, -0.6215, -0.5953, -0.5438, -0.4775, -0.4539, -0.5172, -0.4872, -0.5105]
        )
Patrick von Platen's avatar
Patrick von Platen committed
1030
1031
1032
1033
1034
1035
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

    @slow
    def test_ddim_cifar10(self):
        model_id = "fusing/ddpm-cifar10"

Patrick von Platen's avatar
Patrick von Platen committed
1036
        unet = UNetModel.from_pretrained(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
1037
        noise_scheduler = DDIMScheduler(tensor_format="pt")
Patrick von Platen's avatar
Patrick von Platen committed
1038

Patrick von Platen's avatar
Patrick von Platen committed
1039
        ddim = DDIMPipeline(unet=unet, noise_scheduler=noise_scheduler)
1040
1041

        generator = torch.manual_seed(0)
Patrick von Platen's avatar
Patrick von Platen committed
1042
1043
1044
1045
1046
        image = ddim(generator=generator, eta=0.0)

        image_slice = image[0, -1, -3:, -3:].cpu()

        assert image.shape == (1, 3, 32, 32)
Patrick von Platen's avatar
Patrick von Platen committed
1047
        expected_slice = torch.tensor(
1048
            [-0.6553, -0.6765, -0.6799, -0.6749, -0.7006, -0.6974, -0.6991, -0.7116, -0.7094]
Patrick von Platen's avatar
Patrick von Platen committed
1049
        )
Patrick von Platen's avatar
Patrick von Platen committed
1050
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2
patil-suraj's avatar
patil-suraj committed
1051

Patrick von Platen's avatar
Patrick von Platen committed
1052
1053
1054
1055
1056
1057
1058
    @slow
    def test_pndm_cifar10(self):
        model_id = "fusing/ddpm-cifar10"

        unet = UNetModel.from_pretrained(model_id)
        noise_scheduler = PNDMScheduler(tensor_format="pt")

Patrick von Platen's avatar
Patrick von Platen committed
1059
        pndm = PNDMPipeline(unet=unet, noise_scheduler=noise_scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
1060
        generator = torch.manual_seed(0)
Patrick von Platen's avatar
Patrick von Platen committed
1061
1062
1063
1064
1065
1066
        image = pndm(generator=generator)

        image_slice = image[0, -1, -3:, -3:].cpu()

        assert image.shape == (1, 3, 32, 32)
        expected_slice = torch.tensor(
Patrick von Platen's avatar
Patrick von Platen committed
1067
            [-0.6872, -0.7071, -0.7188, -0.7057, -0.7515, -0.7191, -0.7377, -0.7565, -0.7500]
Patrick von Platen's avatar
Patrick von Platen committed
1068
1069
1070
        )
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

patil-suraj's avatar
patil-suraj committed
1071
    @slow
patil-suraj's avatar
patil-suraj committed
1072
    @unittest.skip("Skipping for now as it takes too long")
patil-suraj's avatar
patil-suraj committed
1073
1074
    def test_ldm_text2img(self):
        model_id = "fusing/latent-diffusion-text2im-large"
Patrick von Platen's avatar
Patrick von Platen committed
1075
        ldm = LatentDiffusionPipeline.from_pretrained(model_id)
patil-suraj's avatar
patil-suraj committed
1076
1077
1078
1079
1080
1081
1082
1083
1084

        prompt = "A painting of a squirrel eating a burger"
        generator = torch.manual_seed(0)
        image = ldm([prompt], generator=generator, num_inference_steps=20)

        image_slice = image[0, -1, -3:, -3:].cpu()

        assert image.shape == (1, 3, 256, 256)
        expected_slice = torch.tensor([0.7295, 0.7358, 0.7256, 0.7435, 0.7095, 0.6884, 0.7325, 0.6921, 0.6458])
Patrick von Platen's avatar
update  
Patrick von Platen committed
1085
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2
1086

patil-suraj's avatar
patil-suraj committed
1087
1088
1089
1090
1091
1092
1093
    @slow
    def test_ldm_text2img_fast(self):
        model_id = "fusing/latent-diffusion-text2im-large"
        ldm = LatentDiffusionPipeline.from_pretrained(model_id)

        prompt = "A painting of a squirrel eating a burger"
        generator = torch.manual_seed(0)
1094
        image = ldm([prompt], generator=generator, num_inference_steps=1)
patil-suraj's avatar
patil-suraj committed
1095
1096
1097
1098

        image_slice = image[0, -1, -3:, -3:].cpu()

        assert image.shape == (1, 3, 256, 256)
patil-suraj's avatar
patil-suraj committed
1099
        expected_slice = torch.tensor([0.3163, 0.8670, 0.6465, 0.1865, 0.6291, 0.5139, 0.2824, 0.3723, 0.4344])
patil-suraj's avatar
patil-suraj committed
1100
1101
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

anton-l's avatar
anton-l committed
1102
1103
1104
    @slow
    def test_glide_text2img(self):
        model_id = "fusing/glide-base"
Patrick von Platen's avatar
Patrick von Platen committed
1105
        glide = GlidePipeline.from_pretrained(model_id)
anton-l's avatar
anton-l committed
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116

        prompt = "a pencil sketch of a corgi"
        generator = torch.manual_seed(0)
        image = glide(prompt, generator=generator, num_inference_steps_upscale=20)

        image_slice = image[0, :3, :3, -1].cpu()

        assert image.shape == (1, 256, 256, 3)
        expected_slice = torch.tensor([0.7119, 0.7073, 0.6460, 0.7780, 0.7423, 0.6926, 0.7378, 0.7189, 0.7784])
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

Patrick von Platen's avatar
Patrick von Platen committed
1117
1118
1119
    @slow
    def test_grad_tts(self):
        model_id = "fusing/grad-tts-libri-tts"
Patrick von Platen's avatar
Patrick von Platen committed
1120
        grad_tts = GradTTSPipeline.from_pretrained(model_id)
1121
1122
        noise_scheduler = GradTTSScheduler()
        grad_tts.noise_scheduler = noise_scheduler
Patrick von Platen's avatar
Patrick von Platen committed
1123
1124

        text = "Hello world, I missed you so much."
Patrick von Platen's avatar
Patrick von Platen committed
1125
        generator = torch.manual_seed(0)
Patrick von Platen's avatar
Patrick von Platen committed
1126
1127

        # generate mel spectograms using text
Patrick von Platen's avatar
Patrick von Platen committed
1128
        mel_spec = grad_tts(text, generator=generator)
Patrick von Platen's avatar
Patrick von Platen committed
1129

Patrick von Platen's avatar
Patrick von Platen committed
1130
1131
        assert mel_spec.shape == (1, 80, 143)
        expected_slice = torch.tensor(
Patrick von Platen's avatar
Patrick von Platen committed
1132
            [-6.7584, -6.8347, -6.3293, -6.6437, -6.7233, -6.4684, -6.1187, -6.3172, -6.6890]
Patrick von Platen's avatar
Patrick von Platen committed
1133
        )
Patrick von Platen's avatar
Patrick von Platen committed
1134
        assert (mel_spec[0, :3, :3].cpu().flatten() - expected_slice).abs().max() < 1e-2
Patrick von Platen's avatar
Patrick von Platen committed
1135

Patrick von Platen's avatar
Patrick von Platen committed
1136
1137
1138
1139
1140
1141
1142
    @slow
    def test_score_sde_ve_pipeline(self):
        model = NCSNpp.from_pretrained("fusing/ffhq_ncsnpp")
        scheduler = ScoreSdeVeScheduler.from_config("fusing/ffhq_ncsnpp")

        sde_ve = ScoreSdeVePipeline(model=model, scheduler=scheduler)

1143
        torch.manual_seed(0)
Patrick von Platen's avatar
Patrick von Platen committed
1144
1145
        image = sde_ve(num_inference_steps=2)

1146
1147
        expected_image_sum = 3382849024.0
        expected_image_mean = 1075.3788
Patrick von Platen's avatar
Patrick von Platen committed
1148
1149
1150
1151

        assert (image.abs().sum() - expected_image_sum).abs().cpu().item() < 1e-2
        assert (image.abs().mean() - expected_image_mean).abs().cpu().item() < 1e-4

Patrick von Platen's avatar
Patrick von Platen committed
1152
1153
    @slow
    def test_score_sde_vp_pipeline(self):
Patrick von Platen's avatar
Patrick von Platen committed
1154
1155
        model = NCSNpp.from_pretrained("fusing/cifar10-ddpmpp-vp")
        scheduler = ScoreSdeVpScheduler.from_config("fusing/cifar10-ddpmpp-vp")
Patrick von Platen's avatar
Patrick von Platen committed
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167

        sde_vp = ScoreSdeVpPipeline(model=model, scheduler=scheduler)

        torch.manual_seed(0)
        image = sde_vp(num_inference_steps=10)

        expected_image_sum = 4183.2012
        expected_image_mean = 1.3617

        assert (image.abs().sum() - expected_image_sum).abs().cpu().item() < 1e-2
        assert (image.abs().mean() - expected_image_mean).abs().cpu().item() < 1e-4

patil-suraj's avatar
patil-suraj committed
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
    @slow
    def test_ldm_uncond(self):
        ldm = LatentDiffusionUncondPipeline.from_pretrained("fusing/latent-diffusion-celeba-256")

        generator = torch.manual_seed(0)
        image = ldm(generator=generator, num_inference_steps=5)

        image_slice = image[0, -1, -3:, -3:].cpu()

        assert image.shape == (1, 3, 256, 256)
patil-suraj's avatar
patil-suraj committed
1178
1179
1180
        expected_slice = torch.tensor(
            [-0.1202, -0.1005, -0.0635, -0.0520, -0.1282, -0.0838, -0.0981, -0.1318, -0.1106]
        )
patil-suraj's avatar
patil-suraj committed
1181
1182
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

1183
1184
1185
1186
    def test_module_from_pipeline(self):
        model = DiffWave(num_res_layers=4)
        noise_scheduler = DDPMScheduler(timesteps=12)

Patrick von Platen's avatar
Patrick von Platen committed
1187
        bddm = BDDMPipeline(model, noise_scheduler)
1188
1189

        # check if the library name for the diffwave moduel is set to pipeline module
1190
        self.assertTrue(bddm.config["diffwave"][0] == "bddm")
1191
1192
1193
1194

        # check if we can save and load the pipeline
        with tempfile.TemporaryDirectory() as tmpdirname:
            bddm.save_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
1195
            _ = BDDMPipeline.from_pretrained(tmpdirname)
1196
            # check if the same works using the DifusionPipeline class
1197
1198
1199
            bddm = DiffusionPipeline.from_pretrained(tmpdirname)

        self.assertTrue(bddm.config["diffwave"][0] == "bddm")