unet_2d_blocks.py 132 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
# Copyright 2023 The HuggingFace Team. All rights reserved.
Patrick von Platen's avatar
Patrick von Platen committed
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
from typing import Any, Dict, Optional, Tuple, Union
15

16
import numpy as np
17
import torch
18
import torch.nn.functional as F
Patrick von Platen's avatar
Patrick von Platen committed
19
20
from torch import nn

21
from ..utils import is_torch_version, logging
22
from ..utils.torch_utils import apply_freeu
23
from .activations import get_activation
24
from .attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0
25
from .dual_transformer_2d import DualTransformer2DModel
26
from .normalization import AdaGroupNorm
27
from .resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, Upsample2D
28
from .transformer_2d import Transformer2DModel
Patrick von Platen's avatar
Patrick von Platen committed
29
30


31
32
33
logger = logging.get_logger(__name__)  # pylint: disable=invalid-name


34
def get_down_block(
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
    down_block_type: str,
    num_layers: int,
    in_channels: int,
    out_channels: int,
    temb_channels: int,
    add_downsample: bool,
    resnet_eps: float,
    resnet_act_fn: str,
    transformer_layers_per_block: int = 1,
    num_attention_heads: Optional[int] = None,
    resnet_groups: Optional[int] = None,
    cross_attention_dim: Optional[int] = None,
    downsample_padding: Optional[int] = None,
    dual_cross_attention: bool = False,
    use_linear_projection: bool = False,
    only_cross_attention: bool = False,
    upcast_attention: bool = False,
    resnet_time_scale_shift: str = "default",
    attention_type: str = "default",
    resnet_skip_time_act: bool = False,
    resnet_out_scale_factor: float = 1.0,
    cross_attention_norm: Optional[str] = None,
    attention_head_dim: Optional[int] = None,
    downsample_type: Optional[str] = None,
    dropout: float = 0.0,
60
):
61
62
63
64
65
66
67
    # If attn head dim is not defined, we default it to the number of heads
    if attention_head_dim is None:
        logger.warn(
            f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
        )
        attention_head_dim = num_attention_heads

Patrick von Platen's avatar
Patrick von Platen committed
68
69
70
    down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
    if down_block_type == "DownBlock2D":
        return DownBlock2D(
71
72
73
74
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
75
            dropout=dropout,
76
77
78
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
79
            resnet_groups=resnet_groups,
Patrick von Platen's avatar
Patrick von Platen committed
80
            downsample_padding=downsample_padding,
Will Berman's avatar
Will Berman committed
81
82
83
84
85
86
87
88
            resnet_time_scale_shift=resnet_time_scale_shift,
        )
    elif down_block_type == "ResnetDownsampleBlock2D":
        return ResnetDownsampleBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
89
            dropout=dropout,
Will Berman's avatar
Will Berman committed
90
91
92
93
94
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
            resnet_groups=resnet_groups,
            resnet_time_scale_shift=resnet_time_scale_shift,
95
96
            skip_time_act=resnet_skip_time_act,
            output_scale_factor=resnet_out_scale_factor,
97
        )
Patrick von Platen's avatar
Patrick von Platen committed
98
    elif down_block_type == "AttnDownBlock2D":
99
100
101
102
        if add_downsample is False:
            downsample_type = None
        else:
            downsample_type = downsample_type or "conv"  # default to 'conv'
Patrick von Platen's avatar
Patrick von Platen committed
103
        return AttnDownBlock2D(
104
105
106
107
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
108
            dropout=dropout,
109
110
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
111
            resnet_groups=resnet_groups,
112
            downsample_padding=downsample_padding,
113
            attention_head_dim=attention_head_dim,
Will Berman's avatar
Will Berman committed
114
            resnet_time_scale_shift=resnet_time_scale_shift,
115
            downsample_type=downsample_type,
116
        )
Patrick von Platen's avatar
Patrick von Platen committed
117
    elif down_block_type == "CrossAttnDownBlock2D":
118
        if cross_attention_dim is None:
119
            raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D")
Patrick von Platen's avatar
Patrick von Platen committed
120
        return CrossAttnDownBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
121
            num_layers=num_layers,
122
            transformer_layers_per_block=transformer_layers_per_block,
Patrick von Platen's avatar
Patrick von Platen committed
123
124
125
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
126
            dropout=dropout,
Patrick von Platen's avatar
Patrick von Platen committed
127
128
129
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
130
            resnet_groups=resnet_groups,
Patrick von Platen's avatar
Patrick von Platen committed
131
            downsample_padding=downsample_padding,
132
            cross_attention_dim=cross_attention_dim,
133
            num_attention_heads=num_attention_heads,
134
            dual_cross_attention=dual_cross_attention,
Suraj Patil's avatar
Suraj Patil committed
135
            use_linear_projection=use_linear_projection,
136
            only_cross_attention=only_cross_attention,
137
            upcast_attention=upcast_attention,
Will Berman's avatar
Will Berman committed
138
            resnet_time_scale_shift=resnet_time_scale_shift,
139
            attention_type=attention_type,
Will Berman's avatar
Will Berman committed
140
141
142
143
144
145
146
147
148
        )
    elif down_block_type == "SimpleCrossAttnDownBlock2D":
        if cross_attention_dim is None:
            raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D")
        return SimpleCrossAttnDownBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
149
            dropout=dropout,
Will Berman's avatar
Will Berman committed
150
151
152
153
154
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
            resnet_groups=resnet_groups,
            cross_attention_dim=cross_attention_dim,
155
            attention_head_dim=attention_head_dim,
Will Berman's avatar
Will Berman committed
156
            resnet_time_scale_shift=resnet_time_scale_shift,
157
158
            skip_time_act=resnet_skip_time_act,
            output_scale_factor=resnet_out_scale_factor,
159
            only_cross_attention=only_cross_attention,
160
            cross_attention_norm=cross_attention_norm,
Patrick von Platen's avatar
Patrick von Platen committed
161
        )
Patrick von Platen's avatar
Patrick von Platen committed
162
163
    elif down_block_type == "SkipDownBlock2D":
        return SkipDownBlock2D(
164
165
166
167
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
168
            dropout=dropout,
169
170
171
172
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
            downsample_padding=downsample_padding,
Will Berman's avatar
Will Berman committed
173
            resnet_time_scale_shift=resnet_time_scale_shift,
174
        )
Patrick von Platen's avatar
Patrick von Platen committed
175
176
    elif down_block_type == "AttnSkipDownBlock2D":
        return AttnSkipDownBlock2D(
177
178
179
180
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
181
            dropout=dropout,
182
183
184
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
185
            attention_head_dim=attention_head_dim,
Will Berman's avatar
Will Berman committed
186
            resnet_time_scale_shift=resnet_time_scale_shift,
187
        )
188
189
190
191
192
    elif down_block_type == "DownEncoderBlock2D":
        return DownEncoderBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
193
            dropout=dropout,
194
195
196
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
197
            resnet_groups=resnet_groups,
198
            downsample_padding=downsample_padding,
Will Berman's avatar
Will Berman committed
199
            resnet_time_scale_shift=resnet_time_scale_shift,
200
        )
Will Berman's avatar
Will Berman committed
201
202
203
204
205
    elif down_block_type == "AttnDownEncoderBlock2D":
        return AttnDownEncoderBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
206
            dropout=dropout,
Will Berman's avatar
Will Berman committed
207
208
209
210
211
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
            resnet_groups=resnet_groups,
            downsample_padding=downsample_padding,
212
            attention_head_dim=attention_head_dim,
Will Berman's avatar
Will Berman committed
213
            resnet_time_scale_shift=resnet_time_scale_shift,
Will Berman's avatar
Will Berman committed
214
        )
215
216
217
218
219
220
    elif down_block_type == "KDownBlock2D":
        return KDownBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
221
            dropout=dropout,
222
223
224
225
226
227
228
229
230
231
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
        )
    elif down_block_type == "KCrossAttnDownBlock2D":
        return KCrossAttnDownBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
232
            dropout=dropout,
233
234
235
236
            add_downsample=add_downsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
            cross_attention_dim=cross_attention_dim,
237
            attention_head_dim=attention_head_dim,
238
239
            add_self_attention=True if not add_downsample else False,
        )
Will Berman's avatar
Will Berman committed
240
    raise ValueError(f"{down_block_type} does not exist.")
241
242
243


def get_up_block(
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
    up_block_type: str,
    num_layers: int,
    in_channels: int,
    out_channels: int,
    prev_output_channel: int,
    temb_channels: int,
    add_upsample: bool,
    resnet_eps: float,
    resnet_act_fn: str,
    resolution_idx: Optional[int] = None,
    transformer_layers_per_block: int = 1,
    num_attention_heads: Optional[int] = None,
    resnet_groups: Optional[int] = None,
    cross_attention_dim: Optional[int] = None,
    dual_cross_attention: bool = False,
    use_linear_projection: bool = False,
    only_cross_attention: bool = False,
    upcast_attention: bool = False,
    resnet_time_scale_shift: str = "default",
    attention_type: str = "default",
    resnet_skip_time_act: bool = False,
    resnet_out_scale_factor: float = 1.0,
    cross_attention_norm: Optional[str] = None,
    attention_head_dim: Optional[int] = None,
    upsample_type: Optional[str] = None,
    dropout: float = 0.0,
) -> nn.Module:
271
272
273
274
275
276
277
    # If attn head dim is not defined, we default it to the number of heads
    if attention_head_dim is None:
        logger.warn(
            f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
        )
        attention_head_dim = num_attention_heads

Patrick von Platen's avatar
Patrick von Platen committed
278
279
280
    up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
    if up_block_type == "UpBlock2D":
        return UpBlock2D(
281
282
            num_layers=num_layers,
            in_channels=in_channels,
Patrick von Platen's avatar
Patrick von Platen committed
283
284
            out_channels=out_channels,
            prev_output_channel=prev_output_channel,
285
            temb_channels=temb_channels,
286
            resolution_idx=resolution_idx,
287
            dropout=dropout,
288
289
290
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
291
            resnet_groups=resnet_groups,
Will Berman's avatar
Will Berman committed
292
293
294
295
296
297
298
299
300
            resnet_time_scale_shift=resnet_time_scale_shift,
        )
    elif up_block_type == "ResnetUpsampleBlock2D":
        return ResnetUpsampleBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            prev_output_channel=prev_output_channel,
            temb_channels=temb_channels,
301
            resolution_idx=resolution_idx,
302
            dropout=dropout,
Will Berman's avatar
Will Berman committed
303
304
305
306
307
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
            resnet_groups=resnet_groups,
            resnet_time_scale_shift=resnet_time_scale_shift,
308
309
            skip_time_act=resnet_skip_time_act,
            output_scale_factor=resnet_out_scale_factor,
310
        )
Patrick von Platen's avatar
Patrick von Platen committed
311
    elif up_block_type == "CrossAttnUpBlock2D":
312
313
        if cross_attention_dim is None:
            raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D")
Patrick von Platen's avatar
Patrick von Platen committed
314
        return CrossAttnUpBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
315
            num_layers=num_layers,
316
            transformer_layers_per_block=transformer_layers_per_block,
Patrick von Platen's avatar
Patrick von Platen committed
317
318
319
320
            in_channels=in_channels,
            out_channels=out_channels,
            prev_output_channel=prev_output_channel,
            temb_channels=temb_channels,
321
            resolution_idx=resolution_idx,
322
            dropout=dropout,
Patrick von Platen's avatar
Patrick von Platen committed
323
324
325
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
326
            resnet_groups=resnet_groups,
327
            cross_attention_dim=cross_attention_dim,
328
            num_attention_heads=num_attention_heads,
329
            dual_cross_attention=dual_cross_attention,
Suraj Patil's avatar
Suraj Patil committed
330
            use_linear_projection=use_linear_projection,
331
            only_cross_attention=only_cross_attention,
332
            upcast_attention=upcast_attention,
Will Berman's avatar
Will Berman committed
333
            resnet_time_scale_shift=resnet_time_scale_shift,
334
            attention_type=attention_type,
