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

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

21
import numpy as np
22
23
import torch

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

Suraj Patil's avatar
Suraj Patil committed
48
49
from ..src.diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipeline

50

Patrick von Platen's avatar
Patrick von Platen committed
51
torch.backends.cuda.matmul.allow_tf32 = False
52
53


54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class SampleObject(ConfigMixin):
    config_name = "config.json"

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


69
70
71
72
73
class ConfigTester(unittest.TestCase):
    def test_load_not_from_mixin(self):
        with self.assertRaises(ValueError):
            ConfigMixin.from_config("dummy_path")

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

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

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

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

    def test_save_load(self):
111
112
113
114
115
116
117
118
119
120
121
122
123
124
        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
125
126
127
128
        # unfreeze configs
        config = dict(config)
        new_config = dict(new_config)

129
130
131
132
133
        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
134
class ModelTesterMixin:
135
    def test_from_pretrained_save_pretrained(self):
patil-suraj's avatar
patil-suraj committed
136
137
138
        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
139
        model.to(torch_device)
patil-suraj's avatar
patil-suraj committed
140
        model.eval()
141
142
143

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

patil-suraj's avatar
patil-suraj committed
147
148
        with torch.no_grad():
            image = model(**inputs_dict)
Patrick von Platen's avatar
Patrick von Platen committed
149
150
151
            if isinstance(image, dict):
                image = image["sample"]

patil-suraj's avatar
patil-suraj committed
152
            new_image = new_model(**inputs_dict)
153

Patrick von Platen's avatar
Patrick von Platen committed
154
155
156
            if isinstance(new_image, dict):
                new_image = new_image["sample"]

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

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

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

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

patil-suraj's avatar
patil-suraj committed
181
    def test_output(self):
patil-suraj's avatar
patil-suraj committed
182
183
184
185
186
187
188
        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)
189

Patrick von Platen's avatar
Patrick von Platen committed
190
191
192
            if isinstance(output, dict):
                output = output["sample"]

patil-suraj's avatar
patil-suraj committed
193
        self.assertIsNotNone(output)
194
        expected_shape = inputs_dict["sample"].shape
patil-suraj's avatar
patil-suraj committed
195
        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
196

patil-suraj's avatar
patil-suraj committed
197
    def test_forward_signature(self):
patil-suraj's avatar
patil-suraj committed
198
199
200
201
202
203
204
        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()]

205
        expected_arg_names = ["sample", "timestep"]
patil-suraj's avatar
patil-suraj committed
206
        self.assertListEqual(arg_names[:2], expected_arg_names)
207

patil-suraj's avatar
patil-suraj committed
208
    def test_model_from_config(self):
patil-suraj's avatar
patil-suraj committed
209
210
211
212
213
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

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

patil-suraj's avatar
patil-suraj committed
215
216
217
218
219
220
221
        # 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()
222

patil-suraj's avatar
patil-suraj committed
223
224
225
226
227
        # 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)
228

patil-suraj's avatar
patil-suraj committed
229
230
        with torch.no_grad():
            output_1 = model(**inputs_dict)
Patrick von Platen's avatar
Patrick von Platen committed
231
232
233
234

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

patil-suraj's avatar
patil-suraj committed
235
            output_2 = new_model(**inputs_dict)
236

Patrick von Platen's avatar
Patrick von Platen committed
237
238
239
            if isinstance(output_2, dict):
                output_2 = output_2["sample"]

patil-suraj's avatar
patil-suraj committed
240
        self.assertEqual(output_1.shape, output_2.shape)
patil-suraj's avatar
patil-suraj committed
241
242

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

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

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

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

257
258
259
260
261
262
263
264
265
    def test_ema_training(self):
        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

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

        output = model(**inputs_dict)
Patrick von Platen's avatar
Patrick von Platen committed
266
267
268
269

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

270
        noise = torch.randn((inputs_dict["sample"].shape[0],) + self.output_shape).to(torch_device)
271
272
273
274
        loss = torch.nn.functional.mse_loss(output, noise)
        loss.backward()
        ema_model.step(model)

