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

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,
patil-suraj's avatar
patil-suraj committed
49
    VQModel,
50
)
51
from diffusers.configuration_utils import ConfigMixin
Patrick von Platen's avatar
Patrick von Platen committed
52
from diffusers.pipeline_utils import DiffusionPipeline
53
from diffusers.pipelines.bddm.pipeline_bddm import DiffWave
Patrick von Platen's avatar
Patrick von Platen committed
54
from diffusers.testing_utils import floats_tensor, slow, torch_device
55
56


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


60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
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],
            ):
77
                self.register_to_config(a=a, b=b, c=c, d=d, e=e)
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

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

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

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

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

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

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

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

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

patil-suraj's avatar
patil-suraj committed
151
    def test_forward_signature(self):
patil-suraj's avatar
patil-suraj committed
152
153
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()]

        expected_arg_names = ["x", "timesteps"]
        self.assertListEqual(arg_names[:2], expected_arg_names)
161

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

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

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

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

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

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

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

patil-suraj's avatar
patil-suraj committed
192
193
194
195
        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
196
        noise = torch.randn((inputs_dict["x"].shape[0],) + self.output_shape).to(torch_device)
patil-suraj's avatar
patil-suraj committed
197
198
        loss = torch.nn.functional.mse_loss(output, noise)
        loss.backward()
199

patil-suraj's avatar
patil-suraj committed
200
201
202
203
204
205
206
207
208
209
210
211
212

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)

patil-suraj's avatar
patil-suraj committed
213
        return {"x": noise, "timesteps": time_step}
214

patil-suraj's avatar
patil-suraj committed
215
    @property
Patrick von Platen's avatar
Patrick von Platen committed
216
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
217
        return (3, 32, 32)
218

patil-suraj's avatar
patil-suraj committed
219
    @property
Patrick von Platen's avatar
Patrick von Platen committed
220
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
221
        return (3, 32, 32)
patil-suraj's avatar
patil-suraj committed
222
223
224
225
226
227
228
229
230
231
232

    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
233

patil-suraj's avatar
patil-suraj committed
234
    def test_from_pretrained_hub(self):
patil-suraj's avatar
patil-suraj committed
235
236
237
        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
238

patil-suraj's avatar
patil-suraj committed
239
        model.to(torch_device)
patil-suraj's avatar
patil-suraj committed
240
241
242
        image = model(**self.dummy_input)

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

patil-suraj's avatar
patil-suraj committed
244
245
246
247
248
249
250
    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)
251

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

patil-suraj's avatar
patil-suraj committed
255
256
        with torch.no_grad():
            output = model(noise, time_step)
257

patil-suraj's avatar
patil-suraj committed
258
259
        output_slice = output[0, -1, -3:, -3:].flatten()
        # fmt: off
Patrick von Platen's avatar
Patrick von Platen committed
260
        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
261
262
263
        # fmt: on
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))

264

Patrick von Platen's avatar
Patrick von Platen committed
265
266
class GlideSuperResUNetTests(ModelTesterMixin, unittest.TestCase):
    model_class = GlideSuperResUNetModel
patil-suraj's avatar
patil-suraj committed
267
268
269
270
271
272
273
274
275
276
277
278
279

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

        return {"x": noise, "timesteps": time_step, "low_res": low_res}
280

patil-suraj's avatar
patil-suraj committed
281
    @property
Patrick von Platen's avatar
Patrick von Platen committed
282
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
283
        return (3, 32, 32)
284

patil-suraj's avatar
patil-suraj committed
285
    @property
Patrick von Platen's avatar
Patrick von Platen committed
286
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
287
        return (6, 32, 32)
288

patil-suraj's avatar
patil-suraj committed
289
290
291
    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "attention_resolutions": (2,),
292
            "channel_mult": (1, 2),
patil-suraj's avatar
patil-suraj committed
293
294
295
296
297
298
299
300
            "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,
301
            "use_scale_shift_norm": True,
patil-suraj's avatar
patil-suraj committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
        }
        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)
316

patil-suraj's avatar
patil-suraj committed
317
318
319
        self.assertIsNotNone(output)
        expected_shape = inputs_dict["x"].shape
        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
320

patil-suraj's avatar
patil-suraj committed
321
    def test_from_pretrained_hub(self):
Patrick von Platen's avatar
Patrick von Platen committed
322
        model, loading_info = GlideSuperResUNetModel.from_pretrained(
323
324
            "fusing/glide-super-res-dummy", output_loading_info=True
        )
patil-suraj's avatar
patil-suraj committed
325
326
327
328
329
330
331
        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"