Will Berman's avatar
Will Berman committed
335
336
337
338
339
340
341
342
343
344
        )
    elif up_block_type == "SimpleCrossAttnUpBlock2D":
        if cross_attention_dim is None:
            raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D")
        return SimpleCrossAttnUpBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            prev_output_channel=prev_output_channel,
            temb_channels=temb_channels,
345
            resolution_idx=resolution_idx,
346
            dropout=dropout,
Will Berman's avatar
Will Berman committed
347
348
349
350
351
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
            resnet_groups=resnet_groups,
            cross_attention_dim=cross_attention_dim,
352
            attention_head_dim=attention_head_dim,
Will Berman's avatar
Will Berman committed
353
            resnet_time_scale_shift=resnet_time_scale_shift,
354
355
            skip_time_act=resnet_skip_time_act,
            output_scale_factor=resnet_out_scale_factor,
356
            only_cross_attention=only_cross_attention,
357
            cross_attention_norm=cross_attention_norm,
Patrick von Platen's avatar
Patrick von Platen committed
358
        )
Patrick von Platen's avatar
Patrick von Platen committed
359
    elif up_block_type == "AttnUpBlock2D":
360
361
362
363
364
        if add_upsample is False:
            upsample_type = None
        else:
            upsample_type = upsample_type or "conv"  # default to 'conv'

Patrick von Platen's avatar
Patrick von Platen committed
365
        return AttnUpBlock2D(
366
367
            num_layers=num_layers,
            in_channels=in_channels,
Patrick von Platen's avatar
Patrick von Platen committed
368
369
            out_channels=out_channels,
            prev_output_channel=prev_output_channel,
370
            temb_channels=temb_channels,
371
            resolution_idx=resolution_idx,
372
            dropout=dropout,
373
374
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
375
            resnet_groups=resnet_groups,
376
            attention_head_dim=attention_head_dim,
Will Berman's avatar
Will Berman committed
377
            resnet_time_scale_shift=resnet_time_scale_shift,
378
            upsample_type=upsample_type,
379
        )
Patrick von Platen's avatar
Patrick von Platen committed
380
381
    elif up_block_type == "SkipUpBlock2D":
        return SkipUpBlock2D(
382
383
384
385
386
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            prev_output_channel=prev_output_channel,
            temb_channels=temb_channels,
387
            resolution_idx=resolution_idx,
388
            dropout=dropout,
389
390
391
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
Will Berman's avatar
Will Berman committed
392
            resnet_time_scale_shift=resnet_time_scale_shift,
393
        )
Patrick von Platen's avatar
Patrick von Platen committed
394
395
    elif up_block_type == "AttnSkipUpBlock2D":
        return AttnSkipUpBlock2D(
396
397
398
399
400
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            prev_output_channel=prev_output_channel,
            temb_channels=temb_channels,
401
            resolution_idx=resolution_idx,
402
            dropout=dropout,
403
404
405
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
406
            attention_head_dim=attention_head_dim,
Will Berman's avatar
Will Berman committed
407
            resnet_time_scale_shift=resnet_time_scale_shift,
408
        )
409
410
411
412
413
    elif up_block_type == "UpDecoderBlock2D":
        return UpDecoderBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
414
            resolution_idx=resolution_idx,
415
            dropout=dropout,
416
417
418
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
419
            resnet_groups=resnet_groups,
Will Berman's avatar
Will Berman committed
420
            resnet_time_scale_shift=resnet_time_scale_shift,
YiYi Xu's avatar
YiYi Xu committed
421
            temb_channels=temb_channels,
422
        )
Will Berman's avatar
Will Berman committed
423
424
425
426
427
    elif up_block_type == "AttnUpDecoderBlock2D":
        return AttnUpDecoderBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
428
            resolution_idx=resolution_idx,
429
            dropout=dropout,
Will Berman's avatar
Will Berman committed
430
431
432
433
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
            resnet_groups=resnet_groups,
434
            attention_head_dim=attention_head_dim,
Will Berman's avatar
Will Berman committed
435
            resnet_time_scale_shift=resnet_time_scale_shift,
YiYi Xu's avatar
YiYi Xu committed
436
            temb_channels=temb_channels,
Will Berman's avatar
Will Berman committed
437
        )
438
439
440
441
442
443
    elif up_block_type == "KUpBlock2D":
        return KUpBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
444
            resolution_idx=resolution_idx,
445
            dropout=dropout,
446
447
448
449
450
451
452
453
454
455
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
        )
    elif up_block_type == "KCrossAttnUpBlock2D":
        return KCrossAttnUpBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            temb_channels=temb_channels,
456
            resolution_idx=resolution_idx,
457
            dropout=dropout,
458
459
460
461
            add_upsample=add_upsample,
            resnet_eps=resnet_eps,
            resnet_act_fn=resnet_act_fn,
            cross_attention_dim=cross_attention_dim,
462
            attention_head_dim=attention_head_dim,
463
464
        )

465
    raise ValueError(f"{up_block_type} does not exist.")
466
467


468
class AutoencoderTinyBlock(nn.Module):
469
    """
Patrick von Platen's avatar
Patrick von Platen committed
470
471
    Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU
    blocks.
472
473
474
475

    Args:
        in_channels (`int`): The number of input channels.
        out_channels (`int`): The number of output channels.
Patrick von Platen's avatar
Patrick von Platen committed
476
477
        act_fn (`str`):
            ` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`.
478
479

    Returns:
Patrick von Platen's avatar
Patrick von Platen committed
480
481
        `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to
        `out_channels`.
482
483
    """

484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
    def __init__(self, in_channels: int, out_channels: int, act_fn: str):
        super().__init__()
        act_fn = get_activation(act_fn)
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            act_fn,
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            act_fn,
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
        )
        self.skip = (
            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
            if in_channels != out_channels
            else nn.Identity()
        )
        self.fuse = nn.ReLU()

501
    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
502
503
504
        return self.fuse(self.conv(x) + self.skip(x))


Patrick von Platen's avatar
Patrick von Platen committed
505
class UNetMidBlock2D(nn.Module):
506
507
508
509
510
511
512
513
514
    """
    A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks.

    Args:
        in_channels (`int`): The number of input channels.
        temb_channels (`int`): The number of temporal embedding channels.
        dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
        num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
        resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
Patrick von Platen's avatar
Patrick von Platen committed
515
516
517
        resnet_time_scale_shift (`str`, *optional*, defaults to `default`):
            The type of normalization to apply to the time embeddings. This can help to improve the performance of the
            model on tasks with long-range temporal dependencies.
518
        resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.
Patrick von Platen's avatar
Patrick von Platen committed
519
520
        resnet_groups (`int`, *optional*, defaults to 32):
            The number of groups to use in the group normalization layers of the resnet blocks.
521
        attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks.
Patrick von Platen's avatar
Patrick von Platen committed
522
523
        resnet_pre_norm (`bool`, *optional*, defaults to `True`):
            Whether to use pre-normalization for the resnet blocks.
524
        add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.
Patrick von Platen's avatar
Patrick von Platen committed
525
526
527
        attention_head_dim (`int`, *optional*, defaults to 1):
            Dimension of a single attention head. The number of attention heads is determined based on this value and
            the number of input channels.
528
529
530
        output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor.

    Returns:
Patrick von Platen's avatar
Patrick von Platen committed
531
532
        `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
        in_channels, height, width)`.
533
534
535

    """

Patrick von Platen's avatar
Patrick von Platen committed
536
537
538
539
    def __init__(
        self,
        in_channels: int,
        temb_channels: int,
540
        dropout: float = 0.0,
541
        num_layers: int = 1,
Patrick von Platen's avatar
Patrick von Platen committed
542
        resnet_eps: float = 1e-6,
YiYi Xu's avatar
YiYi Xu committed
543
        resnet_time_scale_shift: str = "default",  # default, spatial
Patrick von Platen's avatar
Patrick von Platen committed
544
545
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
546
        attn_groups: Optional[int] = None,
547
        resnet_pre_norm: bool = True,
Will Berman's avatar
Will Berman committed
548
        add_attention: bool = True,
549
550
        attention_head_dim: int = 1,
        output_scale_factor: float = 1.0,
Patrick von Platen's avatar
Patrick von Platen committed
551
552
    ):
        super().__init__()
553
        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
Will Berman's avatar
Will Berman committed
554
        self.add_attention = add_attention
Patrick von Platen's avatar
Patrick von Platen committed
555

556
557
558
        if attn_groups is None:
            attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None

559
560
        # there is always at least one resnet
        resnets = [
561
            ResnetBlock2D(
562
563
564
                in_channels=in_channels,
                out_channels=in_channels,
                temb_channels=temb_channels,
565
                eps=resnet_eps,
566
567
568
569
570
571
                groups=resnet_groups,
                dropout=dropout,
                time_embedding_norm=resnet_time_scale_shift,
                non_linearity=resnet_act_fn,
                output_scale_factor=output_scale_factor,
                pre_norm=resnet_pre_norm,
Patrick von Platen's avatar
Patrick von Platen committed
572
            )
573
574
        ]
        attentions = []
Patrick von Platen's avatar
Patrick von Platen committed
575

576
577
578
579
580
581
        if attention_head_dim is None:
            logger.warn(
                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
            )
            attention_head_dim = in_channels

582
        for _ in range(num_layers):
Will Berman's avatar
Will Berman committed
583
584
            if self.add_attention:
                attentions.append(
585
                    Attention(
Will Berman's avatar
Will Berman committed
586
                        in_channels,
587
588
                        heads=in_channels // attention_head_dim,
                        dim_head=attention_head_dim,
Will Berman's avatar
Will Berman committed
589
590
                        rescale_output_factor=output_scale_factor,
                        eps=resnet_eps,
591
                        norm_num_groups=attn_groups,
YiYi Xu's avatar
YiYi Xu committed
592
                        spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
593
594
595
596
                        residual_connection=True,
                        bias=True,
                        upcast_softmax=True,
                        _from_deprecated_attn_block=True,
Will Berman's avatar
Will Berman committed
597
                    )
598
                )
Will Berman's avatar
Will Berman committed
599
600
601
            else:
                attentions.append(None)

602
            resnets.append(
603
                ResnetBlock2D(
604
605
606
                    in_channels=in_channels,
                    out_channels=in_channels,
                    temb_channels=temb_channels,
607
                    eps=resnet_eps,
608
609
610
611
612
613
614
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
Patrick von Platen's avatar
Patrick von Platen committed
615
616
            )

617
618
619
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

620
    def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:
Patrick von Platen's avatar
Patrick von Platen committed
621
        hidden_states = self.resnets[0](hidden_states, temb)
622
        for attn, resnet in zip(self.attentions, self.resnets[1:]):
Will Berman's avatar
Will Berman committed
623
            if attn is not None:
YiYi Xu's avatar
YiYi Xu committed
624
                hidden_states = attn(hidden_states, temb=temb)
Patrick von Platen's avatar
Patrick von Platen committed
625
            hidden_states = resnet(hidden_states, temb)
Patrick von Platen's avatar
Patrick von Platen committed
626

627
        return hidden_states
Patrick von Platen's avatar
Patrick von Platen committed
628

629

Patrick von Platen's avatar
Patrick von Platen committed
630
631
632
633
634
635
636
class UNetMidBlock2DCrossAttn(nn.Module):
    def __init__(
        self,
        in_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
637
        transformer_layers_per_block: Union[int, Tuple[int]] = 1,
Patrick von Platen's avatar
Patrick von Platen committed
638
639
640
641
642
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
643
644
645
646
647
648
649
        num_attention_heads: int = 1,
        output_scale_factor: float = 1.0,
        cross_attention_dim: int = 1280,
        dual_cross_attention: bool = False,
        use_linear_projection: bool = False,
        upcast_attention: bool = False,
        attention_type: str = "default",
Patrick von Platen's avatar
Patrick von Platen committed
650
651
652
    ):
        super().__init__()

653
        self.has_cross_attention = True
654
        self.num_attention_heads = num_attention_heads
Patrick von Platen's avatar
Patrick von Platen committed
655
656
        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)

657
658
659
660
        # support for variable transformer layers per block
        if isinstance(transformer_layers_per_block, int):
            transformer_layers_per_block = [transformer_layers_per_block] * num_layers

