transformer_2d.py 17.1 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
14
#
# 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.
from dataclasses import dataclass
15
from typing import Any, Dict, Optional
16
17
18
19
20
21
22

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
Kashif Rasul's avatar
Kashif Rasul committed
23
from ..utils import BaseOutput, deprecate
24
from .attention import BasicTransformerBlock
Kashif Rasul's avatar
Kashif Rasul committed
25
from .embeddings import PatchEmbed
26
27
28
29
30
31
from .modeling_utils import ModelMixin


@dataclass
class Transformer2DModelOutput(BaseOutput):
    """
Steven Liu's avatar
Steven Liu committed
32
33
    The output of [`Transformer2DModel`].

34
35
    Args:
        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
Steven Liu's avatar
Steven Liu committed
36
37
            The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
            distributions for the unnoised latent pixels.
38
39
40
41
42
43
44
    """

    sample: torch.FloatTensor


class Transformer2DModel(ModelMixin, ConfigMixin):
    """
Steven Liu's avatar
Steven Liu committed
45
    A 2D Transformer model for image-like data.
46
47
48
49
50

    Parameters:
        num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
        attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
        in_channels (`int`, *optional*):
Steven Liu's avatar
Steven Liu committed
51
            The number of channels in the input and output (specify if the input is **continuous**).
52
53
        num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
Steven Liu's avatar
Steven Liu committed
54
55
56
        cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
        sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
            This is fixed during training since it is used to learn a number of position embeddings.
57
        num_vector_embeds (`int`, *optional*):
Steven Liu's avatar
Steven Liu committed
58
            The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
59
            Includes the class for the masked latent pixel.
Steven Liu's avatar
Steven Liu committed
60
61
62
63
64
65
66
        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
        num_embeds_ada_norm ( `int`, *optional*):
            The number of diffusion steps used during training. Pass if at least one of the norm_layers is
            `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
            added to the hidden states.

            During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
67
        attention_bias (`bool`, *optional*):
Steven Liu's avatar
Steven Liu committed
68
            Configure if the `TransformerBlocks` attention should contain a bias parameter.
69
70
71
72
73
74
75
76
    """

    @register_to_config
    def __init__(
        self,
        num_attention_heads: int = 16,
        attention_head_dim: int = 88,
        in_channels: Optional[int] = None,
Kashif Rasul's avatar
Kashif Rasul committed
77
        out_channels: Optional[int] = None,
78
79
80
81
82
83
84
        num_layers: int = 1,
        dropout: float = 0.0,
        norm_num_groups: int = 32,
        cross_attention_dim: Optional[int] = None,
        attention_bias: bool = False,
        sample_size: Optional[int] = None,
        num_vector_embeds: Optional[int] = None,
Kashif Rasul's avatar
Kashif Rasul committed
85
        patch_size: Optional[int] = None,
86
87
88
89
90
        activation_fn: str = "geglu",
        num_embeds_ada_norm: Optional[int] = None,
        use_linear_projection: bool = False,
        only_cross_attention: bool = False,
        upcast_attention: bool = False,
Kashif Rasul's avatar
Kashif Rasul committed
91
92
        norm_type: str = "layer_norm",
        norm_elementwise_affine: bool = True,
93
94
95
96
97
98
99
    ):
        super().__init__()
        self.use_linear_projection = use_linear_projection
        self.num_attention_heads = num_attention_heads
        self.attention_head_dim = attention_head_dim
        inner_dim = num_attention_heads * attention_head_dim

Alexander Pivovarov's avatar
Alexander Pivovarov committed
100
        # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
101
        # Define whether input is continuous or discrete depending on configuration
Kashif Rasul's avatar
Kashif Rasul committed
102
        self.is_input_continuous = (in_channels is not None) and (patch_size is None)
103
        self.is_input_vectorized = num_vector_embeds is not None