patil-suraj's avatar
patil-suraj committed
275
276

class UnetModelTests(ModelTesterMixin, unittest.TestCase):
Patrick von Platen's avatar
Patrick von Platen committed
277
    model_class = UNet2DModel
patil-suraj's avatar
patil-suraj committed
278
279
280
281
282
283
284
285
286
287

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

288
        return {"sample": noise, "timestep": time_step}
289

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

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

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

patil-suraj's avatar
patil-suraj committed
312

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


patil-suraj's avatar
patil-suraj committed
335
class UNetLDMModelTests(ModelTesterMixin, unittest.TestCase):
Patrick von Platen's avatar
Patrick von Platen committed
336
    model_class = UNet2DModel
patil-suraj's avatar
patil-suraj committed
337
338
339
340
341
342
343
344
345
346

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

347
        return {"sample": noise, "timestep": time_step}
patil-suraj's avatar
patil-suraj committed
348
349

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

    @property
Patrick von Platen's avatar
Patrick von Platen committed
354
    def output_shape(self):
patil-suraj's avatar
patil-suraj committed
355
356
357
358
        return (4, 32, 32)

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

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

patil-suraj's avatar
patil-suraj committed
374
        self.assertIsNotNone(model)
Patrick von Platen's avatar
upload  
Patrick von Platen committed
375
        self.assertEqual(len(loading_info["missing_keys"]), 0)
patil-suraj's avatar
patil-suraj committed
376
377

        model.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
378
        image = model(**self.dummy_input)["sample"]
patil-suraj's avatar
patil-suraj committed
379
380
381
382

        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
383
        model = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update")
patil-suraj's avatar
patil-suraj committed
384
385
386
387
388
389
        model.eval()

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

Patrick von Platen's avatar
Patrick von Platen committed
390
        noise = torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
patil-suraj's avatar
patil-suraj committed
391
392
393
        time_step = torch.tensor([10] * noise.shape[0])

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

        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
403

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

patil-suraj's avatar
patil-suraj committed
428

429
class NCSNppModelTests(ModelTesterMixin, unittest.TestCase):
Patrick von Platen's avatar
Patrick von Platen committed
430
    model_class = UNet2DModel
431
432

    @property
Patrick von Platen's avatar
Patrick von Platen committed
433
    def dummy_input(self, sizes=(32, 32)):
434
435
436
437
438
439
        batch_size = 4
        num_channels = 3

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

440
        return {"sample": noise, "timestep": time_step}
441
442

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

    @property
Patrick von Platen's avatar
Patrick von Platen committed
447
    def output_shape(self):
448
449
450
451
        return (3, 32, 32)

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

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

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

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

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

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

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

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

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

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

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

514
    def test_output_pretrained_ve_large(self):
Patrick von Platen's avatar
Patrick von Platen committed
515
        model = UNet2DModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update")
516
517
518
519
520
521
522
523
524
525
        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
526
527
        noise = torch.ones((batch_size, num_channels) + sizes).to(torch_device)
        time_step = torch.tensor(batch_size * [1e-4]).to(torch_device)
528
529

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

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

Patrick von Platen's avatar
Patrick von Platen committed
537
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
538
539


patil-suraj's avatar
patil-suraj committed
540
541
542
543
class VQModelTests(ModelTesterMixin, unittest.TestCase):
    model_class = VQModel

    @property
Patrick von Platen's avatar
Patrick von Platen committed
544
    def dummy_input(self, sizes=(32, 32)):
patil-suraj's avatar
patil-suraj committed
545
546
547
548
549
        batch_size = 4
        num_channels = 3

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

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

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

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

    def prepare_init_args_and_inputs_for_common(self):
        init_dict = {
562
            "block_out_channels": [32, 64],
patil-suraj's avatar
patil-suraj committed
563
            "in_channels": 3,
564
            "out_channels": 3,
565
566
            "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
            "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"],
567
            "latent_channels": 3,