Patrick von Platen's avatar
Patrick von Platen committed
661
662
        # there is always at least one resnet
        resnets = [
663
            ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
664
665
666
667
668
669
670
671
672
673
674
675
676
677
                in_channels=in_channels,
                out_channels=in_channels,
                temb_channels=temb_channels,
                eps=resnet_eps,
                groups=resnet_groups,
                dropout=dropout,
                time_embedding_norm=resnet_time_scale_shift,
                non_linearity=resnet_act_fn,
                output_scale_factor=output_scale_factor,
                pre_norm=resnet_pre_norm,
            )
        ]
        attentions = []

678
        for i in range(num_layers):
679
680
681
            if not dual_cross_attention:
                attentions.append(
                    Transformer2DModel(
682
683
                        num_attention_heads,
                        in_channels // num_attention_heads,
684
                        in_channels=in_channels,
685
                        num_layers=transformer_layers_per_block[i],
686
687
                        cross_attention_dim=cross_attention_dim,
                        norm_num_groups=resnet_groups,
Suraj Patil's avatar
Suraj Patil committed
688
                        use_linear_projection=use_linear_projection,
689
                        upcast_attention=upcast_attention,
690
                        attention_type=attention_type,
691
692
693
694
695
                    )
                )
            else:
                attentions.append(
                    DualTransformer2DModel(
696
697
                        num_attention_heads,
                        in_channels // num_attention_heads,
698
699
700
701
702
                        in_channels=in_channels,
                        num_layers=1,
                        cross_attention_dim=cross_attention_dim,
                        norm_num_groups=resnet_groups,
                    )
Patrick von Platen's avatar
Patrick von Platen committed
703
704
                )
            resnets.append(
705
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
                    in_channels=in_channels,
                    out_channels=in_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )

        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

722
723
        self.gradient_checkpointing = False

724
    def forward(
725
726
727
728
729
730
731
732
        self,
        hidden_states: torch.FloatTensor,
        temb: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
    ) -> torch.FloatTensor:
733
734
        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
        hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)
Patrick von Platen's avatar
Patrick von Platen committed
735
        for attn, resnet in zip(self.attentions, self.resnets[1:]):
736
737
738
739
740
741
742
743
744
745
746
747
            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)

                    return custom_forward

                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
ethansmith2000's avatar
ethansmith2000 committed
748
                hidden_states = attn(
749
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
750
751
752
753
754
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
                    return_dict=False,
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
                )[0]
                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(resnet),
                    hidden_states,
                    temb,
                    **ckpt_kwargs,
                )
            else:
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
                    return_dict=False,
                )[0]
771
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
Will Berman's avatar
Will Berman committed
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787

        return hidden_states


class UNetMidBlock2DSimpleCrossAttn(nn.Module):
    def __init__(
        self,
        in_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
788
789
790
791
792
793
        attention_head_dim: int = 1,
        output_scale_factor: float = 1.0,
        cross_attention_dim: int = 1280,
        skip_time_act: bool = False,
        only_cross_attention: bool = False,
        cross_attention_norm: Optional[str] = None,
Will Berman's avatar
Will Berman committed
794
795
796
797
798
    ):
        super().__init__()

        self.has_cross_attention = True

799
        self.attention_head_dim = attention_head_dim
Will Berman's avatar
Will Berman committed
800
801
        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)

802
        self.num_heads = in_channels // self.attention_head_dim
Will Berman's avatar
Will Berman committed
803
804
805
806
807
808
809
810
811
812
813
814
815
816

        # there is always at least one resnet
        resnets = [
            ResnetBlock2D(
                in_channels=in_channels,
                out_channels=in_channels,
                temb_channels=temb_channels,
                eps=resnet_eps,
                groups=resnet_groups,
                dropout=dropout,
                time_embedding_norm=resnet_time_scale_shift,
                non_linearity=resnet_act_fn,
                output_scale_factor=output_scale_factor,
                pre_norm=resnet_pre_norm,
817
                skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
818
819
820
821
822
            )
        ]
        attentions = []

        for _ in range(num_layers):
823
824
825
826
            processor = (
                AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()
            )

Will Berman's avatar
Will Berman committed
827
            attentions.append(
Patrick von Platen's avatar
Patrick von Platen committed
828
                Attention(
Will Berman's avatar
Will Berman committed
829
830
831
                    query_dim=in_channels,
                    cross_attention_dim=in_channels,
                    heads=self.num_heads,
832
                    dim_head=self.attention_head_dim,
Will Berman's avatar
Will Berman committed
833
834
835
836
                    added_kv_proj_dim=cross_attention_dim,
                    norm_num_groups=resnet_groups,
                    bias=True,
                    upcast_softmax=True,
837
                    only_cross_attention=only_cross_attention,
838
                    cross_attention_norm=cross_attention_norm,
839
                    processor=processor,
Will Berman's avatar
Will Berman committed
840
841
842
843
844
845
846
847
848
849
850
851
852
853
                )
            )
            resnets.append(
                ResnetBlock2D(
                    in_channels=in_channels,
                    out_channels=in_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
854
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
855
856
857
858
859
860
                )
            )

        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

861
    def forward(
862
863
864
865
866
867
868
        self,
        hidden_states: torch.FloatTensor,
        temb: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
869
    ) -> torch.FloatTensor:
870
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
871
        lora_scale = cross_attention_kwargs.get("scale", 1.0)
872
873
874
875
876
877
878
879
880
881
882
883

        if attention_mask is None:
            # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask.
            mask = None if encoder_hidden_states is None else encoder_attention_mask
        else:
            # when attention_mask is defined: we don't even check for encoder_attention_mask.
            # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks.
            # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask.
            #       then we can simplify this whole if/else block to:
            #         mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask
            mask = attention_mask

884
        hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)
Will Berman's avatar
Will Berman committed
885
886
887
888
        for attn, resnet in zip(self.attentions, self.resnets[1:]):
            # attn
            hidden_states = attn(
                hidden_states,
889
                encoder_hidden_states=encoder_hidden_states,
890
                attention_mask=mask,
891
                **cross_attention_kwargs,
Will Berman's avatar
Will Berman committed
892
893
894
            )

            # resnet
895
            hidden_states = resnet(hidden_states, temb, scale=lora_scale)
Patrick von Platen's avatar
Patrick von Platen committed
896
897
898
899

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
900
class AttnDownBlock2D(nn.Module):
901
902
903
904
905
906
907
908
909
910
911
912
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
913
914
915
916
        attention_head_dim: int = 1,
        output_scale_factor: float = 1.0,
        downsample_padding: int = 1,
        downsample_type: str = "conv",
917
918
919
920
    ):
        super().__init__()
        resnets = []
        attentions = []
921
        self.downsample_type = downsample_type
922

923
924
925
926
927
928
        if attention_head_dim is None:
            logger.warn(
                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
            )
            attention_head_dim = out_channels

929
930
931
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
932
                ResnetBlock2D(
933
934
935
936
937
938
939
940
941
942
943
944
945
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )
            attentions.append(
946
                Attention(
947
                    out_channels,
948
949
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
950
                    rescale_output_factor=output_scale_factor,
Patrick von Platen's avatar
Patrick von Platen committed
951
                    eps=resnet_eps,
Will Berman's avatar
Will Berman committed
952
                    norm_num_groups=resnet_groups,
953
954
955
956
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
957
958
959
960
961
962
                )
            )

        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

963
        if downsample_type == "conv":
964
            self.downsamplers = nn.ModuleList(
965
966
                [
                    Downsample2D(
967
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
968
969
                    )
                ]
970
            )
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
        elif downsample_type == "resnet":
            self.downsamplers = nn.ModuleList(
                [
                    ResnetBlock2D(
                        in_channels=out_channels,
                        out_channels=out_channels,
                        temb_channels=temb_channels,
                        eps=resnet_eps,
                        groups=resnet_groups,
                        dropout=dropout,
                        time_embedding_norm=resnet_time_scale_shift,
                        non_linearity=resnet_act_fn,
                        output_scale_factor=output_scale_factor,
                        pre_norm=resnet_pre_norm,
                        down=True,
                    )
                ]
            )
989
990
991
        else:
            self.downsamplers = None

992
993
994
995
996
997
998
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        temb: Optional[torch.FloatTensor] = None,
        upsample_size: Optional[int] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
999
1000
1001
1002
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}

        lora_scale = cross_attention_kwargs.get("scale", 1.0)

1003
1004
1005
        output_states = ()

        for resnet, attn in zip(self.resnets, self.attentions):
1006
1007
1008
            cross_attention_kwargs.update({"scale": lora_scale})
            hidden_states = resnet(hidden_states, temb, scale=lora_scale)
            hidden_states = attn(hidden_states, **cross_attention_kwargs)
1009
            output_states = output_states + (hidden_states,)
1010
1011
1012

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
1013
                if self.downsample_type == "resnet":
1014
                    hidden_states = downsampler(hidden_states, temb=temb, scale=lora_scale)
1015
                else:
1016
                    hidden_states = downsampler(hidden_states, scale=lora_scale)
1017
1018
1019
1020
1021
1022

            output_states += (hidden_states,)

        return hidden_states, output_states


Patrick von Platen's avatar
Patrick von Platen committed
1023
class CrossAttnDownBlock2D(nn.Module):
Patrick von Platen's avatar
Patrick von Platen committed
1024
1025
1026
1027
1028
1029
1030
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
1031
        transformer_layers_per_block: Union[int, Tuple[int]] = 1,
Patrick von Platen's avatar
Patrick von Platen committed
1032
1033
1034
1035
1036
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
        num_attention_heads: int = 1,
        cross_attention_dim: int = 1280,
        output_scale_factor: float = 1.0,
        downsample_padding: int = 1,
        add_downsample: bool = True,
        dual_cross_attention: bool = False,
        use_linear_projection: bool = False,
        only_cross_attention: bool = False,
        upcast_attention: bool = False,
        attention_type: str = "default",
Patrick von Platen's avatar
Patrick von Platen committed
1047
1048
1049
1050
1051
    ):
        super().__init__()
        resnets = []
        attentions = []

1052
        self.has_cross_attention = True
1053
        self.num_attention_heads = num_attention_heads
1054
1055
        if isinstance(transformer_layers_per_block, int):
            transformer_layers_per_block = [transformer_layers_per_block] * num_layers
Patrick von Platen's avatar
Patrick von Platen committed
1056
1057
1058
1059

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
1060
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )
1073
1074
1075
            if not dual_cross_attention:
                attentions.append(
                    Transformer2DModel(
1076
1077
                        num_attention_heads,
                        out_channels // num_attention_heads,
1078
                        in_channels=out_channels,
1079
                        num_layers=transformer_layers_per_block[i],
1080
1081
                        cross_attention_dim=cross_attention_dim,
                        norm_num_groups=resnet_groups,
Suraj Patil's avatar
Suraj Patil committed
1082
                        use_linear_projection=use_linear_projection,
1083
                        only_cross_attention=only_cross_attention,
1084
                        upcast_attention=upcast_attention,
1085
                        attention_type=attention_type,
1086
1087
1088
1089
1090
                    )
                )
            else:
                attentions.append(
                    DualTransformer2DModel(
1091
1092
                        num_attention_heads,
                        out_channels // num_attention_heads,
1093
1094
1095
1096
1097
                        in_channels=out_channels,
                        num_layers=1,
                        cross_attention_dim=cross_attention_dim,
                        norm_num_groups=resnet_groups,
                    )
Patrick von Platen's avatar
Patrick von Platen committed
1098
1099
1100
1101
1102
1103
1104
1105
                )
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

        if add_downsample:
            self.downsamplers = nn.ModuleList(
                [
                    Downsample2D(
1106
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
Patrick von Platen's avatar
Patrick von Platen committed
1107
1108
1109
1110
1111
1112
                    )
                ]
            )
        else:
            self.downsamplers = None

1113
1114
        self.gradient_checkpointing = False