Kashif Rasul's avatar
Kashif Rasul committed
104
105
106
107
108
109
110
111
112
113
114
115
        self.is_input_patches = in_channels is not None and patch_size is not None

        if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
            deprecation_message = (
                f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"
                " incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config."
                " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"
                " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"
                " would be very nice if you could open a Pull request for the `transformer/config.json` file"
            )
            deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)
            norm_type = "ada_norm"
116
117
118
119
120
121

        if self.is_input_continuous and self.is_input_vectorized:
            raise ValueError(
                f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"
                " sure that either `in_channels` or `num_vector_embeds` is None."
            )
Kashif Rasul's avatar
Kashif Rasul committed
122
123
124
125
126
127
        elif self.is_input_vectorized and self.is_input_patches:
            raise ValueError(
                f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"
                " sure that either `num_vector_embeds` or `num_patches` is None."
            )
        elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
128
            raise ValueError(
Kashif Rasul's avatar
Kashif Rasul committed
129
130
                f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"
                f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
            )

        # 2. Define input layers
        if self.is_input_continuous:
            self.in_channels = in_channels

            self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
            if use_linear_projection:
                self.proj_in = nn.Linear(in_channels, inner_dim)
            else:
                self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
        elif self.is_input_vectorized:
            assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
            assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

            self.height = sample_size
            self.width = sample_size
            self.num_vector_embeds = num_vector_embeds
            self.num_latent_pixels = self.height * self.width

            self.latent_image_embedding = ImagePositionalEmbeddings(
                num_embed=num_vector_embeds, embed_dim=inner_dim, height=self.height, width=self.width
            )
Kashif Rasul's avatar
Kashif Rasul committed
154
155
156
157
158
159
160
161
162
163
164
165
166
167
        elif self.is_input_patches:
            assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

            self.height = sample_size
            self.width = sample_size

            self.patch_size = patch_size
            self.pos_embed = PatchEmbed(
                height=sample_size,
                width=sample_size,
                patch_size=patch_size,
                in_channels=in_channels,
                embed_dim=inner_dim,
            )
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182

        # 3. Define transformers blocks
        self.transformer_blocks = nn.ModuleList(
            [
                BasicTransformerBlock(
                    inner_dim,
                    num_attention_heads,
                    attention_head_dim,
                    dropout=dropout,
                    cross_attention_dim=cross_attention_dim,
                    activation_fn=activation_fn,
                    num_embeds_ada_norm=num_embeds_ada_norm,
                    attention_bias=attention_bias,
                    only_cross_attention=only_cross_attention,
                    upcast_attention=upcast_attention,
Kashif Rasul's avatar
Kashif Rasul committed
183
184
                    norm_type=norm_type,
                    norm_elementwise_affine=norm_elementwise_affine,
185
186
187
188
189
190
                )
                for d in range(num_layers)
            ]
        )

        # 4. Define output layers
Kashif Rasul's avatar
Kashif Rasul committed
191
        self.out_channels = in_channels if out_channels is None else out_channels
192
        if self.is_input_continuous:
Alexander Pivovarov's avatar
Alexander Pivovarov committed
193
            # TODO: should use out_channels for continuous projections
194
            if use_linear_projection:
195
                self.proj_out = nn.Linear(inner_dim, in_channels)
196
197
198
199
200
            else:
                self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
        elif self.is_input_vectorized:
            self.norm_out = nn.LayerNorm(inner_dim)
            self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1)
Kashif Rasul's avatar
Kashif Rasul committed
201
202
203
204
        elif self.is_input_patches:
            self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
            self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim)
            self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)