patil-suraj's avatar
patil-suraj committed
568
569
570
571
572
573
574
575
576
577
578
        }
        inputs_dict = self.dummy_input
        return init_dict, inputs_dict

    def test_forward_signature(self):
        pass

    def test_training(self):
        pass

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

        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
589
        model = VQModel.from_pretrained("fusing/vqgan-dummy")
patil-suraj's avatar
patil-suraj committed
590
591
592
593
594
595
        model.eval()

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

596
        image = torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
patil-suraj's avatar
patil-suraj committed
597
598
599
600
601
        with torch.no_grad():
            output = model(image)

        output_slice = output[0, -1, -3:, -3:].flatten()
        # fmt: off
602
        expected_output_slice = torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143])
patil-suraj's avatar
patil-suraj committed
603
        # fmt: on
Patrick von Platen's avatar
up  
Patrick von Platen committed
604
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
patil-suraj's avatar
patil-suraj committed
605
606


Patrick von Platen's avatar
Patrick von Platen committed
607
class AutoencoderKLTests(ModelTesterMixin, unittest.TestCase):
patil-suraj's avatar
patil-suraj committed
608
609
610
611
612
613
614
615
616
617
    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)

618
        return {"sample": image}
patil-suraj's avatar
patil-suraj committed
619
620
621
622
623
624
625
626
627
628
629

    @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 = {
630
            "block_out_channels": [32, 64],
631
632
            "in_channels": 3,
            "out_channels": 3,
633
634
            "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
            "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"],
635
636
            "latent_channels": 4,
        }
patil-suraj's avatar
patil-suraj committed
637
638
639
640
641
642
643
644
        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
645

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

        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
657
        model = AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy")
patil-suraj's avatar
patil-suraj committed
658
659
660
661
662
663
        model.eval()

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

664
        image = torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
patil-suraj's avatar
patil-suraj committed
665
666
667
668
669
        with torch.no_grad():
            output = model(image, sample_posterior=True)

        output_slice = output[0, -1, -3:, -3:].flatten()
        # fmt: off
670
        expected_output_slice = torch.tensor([-4.0078e-01, -3.8304e-04, -1.2681e-01, -1.1462e-01,  2.0095e-01, 1.0893e-01, -8.8248e-02, -3.0361e-01, -9.8646e-03])
patil-suraj's avatar
patil-suraj committed
671
        # fmt: on
Patrick von Platen's avatar
up  
Patrick von Platen committed
672
        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
patil-suraj's avatar
patil-suraj committed
673
674


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

Patrick von Platen's avatar
Patrick von Platen committed
689
        ddpm = DDPMPipeline(model, schedular)
690
691
692

        with tempfile.TemporaryDirectory() as tmpdirname:
            ddpm.save_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
693
            new_ddpm = DDPMPipeline.from_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
694
695

        generator = torch.manual_seed(0)
696

anton-l's avatar
anton-l committed
697
        image = ddpm(generator=generator, output_type="numpy")["sample"]
Patrick von Platen's avatar
Patrick von Platen committed
698
        generator = generator.manual_seed(0)
anton-l's avatar
anton-l committed
699
        new_image = new_ddpm(generator=generator, output_type="numpy")["sample"]
700

anton-l's avatar
anton-l committed
701
        assert np.abs(image - new_image).sum() < 1e-5, "Models don't give the same forward pass"
702
703
704

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

Patrick von Platen's avatar
Patrick von Platen committed
707
        ddpm = DDPMPipeline.from_pretrained(model_path)
708
709
        ddpm_from_hub = DiffusionPipeline.from_pretrained(model_path)

710
711
        ddpm.scheduler.num_timesteps = 10
        ddpm_from_hub.scheduler.num_timesteps = 10
712

Patrick von Platen's avatar
Patrick von Platen committed
713
        generator = torch.manual_seed(0)
714

anton-l's avatar
anton-l committed
715
        image = ddpm(generator=generator, output_type="numpy")["sample"]
Patrick von Platen's avatar
Patrick von Platen committed
716
        generator = generator.manual_seed(0)