1115
    def forward(
1116
1117
1118
1119
1120
1121
1122
        self,
        hidden_states: torch.FloatTensor,
        temb: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
1123
1124
        additional_residuals: Optional[torch.FloatTensor] = None,
    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
Patrick von Platen's avatar
Patrick von Platen committed
1125
1126
        output_states = ()

1127
1128
        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0

Will Berman's avatar
Will Berman committed
1129
1130
1131
        blocks = list(zip(self.resnets, self.attentions))

        for i, (resnet, attn) in enumerate(blocks):
1132
1133
            if self.training and self.gradient_checkpointing:

Will Berman's avatar
Will Berman committed
1134
                def create_custom_forward(module, return_dict=None):
1135
                    def custom_forward(*inputs):
Will Berman's avatar
Will Berman committed
1136
1137
1138
1139
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)
1140
1141
1142

                    return custom_forward

1143
1144
1145
1146
1147
1148
1149
                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(resnet),
                    hidden_states,
                    temb,
                    **ckpt_kwargs,
                )
ethansmith2000's avatar
ethansmith2000 committed
1150
                hidden_states = attn(
1151
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
1152
1153
1154
1155
1156
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
                    return_dict=False,
1157
                )[0]
1158
            else:
1159
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
1160
1161
1162
1163
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
1164
1165
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
1166
1167
                    return_dict=False,
                )[0]
1168

Will Berman's avatar
Will Berman committed
1169
1170
1171
1172
            # apply additional residuals to the output of the last pair of resnet and attention blocks
            if i == len(blocks) - 1 and additional_residuals is not None:
                hidden_states = hidden_states + additional_residuals

1173
            output_states = output_states + (hidden_states,)
Patrick von Platen's avatar
Patrick von Platen committed
1174
1175
1176

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
1177
                hidden_states = downsampler(hidden_states, scale=lora_scale)
Patrick von Platen's avatar
Patrick von Platen committed
1178

1179
            output_states = output_states + (hidden_states,)
Patrick von Platen's avatar
Patrick von Platen committed
1180
1181
1182
1183

        return hidden_states, output_states


Patrick von Platen's avatar
Patrick von Platen committed
1184
class DownBlock2D(nn.Module):
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
1197
1198
1199
        output_scale_factor: float = 1.0,
        add_downsample: bool = True,
        downsample_padding: int = 1,
1200
1201
1202
1203
1204
1205
1206
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
1207
                ResnetBlock2D(
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )

        self.resnets = nn.ModuleList(resnets)

        if add_downsample:
            self.downsamplers = nn.ModuleList(
Patrick von Platen's avatar
Patrick von Platen committed
1225
1226
                [
                    Downsample2D(
1227
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
Patrick von Platen's avatar
Patrick von Platen committed
1228
1229
                    )
                ]
1230
1231
1232
1233
            )
        else:
            self.downsamplers = None

1234
1235
        self.gradient_checkpointing = False

1236
1237
1238
    def forward(
        self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
1239
1240
1241
        output_states = ()

        for resnet in self.resnets:
1242
1243
1244
1245
1246
1247
1248
1249
            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)

                    return custom_forward

1250
1251
1252
1253
1254
1255
1256
1257
                if is_torch_version(">=", "1.11.0"):
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
                    )
                else:
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb
                    )
1258
            else:
1259
                hidden_states = resnet(hidden_states, temb, scale=scale)
1260

1261
            output_states = output_states + (hidden_states,)
1262
1263
1264

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
1265
                hidden_states = downsampler(hidden_states, scale=scale)
1266

1267
            output_states = output_states + (hidden_states,)
1268
1269
1270
1271

        return hidden_states, output_states


1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
class DownEncoderBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
1284
1285
1286
        output_scale_factor: float = 1.0,
        add_downsample: bool = True,
        downsample_padding: int = 1,
1287
1288
1289
1290
1291
1292
1293
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
1294
                ResnetBlock2D(
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=None,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )

        self.resnets = nn.ModuleList(resnets)

        if add_downsample:
            self.downsamplers = nn.ModuleList(
                [
                    Downsample2D(
1314
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
1315
1316
1317
1318
1319
1320
                    )
                ]
            )
        else:
            self.downsamplers = None

1321
    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
1322
        for resnet in self.resnets:
1323
            hidden_states = resnet(hidden_states, temb=None, scale=scale)
1324
1325
1326

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
1327
                hidden_states = downsampler(hidden_states, scale)
1328
1329
1330
1331

        return hidden_states


1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
class AttnDownEncoderBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
1344
1345
1346
1347
        attention_head_dim: int = 1,
        output_scale_factor: float = 1.0,
        add_downsample: bool = True,
        downsample_padding: int = 1,
1348
1349
1350
1351
1352
    ):
        super().__init__()
        resnets = []
        attentions = []

1353
1354
1355
1356
1357
1358
        if attention_head_dim is None:
            logger.warn(
                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
            )
            attention_head_dim = out_channels

1359
1360
1361
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
1362
                ResnetBlock2D(
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=None,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )
            attentions.append(
1376
                Attention(
1377
                    out_channels,
1378
1379
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
1380
1381
                    rescale_output_factor=output_scale_factor,
                    eps=resnet_eps,
Will Berman's avatar
Will Berman committed
1382
                    norm_num_groups=resnet_groups,
1383
1384
1385
1386
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
                )
            )

        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

        if add_downsample:
            self.downsamplers = nn.ModuleList(
                [
                    Downsample2D(
1397
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
1398
1399
1400
1401
1402
1403
                    )
                ]
            )
        else:
            self.downsamplers = None

1404
    def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
1405
        for resnet, attn in zip(self.resnets, self.attentions):
1406
1407
1408
            hidden_states = resnet(hidden_states, temb=None, scale=scale)
            cross_attention_kwargs = {"scale": scale}
            hidden_states = attn(hidden_states, **cross_attention_kwargs)
1409
1410
1411

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
1412
                hidden_states = downsampler(hidden_states, scale)
1413
1414
1415
1416

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
1417
class AttnSkipDownBlock2D(nn.Module):
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_pre_norm: bool = True,
1429
1430
1431
        attention_head_dim: int = 1,
        output_scale_factor: float = np.sqrt(2.0),
        add_downsample: bool = True,
1432
1433
1434
1435
1436
    ):
        super().__init__()
        self.attentions = nn.ModuleList([])
        self.resnets = nn.ModuleList([])

1437
1438
1439
1440
1441
1442
        if attention_head_dim is None:
            logger.warn(
                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
            )
            attention_head_dim = out_channels

1443
1444
1445
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            self.resnets.append(
1446
                ResnetBlock2D(
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=min(in_channels // 4, 32),
                    groups_out=min(out_channels // 4, 32),
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )
            self.attentions.append(
1461
                Attention(
1462
                    out_channels,
1463
1464
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
1465
1466
                    rescale_output_factor=output_scale_factor,
                    eps=resnet_eps,
1467
1468
1469
1470
1471
                    norm_num_groups=32,
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
1472
1473
1474
1475
                )
            )

        if add_downsample:
1476
            self.resnet_down = ResnetBlock2D(
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
                in_channels=out_channels,
                out_channels=out_channels,
                temb_channels=temb_channels,
                eps=resnet_eps,
                groups=min(out_channels // 4, 32),
                dropout=dropout,
                time_embedding_norm=resnet_time_scale_shift,
                non_linearity=resnet_act_fn,
                output_scale_factor=output_scale_factor,
                pre_norm=resnet_pre_norm,
1487
                use_in_shortcut=True,
1488
1489
1490
                down=True,
                kernel="fir",
            )
1491
            self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
1492
1493
1494
1495
1496
1497
            self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1))
        else:
            self.resnet_down = None
            self.downsamplers = None
            self.skip_conv = None

1498
1499
1500
1501
1502
1503
1504
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        temb: Optional[torch.FloatTensor] = None,
        skip_sample: Optional[torch.FloatTensor] = None,
        scale: float = 1.0,
    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...], torch.FloatTensor]:
1505
1506
1507
        output_states = ()

        for resnet, attn in zip(self.resnets, self.attentions):
1508
1509
1510
            hidden_states = resnet(hidden_states, temb, scale=scale)
            cross_attention_kwargs = {"scale": scale}
            hidden_states = attn(hidden_states, **cross_attention_kwargs)
1511
1512
1513
            output_states += (hidden_states,)

        if self.downsamplers is not None:
1514
            hidden_states = self.resnet_down(hidden_states, temb, scale=scale)
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
            for downsampler in self.downsamplers:
                skip_sample = downsampler(skip_sample)

            hidden_states = self.skip_conv(skip_sample) + hidden_states

            output_states += (hidden_states,)

        return hidden_states, output_states, skip_sample


Patrick von Platen's avatar
Patrick von Platen committed
1525
class SkipDownBlock2D(nn.Module):
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_pre_norm: bool = True,
1537
1538
1539
        output_scale_factor: float = np.sqrt(2.0),
        add_downsample: bool = True,
        downsample_padding: int = 1,
1540
1541
1542
1543
1544
1545
1546
    ):
        super().__init__()
        self.resnets = nn.ModuleList([])

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            self.resnets.append(
1547
                ResnetBlock2D(
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=min(in_channels // 4, 32),
                    groups_out=min(out_channels // 4, 32),
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )

        if add_downsample:
1563
            self.resnet_down = ResnetBlock2D(
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
                in_channels=out_channels,
                out_channels=out_channels,
                temb_channels=temb_channels,
                eps=resnet_eps,
                groups=min(out_channels // 4, 32),
                dropout=dropout,
                time_embedding_norm=resnet_time_scale_shift,
                non_linearity=resnet_act_fn,
                output_scale_factor=output_scale_factor,
                pre_norm=resnet_pre_norm,
1574
                use_in_shortcut=True,
1575
1576
1577
                down=True,
                kernel="fir",
            )
1578
            self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
1579
1580
1581
1582
1583
1584
            self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1))
        else:
            self.resnet_down = None
            self.downsamplers = None
            self.skip_conv = None

1585
1586
1587
1588
1589
1590
1591
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        temb: Optional[torch.FloatTensor] = None,
        skip_sample: Optional[torch.FloatTensor] = None,
        scale: float = 1.0,
    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...], torch.FloatTensor]:
1592
1593
1594
        output_states = ()

        for resnet in self.resnets:
1595
            hidden_states = resnet(hidden_states, temb, scale)
1596
1597
1598
            output_states += (hidden_states,)

        if self.downsamplers is not None:
1599
            hidden_states = self.resnet_down(hidden_states, temb, scale)
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
            for downsampler in self.downsamplers:
                skip_sample = downsampler(skip_sample)

            hidden_states = self.skip_conv(skip_sample) + hidden_states

            output_states += (hidden_states,)

        return hidden_states, output_states, skip_sample


Will Berman's avatar
Will Berman committed
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
class ResnetDownsampleBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
1623
1624
1625
        output_scale_factor: float = 1.0,
        add_downsample: bool = True,
        skip_time_act: bool = False,
Will Berman's avatar
Will Berman committed
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
                ResnetBlock2D(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
1644
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
                )
            )

        self.resnets = nn.ModuleList(resnets)

        if add_downsample:
            self.downsamplers = nn.ModuleList(
                [
                    ResnetBlock2D(
                        in_channels=out_channels,
                        out_channels=out_channels,
                        temb_channels=temb_channels,
                        eps=resnet_eps,
                        groups=resnet_groups,
                        dropout=dropout,
                        time_embedding_norm=resnet_time_scale_shift,
                        non_linearity=resnet_act_fn,
                        output_scale_factor=output_scale_factor,
                        pre_norm=resnet_pre_norm,
1664
                        skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
1665
1666
1667
1668
1669
1670
1671
1672
1673
                        down=True,
                    )
                ]
            )
        else:
            self.downsamplers = None

        self.gradient_checkpointing = False

1674
1675
1676
    def forward(
        self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
Will Berman's avatar
Will Berman committed
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
        output_states = ()

        for resnet in self.resnets:
            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)

                    return custom_forward

1688
1689
1690
1691
1692
1693
1694
1695
                if is_torch_version(">=", "1.11.0"):
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
                    )
                else:
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb
                    )
Will Berman's avatar
Will Berman committed
1696
            else:
1697
                hidden_states = resnet(hidden_states, temb, scale)
Will Berman's avatar
Will Berman committed
1698

1699
            output_states = output_states + (hidden_states,)
