unet_2d.py 15.3 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
# Copyright 2023 The HuggingFace Team. All rights reserved.
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
14
15
from dataclasses import dataclass
from typing import Optional, Tuple, Union
Patrick von Platen's avatar
Patrick von Platen committed
16
17
18
19
20

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
21
from ..utils import BaseOutput
Patrick von Platen's avatar
Patrick von Platen committed
22
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
23
from .modeling_utils import ModelMixin
24
from .unet_2d_blocks import UNetMidBlock2D, get_down_block, get_up_block
Patrick von Platen's avatar
Patrick von Platen committed
25
26


27
28
29
@dataclass
class UNet2DOutput(BaseOutput):
    """
Steven Liu's avatar
Steven Liu committed
30
31
    The output of [`UNet2DModel`].

32
33
    Args:
        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
Steven Liu's avatar
Steven Liu committed
34
            The hidden states output from the last layer of the model.
35
36
37
38
39
    """

    sample: torch.FloatTensor


Patrick von Platen's avatar
Patrick von Platen committed
40
class UNet2DModel(ModelMixin, ConfigMixin):
Kashif Rasul's avatar
Kashif Rasul committed
41
    r"""
Steven Liu's avatar
Steven Liu committed
42
    A 2D UNet model that takes a noisy sample and a timestep and returns a sample shaped output.
Kashif Rasul's avatar
Kashif Rasul committed
43

Steven Liu's avatar
Steven Liu committed
44
45
    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
    for all models (such as downloading or saving).
Kashif Rasul's avatar
Kashif Rasul committed
46
47

    Parameters:
48
        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
George Ogden's avatar
George Ogden committed
49
50
            Height and width of input/output sample. Dimensions must be a multiple of `2 ** (len(block_out_channels) -
            1)`.
Steven Liu's avatar
Steven Liu committed
51
        in_channels (`int`, *optional*, defaults to 3): Number of channels in the input sample.
Kashif Rasul's avatar
Kashif Rasul committed
52
53
54
        out_channels (`int`, *optional*, defaults to 3): Number of channels in the output.
        center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
        time_embedding_type (`str`, *optional*, defaults to `"positional"`): Type of time embedding to use.
Steven Liu's avatar
Steven Liu committed
55
56
57
58
59
        freq_shift (`int`, *optional*, defaults to 0): Frequency shift for Fourier time embedding.
        flip_sin_to_cos (`bool`, *optional*, defaults to `True`):
            Whether to flip sin to cos for Fourier time embedding.
        down_block_types (`Tuple[str]`, *optional*, defaults to `("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")`):
            Tuple of downsample block types.
Will Berman's avatar
Will Berman committed
60
        mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2D"`):
Steven Liu's avatar
Steven Liu committed
61
62
63
64
65
            Block type for middle of UNet, it can be either `UNetMidBlock2D` or `UnCLIPUNetMidBlock2D`.
        up_block_types (`Tuple[str]`, *optional*, defaults to `("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")`):
            Tuple of upsample block types.
        block_out_channels (`Tuple[int]`, *optional*, defaults to `(224, 448, 672, 896)`):
            Tuple of block output channels.
Kashif Rasul's avatar
Kashif Rasul committed
66
67
68
        layers_per_block (`int`, *optional*, defaults to `2`): The number of layers per block.
        mid_block_scale_factor (`float`, *optional*, defaults to `1`): The scale factor for the mid block.
        downsample_padding (`int`, *optional*, defaults to `1`): The padding for the downsample convolution.
69
70
71
72
        downsample_type (`str`, *optional*, defaults to `conv`):
            The downsample type for downsampling layers. Choose between "conv" and "resnet"
        upsample_type (`str`, *optional*, defaults to `conv`):
            The upsample type for upsampling layers. Choose between "conv" and "resnet"
73
        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
Kashif Rasul's avatar
Kashif Rasul committed
74
75
        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
        attention_head_dim (`int`, *optional*, defaults to `8`): The attention head dimension.
Steven Liu's avatar
Steven Liu committed
76
77
        norm_num_groups (`int`, *optional*, defaults to `32`): The number of groups for normalization.
        norm_eps (`float`, *optional*, defaults to `1e-5`): The epsilon for normalization.
Will Berman's avatar
Will Berman committed
78
        resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
Steven Liu's avatar
Steven Liu committed
79
80
            for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
        class_embed_type (`str`, *optional*, defaults to `None`):
81
82
            The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
            `"timestep"`, or `"identity"`.
Steven Liu's avatar
Steven Liu committed
83
84
85
        num_class_embeds (`int`, *optional*, defaults to `None`):
            Input dimension of the learnable embedding matrix to be projected to `time_embed_dim` when performing class
            conditioning with `class_embed_type` equal to `None`.
Kashif Rasul's avatar
Kashif Rasul committed
86
87
    """