anton-l's avatar
anton-l committed
717
        new_image = ddpm_from_hub(generator=generator, output_type="numpy")["sample"]
718

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

721
722
    @slow
    def test_output_format(self):
Patrick von Platen's avatar
Patrick von Platen committed
723
        model_path = "google/ddpm-cifar10-32"
724
725
726
727
728
729
730
731
732
733
734
735
736

        pipe = DDIMPipeline.from_pretrained(model_path)

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

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

anton-l's avatar
anton-l committed
737
738
739
740
741
        # use PIL by default
        images = pipe(generator=generator)["sample"]
        assert isinstance(images, list)
        assert isinstance(images[0], PIL.Image.Image)

Patrick von Platen's avatar
Patrick von Platen committed
742
743
    @slow
    def test_ddpm_cifar10(self):
Patrick von Platen's avatar
Patrick von Platen committed
744
        model_id = "google/ddpm-cifar10-32"
Patrick von Platen's avatar
Patrick von Platen committed
745

Patrick von Platen's avatar
Patrick von Platen committed
746
        unet = UNet2DModel.from_pretrained(model_id)
747
748
        scheduler = DDPMScheduler.from_config(model_id)
        scheduler = scheduler.set_format("pt")
Patrick von Platen's avatar
Patrick von Platen committed
749