Will Berman's avatar
Will Berman committed
1700
1701
1702

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
1703
                hidden_states = downsampler(hidden_states, temb, scale)
Will Berman's avatar
Will Berman committed
1704

1705
            output_states = output_states + (hidden_states,)
Will Berman's avatar
Will Berman committed
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722

        return hidden_states, output_states


class SimpleCrossAttnDownBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
1723
1724
1725
1726
1727
1728
1729
        attention_head_dim: int = 1,
        cross_attention_dim: int = 1280,
        output_scale_factor: float = 1.0,
        add_downsample: bool = True,
        skip_time_act: bool = False,
        only_cross_attention: bool = False,
        cross_attention_norm: Optional[str] = None,
Will Berman's avatar
Will Berman committed
1730
1731
1732
1733
1734
1735
1736
1737
    ):
        super().__init__()

        self.has_cross_attention = True

        resnets = []
        attentions = []

1738
1739
        self.attention_head_dim = attention_head_dim
        self.num_heads = out_channels // self.attention_head_dim
Will Berman's avatar
Will Berman committed
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
                ResnetBlock2D(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
1755
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
1756
1757
                )
            )
1758
1759
1760
1761
1762

            processor = (
                AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()
            )

Will Berman's avatar
Will Berman committed
1763
            attentions.append(
Patrick von Platen's avatar
Patrick von Platen committed
1764
                Attention(
Will Berman's avatar
Will Berman committed
1765
1766
1767
                    query_dim=out_channels,
                    cross_attention_dim=out_channels,
                    heads=self.num_heads,
1768
                    dim_head=attention_head_dim,
Will Berman's avatar
Will Berman committed
1769
1770
1771
1772
                    added_kv_proj_dim=cross_attention_dim,
                    norm_num_groups=resnet_groups,
                    bias=True,
                    upcast_softmax=True,
1773
                    only_cross_attention=only_cross_attention,
1774
                    cross_attention_norm=cross_attention_norm,
1775
                    processor=processor,
Will Berman's avatar
Will Berman committed
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
                )
            )
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

        if add_downsample:
            self.downsamplers = nn.ModuleList(
                [
                    ResnetBlock2D(
                        in_channels=out_channels,
                        out_channels=out_channels,
                        temb_channels=temb_channels,
                        eps=resnet_eps,
                        groups=resnet_groups,
                        dropout=dropout,
                        time_embedding_norm=resnet_time_scale_shift,
                        non_linearity=resnet_act_fn,
                        output_scale_factor=output_scale_factor,
                        pre_norm=resnet_pre_norm,
1795
                        skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
1796
1797
1798
1799
1800
1801
1802
1803
1804
                        down=True,
                    )
                ]
            )
        else:
            self.downsamplers = None

        self.gradient_checkpointing = False

1805
    def forward(
1806
1807
1808
1809
1810
1811
1812
        self,
        hidden_states: torch.FloatTensor,
        temb: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
1813
    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
Will Berman's avatar
Will Berman committed
1814
        output_states = ()
1815
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
Will Berman's avatar
Will Berman committed
1816

1817
1818
        lora_scale = cross_attention_kwargs.get("scale", 1.0)

1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
        if attention_mask is None:
            # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask.
            mask = None if encoder_hidden_states is None else encoder_attention_mask
        else:
            # when attention_mask is defined: we don't even check for encoder_attention_mask.
            # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks.
            # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask.
            #       then we can simplify this whole if/else block to:
            #         mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask
            mask = attention_mask

Will Berman's avatar
Will Berman committed
1830
        for resnet, attn in zip(self.resnets, self.attentions):
1831
            if self.training and self.gradient_checkpointing:
Will Berman's avatar
Will Berman committed
1832

1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)

                    return custom_forward

                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
ethansmith2000's avatar
ethansmith2000 committed
1843
                hidden_states = attn(
1844
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
1845
1846
1847
1848
                    encoder_hidden_states=encoder_hidden_states,
                    attention_mask=mask,
                    **cross_attention_kwargs,
                )
1849
            else:
1850
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
1851
1852
1853
1854

                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
1855
                    attention_mask=mask,
1856
1857
                    **cross_attention_kwargs,
                )
Will Berman's avatar
Will Berman committed
1858

1859
            output_states = output_states + (hidden_states,)
Will Berman's avatar
Will Berman committed
1860
1861
1862

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
1863
                hidden_states = downsampler(hidden_states, temb, scale=lora_scale)
Will Berman's avatar
Will Berman committed
1864

1865
            output_states = output_states + (hidden_states,)
Will Berman's avatar
Will Berman committed
1866
1867
1868
1869

        return hidden_states, output_states


1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
class KDownBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 4,
        resnet_eps: float = 1e-5,
        resnet_act_fn: str = "gelu",
        resnet_group_size: int = 32,
1881
        add_downsample: bool = False,
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            groups = in_channels // resnet_group_size
            groups_out = out_channels // resnet_group_size

            resnets.append(
                ResnetBlock2D(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    dropout=dropout,
                    temb_channels=temb_channels,
                    groups=groups,
                    groups_out=groups_out,
                    eps=resnet_eps,
                    non_linearity=resnet_act_fn,
                    time_embedding_norm="ada_group",
                    conv_shortcut_bias=False,
                )
            )

        self.resnets = nn.ModuleList(resnets)

        if add_downsample:
            # YiYi's comments- might be able to use FirDownsample2D, look into details later
            self.downsamplers = nn.ModuleList([KDownsample2D()])
        else:
            self.downsamplers = None

        self.gradient_checkpointing = False

1916
1917
1918
    def forward(
        self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
        output_states = ()

        for resnet in self.resnets:
            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)

                    return custom_forward

1930
1931
1932
1933
1934
1935
1936
1937
                if is_torch_version(">=", "1.11.0"):
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
                    )
                else:
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb
                    )
1938
            else:
1939
                hidden_states = resnet(hidden_states, temb, scale)
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959

            output_states += (hidden_states,)

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
                hidden_states = downsampler(hidden_states)

        return hidden_states, output_states


class KCrossAttnDownBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        cross_attention_dim: int,
        dropout: float = 0.0,
        num_layers: int = 4,
        resnet_group_size: int = 32,
1960
        add_downsample: bool = True,
1961
        attention_head_dim: int = 64,
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
        add_self_attention: bool = False,
        resnet_eps: float = 1e-5,
        resnet_act_fn: str = "gelu",
    ):
        super().__init__()
        resnets = []
        attentions = []

        self.has_cross_attention = True

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            groups = in_channels // resnet_group_size
            groups_out = out_channels // resnet_group_size

            resnets.append(
                ResnetBlock2D(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    dropout=dropout,
                    temb_channels=temb_channels,
                    groups=groups,
                    groups_out=groups_out,
                    eps=resnet_eps,
                    non_linearity=resnet_act_fn,
                    time_embedding_norm="ada_group",
                    conv_shortcut_bias=False,
                )
            )
            attentions.append(
                KAttentionBlock(
                    out_channels,
1994
1995
                    out_channels // attention_head_dim,
                    attention_head_dim,
1996
1997
1998
1999
                    cross_attention_dim=cross_attention_dim,
                    temb_channels=temb_channels,
                    attention_bias=True,
                    add_self_attention=add_self_attention,
2000
                    cross_attention_norm="layer_norm",
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
                    group_size=resnet_group_size,
                )
            )

        self.resnets = nn.ModuleList(resnets)
        self.attentions = nn.ModuleList(attentions)

        if add_downsample:
            self.downsamplers = nn.ModuleList([KDownsample2D()])
        else:
            self.downsamplers = None

        self.gradient_checkpointing = False

    def forward(
2016
2017
2018
2019
2020
2021
2022
        self,
        hidden_states: torch.FloatTensor,
        temb: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
2023
    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
2024
        output_states = ()
2025
        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038

        for resnet, attn in zip(self.resnets, self.attentions):
            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)

                    return custom_forward

2039
2040
2041
2042
2043
2044
2045
                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(resnet),
                    hidden_states,
                    temb,
                    **ckpt_kwargs,
                )
ethansmith2000's avatar
ethansmith2000 committed
2046
                hidden_states = attn(
2047
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
2048
2049
2050
2051
2052
                    encoder_hidden_states=encoder_hidden_states,
                    emb=temb,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
                    encoder_attention_mask=encoder_attention_mask,
2053
                )
2054
            else:
2055
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
2056
2057
2058
2059
2060
2061
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    emb=temb,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
2062
                    encoder_attention_mask=encoder_attention_mask,
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
                )

            if self.downsamplers is None:
                output_states += (None,)
            else:
                output_states += (hidden_states,)

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
                hidden_states = downsampler(hidden_states)

        return hidden_states, output_states


Patrick von Platen's avatar
Patrick von Platen committed
2077
class AttnUpBlock2D(nn.Module):
2078
2079
2080
    def __init__(
        self,
        in_channels: int,
Patrick von Platen's avatar
Patrick von Platen committed
2081
2082
        prev_output_channel: int,
        out_channels: int,
2083
        temb_channels: int,
2084
        resolution_idx: int = None,
2085
2086
2087
2088
2089
2090
2091
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
2092
2093
2094
        attention_head_dim: int = 1,
        output_scale_factor: float = 1.0,
        upsample_type: str = "conv",
2095
2096
2097
2098
2099
    ):
        super().__init__()
        resnets = []
        attentions = []

2100
2101
        self.upsample_type = upsample_type

2102
2103
2104
2105
2106
2107
        if attention_head_dim is None:
            logger.warn(
                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
            )
            attention_head_dim = out_channels

2108
        for i in range(num_layers):
Patrick von Platen's avatar
Patrick von Platen committed
2109
2110
2111
            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels

2112
            resnets.append(
2113
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
2114
2115
                    in_channels=resnet_in_channels + res_skip_channels,
                    out_channels=out_channels,
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )
            attentions.append(
2127
                Attention(
Patrick von Platen's avatar
Patrick von Platen committed
2128
                    out_channels,
2129
2130
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
2131
                    rescale_output_factor=output_scale_factor,
Patrick von Platen's avatar
Patrick von Platen committed
2132
                    eps=resnet_eps,
Will Berman's avatar
Will Berman committed
2133
                    norm_num_groups=resnet_groups,
2134
2135
2136
2137
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
2138
2139
2140
2141
2142
2143
                )
            )

        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

2144
        if upsample_type == "conv":
Patrick von Platen's avatar
Patrick von Platen committed
2145
            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
        elif upsample_type == "resnet":
            self.upsamplers = nn.ModuleList(
                [
                    ResnetBlock2D(
                        in_channels=out_channels,
                        out_channels=out_channels,
                        temb_channels=temb_channels,
                        eps=resnet_eps,
                        groups=resnet_groups,
                        dropout=dropout,
                        time_embedding_norm=resnet_time_scale_shift,
                        non_linearity=resnet_act_fn,
                        output_scale_factor=output_scale_factor,
                        pre_norm=resnet_pre_norm,
                        up=True,
                    )
                ]
            )
2164
2165
2166
        else:
            self.upsamplers = None

2167
2168
        self.resolution_idx = resolution_idx

2169
2170
2171
2172
2173
2174
2175
2176
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
        temb: Optional[torch.FloatTensor] = None,
        upsample_size: Optional[int] = None,
        scale: float = 1.0,
    ) -> torch.FloatTensor:
2177
2178
2179
2180
2181
2182
        for resnet, attn in zip(self.resnets, self.attentions):
            # pop res hidden states
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

2183
2184
2185
            hidden_states = resnet(hidden_states, temb, scale=scale)
            cross_attention_kwargs = {"scale": scale}
            hidden_states = attn(hidden_states, **cross_attention_kwargs)
2186
2187
2188

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2189
                if self.upsample_type == "resnet":
2190
                    hidden_states = upsampler(hidden_states, temb=temb, scale=scale)
2191
                else:
2192
                    hidden_states = upsampler(hidden_states, scale=scale)
2193
2194
2195
2196

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
2197
class CrossAttnUpBlock2D(nn.Module):
Patrick von Platen's avatar
Patrick von Platen committed
2198
2199
2200
2201
2202
2203
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        prev_output_channel: int,
        temb_channels: int,
