transformer_2d.py 17.9 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
from .lora import LoRACompatibleConv, LoRACompatibleLinear
27
28
29
30
31
32
from .modeling_utils import ModelMixin


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

35
36
    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
37
38
            The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
            distributions for the unnoised latent pixels.
39
40
41
42
43
44
45
    """

    sample: torch.FloatTensor


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

    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
52
            The number of channels in the input and output (specify if the input is **continuous**).
53
54
        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
55
56
57
        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.
58
        num_vector_embeds (`int`, *optional*):
Steven Liu's avatar
Steven Liu committed
59
            The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
60
            Includes the class for the masked latent pixel.
Steven Liu's avatar
Steven Liu committed
61
62
63
64
65
66
67
        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`.
68
        attention_bias (`bool`, *optional*):
Steven Liu's avatar
Steven Liu committed
69
            Configure if the `TransformerBlocks` attention should contain a bias parameter.
70
71
72
73
74
75
76
77
    """

    @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
78
        out_channels: Optional[int] = None,
79
80
81
82
83
84
85
        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
86
        patch_size: Optional[int] = None,
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,
Sanchit Gandhi's avatar
Sanchit Gandhi committed
91
        double_self_attention: bool = False,
92
        upcast_attention: bool = False,
Kashif Rasul's avatar
Kashif Rasul committed
93
94
        norm_type: str = "layer_norm",
        norm_elementwise_affine: bool = True,
95
        attention_type: str = "default",
96
97
98
99
100
101
102
    ):
        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
103
        # 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)`
104
        # Define whether input is continuous or discrete depending on configuration
Kashif Rasul's avatar
Kashif Rasul committed
105
        self.is_input_continuous = (in_channels is not None) and (patch_size is None)
106
        self.is_input_vectorized = num_vector_embeds is not None
Kashif Rasul's avatar
Kashif Rasul committed
107
108
109
110
111
112
113
114
115
116
117
118
        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"
119
120
121
122
123
124

        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
125
126
127
128
129
130
        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:
131
            raise ValueError(
Kashif Rasul's avatar
Kashif Rasul committed
132
133
                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."
134
135
136
137
138
139
140
141
            )

        # 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:
142
                self.proj_in = LoRACompatibleLinear(in_channels, inner_dim)
143
            else:
144
                self.proj_in = LoRACompatibleConv(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
145
146
147
148
149
150
151
152
153
154
155
156
        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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
        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,
            )
171
172
173
174
175
176
177
178
179
180
181
182
183
184

        # 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,
Sanchit Gandhi's avatar
Sanchit Gandhi committed
185
                    double_self_attention=double_self_attention,
186
                    upcast_attention=upcast_attention,
Kashif Rasul's avatar
Kashif Rasul committed
187
188
                    norm_type=norm_type,
                    norm_elementwise_affine=norm_elementwise_affine,
189
                    attention_type=attention_type,
190
191
192
193
194
195
                )
                for d in range(num_layers)
            ]
        )

        # 4. Define output layers
Kashif Rasul's avatar
Kashif Rasul committed
196
        self.out_channels = in_channels if out_channels is None else out_channels
197
        if self.is_input_continuous:
Alexander Pivovarov's avatar
Alexander Pivovarov committed
198
            # TODO: should use out_channels for continuous projections
199
            if use_linear_projection:
200
                self.proj_out = LoRACompatibleLinear(inner_dim, in_channels)
201
            else:
202
                self.proj_out = LoRACompatibleConv(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
203
204
205
        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
206
207
208
209
        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)
210

211
212
        self.gradient_checkpointing = False

213
214
    def forward(
        self,
215
216
217
218
219
220
221
        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,
222
223
224
        return_dict: bool = True,
    ):
        """
Steven Liu's avatar
Steven Liu committed
225
226
        The [`Transformer2DModel`] forward method.

227
        Args:
Steven Liu's avatar
Steven Liu committed
228
229
            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`.
230
            encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
231
232
                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
                self-attention.
233
            timestep ( `torch.LongTensor`, *optional*):
Steven Liu's avatar
Steven Liu committed
234
                Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
Kashif Rasul's avatar
Kashif Rasul committed
235
            class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
Steven Liu's avatar
Steven Liu committed
236
237
238
239
240
241
242
243
244
                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
245
                above. This bias will be added to the cross-attention scores.
246
            return_dict (`bool`, *optional*, defaults to `True`):
Steven Liu's avatar
Steven Liu committed
247
248
                Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
                tuple.
249
250

        Returns:
Steven Liu's avatar
Steven Liu committed
251
252
            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
            `tuple` where the first element is the sample tensor.
253
        """
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
        # 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)

277
278
        # 1. Input
        if self.is_input_continuous:
Kashif Rasul's avatar
Kashif Rasul committed
279
            batch, _, height, width = hidden_states.shape
280
281
282
283
284
285
286
287
288
289
290
291
292
            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
293
294
        elif self.is_input_patches:
            hidden_states = self.pos_embed(hidden_states)
295
296
297

        # 2. Blocks
        for block in self.transformer_blocks:
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
            if self.training and self.gradient_checkpointing:
                hidden_states = torch.utils.checkpoint.checkpoint(
                    block,
                    hidden_states,
                    attention_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    timestep,
                    cross_attention_kwargs,
                    class_labels,
                    use_reentrant=False,
                )
            else:
                hidden_states = block(
                    hidden_states,
                    attention_mask=attention_mask,
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_attention_mask,
                    timestep=timestep,
                    cross_attention_kwargs=cross_attention_kwargs,
                    class_labels=class_labels,
                )
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338

        # 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
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
        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)
            )
357
358
359
360
361

        if not return_dict:
            return (output,)

        return Transformer2DModelOutput(sample=output)