Patrick von Platen's avatar
Patrick von Platen committed
88
89
90
    @register_to_config
    def __init__(
        self,
91
        sample_size: Optional[Union[int, Tuple[int, int]]] = None,
Sid Sahai's avatar
Sid Sahai committed
92
93
94
95
96
97
98
99
100
101
102
103
        in_channels: int = 3,
        out_channels: int = 3,
        center_input_sample: bool = False,
        time_embedding_type: str = "positional",
        freq_shift: int = 0,
        flip_sin_to_cos: bool = True,
        down_block_types: Tuple[str] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),
        up_block_types: Tuple[str] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
        block_out_channels: Tuple[int] = (224, 448, 672, 896),
        layers_per_block: int = 2,
        mid_block_scale_factor: float = 1,
        downsample_padding: int = 1,
104
105
        downsample_type: str = "conv",
        upsample_type: str = "conv",
106
        dropout: float = 0.0,
Sid Sahai's avatar
Sid Sahai committed
107
        act_fn: str = "silu",
Will Berman's avatar
Will Berman committed
108
        attention_head_dim: Optional[int] = 8,
Sid Sahai's avatar
Sid Sahai committed
109
110
        norm_num_groups: int = 32,
        norm_eps: float = 1e-5,
Will Berman's avatar
Will Berman committed
111
112
        resnet_time_scale_shift: str = "default",
        add_attention: bool = True,
113
114
        class_embed_type: Optional[str] = None,
        num_class_embeds: Optional[int] = None,
Patrick von Platen's avatar
Patrick von Platen committed
115
116
117
118
119
120
    ):
        super().__init__()

        self.sample_size = sample_size
        time_embed_dim = block_out_channels[0] * 4

Will Berman's avatar
Will Berman committed
121
122
123
124
125
126
127
128
129
130
131
        # Check inputs
        if len(down_block_types) != len(up_block_types):
            raise ValueError(
                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
            )

        if len(block_out_channels) != len(down_block_types):
            raise ValueError(
                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
            )

Patrick von Platen's avatar
Patrick von Platen committed
132
133
134
135
136
137
138
139
140
141
142
143
144
        # input
        self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1))

        # time
        if time_embedding_type == "fourier":
            self.time_proj = GaussianFourierProjection(embedding_size=block_out_channels[0], scale=16)
            timestep_input_dim = 2 * block_out_channels[0]
        elif time_embedding_type == "positional":
            self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
            timestep_input_dim = block_out_channels[0]

        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)

145
146
147
148
149
150
151
152
153
154
        # class embedding
        if class_embed_type is None and num_class_embeds is not None:
            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
        elif class_embed_type == "timestep":
            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
        elif class_embed_type == "identity":
            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
        else:
            self.class_embedding = None

Patrick von Platen's avatar
Patrick von Platen committed
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
        self.down_blocks = nn.ModuleList([])
        self.mid_block = None
        self.up_blocks = nn.ModuleList([])

        # down
        output_channel = block_out_channels[0]
        for i, down_block_type in enumerate(down_block_types):
            input_channel = output_channel
            output_channel = block_out_channels[i]
            is_final_block = i == len(block_out_channels) - 1

            down_block = get_down_block(
                down_block_type,
                num_layers=layers_per_block,
                in_channels=input_channel,
                out_channels=output_channel,
                temb_channels=time_embed_dim,
                add_downsample=not is_final_block,
                resnet_eps=norm_eps,
                resnet_act_fn=act_fn,
175
                resnet_groups=norm_num_groups,
176
                attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel,
Patrick von Platen's avatar
Patrick von Platen committed
177
                downsample_padding=downsample_padding,
Will Berman's avatar
Will Berman committed
178
                resnet_time_scale_shift=resnet_time_scale_shift,
179
                downsample_type=downsample_type,
180
                dropout=dropout,
Patrick von Platen's avatar
Patrick von Platen committed
181
182
183
184
185
186
187
            )
            self.down_blocks.append(down_block)

        # mid
        self.mid_block = UNetMidBlock2D(
            in_channels=block_out_channels[-1],
            temb_channels=time_embed_dim,
188
            dropout=dropout,
Patrick von Platen's avatar
Patrick von Platen committed
189
190
191
            resnet_eps=norm_eps,
            resnet_act_fn=act_fn,
            output_scale_factor=mid_block_scale_factor,
Will Berman's avatar
Will Berman committed
192
            resnet_time_scale_shift=resnet_time_scale_shift,
193
            attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1],
Patrick von Platen's avatar
Patrick von Platen committed
194
            resnet_groups=norm_num_groups,
Will Berman's avatar
Will Berman committed
195
            add_attention=add_attention,