2204
        resolution_idx: Optional[int] = None,
Patrick von Platen's avatar
Patrick von Platen committed
2205
2206
        dropout: float = 0.0,
        num_layers: int = 1,
2207
        transformer_layers_per_block: Union[int, Tuple[int]] = 1,
Patrick von Platen's avatar
Patrick von Platen committed
2208
2209
2210
2211
2212
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
2213
2214
2215
2216
2217
2218
2219
2220
2221
        num_attention_heads: int = 1,
        cross_attention_dim: int = 1280,
        output_scale_factor: float = 1.0,
        add_upsample: bool = True,
        dual_cross_attention: bool = False,
        use_linear_projection: bool = False,
        only_cross_attention: bool = False,
        upcast_attention: bool = False,
        attention_type: str = "default",
Patrick von Platen's avatar
Patrick von Platen committed
2222
2223
2224
2225
2226
    ):
        super().__init__()
        resnets = []
        attentions = []

2227
        self.has_cross_attention = True
2228
        self.num_attention_heads = num_attention_heads
Patrick von Platen's avatar
Patrick von Platen committed
2229

2230
2231
2232
        if isinstance(transformer_layers_per_block, int):
            transformer_layers_per_block = [transformer_layers_per_block] * num_layers

Patrick von Platen's avatar
Patrick von Platen committed
2233
2234
2235
2236
2237
        for i in range(num_layers):
            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels

            resnets.append(
2238
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
                    in_channels=resnet_in_channels + res_skip_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )
2251
2252
2253
            if not dual_cross_attention:
                attentions.append(
                    Transformer2DModel(
2254
2255
                        num_attention_heads,
                        out_channels // num_attention_heads,
2256
                        in_channels=out_channels,
2257
                        num_layers=transformer_layers_per_block[i],
2258
2259
                        cross_attention_dim=cross_attention_dim,
                        norm_num_groups=resnet_groups,
Suraj Patil's avatar
Suraj Patil committed
2260
                        use_linear_projection=use_linear_projection,
2261
                        only_cross_attention=only_cross_attention,
2262
                        upcast_attention=upcast_attention,
2263
                        attention_type=attention_type,
2264
2265
2266
2267
2268
                    )
                )
            else:
                attentions.append(
                    DualTransformer2DModel(
2269
2270
                        num_attention_heads,
                        out_channels // num_attention_heads,
2271
2272
2273
2274
2275
                        in_channels=out_channels,
                        num_layers=1,
                        cross_attention_dim=cross_attention_dim,
                        norm_num_groups=resnet_groups,
                    )
Patrick von Platen's avatar
Patrick von Platen committed
2276
2277
2278
2279
2280
2281
2282
2283
2284
                )
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

        if add_upsample:
            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
        else:
            self.upsamplers = None

2285
        self.gradient_checkpointing = False
2286
        self.resolution_idx = resolution_idx
2287
2288
2289

    def forward(
        self,
2290
2291
2292
2293
2294
2295
2296
2297
        hidden_states: torch.FloatTensor,
        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
        temb: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        upsample_size: Optional[int] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
2298
    ) -> torch.FloatTensor:
2299
        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
2300
2301
2302
2303
2304
2305
        is_freeu_enabled = (
            getattr(self, "s1", None)
            and getattr(self, "s2", None)
            and getattr(self, "b1", None)
            and getattr(self, "b2", None)
        )
2306

Patrick von Platen's avatar
Patrick von Platen committed
2307
2308
2309
2310
        for resnet, attn in zip(self.resnets, self.attentions):
            # pop res hidden states
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323

            # FreeU: Only operate on the first two stages
            if is_freeu_enabled:
                hidden_states, res_hidden_states = apply_freeu(
                    self.resolution_idx,
                    hidden_states,
                    res_hidden_states,
                    s1=self.s1,
                    s2=self.s2,
                    b1=self.b1,
                    b2=self.b2,
                )

Patrick von Platen's avatar
Patrick von Platen committed
2324
2325
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

2326
2327
            if self.training and self.gradient_checkpointing:

Will Berman's avatar
Will Berman committed
2328
                def create_custom_forward(module, return_dict=None):
2329
                    def custom_forward(*inputs):
Will Berman's avatar
Will Berman committed
2330
2331
2332
2333
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)
2334
2335
2336

                    return custom_forward

2337
2338
2339
2340
2341
2342
2343
                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(resnet),
                    hidden_states,
                    temb,
                    **ckpt_kwargs,
                )
ethansmith2000's avatar
ethansmith2000 committed
2344
                hidden_states = attn(
2345
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
2346
2347
2348
2349
2350
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
                    return_dict=False,
2351
                )[0]
2352
            else:
2353
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
2354
2355
2356
2357
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
2358
2359
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
2360
2361
                    return_dict=False,
                )[0]
Patrick von Platen's avatar
Patrick von Platen committed
2362
2363
2364

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2365
                hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale)
Patrick von Platen's avatar
Patrick von Platen committed
2366
2367
2368
2369

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
2370
class UpBlock2D(nn.Module):
2371
2372
2373
    def __init__(
        self,
        in_channels: int,
Patrick von Platen's avatar
Patrick von Platen committed
2374
2375
        prev_output_channel: int,
        out_channels: int,
2376
        temb_channels: int,
2377
        resolution_idx: Optional[int] = None,
2378
2379
2380
2381
2382
2383
2384
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
2385
2386
        output_scale_factor: float = 1.0,
        add_upsample: bool = True,
2387
2388
2389
2390
2391
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
Patrick von Platen's avatar
Patrick von Platen committed
2392
2393
2394
            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels

2395
            resnets.append(
2396
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
2397
2398
                    in_channels=resnet_in_channels + res_skip_channels,
                    out_channels=out_channels,
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )

        self.resnets = nn.ModuleList(resnets)

        if add_upsample:
Patrick von Platen's avatar
Patrick von Platen committed
2413
            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
2414
2415
2416
        else:
            self.upsamplers = None

2417
        self.gradient_checkpointing = False
2418
        self.resolution_idx = resolution_idx
2419

2420
2421
2422
2423
2424
2425
2426
2427
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
        temb: Optional[torch.FloatTensor] = None,
        upsample_size: Optional[int] = None,
        scale: float = 1.0,
    ) -> torch.FloatTensor:
2428
2429
2430
2431
2432
2433
2434
        is_freeu_enabled = (
            getattr(self, "s1", None)
            and getattr(self, "s2", None)
            and getattr(self, "b1", None)
            and getattr(self, "b2", None)
        )

2435
2436
2437
2438
        for resnet in self.resnets:
            # pop res hidden states
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451

            # FreeU: Only operate on the first two stages
            if is_freeu_enabled:
                hidden_states, res_hidden_states = apply_freeu(
                    self.resolution_idx,
                    hidden_states,
                    res_hidden_states,
                    s1=self.s1,
                    s2=self.s2,
                    b1=self.b1,
                    b2=self.b2,
                )

2452
2453
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

2454
2455
2456
2457
2458
2459
2460
2461
            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)

                    return custom_forward

2462
2463
2464
2465
2466
2467
2468
2469
                if is_torch_version(">=", "1.11.0"):
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
                    )
                else:
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb
                    )
2470
            else:
2471
                hidden_states = resnet(hidden_states, temb, scale=scale)
2472
2473
2474

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2475
                hidden_states = upsampler(hidden_states, upsample_size, scale=scale)
2476
2477

        return hidden_states
2478
2479


2480
2481
2482
2483
2484
class UpDecoderBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
2485
        resolution_idx: Optional[int] = None,
2486
2487
2488
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
YiYi Xu's avatar
YiYi Xu committed
2489
        resnet_time_scale_shift: str = "default",  # default, spatial
2490
2491
2492
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
2493
2494
2495
        output_scale_factor: float = 1.0,
        add_upsample: bool = True,
        temb_channels: Optional[int] = None,
2496
2497
2498
2499
2500
2501
2502
2503
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
            input_channels = in_channels if i == 0 else out_channels

            resnets.append(
2504
                ResnetBlock2D(
2505
2506
                    in_channels=input_channels,
                    out_channels=out_channels,
YiYi Xu's avatar
YiYi Xu committed
2507
                    temb_channels=temb_channels,
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )

        self.resnets = nn.ModuleList(resnets)

        if add_upsample:
            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
        else:
            self.upsamplers = None

2525
2526
        self.resolution_idx = resolution_idx

2527
2528
2529
    def forward(
        self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
    ) -> torch.FloatTensor:
2530
        for resnet in self.resnets:
2531
            hidden_states = resnet(hidden_states, temb=temb, scale=scale)
2532
2533
2534
2535
2536
2537
2538
2539

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
                hidden_states = upsampler(hidden_states)

        return hidden_states


2540
2541
2542
2543
2544
class AttnUpDecoderBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
2545
        resolution_idx: Optional[int] = None,
2546
2547
2548
2549
2550
2551
2552
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
2553
2554
2555
2556
        attention_head_dim: int = 1,
        output_scale_factor: float = 1.0,
        add_upsample: bool = True,
        temb_channels: Optional[int] = None,
2557
2558
2559
2560
2561
    ):
        super().__init__()
        resnets = []
        attentions = []

2562
2563
2564
2565
2566
2567
        if attention_head_dim is None:
            logger.warn(
                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}."
            )
            attention_head_dim = out_channels

