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

17
18
19
20
21
import tempfile
import unittest

import torch

Patrick von Platen's avatar
Patrick von Platen committed
22
from diffusers import DDIM, DDPM, DDIMScheduler, DDPMScheduler, LatentDiffusion, UNetModel
23
from diffusers.configuration_utils import ConfigMixin
Patrick von Platen's avatar
Patrick von Platen committed
24
from diffusers.pipeline_utils import DiffusionPipeline
Patrick von Platen's avatar
Patrick von Platen committed
25
from diffusers.testing_utils import floats_tensor, slow, torch_device
26
27


Patrick von Platen's avatar
Patrick von Platen committed
28
torch.backends.cuda.matmul.allow_tf32 = False
29
30


31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
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],
            ):
                self.register(a=a, b=b, c=c, d=d, e=e)

        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

        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


69
class ModelTesterMixin(unittest.TestCase):
Patrick von Platen's avatar
Patrick von Platen committed
70
71
    @property
    def dummy_input(self):
Patrick von Platen's avatar
up  
Patrick von Platen committed
72
        batch_size = 4
Patrick von Platen's avatar
Patrick von Platen committed
73
74
75
        num_channels = 3
        sizes = (32, 32)

Patrick von Platen's avatar
Patrick von Platen committed
76
        noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
77
        time_step = torch.tensor([10]).to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
78
79
80

        return (noise, time_step)

81
    def test_from_pretrained_save_pretrained(self):
Patrick von Platen's avatar
improve  
Patrick von Platen committed
82
        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
83
        model.to(torch_device)
84
85
86
87

        with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_pretrained(tmpdirname)
            new_model = UNetModel.from_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
88
            new_model.to(torch_device)
89

Patrick von Platen's avatar
Patrick von Platen committed
90
        dummy_input = self.dummy_input
91

Patrick von Platen's avatar
Patrick von Platen committed
92
93
        image = model(*dummy_input)
        new_image = new_model(*dummy_input)
94
95

        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
96
97
98

    def test_from_pretrained_hub(self):
        model = UNetModel.from_pretrained("fusing/ddpm_dummy")
Patrick von Platen's avatar
Patrick von Platen committed
99
        model.to(torch_device)
Patrick von Platen's avatar
Patrick von Platen committed
100
101
102
103

        image = model(*self.dummy_input)

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


106
107
108
109
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
110
        schedular = DDPMScheduler(timesteps=10)
111
112
113
114
115
116

        ddpm = DDPM(model, schedular)

        with tempfile.TemporaryDirectory() as tmpdirname:
            ddpm.save_pretrained(tmpdirname)
            new_ddpm = DDPM.from_pretrained(tmpdirname)
Patrick von Platen's avatar
Patrick von Platen committed
117
118

        generator = torch.manual_seed(0)
119

patil-suraj's avatar
patil-suraj committed
120
        image = ddpm(generator=generator)
Patrick von Platen's avatar
Patrick von Platen committed
121
        generator = generator.manual_seed(0)
patil-suraj's avatar
patil-suraj committed
122
        new_image = new_ddpm(generator=generator)
123
124
125
126
127
128
129
130
131
132
133
134
135

        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"

        ddpm = DDPM.from_pretrained(model_path)
        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
136
        generator = torch.manual_seed(0)
137

patil-suraj's avatar
patil-suraj committed
138
        image = ddpm(generator=generator)
Patrick von Platen's avatar
Patrick von Platen committed
139
        generator = generator.manual_seed(0)
patil-suraj's avatar
patil-suraj committed
140
        new_image = ddpm_from_hub(generator=generator)
141
142

        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
143
144
145
146
147
148

    @slow
    def test_ddpm_cifar10(self):
        generator = torch.manual_seed(0)
        model_id = "fusing/ddpm-cifar10"

Patrick von Platen's avatar
Patrick von Platen committed
149
        unet = UNetModel.from_pretrained(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
150
        noise_scheduler = DDPMScheduler.from_config(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
151
        noise_scheduler = noise_scheduler.set_format("pt")
Patrick von Platen's avatar
Patrick von Platen committed
152
153

        ddpm = DDPM(unet=unet, noise_scheduler=noise_scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
154
155
156
157
158
159
160
161
162
163
164
165
166
        image = ddpm(generator=generator)

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

        assert image.shape == (1, 3, 32, 32)
        expected_slice = torch.tensor([0.2250, 0.3375, 0.2360, 0.0930, 0.3440, 0.3156, 0.1937, 0.3585, 0.1761])
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2

    @slow
    def test_ddim_cifar10(self):
        generator = torch.manual_seed(0)
        model_id = "fusing/ddpm-cifar10"

Patrick von Platen's avatar
Patrick von Platen committed
167
        unet = UNetModel.from_pretrained(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
168
        noise_scheduler = DDIMScheduler(tensor_format="pt")
Patrick von Platen's avatar
Patrick von Platen committed
169
170

        ddim = DDIM(unet=unet, noise_scheduler=noise_scheduler)
Patrick von Platen's avatar
Patrick von Platen committed
171
172
173
174
175
        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
176
177
178
        expected_slice = torch.tensor(
            [-0.7383, -0.7385, -0.7298, -0.7364, -0.7414, -0.7239, -0.6737, -0.6813, -0.7068]
        )
Patrick von Platen's avatar
Patrick von Platen committed
179
        assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2
patil-suraj's avatar
patil-suraj committed
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194

    @slow
    def test_ldm_text2img(self):
        model_id = "fusing/latent-diffusion-text2im-large"
        ldm = LatentDiffusion.from_pretrained(model_id)

        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()
        print(image_slice.shape)

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