332

patil-suraj's avatar
patil-suraj committed
333
    def test_output_pretrained(self):
Patrick von Platen's avatar
Patrick von Platen committed
334
        model = GlideSuperResUNetModel.from_pretrained("fusing/glide-super-res-dummy")
patil-suraj's avatar
patil-suraj committed
335
336
337
338

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

340
        noise = torch.randn(1, 3, 64, 64)
patil-suraj's avatar
patil-suraj committed
341
342
        low_res = torch.randn(1, 3, 4, 4)
        time_step = torch.tensor([42] * noise.shape[0])
343

patil-suraj's avatar
patil-suraj committed
344
345
        with torch.no_grad():
            output = model(noise, time_step, low_res)
346

patil-suraj's avatar
patil-suraj committed
347
348
349
        output, _ = torch.split(output, 3, dim=1)
        output_slice = output[0, -1, -3:, -3:].flatten()
        # fmt: off
350
        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
351
352
        # fmt: on
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
patil-suraj's avatar
patil-suraj committed
353

anton-l's avatar
anton-l committed
354

Patrick von Platen's avatar
Patrick von Platen committed
355
356
class GlideTextToImageUNetModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = GlideTextToImageUNetModel
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372

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

        return {"x": noise, "timesteps": time_step, "transformer_out": emb}

    @property
Patrick von Platen's avatar
Patrick von Platen committed
373
    def input_shape(self):
374
375
376
        return (3, 32, 32)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
377
    def output_shape(self):
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
        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)
        expected_shape = inputs_dict["x"].shape
        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
414
        model, loading_info = GlideTextToImageUNetModel.from_pretrained(
415
416
417
418
419
420
421
422
423
424
425
            "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
426
        model = GlideTextToImageUNetModel.from_pretrained("fusing/unet-glide-text2im-dummy")
427
428
429
430
431
432
433
434
435
436
437

        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
438
        model.to(torch_device)
439
440
441
442
        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
443
        output_slice = output[0, -1, -3:, -3:].cpu().flatten()
444
        # fmt: off
Patrick von Platen's avatar
Patrick von Platen committed
445
        expected_output_slice = torch.tensor([2.7766, -10.3558, -14.9149, -0.9376, -14.9175, -17.7679, -5.5565, -12.9521, -12.9845])
446
447
448
449
        # fmt: on
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))


patil-suraj's avatar
patil-suraj committed
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
class UNetLDMModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = UNetLDMModel

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

        return {"x": noise, "timesteps": time_step}

    @property
Patrick von Platen's avatar
Patrick von Platen committed
465
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
466
467
468
        return (4, 32, 32)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
469
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
        return (4, 32, 32)

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
            "image_size": 32,
            "in_channels": 4,
            "out_channels": 4,
            "model_channels": 32,
            "num_res_blocks": 2,
            "attention_resolutions": (16,),
            "channel_mult": (1, 2),
            "num_heads": 2,
            "conv_resample": True,
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict
anton-l's avatar
anton-l committed
486

patil-suraj's avatar
patil-suraj committed
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
    def test_from_pretrained_hub(self):
        model, loading_info = UNetLDMModel.from_pretrained("fusing/unet-ldm-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 = UNetLDMModel.from_pretrained("fusing/unet-ldm-dummy")
        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
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
    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
540

patil-suraj's avatar
patil-suraj committed
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
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)

        return {"x": noise, "timesteps": time_step, "mu": condition, "mask": mask}

    @property
Patrick von Platen's avatar
Patrick von Platen committed
558
    def input_shape(self):
patil-suraj's avatar
patil-suraj committed
559
560
561
        return (4, 32, 16)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
562
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
563
564
565
566
567
568
569
570
571
572
573
574
575
        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
576

patil-suraj's avatar
patil-suraj committed
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
    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
594

patil-suraj's avatar
patil-suraj committed
595
596
597
598
599
600
601
602
603
604
605
606
        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
607
        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
608
609
610
        # fmt: on

        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
611
612
613
614
615


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

Patrick von Platen's avatar
Patrick von Platen committed
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
    @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)

        return {"x": noise, "timesteps": time_step}

    @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

648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
    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
678
        expected_output_slice = torch.tensor([-0.2714, 0.1042, -0.0794, -0.2820, 0.0803, -0.0811, -0.2345, 0.0580, -0.0584])
679
680
681
        # fmt: on

        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
patil-suraj's avatar
patil-suraj committed
682
683


684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
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)

        return {"x": noise, "timesteps": time_step}

    @property