2568
2569
2570
2571
        for i in range(num_layers):
            input_channels = in_channels if i == 0 else out_channels

            resnets.append(
2572
                ResnetBlock2D(
2573
2574
                    in_channels=input_channels,
                    out_channels=out_channels,
YiYi Xu's avatar
YiYi Xu committed
2575
                    temb_channels=temb_channels,
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )
            attentions.append(
2586
                Attention(
2587
                    out_channels,
2588
2589
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
2590
2591
                    rescale_output_factor=output_scale_factor,
                    eps=resnet_eps,
2592
                    norm_num_groups=resnet_groups if resnet_time_scale_shift != "spatial" else None,
YiYi Xu's avatar
YiYi Xu committed
2593
                    spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
2594
2595
2596
2597
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
                )
            )

        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

        if add_upsample:
            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
        else:
            self.upsamplers = None

2609
2610
        self.resolution_idx = resolution_idx

2611
2612
2613
    def forward(
        self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
    ) -> torch.FloatTensor:
2614
        for resnet, attn in zip(self.resnets, self.attentions):
2615
2616
2617
            hidden_states = resnet(hidden_states, temb=temb, scale=scale)
            cross_attention_kwargs = {"scale": scale}
            hidden_states = attn(hidden_states, temb=temb, **cross_attention_kwargs)
2618
2619
2620

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2621
                hidden_states = upsampler(hidden_states, scale=scale)
2622
2623
2624
2625

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
2626
class AttnSkipUpBlock2D(nn.Module):
2627
2628
2629
2630
2631
2632
    def __init__(
        self,
        in_channels: int,
        prev_output_channel: int,
        out_channels: int,
        temb_channels: int,
2633
        resolution_idx: Optional[int] = None,
2634
2635
2636
2637
2638
2639
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_pre_norm: bool = True,
2640
2641
2642
        attention_head_dim: int = 1,
        output_scale_factor: float = np.sqrt(2.0),
        add_upsample: bool = True,
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
    ):
        super().__init__()
        self.attentions = nn.ModuleList([])
        self.resnets = nn.ModuleList([])

        for i in range(num_layers):
            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels

            self.resnets.append(
2653
                ResnetBlock2D(
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
                    in_channels=resnet_in_channels + res_skip_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=min(resnet_in_channels + res_skip_channels // 4, 32),
                    groups_out=min(out_channels // 4, 32),
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )

2668
2669
2670
2671
2672
2673
        if attention_head_dim is None:
            logger.warn(
                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}."
            )
            attention_head_dim = out_channels

2674
        self.attentions.append(
2675
            Attention(
2676
                out_channels,
2677
2678
                heads=out_channels // attention_head_dim,
                dim_head=attention_head_dim,
2679
2680
                rescale_output_factor=output_scale_factor,
                eps=resnet_eps,
2681
2682
2683
2684
2685
                norm_num_groups=32,
                residual_connection=True,
                bias=True,
                upcast_softmax=True,
                _from_deprecated_attn_block=True,
2686
2687
2688
2689
2690
            )
        )

        self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels)
        if add_upsample:
2691
            self.resnet_up = ResnetBlock2D(
2692
2693
2694
2695
2696
2697
2698
2699
2700
2701
2702
                in_channels=out_channels,
                out_channels=out_channels,
                temb_channels=temb_channels,
                eps=resnet_eps,
                groups=min(out_channels // 4, 32),
                groups_out=min(out_channels // 4, 32),
                dropout=dropout,
                time_embedding_norm=resnet_time_scale_shift,
                non_linearity=resnet_act_fn,
                output_scale_factor=output_scale_factor,
                pre_norm=resnet_pre_norm,
2703
                use_in_shortcut=True,
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
2716
2717
                up=True,
                kernel="fir",
            )
            self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            self.skip_norm = torch.nn.GroupNorm(
                num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True
            )
            self.act = nn.SiLU()
        else:
            self.resnet_up = None
            self.skip_conv = None
            self.skip_norm = None
            self.act = None

2718
2719
        self.resolution_idx = resolution_idx

2720
2721
2722
2723
2724
2725
2726
2727
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
        temb: Optional[torch.FloatTensor] = None,
        skip_sample=None,
        scale: float = 1.0,
    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
2728
2729
2730
2731
2732
2733
        for resnet in self.resnets:
            # pop res hidden states
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

2734
            hidden_states = resnet(hidden_states, temb, scale=scale)
2735

2736
2737
        cross_attention_kwargs = {"scale": scale}
        hidden_states = self.attentions[0](hidden_states, **cross_attention_kwargs)
2738
2739
2740
2741
2742
2743
2744
2745
2746
2747
2748
2749
2750

        if skip_sample is not None:
            skip_sample = self.upsampler(skip_sample)
        else:
            skip_sample = 0

        if self.resnet_up is not None:
            skip_sample_states = self.skip_norm(hidden_states)
            skip_sample_states = self.act(skip_sample_states)
            skip_sample_states = self.skip_conv(skip_sample_states)

            skip_sample = skip_sample + skip_sample_states

2751
            hidden_states = self.resnet_up(hidden_states, temb, scale=scale)
2752
2753
2754
2755

        return hidden_states, skip_sample


Patrick von Platen's avatar
Patrick von Platen committed
2756
class SkipUpBlock2D(nn.Module):
2757
2758
2759
2760
2761
2762
    def __init__(
        self,
        in_channels: int,
        prev_output_channel: int,
        out_channels: int,
        temb_channels: int,
2763
        resolution_idx: Optional[int] = None,
2764
2765
2766
2767
2768
2769
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_pre_norm: bool = True,
2770
2771
2772
        output_scale_factor: float = np.sqrt(2.0),
        add_upsample: bool = True,
        upsample_padding: int = 1,
2773
2774
2775
2776
2777
2778
2779
2780
2781
    ):
        super().__init__()
        self.resnets = nn.ModuleList([])

        for i in range(num_layers):
            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels

            self.resnets.append(
2782
                ResnetBlock2D(
2783
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
                    in_channels=resnet_in_channels + res_skip_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=min((resnet_in_channels + res_skip_channels) // 4, 32),
                    groups_out=min(out_channels // 4, 32),
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
                )
            )

        self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels)
        if add_upsample:
2799
            self.resnet_up = ResnetBlock2D(
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
                in_channels=out_channels,
                out_channels=out_channels,
                temb_channels=temb_channels,
                eps=resnet_eps,
                groups=min(out_channels // 4, 32),
                groups_out=min(out_channels // 4, 32),
                dropout=dropout,
                time_embedding_norm=resnet_time_scale_shift,
                non_linearity=resnet_act_fn,
                output_scale_factor=output_scale_factor,
                pre_norm=resnet_pre_norm,
2811
                use_in_shortcut=True,
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
2824
2825
                up=True,
                kernel="fir",
            )
            self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            self.skip_norm = torch.nn.GroupNorm(
                num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True
            )
            self.act = nn.SiLU()
        else:
            self.resnet_up = None
            self.skip_conv = None
            self.skip_norm = None
            self.act = None

2826
2827
        self.resolution_idx = resolution_idx

2828
2829
2830
2831
2832
2833
2834
2835
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
        temb: Optional[torch.FloatTensor] = None,
        skip_sample=None,
        scale: float = 1.0,
    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
2836
2837
2838
2839
2840
2841
        for resnet in self.resnets:
            # pop res hidden states
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

2842
            hidden_states = resnet(hidden_states, temb, scale=scale)
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
2855

        if skip_sample is not None:
            skip_sample = self.upsampler(skip_sample)
        else:
            skip_sample = 0

        if self.resnet_up is not None:
            skip_sample_states = self.skip_norm(hidden_states)
            skip_sample_states = self.act(skip_sample_states)
            skip_sample_states = self.skip_conv(skip_sample_states)

            skip_sample = skip_sample + skip_sample_states

2856
            hidden_states = self.resnet_up(hidden_states, temb, scale=scale)
2857
2858

        return hidden_states, skip_sample
Will Berman's avatar
Will Berman committed
2859
2860
2861
2862
2863
2864
2865
2866
2867


class ResnetUpsampleBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        prev_output_channel: int,
        out_channels: int,
        temb_channels: int,
2868
        resolution_idx: Optional[int] = None,
Will Berman's avatar
Will Berman committed
2869
2870
2871
2872
2873
2874
2875
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
2876
2877
2878
        output_scale_factor: float = 1.0,
        add_upsample: bool = True,
        skip_time_act: bool = False,
Will Berman's avatar
Will Berman committed
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels

            resnets.append(
                ResnetBlock2D(
                    in_channels=resnet_in_channels + res_skip_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
2899
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
                )
            )

        self.resnets = nn.ModuleList(resnets)

        if add_upsample:
            self.upsamplers = nn.ModuleList(
                [
                    ResnetBlock2D(
                        in_channels=out_channels,
                        out_channels=out_channels,
                        temb_channels=temb_channels,
                        eps=resnet_eps,
                        groups=resnet_groups,
                        dropout=dropout,
                        time_embedding_norm=resnet_time_scale_shift,
                        non_linearity=resnet_act_fn,
                        output_scale_factor=output_scale_factor,
                        pre_norm=resnet_pre_norm,
2919
                        skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
2920
2921
2922
2923
2924
2925
2926
2927
                        up=True,
                    )
                ]
            )
        else:
            self.upsamplers = None

        self.gradient_checkpointing = False
2928
        self.resolution_idx = resolution_idx
Will Berman's avatar
Will Berman committed
2929

2930
2931
2932
2933
2934
2935
2936
2937
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
        temb: Optional[torch.FloatTensor] = None,
        upsample_size: Optional[int] = None,
        scale: float = 1.0,
    ) -> torch.FloatTensor:
Will Berman's avatar
Will Berman committed
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
        for resnet in self.resnets:
            # pop res hidden states
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)

                    return custom_forward

2952
2953
2954
2955
2956
2957
2958
2959
                if is_torch_version(">=", "1.11.0"):
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
                    )
                else:
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb
                    )
Will Berman's avatar
Will Berman committed
2960
            else:
2961
                hidden_states = resnet(hidden_states, temb, scale=scale)
Will Berman's avatar
Will Berman committed
2962
2963
2964

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2965
                hidden_states = upsampler(hidden_states, temb, scale=scale)
Will Berman's avatar
Will Berman committed
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976

        return hidden_states


class SimpleCrossAttnUpBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        prev_output_channel: int,
        temb_channels: int,
2977
        resolution_idx: Optional[int] = None,
Will Berman's avatar
Will Berman committed
2978
2979
2980
2981
2982
2983
2984
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
        resnet_time_scale_shift: str = "default",
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
2985
2986
2987
2988
2989
2990
2991
        attention_head_dim: int = 1,
        cross_attention_dim: int = 1280,
        output_scale_factor: float = 1.0,
        add_upsample: bool = True,
        skip_time_act: bool = False,
        only_cross_attention: bool = False,
        cross_attention_norm: Optional[str] = None,
Will Berman's avatar
Will Berman committed
2992
2993
2994
2995
2996
2997
    ):
        super().__init__()
        resnets = []
        attentions = []

        self.has_cross_attention = True
2998
        self.attention_head_dim = attention_head_dim
Will Berman's avatar
Will Berman committed
2999

3000
        self.num_heads = out_channels // self.attention_head_dim
Will Berman's avatar
Will Berman committed
3001
3002
3003
3004
3005
3006
3007
3008
3009
3010
3011
3012
3013
3014
3015
3016
3017

        for i in range(num_layers):
            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels

            resnets.append(
                ResnetBlock2D(
                    in_channels=resnet_in_channels + res_skip_channels,
                    out_channels=out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=resnet_groups,
                    dropout=dropout,
                    time_embedding_norm=resnet_time_scale_shift,
                    non_linearity=resnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=resnet_pre_norm,
3018
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
3019
3020
                )
            )
3021
3022
3023
3024
3025

            processor = (
                AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()
            )

Will Berman's avatar
Will Berman committed
3026
            attentions.append(
Patrick von Platen's avatar
Patrick von Platen committed
3027
                Attention(
Will Berman's avatar
Will Berman committed
3028
3029
3030
                    query_dim=out_channels,
                    cross_attention_dim=out_channels,
                    heads=self.num_heads,
3031
                    dim_head=self.attention_head_dim,
Will Berman's avatar
Will Berman committed
3032
3033
3034
3035
                    added_kv_proj_dim=cross_attention_dim,
                    norm_num_groups=resnet_groups,
                    bias=True,
                    upcast_softmax=True,
3036
                    only_cross_attention=only_cross_attention,
3037
                    cross_attention_norm=cross_attention_norm,
3038
                    processor=processor,
Will Berman's avatar
Will Berman committed
3039
3040
3041
3042
3043
3044
3045
3046
3047
3048
3049
3050
3051
3052
3053
3054
3055
3056
3057
                )
            )
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

        if add_upsample:
            self.upsamplers = nn.ModuleList(
                [
                    ResnetBlock2D(
                        in_channels=out_channels,
                        out_channels=out_channels,
                        temb_channels=temb_channels,
                        eps=resnet_eps,
                        groups=resnet_groups,
                        dropout=dropout,
                        time_embedding_norm=resnet_time_scale_shift,
                        non_linearity=resnet_act_fn,
                        output_scale_factor=output_scale_factor,
                        pre_norm=resnet_pre_norm,
3058
                        skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
3059
3060
3061
3062
3063
3064
3065
3066
                        up=True,
                    )
                ]
            )
        else:
            self.upsamplers = None

        self.gradient_checkpointing = False
3067
        self.resolution_idx = resolution_idx
Will Berman's avatar
Will Berman committed
3068
3069
3070

    def forward(
        self,
3071
3072
3073
3074
3075
3076
3077
3078
        hidden_states: torch.FloatTensor,
        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
        temb: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        upsample_size: Optional[int] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
3079
    ) -> torch.FloatTensor:
3080
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
3081

3082
        lora_scale = cross_attention_kwargs.get("scale", 1.0)
3083
3084
3085
3086
3087
3088
3089
3090
3091
3092
3093
        if attention_mask is None:
            # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask.
            mask = None if encoder_hidden_states is None else encoder_attention_mask
        else:
            # when attention_mask is defined: we don't even check for encoder_attention_mask.
            # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks.
            # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask.
            #       then we can simplify this whole if/else block to:
            #         mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask
            mask = attention_mask

Will Berman's avatar
Will Berman committed
3094
3095
3096
3097
3098
3099
3100
        for resnet, attn in zip(self.resnets, self.attentions):
            # resnet
            # pop res hidden states
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

3101
            if self.training and self.gradient_checkpointing:
Will Berman's avatar
Will Berman committed
3102

3103
3104
3105
3106
3107
3108
3109
3110
3111
3112
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)

                    return custom_forward

                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
ethansmith2000's avatar
ethansmith2000 committed
3113
                hidden_states = attn(
3114
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
3115
3116
3117
3118
                    encoder_hidden_states=encoder_hidden_states,
                    attention_mask=mask,
                    **cross_attention_kwargs,
                )
3119
            else:
3120
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
3121
3122
3123
3124

                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
3125
                    attention_mask=mask,
3126
3127
                    **cross_attention_kwargs,
                )