750
        ddpm = DDPMPipeline(unet=unet, scheduler=scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
751
752

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

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

757
758
759
        assert image.shape == (1, 32, 32, 3)
        expected_slice = np.array([0.41995, 0.35885, 0.19385, 0.38475, 0.3382, 0.2647, 0.41545, 0.3582, 0.33845])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
760
761
762

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

Patrick von Platen's avatar
Patrick von Platen committed
765
        unet = UNet2DModel.from_pretrained(model_id)
766
        scheduler = DDIMScheduler.from_config(model_id)
767

768
        ddpm = DDIMPipeline(unet=unet, scheduler=scheduler)
769
770

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

773
        image_slice = image[0, -3:, -3:, -1]
774

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

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

Patrick von Platen's avatar
Patrick von Platen committed
783
        unet = UNet2DModel.from_pretrained(model_id)
784
        scheduler = DDIMScheduler(tensor_format="pt")
Patrick von Platen's avatar
Patrick von Platen committed
785

786
        ddim = DDIMPipeline(unet=unet, scheduler=scheduler)
787
788

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

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

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

Patrick von Platen's avatar
Patrick von Platen committed
797
798
    @slow
    def test_pndm_cifar10(self):
Patrick von Platen's avatar
Patrick von Platen committed
799
        model_id = "google/ddpm-cifar10-32"
Patrick von Platen's avatar
Patrick von Platen committed
800

Patrick von Platen's avatar
Patrick von Platen committed
801
        unet = UNet2DModel.from_pretrained(model_id)
802
        scheduler = PNDMScheduler(tensor_format="pt")
Patrick von Platen's avatar
Patrick von Platen committed
803

804
        pndm = PNDMPipeline(unet=unet, scheduler=scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
805
        generator = torch.manual_seed(0)
anton-l's avatar
anton-l committed
806
        image = pndm(generator=generator, output_type="numpy")["sample"]
Patrick von Platen's avatar
Patrick von Platen committed
807

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

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

patil-suraj's avatar
patil-suraj committed
814
815
    @slow
    def test_ldm_text2img(self):
Patrick von Platen's avatar
Patrick von Platen committed
816
        ldm = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")
patil-suraj's avatar
patil-suraj committed
817
818
819

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

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

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

patil-suraj's avatar
patil-suraj committed
830
831
    @slow
    def test_ldm_text2img_fast(self):
Patrick von Platen's avatar
Patrick von Platen committed
832
        ldm = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")
patil-suraj's avatar
patil-suraj committed
833
834
835

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

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

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

Suraj Patil's avatar
Suraj Patil committed
844
845
    @slow
    def test_stable_diffusion(self):
846
        pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-1-diffusers")
Suraj Patil's avatar
Suraj Patil committed
847
848
849

        prompt = "A painting of a squirrel eating a burger"
        generator = torch.manual_seed(0)
850
        image = pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=20, output_type="numpy")[
Suraj Patil's avatar
Suraj Patil committed
851
852
853
854
855
856
            "sample"
        ]

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

        assert image.shape == (1, 512, 512, 3)
857
858
859
        # fmt: off
        expected_slice = np.array([0.09609553, 0.09020892, 0.07902172, 0.07634321, 0.08755809, 0.06491277, 0.07687345, 0.07173461, 0.07374045])
        # fmt: on
Suraj Patil's avatar
Suraj Patil committed
860
861
862
863
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2

    @slow
    def test_stable_diffusion_fast(self):
864
        pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-1-diffusers")
Suraj Patil's avatar
Suraj Patil committed
865
866
867

        prompt = "A painting of a squirrel eating a burger"
        generator = torch.manual_seed(0)
868
        image = pipe([prompt], generator=generator, num_inference_steps=5, output_type="numpy")["sample"]
Suraj Patil's avatar
Suraj Patil committed
869
870
871
872

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

        assert image.shape == (1, 512, 512, 3)
873
874
875
        # fmt: off
        expected_slice = np.array([0.16537648, 0.17572534, 0.14657784, 0.20084214, 0.19819549, 0.16032678, 0.30438453, 0.22730353, 0.21307352])
        # fmt: on
Suraj Patil's avatar
Suraj Patil committed
876
877
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2

Patrick von Platen's avatar
Patrick von Platen committed
878
879
    @slow
    def test_score_sde_ve_pipeline(self):
Patrick von Platen's avatar
Patrick von Platen committed
880
881
        model_id = "google/ncsnpp-church-256"
        model = UNet2DModel.from_pretrained(model_id)
882

Patrick von Platen's avatar
Patrick von Platen committed
883
        scheduler = ScoreSdeVeScheduler.from_config(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
884

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

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

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

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

patil-suraj's avatar
patil-suraj committed
896
897
    @slow
    def test_ldm_uncond(self):
Patrick von Platen's avatar
Patrick von Platen committed
898
        ldm = LDMPipeline.from_pretrained("CompVis/ldm-celebahq-256")
patil-suraj's avatar
patil-suraj committed
899
900

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

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

905
906
907
        assert image.shape == (1, 256, 256, 3)
        expected_slice = np.array([0.4399, 0.44975, 0.46825, 0.474, 0.4359, 0.4581, 0.45095, 0.4341, 0.4447])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925

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

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

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

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

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

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

929
    @unittest.skip("(Anton) The test is failing for large batch sizes, needs investigation")
930
931
932
933
934
935
936
937
938
939
940
    def test_ddpm_ddim_equality_batched(self):
        model_id = "google/ddpm-cifar10-32"

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

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

        generator = torch.manual_seed(0)
941
        ddpm_images = ddpm(batch_size=4, generator=generator, output_type="numpy")["sample"]
942
943

        generator = torch.manual_seed(0)
944
        ddim_images = ddim(batch_size=4, generator=generator, num_inference_steps=1000, eta=1.0, output_type="numpy")[
945
946
947
            "sample"
        ]

948
        # the values aren't exactly equal, but the images look the same visually
949
        assert np.abs(ddpm_images - ddim_images).max() < 1e-1
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965

    @slow
    def test_karras_ve_pipeline(self):
        model_id = "google/ncsnpp-celebahq-256"
        model = UNet2DModel.from_pretrained(model_id)
        scheduler = KarrasVeScheduler(tensor_format="pt")

        pipe = KarrasVePipeline(unet=model, scheduler=scheduler)

        generator = torch.manual_seed(0)
        image = pipe(num_inference_steps=20, generator=generator, output_type="numpy")["sample"]

        image_slice = image[0, -3:, -3:, -1]
        assert image.shape == (1, 256, 256, 3)
        expected_slice = np.array([0.26815, 0.1581, 0.2658, 0.23248, 0.1550, 0.2539, 0.1131, 0.1024, 0.0837])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2