Patrick von Platen's avatar
Patrick von Platen committed
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
        )

        # up
        reversed_block_out_channels = list(reversed(block_out_channels))
        output_channel = reversed_block_out_channels[0]
        for i, up_block_type in enumerate(up_block_types):
            prev_output_channel = output_channel
            output_channel = reversed_block_out_channels[i]
            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]

            is_final_block = i == len(block_out_channels) - 1

            up_block = get_up_block(
                up_block_type,
                num_layers=layers_per_block + 1,
                in_channels=input_channel,
                out_channels=output_channel,
                prev_output_channel=prev_output_channel,
                temb_channels=time_embed_dim,
                add_upsample=not is_final_block,
                resnet_eps=norm_eps,
                resnet_act_fn=act_fn,
218
                resnet_groups=norm_num_groups,
219
                attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel,
Will Berman's avatar
Will Berman committed
220
                resnet_time_scale_shift=resnet_time_scale_shift,
221
                upsample_type=upsample_type,
222
                dropout=dropout,
Patrick von Platen's avatar
Patrick von Platen committed
223
224
225
226
227
228
229
230
            )
            self.up_blocks.append(up_block)
            prev_output_channel = output_channel

        # out
        num_groups_out = norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4, 32)
        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=num_groups_out, eps=norm_eps)
        self.conv_act = nn.SiLU()
231
        self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, padding=1)
Patrick von Platen's avatar
Patrick von Platen committed
232
233

    def forward(
234
235
236
        self,
        sample: torch.FloatTensor,
        timestep: Union[torch.Tensor, float, int],
237
        class_labels: Optional[torch.Tensor] = None,
238
239
        return_dict: bool = True,
    ) -> Union[UNet2DOutput, Tuple]:
240
        r"""
Steven Liu's avatar
Steven Liu committed
241
242
        The [`UNet2DModel`] forward method.

Kashif Rasul's avatar
Kashif Rasul committed
243
        Args:
Steven Liu's avatar
Steven Liu committed
244
245
246
            sample (`torch.FloatTensor`):
                The noisy input tensor with the following shape `(batch, channel, height, width)`.
            timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
247
248
            class_labels (`torch.FloatTensor`, *optional*, defaults to `None`):
                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
Kashif Rasul's avatar
Kashif Rasul committed
249
250
251
252
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~models.unet_2d.UNet2DOutput`] instead of a plain tuple.

        Returns:
Steven Liu's avatar
Steven Liu committed
253
254
255
            [`~models.unet_2d.UNet2DOutput`] or `tuple`:
                If `return_dict` is True, an [`~models.unet_2d.UNet2DOutput`] is returned, otherwise a `tuple` is
                returned where the first element is the sample tensor.
Kashif Rasul's avatar
Kashif Rasul committed
256
        """
Patrick von Platen's avatar
Patrick von Platen committed
257
258
259
260
261
262
263
264
265
266
267
        # 0. center input if necessary
        if self.config.center_input_sample:
            sample = 2 * sample - 1.0

        # 1. time
        timesteps = timestep
        if not torch.is_tensor(timesteps):
            timesteps = torch.tensor([timesteps], dtype=torch.long, device=sample.device)
        elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:
            timesteps = timesteps[None].to(sample.device)

268
269
        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
        timesteps = timesteps * torch.ones(sample.shape[0], dtype=timesteps.dtype, device=timesteps.device)
270

Patrick von Platen's avatar
Patrick von Platen committed
271
        t_emb = self.time_proj(timesteps)
272
273
274
275
276

        # timesteps does not contain any weights and will always return f32 tensors
        # but time_embedding might actually be running in fp16. so we need to cast here.
        # there might be better ways to encapsulate this.
        t_emb = t_emb.to(dtype=self.dtype)
Patrick von Platen's avatar
Patrick von Platen committed
277
278
        emb = self.time_embedding(t_emb)

279
280
281
282
283
284
285
286
287
288
        if self.class_embedding is not None:
            if class_labels is None:
                raise ValueError("class_labels should be provided when doing class conditioning")

            if self.config.class_embed_type == "timestep":
                class_labels = self.time_proj(class_labels)

            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
            emb = emb + class_emb

Patrick von Platen's avatar
Patrick von Platen committed
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        # 2. pre-process
        skip_sample = sample
        sample = self.conv_in(sample)

        # 3. down
        down_block_res_samples = (sample,)
        for downsample_block in self.down_blocks:
            if hasattr(downsample_block, "skip_conv"):
                sample, res_samples, skip_sample = downsample_block(
                    hidden_states=sample, temb=emb, skip_sample=skip_sample
                )
            else:
                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)

            down_block_res_samples += res_samples

        # 4. mid
        sample = self.mid_block(sample, emb)

        # 5. up
        skip_sample = None
        for upsample_block in self.up_blocks:
            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]

            if hasattr(upsample_block, "skip_conv"):
                sample, skip_sample = upsample_block(sample, res_samples, emb, skip_sample)
            else:
                sample = upsample_block(sample, res_samples, emb)

        # 6. post-process
320
        sample = self.conv_norm_out(sample)
Patrick von Platen's avatar
Patrick von Platen committed
321
322
323
324
325
326
327
328
329
330
        sample = self.conv_act(sample)
        sample = self.conv_out(sample)

        if skip_sample is not None:
            sample += skip_sample

        if self.config.time_embedding_type == "fourier":
            timesteps = timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]))))
            sample = sample / timesteps

331
332
        if not return_dict:
            return (sample,)
Patrick von Platen's avatar
Patrick von Platen committed
333

334
        return UNet2DOutput(sample=sample)