Will Berman's avatar
Will Berman committed
3128
3129
3130

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
3131
                hidden_states = upsampler(hidden_states, temb, scale=lora_scale)
Will Berman's avatar
Will Berman committed
3132
3133

        return hidden_states
3134
3135
3136
3137
3138
3139
3140
3141


class KUpBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
3142
        resolution_idx: int,
3143
3144
3145
3146
3147
        dropout: float = 0.0,
        num_layers: int = 5,
        resnet_eps: float = 1e-5,
        resnet_act_fn: str = "gelu",
        resnet_group_size: Optional[int] = 32,
3148
        add_upsample: bool = True,
3149
3150
3151
3152
3153
3154
3155
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
3166
3167
3168
3169
3170
3171
3172
3173
3174
3175
3176
3177
3178
3179
3180
3181
3182
3183
    ):
        super().__init__()
        resnets = []
        k_in_channels = 2 * out_channels
        k_out_channels = in_channels
        num_layers = num_layers - 1

        for i in range(num_layers):
            in_channels = k_in_channels if i == 0 else out_channels
            groups = in_channels // resnet_group_size
            groups_out = out_channels // resnet_group_size

            resnets.append(
                ResnetBlock2D(
                    in_channels=in_channels,
                    out_channels=k_out_channels if (i == num_layers - 1) else out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=groups,
                    groups_out=groups_out,
                    dropout=dropout,
                    non_linearity=resnet_act_fn,
                    time_embedding_norm="ada_group",
                    conv_shortcut_bias=False,
                )
            )

        self.resnets = nn.ModuleList(resnets)

        if add_upsample:
            self.upsamplers = nn.ModuleList([KUpsample2D()])
        else:
            self.upsamplers = None

        self.gradient_checkpointing = False
3184
        self.resolution_idx = resolution_idx
3185

3186
3187
3188
3189
3190
3191
3192
3193
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
        temb: Optional[torch.FloatTensor] = None,
        upsample_size: Optional[int] = None,
        scale: float = 1.0,
    ) -> torch.FloatTensor:
3194
3195
3196
3197
3198
3199
3200
3201
3202
3203
3204
3205
3206
        res_hidden_states_tuple = res_hidden_states_tuple[-1]
        if res_hidden_states_tuple is not None:
            hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1)

        for resnet in self.resnets:
            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)

                    return custom_forward

3207
3208
3209
3210
3211
3212
3213
3214
                if is_torch_version(">=", "1.11.0"):
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
                    )
                else:
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(resnet), hidden_states, temb
                    )
3215
            else:
3216
                hidden_states = resnet(hidden_states, temb, scale=scale)
3217
3218
3219
3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
3230

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
                hidden_states = upsampler(hidden_states)

        return hidden_states


class KCrossAttnUpBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
3231
        resolution_idx: int,
3232
3233
3234
3235
3236
        dropout: float = 0.0,
        num_layers: int = 4,
        resnet_eps: float = 1e-5,
        resnet_act_fn: str = "gelu",
        resnet_group_size: int = 32,
3237
        attention_head_dim: int = 1,  # attention dim_head
3238
3239
3240
3241
3242
3243
3244
3245
3246
3247
3248
3249
3250
        cross_attention_dim: int = 768,
        add_upsample: bool = True,
        upcast_attention: bool = False,
    ):
        super().__init__()
        resnets = []
        attentions = []

        is_first_block = in_channels == out_channels == temb_channels
        is_middle_block = in_channels != out_channels
        add_self_attention = True if is_first_block else False

        self.has_cross_attention = True
3251
        self.attention_head_dim = attention_head_dim
3252
3253
3254
3255
3256
3257
3258
3259
3260
3261
3262
3263
3264
3265
3266
3267
3268
3269
3270
3271
3272
3273
3274
3275
3276
3277
3278
3279
3280
3281
3282
3283
3284
3285
3286

        # in_channels, and out_channels for the block (k-unet)
        k_in_channels = out_channels if is_first_block else 2 * out_channels
        k_out_channels = in_channels

        num_layers = num_layers - 1

        for i in range(num_layers):
            in_channels = k_in_channels if i == 0 else out_channels
            groups = in_channels // resnet_group_size
            groups_out = out_channels // resnet_group_size

            if is_middle_block and (i == num_layers - 1):
                conv_2d_out_channels = k_out_channels
            else:
                conv_2d_out_channels = None

            resnets.append(
                ResnetBlock2D(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    conv_2d_out_channels=conv_2d_out_channels,
                    temb_channels=temb_channels,
                    eps=resnet_eps,
                    groups=groups,
                    groups_out=groups_out,
                    dropout=dropout,
                    non_linearity=resnet_act_fn,
                    time_embedding_norm="ada_group",
                    conv_shortcut_bias=False,
                )
            )
            attentions.append(
                KAttentionBlock(
                    k_out_channels if (i == num_layers - 1) else out_channels,
3287
                    k_out_channels // attention_head_dim
3288
                    if (i == num_layers - 1)
3289
3290
                    else out_channels // attention_head_dim,
                    attention_head_dim,
3291
3292
3293
3294
                    cross_attention_dim=cross_attention_dim,
                    temb_channels=temb_channels,
                    attention_bias=True,
                    add_self_attention=add_self_attention,
3295
                    cross_attention_norm="layer_norm",
3296
3297
3298
3299
3300
3301
3302
3303
3304
3305
3306
3307
3308
                    upcast_attention=upcast_attention,
                )
            )

        self.resnets = nn.ModuleList(resnets)
        self.attentions = nn.ModuleList(attentions)

        if add_upsample:
            self.upsamplers = nn.ModuleList([KUpsample2D()])
        else:
            self.upsamplers = None

        self.gradient_checkpointing = False
3309
        self.resolution_idx = resolution_idx
3310
3311
3312

    def forward(
        self,
3313
3314
3315
3316
3317
3318
3319
3320
        hidden_states: torch.FloatTensor,
        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
        temb: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        upsample_size: Optional[int] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
3321
    ) -> torch.FloatTensor:
3322
3323
3324
3325
        res_hidden_states_tuple = res_hidden_states_tuple[-1]
        if res_hidden_states_tuple is not None:
            hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1)

3326
        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
3327
3328
3329
3330
3331
3332
3333
3334
3335
3336
3337
3338
        for resnet, attn in zip(self.resnets, self.attentions):
            if self.training and self.gradient_checkpointing:

                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)

                    return custom_forward

3339
3340
3341
3342
3343
3344
3345
                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(resnet),
                    hidden_states,
                    temb,
                    **ckpt_kwargs,
                )
ethansmith2000's avatar
ethansmith2000 committed
3346
                hidden_states = attn(
3347
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
3348
3349
3350
3351
3352
3353
                    encoder_hidden_states=encoder_hidden_states,
                    emb=temb,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
                    encoder_attention_mask=encoder_attention_mask,
                )
3354
            else:
3355
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
3356
3357
3358
3359
3360
3361
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    emb=temb,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
3362
                    encoder_attention_mask=encoder_attention_mask,
3363
3364
3365
3366
3367
3368
3369
3370
3371
3372
3373
3374
3375
3376
3377
3378
3379
3380
3381
3382
                )

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
                hidden_states = upsampler(hidden_states)

        return hidden_states


# can potentially later be renamed to `No-feed-forward` attention
class KAttentionBlock(nn.Module):
    r"""
    A basic Transformer block.

    Parameters:
        dim (`int`): The number of channels in the input and output.
        num_attention_heads (`int`): The number of heads to use for multi-head attention.
        attention_head_dim (`int`): The number of channels in each head.
        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
3383
3384
3385
3386
3387
3388
3389
3390
3391
3392
3393
3394
        attention_bias (`bool`, *optional*, defaults to `False`):
            Configure if the attention layers should contain a bias parameter.
        upcast_attention (`bool`, *optional*, defaults to `False`):
            Set to `True` to upcast the attention computation to `float32`.
        temb_channels (`int`, *optional*, defaults to 768):
            The number of channels in the token embedding.
        add_self_attention (`bool`, *optional*, defaults to `False`):
            Set to `True` to add self-attention to the block.
        cross_attention_norm (`str`, *optional*, defaults to `None`):
            The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.
        group_size (`int`, *optional*, defaults to 32):
            The number of groups to separate the channels into for group normalization.
3395
3396
3397
3398
3399
3400
3401
3402
3403
3404
3405
3406
3407
    """

    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        dropout: float = 0.0,
        cross_attention_dim: Optional[int] = None,
        attention_bias: bool = False,
        upcast_attention: bool = False,
        temb_channels: int = 768,  # for ada_group_norm
        add_self_attention: bool = False,
3408
        cross_attention_norm: Optional[str] = None,
3409
3410
3411
3412
3413
3414
3415
3416
        group_size: int = 32,
    ):
        super().__init__()
        self.add_self_attention = add_self_attention

        # 1. Self-Attn
        if add_self_attention:
            self.norm1 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size))
Patrick von Platen's avatar
Patrick von Platen committed
3417
            self.attn1 = Attention(
3418
3419
3420
3421
3422
3423
                query_dim=dim,
                heads=num_attention_heads,
                dim_head=attention_head_dim,
                dropout=dropout,
                bias=attention_bias,
                cross_attention_dim=None,
3424
                cross_attention_norm=None,
3425
3426
3427
3428
            )

        # 2. Cross-Attn
        self.norm2 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size))
Patrick von Platen's avatar
Patrick von Platen committed
3429
        self.attn2 = Attention(
3430
3431
3432
3433
3434
3435
3436
3437
3438
3439
            query_dim=dim,
            cross_attention_dim=cross_attention_dim,
            heads=num_attention_heads,
            dim_head=attention_head_dim,
            dropout=dropout,
            bias=attention_bias,
            upcast_attention=upcast_attention,
            cross_attention_norm=cross_attention_norm,
        )

3440
    def _to_3d(self, hidden_states: torch.FloatTensor, height: int, weight: int) -> torch.FloatTensor:
3441
3442
        return hidden_states.permute(0, 2, 3, 1).reshape(hidden_states.shape[0], height * weight, -1)

3443
    def _to_4d(self, hidden_states: torch.FloatTensor, height: int, weight: int) -> torch.FloatTensor:
3444
3445
3446
3447
        return hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, height, weight)

    def forward(
        self,
3448
3449
3450
3451
3452
3453
3454
3455
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        # TODO: mark emb as non-optional (self.norm2 requires it).
        #       requires assessing impact of change to positional param interface.
        emb: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
3456
    ) -> torch.FloatTensor:
3457
3458
3459
3460
3461
3462
3463
3464
3465
3466
3467
3468
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}

        # 1. Self-Attention
        if self.add_self_attention:
            norm_hidden_states = self.norm1(hidden_states, emb)

            height, weight = norm_hidden_states.shape[2:]
            norm_hidden_states = self._to_3d(norm_hidden_states, height, weight)

            attn_output = self.attn1(
                norm_hidden_states,
                encoder_hidden_states=None,
3469
                attention_mask=attention_mask,
3470
3471
3472
3473
3474
3475
3476
3477
3478
3479
3480
3481
3482
3483
                **cross_attention_kwargs,
            )
            attn_output = self._to_4d(attn_output, height, weight)

            hidden_states = attn_output + hidden_states

        # 2. Cross-Attention/None
        norm_hidden_states = self.norm2(hidden_states, emb)

        height, weight = norm_hidden_states.shape[2:]
        norm_hidden_states = self._to_3d(norm_hidden_states, height, weight)
        attn_output = self.attn2(
            norm_hidden_states,
            encoder_hidden_states=encoder_hidden_states,
3484
            attention_mask=attention_mask if encoder_hidden_states is None else encoder_attention_mask,
3485
3486
3487
3488
3489
3490
3491
            **cross_attention_kwargs,
        )
        attn_output = self._to_4d(attn_output, height, weight)

        hidden_states = attn_output + hidden_states

        return hidden_states