Patrick von Platen's avatar
Patrick von Platen committed
699
    def input_shape(self):
700
701
702
        return (3, 32, 32)

    @property
Patrick von Platen's avatar
Patrick von Platen committed
703
    def output_shape(self):
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
        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)

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

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

        output_slice = output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        expected_output_slice = torch.tensor([3.1909e-07, -8.5393e-08, 4.8460e-07, -4.5550e-07, -1.3205e-06, -6.3475e-07, 9.7837e-07, 2.9974e-07, 1.2345e-06])
        # fmt: on

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

    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)

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

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

        output_slice = output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        expected_output_slice = torch.tensor([-8.3299e-07, -9.0431e-07, 4.0585e-08, 9.7563e-07, 1.0280e-06, 1.0133e-06, 1.4979e-06, -2.9716e-07, -6.1817e-07])
        # fmt: on

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

    def test_output_pretrained_vp(self):
        model = NCSNpp.from_pretrained("fusing/ddpm-cifar10-vp-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)

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

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

        output_slice = output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        expected_output_slice = torch.tensor([-3.9086e-07, -1.1001e-05, 1.8881e-06, 1.1106e-05, 1.6629e-06, 2.9820e-06, 8.4978e-06, 8.0253e-07, 1.5435e-06])
        # fmt: on

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


patil-suraj's avatar
patil-suraj committed
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
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
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)

        return {"x": image}

    @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
        expected_output_slice = torch.tensor([-1.1321,  0.1056,  0.3505, -0.6461, -0.2014,  0.0419, -0.5763, -0.8462,
        -0.4218])
        # fmt: on
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))


patil-suraj's avatar
patil-suraj committed
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
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)

        return {"x": image}

    @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
918
            "attn_resolutions": [],
patil-suraj's avatar
patil-suraj committed
919
920
921
922
923
924
925
926
927
        }
        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
928

patil-suraj's avatar
patil-suraj committed
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
    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
        expected_output_slice = torch.tensor([-0.0814, -0.0229, -0.1320, -0.4123, -0.0366, -0.3473,  0.0438, -0.1662,
         0.1750])
        # fmt: on
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))


959
960
961
962
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
963
        schedular = DDPMScheduler(timesteps=10)
964

Patrick von Platen's avatar
Patrick von Platen committed
965
        ddpm = DDPMPipeline(model, schedular)
966
967
968

        with tempfile.TemporaryDirectory() as tmpdirname:
            ddpm.save_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
969
            new_ddpm = DDPMPipeline.from_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
970
971

        generator = torch.manual_seed(0)
972

patil-suraj's avatar
patil-suraj committed
973
        image = ddpm(generator=generator)
Patrick von Platen's avatar
Patrick von Platen committed
974
        generator = generator.manual_seed(0)
patil-suraj's avatar
patil-suraj committed
975
        new_image = new_ddpm(generator=generator)
976
977
978
979
980
981
982

        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
983
        ddpm = DDPMPipeline.from_pretrained(model_path)
984
985
986
987
988
        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
989
        generator = torch.manual_seed(0)
990

patil-suraj's avatar
patil-suraj committed
991
        image = ddpm(generator=generator)
Patrick von Platen's avatar
Patrick von Platen committed
992
        generator = generator.manual_seed(0)
patil-suraj's avatar
patil-suraj committed
993
        new_image = ddpm_from_hub(generator=generator)
994
995

        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
996
997
998
999
1000

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