205
206
207

    def forward(
        self,
208
209
210
211
212
213
214
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        timestep: Optional[torch.LongTensor] = None,
        class_labels: Optional[torch.LongTensor] = None,
        cross_attention_kwargs: Dict[str, Any] = None,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
215
216
217
        return_dict: bool = True,
    ):
        """
Steven Liu's avatar
Steven Liu committed
218
219
        The [`Transformer2DModel`] forward method.

220
        Args:
Steven Liu's avatar
Steven Liu committed
221
222
            hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
                Input `hidden_states`.
223
            encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
224
225
                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
                self-attention.
226
            timestep ( `torch.LongTensor`, *optional*):
Steven Liu's avatar
Steven Liu committed
227
                Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
Kashif Rasul's avatar
Kashif Rasul committed
228
            class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
Steven Liu's avatar
Steven Liu committed
229
230
231
232
233
234
235
236
237
                Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
                `AdaLayerZeroNorm`.
            encoder_attention_mask ( `torch.Tensor`, *optional*):
                Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:

                    * Mask `(batch, sequence_length)` True = keep, False = discard.
                    * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.

                If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
238
                above. This bias will be added to the cross-attention scores.
239
            return_dict (`bool`, *optional*, defaults to `True`):
Steven Liu's avatar
Steven Liu committed
240
241
                Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
                tuple.
242
243

        Returns:
Steven Liu's avatar
Steven Liu committed
244
245
            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
            `tuple` where the first element is the sample tensor.
246
        """
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
        #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
        #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
        # expects mask of shape:
        #   [batch, key_tokens]
        # adds singleton query_tokens dimension:
        #   [batch,                    1, key_tokens]
        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
        if attention_mask is not None and attention_mask.ndim == 2:
            # assume that mask is expressed as:
            #   (1 = keep,      0 = discard)
            # convert mask into a bias that can be added to attention scores:
            #       (keep = +0,     discard = -10000.0)
            attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
            attention_mask = attention_mask.unsqueeze(1)

        # convert encoder_attention_mask to a bias the same way we do for attention_mask
        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
            encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)

270
271
        # 1. Input
        if self.is_input_continuous:
Kashif Rasul's avatar
Kashif Rasul committed
272
            batch, _, height, width = hidden_states.shape
273
274
275
276
277
278
279
280
281
282
283
284
285
            residual = hidden_states

            hidden_states = self.norm(hidden_states)
            if not self.use_linear_projection:
                hidden_states = self.proj_in(hidden_states)
                inner_dim = hidden_states.shape[1]
                hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
            else:
                inner_dim = hidden_states.shape[1]
                hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
                hidden_states = self.proj_in(hidden_states)
        elif self.is_input_vectorized:
            hidden_states = self.latent_image_embedding(hidden_states)
Kashif Rasul's avatar
Kashif Rasul committed
286
287
        elif self.is_input_patches:
            hidden_states = self.pos_embed(hidden_states)
288
289
290
291
292

        # 2. Blocks
        for block in self.transformer_blocks:
            hidden_states = block(
                hidden_states,
293
                attention_mask=attention_mask,
294
                encoder_hidden_states=encoder_hidden_states,
295
                encoder_attention_mask=encoder_attention_mask,
296
297
                timestep=timestep,
                cross_attention_kwargs=cross_attention_kwargs,
Kashif Rasul's avatar
Kashif Rasul committed
298
                class_labels=class_labels,
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
            )

        # 3. Output
        if self.is_input_continuous:
            if not self.use_linear_projection:
                hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
                hidden_states = self.proj_out(hidden_states)
            else:
                hidden_states = self.proj_out(hidden_states)
                hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()

            output = hidden_states + residual
        elif self.is_input_vectorized:
            hidden_states = self.norm_out(hidden_states)
            logits = self.out(hidden_states)
            # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
            logits = logits.permute(0, 2, 1)

            # log(p(x_0))
            output = F.log_softmax(logits.double(), dim=1).float()
Kashif Rasul's avatar
Kashif Rasul committed
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
        elif self.is_input_patches:
            # TODO: cleanup!
            conditioning = self.transformer_blocks[0].norm1.emb(
                timestep, class_labels, hidden_dtype=hidden_states.dtype
            )
            shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
            hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
            hidden_states = self.proj_out_2(hidden_states)

            # unpatchify
            height = width = int(hidden_states.shape[1] ** 0.5)
            hidden_states = hidden_states.reshape(
                shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)
            )
            hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
            output = hidden_states.reshape(
                shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)
            )
337
338
339
340
341

        if not return_dict:
            return (output,)

        return Transformer2DModelOutput(sample=output)