Patrick von Platen's avatar
Patrick von Platen committed
1001
        unet = UNetModel.from_pretrained(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
1002
        noise_scheduler = DDPMScheduler.from_config(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
1003
        noise_scheduler = noise_scheduler.set_format("pt")
Patrick von Platen's avatar
Patrick von Platen committed
1004

Patrick von Platen's avatar
Patrick von Platen committed
1005
        ddpm = DDPMPipeline(unet=unet, noise_scheduler=noise_scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
1006
1007

        generator = torch.manual_seed(0)
Patrick von Platen's avatar
Patrick von Platen committed
1008
1009
1010
1011
1012
        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
1013
1014
1015
        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
1016
1017
1018
1019
1020
1021
        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
1022
        unet = UNetModel.from_pretrained(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
1023
        noise_scheduler = DDIMScheduler(tensor_format="pt")
Patrick von Platen's avatar
Patrick von Platen committed
1024

Patrick von Platen's avatar
Patrick von Platen committed
1025
        ddim = DDIMPipeline(unet=unet, noise_scheduler=noise_scheduler)
1026
1027

        generator = torch.manual_seed(0)
Patrick von Platen's avatar
Patrick von Platen committed
1028
1029
1030
1031
1032
        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
1033
        expected_slice = torch.tensor(
1034
            [-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
1035
        )
Patrick von Platen's avatar
Patrick von Platen committed
1036
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2
patil-suraj's avatar
patil-suraj committed
1037

Patrick von Platen's avatar
Patrick von Platen committed
1038
1039
1040
1041
1042
1043
1044
    @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
1045
        pndm = PNDMPipeline(unet=unet, noise_scheduler=noise_scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
1046
        generator = torch.manual_seed(0)
Patrick von Platen's avatar
Patrick von Platen committed
1047
1048
1049
1050
1051
1052
        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
1053
            [-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
1054
1055
1056
        )
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

patil-suraj's avatar
patil-suraj committed
1057
    @slow
patil-suraj's avatar
patil-suraj committed
1058
    @unittest.skip("Skipping for now as it takes too long")
patil-suraj's avatar
patil-suraj committed
1059
1060
    def test_ldm_text2img(self):
        model_id = "fusing/latent-diffusion-text2im-large"
Patrick von Platen's avatar
Patrick von Platen committed
1061
        ldm = LatentDiffusionPipeline.from_pretrained(model_id)
patil-suraj's avatar
patil-suraj committed
1062
1063
1064
1065
1066
1067
1068
1069
1070

        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
1071
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2
1072

patil-suraj's avatar
patil-suraj committed
1073
1074
1075
1076
1077
1078
1079
    @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)
1080
        image = ldm([prompt], generator=generator, num_inference_steps=1)
patil-suraj's avatar
patil-suraj committed
1081
1082
1083
1084

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

        assert image.shape == (1, 3, 256, 256)
patil-suraj's avatar
patil-suraj committed
1085
        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
1086
1087
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

anton-l's avatar
anton-l committed
1088
1089
1090
    @slow
    def test_glide_text2img(self):
        model_id = "fusing/glide-base"
Patrick von Platen's avatar
Patrick von Platen committed
1091
        glide = GlidePipeline.from_pretrained(model_id)
anton-l's avatar
anton-l committed
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102

        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
1103
1104
1105
    @slow
    def test_grad_tts(self):
        model_id = "fusing/grad-tts-libri-tts"
Patrick von Platen's avatar
Patrick von Platen committed
1106
        grad_tts = GradTTSPipeline.from_pretrained(model_id)
1107
1108
        noise_scheduler = GradTTSScheduler()
        grad_tts.noise_scheduler = noise_scheduler
Patrick von Platen's avatar
Patrick von Platen committed
1109
1110

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

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

Patrick von Platen's avatar
Patrick von Platen committed
1116
1117
        assert mel_spec.shape == (1, 80, 143)
        expected_slice = torch.tensor(
Patrick von Platen's avatar
Patrick von Platen committed
1118
            [-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
1119
        )
Patrick von Platen's avatar
Patrick von Platen committed
1120
        assert (mel_spec[0, :3, :3].cpu().flatten() - expected_slice).abs().max() < 1e-2
Patrick von Platen's avatar
Patrick von Platen committed
1121

Patrick von Platen's avatar
Patrick von Platen committed
1122
1123
1124
1125
1126
1127
1128
    @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)

1129
        torch.manual_seed(0)
Patrick von Platen's avatar
Patrick von Platen committed
1130
1131
        image = sde_ve(num_inference_steps=2)

1132
1133
        expected_image_sum = 3382849024.0
        expected_image_mean = 1075.3788
Patrick von Platen's avatar
Patrick von Platen committed
1134
1135
1136
1137

        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
1138
1139
    @slow
    def test_score_sde_vp_pipeline(self):
Patrick von Platen's avatar
Patrick von Platen committed
1140
1141
        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
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153

        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
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
    @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)
        expected_slice = torch.tensor([0.5025, 0.4121, 0.3851, 0.4806, 0.3996, 0.3745, 0.4839, 0.4559, 0.4293])
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

1167
1168
1169
1170
    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
1171
        bddm = BDDMPipeline(model, noise_scheduler)
1172
1173

        # check if the library name for the diffwave moduel is set to pipeline module
1174
        self.assertTrue(bddm.config["diffwave"][0] == "bddm")
1175
1176
1177
1178

        # 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
1179
            _ = BDDMPipeline.from_pretrained(tmpdirname)
1180
            # check if the same works using the DifusionPipeline class
1181
1182
1183
            bddm = DiffusionPipeline.from_pretrained(tmpdirname)

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