unet_2d_blocks.py 125 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
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 import AdaGroupNorm
25
from .attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0
26
from .dual_transformer_2d import DualTransformer2DModel
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
35
36
37
38
39
40
41
42
def get_down_block(
    down_block_type,
    num_layers,
    in_channels,
    out_channels,
    temb_channels,
    add_downsample,
    resnet_eps,
    resnet_act_fn,
43
    transformer_layers_per_block=1,
44
    num_attention_heads=None,
45
    resnet_groups=None,
46
    cross_attention_dim=None,
Patrick von Platen's avatar
Patrick von Platen committed
47
    downsample_padding=None,
48
    dual_cross_attention=False,
Suraj Patil's avatar
Suraj Patil committed
49
    use_linear_projection=False,
50
    only_cross_attention=False,
51
    upcast_attention=False,
Will Berman's avatar
Will Berman committed
52
    resnet_time_scale_shift="default",
53
    attention_type="default",
54
55
    resnet_skip_time_act=False,
    resnet_out_scale_factor=1.0,
56
    cross_attention_norm=None,
57
    attention_head_dim=None,
58
    downsample_type=None,
59
    dropout=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
244
245
246


def get_up_block(
    up_block_type,
    num_layers,
    in_channels,
Patrick von Platen's avatar
Patrick von Platen committed
247
248
    out_channels,
    prev_output_channel,
249
250
251
252
    temb_channels,
    add_upsample,
    resnet_eps,
    resnet_act_fn,
253
    resolution_idx=None,
254
    transformer_layers_per_block=1,
255
    num_attention_heads=None,
256
    resnet_groups=None,
257
    cross_attention_dim=None,
258
    dual_cross_attention=False,
Suraj Patil's avatar
Suraj Patil committed
259
    use_linear_projection=False,
260
    only_cross_attention=False,
261
    upcast_attention=False,
Will Berman's avatar
Will Berman committed
262
    resnet_time_scale_shift="default",
263
    attention_type="default",
264
265
    resnet_skip_time_act=False,
    resnet_out_scale_factor=1.0,
266
    cross_attention_norm=None,
267
    attention_head_dim=None,
268
    upsample_type=None,
269
    dropout=0.0,
270
):
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
501
502
503
504
    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()

    def forward(self, x):
        return self.fuse(self.conv(x) + self.skip(x))


Patrick von Platen's avatar
Patrick von Platen committed
505
506
507
508
509
class UNetMidBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        temb_channels: int,
510
        dropout: float = 0.0,
511
        num_layers: int = 1,
Patrick von Platen's avatar
Patrick von Platen committed
512
        resnet_eps: float = 1e-6,
YiYi Xu's avatar
YiYi Xu committed
513
        resnet_time_scale_shift: str = "default",  # default, spatial
Patrick von Platen's avatar
Patrick von Platen committed
514
515
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
516
        attn_groups: Optional[int] = None,
517
        resnet_pre_norm: bool = True,
Will Berman's avatar
Will Berman committed
518
        add_attention: bool = True,
519
        attention_head_dim=1,
Patrick von Platen's avatar
Patrick von Platen committed
520
521
522
        output_scale_factor=1.0,
    ):
        super().__init__()
523
        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
Will Berman's avatar
Will Berman committed
524
        self.add_attention = add_attention
Patrick von Platen's avatar
Patrick von Platen committed
525

526
527
528
        if attn_groups is None:
            attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None

529
530
        # there is always at least one resnet
        resnets = [
531
            ResnetBlock2D(
532
533
534
                in_channels=in_channels,
                out_channels=in_channels,
                temb_channels=temb_channels,
535
                eps=resnet_eps,
536
537
538
539
540
541
                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
542
            )
543
544
        ]
        attentions = []
Patrick von Platen's avatar
Patrick von Platen committed
545

546
547
548
549
550
551
        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

552
        for _ in range(num_layers):
Will Berman's avatar
Will Berman committed
553
554
            if self.add_attention:
                attentions.append(
555
                    Attention(
Will Berman's avatar
Will Berman committed
556
                        in_channels,
557
558
                        heads=in_channels // attention_head_dim,
                        dim_head=attention_head_dim,
Will Berman's avatar
Will Berman committed
559
560
                        rescale_output_factor=output_scale_factor,
                        eps=resnet_eps,
561
                        norm_num_groups=attn_groups,
YiYi Xu's avatar
YiYi Xu committed
562
                        spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
563
564
565
566
                        residual_connection=True,
                        bias=True,
                        upcast_softmax=True,
                        _from_deprecated_attn_block=True,
Will Berman's avatar
Will Berman committed
567
                    )
568
                )
Will Berman's avatar
Will Berman committed
569
570
571
            else:
                attentions.append(None)

572
            resnets.append(
573
                ResnetBlock2D(
574
575
576
                    in_channels=in_channels,
                    out_channels=in_channels,
                    temb_channels=temb_channels,
577
                    eps=resnet_eps,
578
579
580
581
582
583
584
                    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
585
586
            )

587
588
589
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

Will Berman's avatar
Will Berman committed
590
    def forward(self, hidden_states, temb=None):
Patrick von Platen's avatar
Patrick von Platen committed
591
        hidden_states = self.resnets[0](hidden_states, temb)
592
        for attn, resnet in zip(self.attentions, self.resnets[1:]):
Will Berman's avatar
Will Berman committed
593
            if attn is not None:
YiYi Xu's avatar
YiYi Xu committed
594
                hidden_states = attn(hidden_states, temb=temb)
Patrick von Platen's avatar
Patrick von Platen committed
595
            hidden_states = resnet(hidden_states, temb)
Patrick von Platen's avatar
Patrick von Platen committed
596

597
        return hidden_states
Patrick von Platen's avatar
Patrick von Platen committed
598

599

Patrick von Platen's avatar
Patrick von Platen committed
600
601
602
603
604
605
606
class UNetMidBlock2DCrossAttn(nn.Module):
    def __init__(
        self,
        in_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
607
        transformer_layers_per_block: int = 1,
Patrick von Platen's avatar
Patrick von Platen committed
608
609
610
611
612
        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,
613
        num_attention_heads=1,
Patrick von Platen's avatar
Patrick von Platen committed
614
615
        output_scale_factor=1.0,
        cross_attention_dim=1280,
616
        dual_cross_attention=False,
Suraj Patil's avatar
Suraj Patil committed
617
        use_linear_projection=False,
618
        upcast_attention=False,
619
        attention_type="default",
Patrick von Platen's avatar
Patrick von Platen committed
620
621
622
    ):
        super().__init__()

623
        self.has_cross_attention = True
624
        self.num_attention_heads = num_attention_heads
Patrick von Platen's avatar
Patrick von Platen committed
625
626
627
628
        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)

        # there is always at least one resnet
        resnets = [
629
            ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
                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 = []

        for _ in range(num_layers):
645
646
647
            if not dual_cross_attention:
                attentions.append(
                    Transformer2DModel(
648
649
                        num_attention_heads,
                        in_channels // num_attention_heads,
650
                        in_channels=in_channels,
651
                        num_layers=transformer_layers_per_block,
652
653
                        cross_attention_dim=cross_attention_dim,
                        norm_num_groups=resnet_groups,
Suraj Patil's avatar
Suraj Patil committed
654
                        use_linear_projection=use_linear_projection,
655
                        upcast_attention=upcast_attention,
656
                        attention_type=attention_type,
657
658
659
660
661
                    )
                )
            else:
                attentions.append(
                    DualTransformer2DModel(
662
663
                        num_attention_heads,
                        in_channels // num_attention_heads,
664
665
666
667
668
                        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
669
670
                )
            resnets.append(
671
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
                    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)

688
689
        self.gradient_checkpointing = False

690
    def forward(
691
692
693
694
695
696
697
698
        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:
699
700
        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
701
        for attn, resnet in zip(self.attentions, self.resnets[1:]):
702
703
704
705
706
707
708
709
710
711
712
713
            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
714
                hidden_states = attn(
715
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
716
717
718
719
720
                    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,
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
                )[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]
737
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
Will Berman's avatar
Will Berman committed
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753

        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,
754
        attention_head_dim=1,
Will Berman's avatar
Will Berman committed
755
756
        output_scale_factor=1.0,
        cross_attention_dim=1280,
757
        skip_time_act=False,
758
        only_cross_attention=False,
759
        cross_attention_norm=None,
Will Berman's avatar
Will Berman committed
760
761
762
763
764
    ):
        super().__init__()

        self.has_cross_attention = True

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

768
        self.num_heads = in_channels // self.attention_head_dim
Will Berman's avatar
Will Berman committed
769
770
771
772
773
774
775
776
777
778
779
780
781
782

        # 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,
783
                skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
784
785
786
787
788
            )
        ]
        attentions = []

        for _ in range(num_layers):
789
790
791
792
            processor = (
                AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()
            )

Will Berman's avatar
Will Berman committed
793
            attentions.append(
Patrick von Platen's avatar
Patrick von Platen committed
794
                Attention(
Will Berman's avatar
Will Berman committed
795
796
797
                    query_dim=in_channels,
                    cross_attention_dim=in_channels,
                    heads=self.num_heads,
798
                    dim_head=self.attention_head_dim,
Will Berman's avatar
Will Berman committed
799
800
801
802
                    added_kv_proj_dim=cross_attention_dim,
                    norm_num_groups=resnet_groups,
                    bias=True,
                    upcast_softmax=True,
803
                    only_cross_attention=only_cross_attention,
804
                    cross_attention_norm=cross_attention_norm,
805
                    processor=processor,
Will Berman's avatar
Will Berman committed
806
807
808
809
810
811
812
813
814
815
816
817
818
819
                )
            )
            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,
820
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
821
822
823
824
825
826
                )
            )

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

827
    def forward(
828
829
830
831
832
833
834
        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,
835
836
    ):
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
837
        lora_scale = cross_attention_kwargs.get("scale", 1.0)
838
839
840
841
842
843
844
845
846
847
848
849

        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

850
        hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)
Will Berman's avatar
Will Berman committed
851
852
853
854
        for attn, resnet in zip(self.attentions, self.resnets[1:]):
            # attn
            hidden_states = attn(
                hidden_states,
855
                encoder_hidden_states=encoder_hidden_states,
856
                attention_mask=mask,
857
                **cross_attention_kwargs,
Will Berman's avatar
Will Berman committed
858
859
860
            )

            # resnet
861
            hidden_states = resnet(hidden_states, temb, scale=lora_scale)
Patrick von Platen's avatar
Patrick von Platen committed
862
863
864
865

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
866
class AttnDownBlock2D(nn.Module):
867
868
869
870
871
872
873
874
875
876
877
878
    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,
879
        attention_head_dim=1,
880
        output_scale_factor=1.0,
881
        downsample_padding=1,
882
        downsample_type="conv",
883
884
885
886
    ):
        super().__init__()
        resnets = []
        attentions = []
887
        self.downsample_type = downsample_type
888

889
890
891
892
893
894
        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

895
896
897
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
898
                ResnetBlock2D(
899
900
901
902
903
904
905
906
907
908
909
910
911
                    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(
912
                Attention(
913
                    out_channels,
914
915
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
916
                    rescale_output_factor=output_scale_factor,
Patrick von Platen's avatar
Patrick von Platen committed
917
                    eps=resnet_eps,
Will Berman's avatar
Will Berman committed
918
                    norm_num_groups=resnet_groups,
919
920
921
922
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
923
924
925
926
927
928
                )
            )

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

929
        if downsample_type == "conv":
930
            self.downsamplers = nn.ModuleList(
931
932
                [
                    Downsample2D(
933
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
934
935
                    )
                ]
936
            )
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
        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,
                    )
                ]
            )
955
956
957
        else:
            self.downsamplers = None

958
959
960
961
962
    def forward(self, hidden_states, temb=None, upsample_size=None, cross_attention_kwargs=None):
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}

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

963
964
965
        output_states = ()

        for resnet, attn in zip(self.resnets, self.attentions):
966
967
968
            cross_attention_kwargs.update({"scale": lora_scale})
            hidden_states = resnet(hidden_states, temb, scale=lora_scale)
            hidden_states = attn(hidden_states, **cross_attention_kwargs)
969
            output_states = output_states + (hidden_states,)
970
971
972

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
973
                if self.downsample_type == "resnet":
974
                    hidden_states = downsampler(hidden_states, temb=temb, scale=lora_scale)
975
                else:
976
                    hidden_states = downsampler(hidden_states, scale=lora_scale)
977
978
979
980
981
982

            output_states += (hidden_states,)

        return hidden_states, output_states


Patrick von Platen's avatar
Patrick von Platen committed
983
class CrossAttnDownBlock2D(nn.Module):
Patrick von Platen's avatar
Patrick von Platen committed
984
985
986
987
988
989
990
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
991
        transformer_layers_per_block: int = 1,
Patrick von Platen's avatar
Patrick von Platen committed
992
993
994
995
996
        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,
997
        num_attention_heads=1,
Patrick von Platen's avatar
Patrick von Platen committed
998
999
1000
1001
        cross_attention_dim=1280,
        output_scale_factor=1.0,
        downsample_padding=1,
        add_downsample=True,
1002
        dual_cross_attention=False,
Suraj Patil's avatar
Suraj Patil committed
1003
        use_linear_projection=False,
1004
        only_cross_attention=False,
1005
        upcast_attention=False,
1006
        attention_type="default",
Patrick von Platen's avatar
Patrick von Platen committed
1007
1008
1009
1010
1011
    ):
        super().__init__()
        resnets = []
        attentions = []

1012
        self.has_cross_attention = True
1013
        self.num_attention_heads = num_attention_heads
Patrick von Platen's avatar
Patrick von Platen committed
1014
1015
1016
1017

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
1018
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
                    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,
                )
            )
1031
1032
1033
            if not dual_cross_attention:
                attentions.append(
                    Transformer2DModel(
1034
1035
                        num_attention_heads,
                        out_channels // num_attention_heads,
1036
                        in_channels=out_channels,
1037
                        num_layers=transformer_layers_per_block,
1038
1039
                        cross_attention_dim=cross_attention_dim,
                        norm_num_groups=resnet_groups,
Suraj Patil's avatar
Suraj Patil committed
1040
                        use_linear_projection=use_linear_projection,
1041
                        only_cross_attention=only_cross_attention,
1042
                        upcast_attention=upcast_attention,
1043
                        attention_type=attention_type,
1044
1045
1046
1047
1048
                    )
                )
            else:
                attentions.append(
                    DualTransformer2DModel(
1049
1050
                        num_attention_heads,
                        out_channels // num_attention_heads,
1051
1052
1053
1054
1055
                        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
1056
1057
1058
1059
1060
1061
1062
1063
                )
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)

        if add_downsample:
            self.downsamplers = nn.ModuleList(
                [
                    Downsample2D(
1064
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
Patrick von Platen's avatar
Patrick von Platen committed
1065
1066
1067
1068
1069
1070
                    )
                ]
            )
        else:
            self.downsamplers = None

1071
1072
        self.gradient_checkpointing = False

1073
    def forward(
1074
1075
1076
1077
1078
1079
1080
        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,
Will Berman's avatar
Will Berman committed
1081
        additional_residuals=None,
1082
    ):
Patrick von Platen's avatar
Patrick von Platen committed
1083
1084
        output_states = ()

1085
1086
        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
1087
1088
1089
        blocks = list(zip(self.resnets, self.attentions))

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

Will Berman's avatar
Will Berman committed
1092
                def create_custom_forward(module, return_dict=None):
1093
                    def custom_forward(*inputs):
Will Berman's avatar
Will Berman committed
1094
1095
1096
1097
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)
1098
1099
1100

                    return custom_forward

1101
1102
1103
1104
1105
1106
1107
                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
1108
                hidden_states = attn(
1109
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
1110
1111
1112
1113
1114
                    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,
1115
                )[0]
1116
            else:
1117
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
1118
1119
1120
1121
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
1122
1123
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
1124
1125
                    return_dict=False,
                )[0]
1126

Will Berman's avatar
Will Berman committed
1127
1128
1129
1130
            # 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

1131
            output_states = output_states + (hidden_states,)
Patrick von Platen's avatar
Patrick von Platen committed
1132
1133
1134

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

1137
            output_states = output_states + (hidden_states,)
Patrick von Platen's avatar
Patrick von Platen committed
1138
1139
1140
1141

        return hidden_states, output_states


Patrick von Platen's avatar
Patrick von Platen committed
1142
class DownBlock2D(nn.Module):
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
    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,
        output_scale_factor=1.0,
        add_downsample=True,
Patrick von Platen's avatar
Patrick von Platen committed
1157
        downsample_padding=1,
1158
1159
1160
1161
1162
1163
1164
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
1165
                ResnetBlock2D(
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
                    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
1183
1184
                [
                    Downsample2D(
1185
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
Patrick von Platen's avatar
Patrick von Platen committed
1186
1187
                    )
                ]
1188
1189
1190
1191
            )
        else:
            self.downsamplers = None

1192
1193
        self.gradient_checkpointing = False

1194
    def forward(self, hidden_states, temb=None, scale: float = 1.0):
1195
1196
1197
        output_states = ()

        for resnet in self.resnets:
1198
1199
1200
1201
1202
1203
1204
1205
            if self.training and self.gradient_checkpointing:

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

                    return custom_forward

1206
1207
1208
1209
1210
1211
1212
1213
                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
                    )
1214
            else:
1215
                hidden_states = resnet(hidden_states, temb, scale=scale)
1216

1217
            output_states = output_states + (hidden_states,)
1218
1219
1220

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

1223
            output_states = output_states + (hidden_states,)
1224
1225
1226
1227

        return hidden_states, output_states


1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
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,
        output_scale_factor=1.0,
        add_downsample=True,
        downsample_padding=1,
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
1250
                ResnetBlock2D(
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
                    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(
1270
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
1271
1272
1273
1274
1275
1276
                    )
                ]
            )
        else:
            self.downsamplers = None

1277
    def forward(self, hidden_states, scale: float = 1.0):
1278
        for resnet in self.resnets:
1279
            hidden_states = resnet(hidden_states, temb=None, scale=scale)
1280
1281
1282

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
1283
                hidden_states = downsampler(hidden_states, scale)
1284
1285
1286
1287

        return hidden_states


1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
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,
1300
        attention_head_dim=1,
1301
1302
1303
1304
1305
1306
1307
1308
        output_scale_factor=1.0,
        add_downsample=True,
        downsample_padding=1,
    ):
        super().__init__()
        resnets = []
        attentions = []

1309
1310
1311
1312
1313
1314
        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

1315
1316
1317
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(
1318
                ResnetBlock2D(
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
                    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(
1332
                Attention(
1333
                    out_channels,
1334
1335
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
1336
1337
                    rescale_output_factor=output_scale_factor,
                    eps=resnet_eps,
Will Berman's avatar
Will Berman committed
1338
                    norm_num_groups=resnet_groups,
1339
1340
1341
1342
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
                )
            )

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

        if add_downsample:
            self.downsamplers = nn.ModuleList(
                [
                    Downsample2D(
1353
                        out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
1354
1355
1356
1357
1358
1359
                    )
                ]
            )
        else:
            self.downsamplers = None

1360
    def forward(self, hidden_states, scale: float = 1.0):
1361
        for resnet, attn in zip(self.resnets, self.attentions):
1362
1363
1364
            hidden_states = resnet(hidden_states, temb=None, scale=scale)
            cross_attention_kwargs = {"scale": scale}
            hidden_states = attn(hidden_states, **cross_attention_kwargs)
1365
1366
1367

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
1368
                hidden_states = downsampler(hidden_states, scale)
1369
1370
1371
1372

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
1373
class AttnSkipDownBlock2D(nn.Module):
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
    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,
1385
        attention_head_dim=1,
1386
1387
1388
1389
1390
1391
1392
        output_scale_factor=np.sqrt(2.0),
        add_downsample=True,
    ):
        super().__init__()
        self.attentions = nn.ModuleList([])
        self.resnets = nn.ModuleList([])

1393
1394
1395
1396
1397
1398
        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

1399
1400
1401
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            self.resnets.append(
1402
                ResnetBlock2D(
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
                    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(
1417
                Attention(
1418
                    out_channels,
1419
1420
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
1421
1422
                    rescale_output_factor=output_scale_factor,
                    eps=resnet_eps,
1423
1424
1425
1426
1427
                    norm_num_groups=32,
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
1428
1429
1430
1431
                )
            )

        if add_downsample:
1432
            self.resnet_down = ResnetBlock2D(
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
                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,
1443
                use_in_shortcut=True,
1444
1445
1446
                down=True,
                kernel="fir",
            )
1447
            self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
1448
1449
1450
1451
1452
1453
            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

1454
    def forward(self, hidden_states, temb=None, skip_sample=None, scale: float = 1.0):
1455
1456
1457
        output_states = ()

        for resnet, attn in zip(self.resnets, self.attentions):
1458
1459
1460
            hidden_states = resnet(hidden_states, temb, scale=scale)
            cross_attention_kwargs = {"scale": scale}
            hidden_states = attn(hidden_states, **cross_attention_kwargs)
1461
1462
1463
            output_states += (hidden_states,)

        if self.downsamplers is not None:
1464
            hidden_states = self.resnet_down(hidden_states, temb, scale=scale)
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
            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
1475
class SkipDownBlock2D(nn.Module):
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
    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,
        output_scale_factor=np.sqrt(2.0),
        add_downsample=True,
        downsample_padding=1,
    ):
        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(
1497
                ResnetBlock2D(
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
                    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:
1513
            self.resnet_down = ResnetBlock2D(
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
                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,
1524
                use_in_shortcut=True,
1525
1526
1527
                down=True,
                kernel="fir",
            )
1528
            self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
1529
1530
1531
1532
1533
1534
            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

1535
    def forward(self, hidden_states, temb=None, skip_sample=None, scale: float = 1.0):
1536
1537
1538
        output_states = ()

        for resnet in self.resnets:
1539
            hidden_states = resnet(hidden_states, temb, scale)
1540
1541
1542
            output_states += (hidden_states,)

        if self.downsamplers is not None:
1543
            hidden_states = self.resnet_down(hidden_states, temb, scale)
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
            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
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
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,
        output_scale_factor=1.0,
        add_downsample=True,
1569
        skip_time_act=False,
Will Berman's avatar
Will Berman committed
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
    ):
        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,
1588
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
                )
            )

        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,
1608
                        skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
1609
1610
1611
1612
1613
1614
1615
1616
1617
                        down=True,
                    )
                ]
            )
        else:
            self.downsamplers = None

        self.gradient_checkpointing = False

1618
    def forward(self, hidden_states, temb=None, scale: float = 1.0):
Will Berman's avatar
Will Berman committed
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
        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

1630
1631
1632
1633
1634
1635
1636
1637
                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
1638
            else:
1639
                hidden_states = resnet(hidden_states, temb, scale)
Will Berman's avatar
Will Berman committed
1640

1641
            output_states = output_states + (hidden_states,)
Will Berman's avatar
Will Berman committed
1642
1643
1644

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

1647
            output_states = output_states + (hidden_states,)
Will Berman's avatar
Will Berman committed
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664

        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,
1665
        attention_head_dim=1,
Will Berman's avatar
Will Berman committed
1666
1667
1668
        cross_attention_dim=1280,
        output_scale_factor=1.0,
        add_downsample=True,
1669
        skip_time_act=False,
1670
        only_cross_attention=False,
1671
        cross_attention_norm=None,
Will Berman's avatar
Will Berman committed
1672
1673
1674
1675
1676
1677
1678
1679
    ):
        super().__init__()

        self.has_cross_attention = True

        resnets = []
        attentions = []

1680
1681
        self.attention_head_dim = attention_head_dim
        self.num_heads = out_channels // self.attention_head_dim
Will Berman's avatar
Will Berman committed
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696

        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,
1697
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
1698
1699
                )
            )
1700
1701
1702
1703
1704

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

Will Berman's avatar
Will Berman committed
1705
            attentions.append(
Patrick von Platen's avatar
Patrick von Platen committed
1706
                Attention(
Will Berman's avatar
Will Berman committed
1707
1708
1709
                    query_dim=out_channels,
                    cross_attention_dim=out_channels,
                    heads=self.num_heads,
1710
                    dim_head=attention_head_dim,
Will Berman's avatar
Will Berman committed
1711
1712
1713
1714
                    added_kv_proj_dim=cross_attention_dim,
                    norm_num_groups=resnet_groups,
                    bias=True,
                    upcast_softmax=True,
1715
                    only_cross_attention=only_cross_attention,
1716
                    cross_attention_norm=cross_attention_norm,
1717
                    processor=processor,
Will Berman's avatar
Will Berman committed
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
                )
            )
        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,
1737
                        skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
1738
1739
1740
1741
1742
1743
1744
1745
1746
                        down=True,
                    )
                ]
            )
        else:
            self.downsamplers = None

        self.gradient_checkpointing = False

1747
    def forward(
1748
1749
1750
1751
1752
1753
1754
        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,
1755
    ):
Will Berman's avatar
Will Berman committed
1756
        output_states = ()
1757
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
Will Berman's avatar
Will Berman committed
1758

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

1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
        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
1772
        for resnet, attn in zip(self.resnets, self.attentions):
1773
            if self.training and self.gradient_checkpointing:
Will Berman's avatar
Will Berman committed
1774

1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
                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
1785
                hidden_states = attn(
1786
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
1787
1788
1789
1790
                    encoder_hidden_states=encoder_hidden_states,
                    attention_mask=mask,
                    **cross_attention_kwargs,
                )
1791
            else:
1792
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
1793
1794
1795
1796

                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
1797
                    attention_mask=mask,
1798
1799
                    **cross_attention_kwargs,
                )
Will Berman's avatar
Will Berman committed
1800

1801
            output_states = output_states + (hidden_states,)
Will Berman's avatar
Will Berman committed
1802
1803
1804

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

1807
            output_states = output_states + (hidden_states,)
Will Berman's avatar
Will Berman committed
1808
1809
1810
1811

        return hidden_states, output_states


1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
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,
        add_downsample=False,
    ):
        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

1858
    def forward(self, hidden_states, temb=None, scale: float = 1.0):
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
        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

1870
1871
1872
1873
1874
1875
1876
1877
                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
                    )
1878
            else:
1879
                hidden_states = resnet(hidden_states, temb, scale)
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900

            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,
        add_downsample=True,
1901
        attention_head_dim: int = 64,
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
        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,
1934
1935
                    out_channels // attention_head_dim,
                    attention_head_dim,
1936
1937
1938
1939
                    cross_attention_dim=cross_attention_dim,
                    temb_channels=temb_channels,
                    attention_bias=True,
                    add_self_attention=add_self_attention,
1940
                    cross_attention_norm="layer_norm",
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
                    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(
1956
1957
1958
1959
1960
1961
1962
        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,
1963
1964
    ):
        output_states = ()
1965
        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978

        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

1979
1980
1981
1982
1983
1984
1985
                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
1986
                hidden_states = attn(
1987
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
1988
1989
1990
1991
1992
                    encoder_hidden_states=encoder_hidden_states,
                    emb=temb,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
                    encoder_attention_mask=encoder_attention_mask,
1993
                )
1994
            else:
1995
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
1996
1997
1998
1999
2000
2001
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    emb=temb,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
2002
                    encoder_attention_mask=encoder_attention_mask,
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
                )

            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
2017
class AttnUpBlock2D(nn.Module):
2018
2019
2020
    def __init__(
        self,
        in_channels: int,
Patrick von Platen's avatar
Patrick von Platen committed
2021
2022
        prev_output_channel: int,
        out_channels: int,
2023
        temb_channels: int,
2024
        resolution_idx: int = None,
2025
2026
2027
2028
2029
2030
2031
        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,
2032
        attention_head_dim=1,
2033
        output_scale_factor=1.0,
2034
        upsample_type="conv",
2035
2036
2037
2038
2039
    ):
        super().__init__()
        resnets = []
        attentions = []

2040
2041
        self.upsample_type = upsample_type

2042
2043
2044
2045
2046
2047
        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

2048
        for i in range(num_layers):
Patrick von Platen's avatar
Patrick von Platen committed
2049
2050
2051
            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

2052
            resnets.append(
2053
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
2054
2055
                    in_channels=resnet_in_channels + res_skip_channels,
                    out_channels=out_channels,
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
                    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(
2067
                Attention(
Patrick von Platen's avatar
Patrick von Platen committed
2068
                    out_channels,
2069
2070
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
2071
                    rescale_output_factor=output_scale_factor,
Patrick von Platen's avatar
Patrick von Platen committed
2072
                    eps=resnet_eps,
Will Berman's avatar
Will Berman committed
2073
                    norm_num_groups=resnet_groups,
2074
2075
2076
2077
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
2078
2079
2080
2081
2082
2083
                )
            )

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

2084
        if upsample_type == "conv":
Patrick von Platen's avatar
Patrick von Platen committed
2085
            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
        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,
                    )
                ]
            )
2104
2105
2106
        else:
            self.upsamplers = None

2107
2108
        self.resolution_idx = resolution_idx

2109
    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0):
2110
2111
2112
2113
2114
2115
        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)

2116
2117
2118
            hidden_states = resnet(hidden_states, temb, scale=scale)
            cross_attention_kwargs = {"scale": scale}
            hidden_states = attn(hidden_states, **cross_attention_kwargs)
2119
2120
2121

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2122
                if self.upsample_type == "resnet":
2123
                    hidden_states = upsampler(hidden_states, temb=temb, scale=scale)
2124
                else:
2125
                    hidden_states = upsampler(hidden_states, scale=scale)
2126
2127
2128
2129

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
2130
class CrossAttnUpBlock2D(nn.Module):
Patrick von Platen's avatar
Patrick von Platen committed
2131
2132
2133
2134
2135
2136
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        prev_output_channel: int,
        temb_channels: int,
2137
        resolution_idx: int = None,
Patrick von Platen's avatar
Patrick von Platen committed
2138
2139
        dropout: float = 0.0,
        num_layers: int = 1,
2140
        transformer_layers_per_block: int = 1,
Patrick von Platen's avatar
Patrick von Platen committed
2141
2142
2143
2144
2145
        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,
2146
        num_attention_heads=1,
Patrick von Platen's avatar
Patrick von Platen committed
2147
2148
2149
        cross_attention_dim=1280,
        output_scale_factor=1.0,
        add_upsample=True,
2150
        dual_cross_attention=False,
Suraj Patil's avatar
Suraj Patil committed
2151
        use_linear_projection=False,
2152
        only_cross_attention=False,
2153
        upcast_attention=False,
2154
        attention_type="default",
Patrick von Platen's avatar
Patrick von Platen committed
2155
2156
2157
2158
2159
    ):
        super().__init__()
        resnets = []
        attentions = []

2160
        self.has_cross_attention = True
2161
        self.num_attention_heads = num_attention_heads
Patrick von Platen's avatar
Patrick von Platen committed
2162
2163
2164
2165
2166
2167

        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(
2168
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
                    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,
                )
            )
2181
2182
2183
            if not dual_cross_attention:
                attentions.append(
                    Transformer2DModel(
2184
2185
                        num_attention_heads,
                        out_channels // num_attention_heads,
2186
                        in_channels=out_channels,
2187
                        num_layers=transformer_layers_per_block,
2188
2189
                        cross_attention_dim=cross_attention_dim,
                        norm_num_groups=resnet_groups,
Suraj Patil's avatar
Suraj Patil committed
2190
                        use_linear_projection=use_linear_projection,
2191
                        only_cross_attention=only_cross_attention,
2192
                        upcast_attention=upcast_attention,
2193
                        attention_type=attention_type,
2194
2195
2196
2197
2198
                    )
                )
            else:
                attentions.append(
                    DualTransformer2DModel(
2199
2200
                        num_attention_heads,
                        out_channels // num_attention_heads,
2201
2202
2203
2204
2205
                        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
2206
2207
2208
2209
2210
2211
2212
2213
2214
                )
        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

2215
        self.gradient_checkpointing = False
2216
        self.resolution_idx = resolution_idx
2217
2218
2219

    def forward(
        self,
2220
2221
2222
2223
2224
2225
2226
2227
        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,
2228
    ):
2229
        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
2230
2231
2232
2233
2234
2235
        is_freeu_enabled = (
            getattr(self, "s1", None)
            and getattr(self, "s2", None)
            and getattr(self, "b1", None)
            and getattr(self, "b2", None)
        )
2236

Patrick von Platen's avatar
Patrick von Platen committed
2237
2238
2239
2240
        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]
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253

            # 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
2254
2255
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

2256
2257
            if self.training and self.gradient_checkpointing:

Will Berman's avatar
Will Berman committed
2258
                def create_custom_forward(module, return_dict=None):
2259
                    def custom_forward(*inputs):
Will Berman's avatar
Will Berman committed
2260
2261
2262
2263
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)
2264
2265
2266

                    return custom_forward

2267
2268
2269
2270
2271
2272
2273
                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
2274
                hidden_states = attn(
2275
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
2276
2277
2278
2279
2280
                    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,
2281
                )[0]
2282
            else:
2283
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
2284
2285
2286
2287
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    cross_attention_kwargs=cross_attention_kwargs,
2288
2289
                    attention_mask=attention_mask,
                    encoder_attention_mask=encoder_attention_mask,
2290
2291
                    return_dict=False,
                )[0]
Patrick von Platen's avatar
Patrick von Platen committed
2292
2293
2294

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2295
                hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale)
Patrick von Platen's avatar
Patrick von Platen committed
2296
2297
2298
2299

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
2300
class UpBlock2D(nn.Module):
2301
2302
2303
    def __init__(
        self,
        in_channels: int,
Patrick von Platen's avatar
Patrick von Platen committed
2304
2305
        prev_output_channel: int,
        out_channels: int,
2306
        temb_channels: int,
2307
        resolution_idx: int = None,
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
        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,
        output_scale_factor=1.0,
        add_upsample=True,
    ):
        super().__init__()
        resnets = []

        for i in range(num_layers):
Patrick von Platen's avatar
Patrick von Platen committed
2322
2323
2324
            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

2325
            resnets.append(
2326
                ResnetBlock2D(
Patrick von Platen's avatar
Patrick von Platen committed
2327
2328
                    in_channels=resnet_in_channels + res_skip_channels,
                    out_channels=out_channels,
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
                    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
2343
            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
2344
2345
2346
        else:
            self.upsamplers = None

2347
        self.gradient_checkpointing = False
2348
        self.resolution_idx = resolution_idx
2349

2350
    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0):
2351
2352
2353
2354
2355
2356
2357
        is_freeu_enabled = (
            getattr(self, "s1", None)
            and getattr(self, "s2", None)
            and getattr(self, "b1", None)
            and getattr(self, "b2", None)
        )

2358
2359
2360
2361
        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]
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374

            # 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,
                )

2375
2376
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

2377
2378
2379
2380
2381
2382
2383
2384
            if self.training and self.gradient_checkpointing:

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

                    return custom_forward

2385
2386
2387
2388
2389
2390
2391
2392
                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
                    )
2393
            else:
2394
                hidden_states = resnet(hidden_states, temb, scale=scale)
2395
2396
2397

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2398
                hidden_states = upsampler(hidden_states, upsample_size, scale=scale)
2399
2400

        return hidden_states
2401
2402


2403
2404
2405
2406
2407
class UpDecoderBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
2408
        resolution_idx: int = None,
2409
2410
2411
        dropout: float = 0.0,
        num_layers: int = 1,
        resnet_eps: float = 1e-6,
YiYi Xu's avatar
YiYi Xu committed
2412
        resnet_time_scale_shift: str = "default",  # default, spatial
2413
2414
2415
2416
2417
        resnet_act_fn: str = "swish",
        resnet_groups: int = 32,
        resnet_pre_norm: bool = True,
        output_scale_factor=1.0,
        add_upsample=True,
YiYi Xu's avatar
YiYi Xu committed
2418
        temb_channels=None,
2419
2420
2421
2422
2423
2424
2425
2426
    ):
        super().__init__()
        resnets = []

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

            resnets.append(
2427
                ResnetBlock2D(
2428
2429
                    in_channels=input_channels,
                    out_channels=out_channels,
YiYi Xu's avatar
YiYi Xu committed
2430
                    temb_channels=temb_channels,
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
                    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

2448
2449
        self.resolution_idx = resolution_idx

2450
    def forward(self, hidden_states, temb=None, scale: float = 1.0):
2451
        for resnet in self.resnets:
2452
            hidden_states = resnet(hidden_states, temb=temb, scale=scale)
2453
2454
2455
2456
2457
2458
2459
2460

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

        return hidden_states


2461
2462
2463
2464
2465
class AttnUpDecoderBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
2466
        resolution_idx: int = None,
2467
2468
2469
2470
2471
2472
2473
        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,
2474
        attention_head_dim=1,
2475
2476
        output_scale_factor=1.0,
        add_upsample=True,
YiYi Xu's avatar
YiYi Xu committed
2477
        temb_channels=None,
2478
2479
2480
2481
2482
    ):
        super().__init__()
        resnets = []
        attentions = []

2483
2484
2485
2486
2487
2488
        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

2489
2490
2491
2492
        for i in range(num_layers):
            input_channels = in_channels if i == 0 else out_channels

            resnets.append(
2493
                ResnetBlock2D(
2494
2495
                    in_channels=input_channels,
                    out_channels=out_channels,
YiYi Xu's avatar
YiYi Xu committed
2496
                    temb_channels=temb_channels,
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
                    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(
2507
                Attention(
2508
                    out_channels,
2509
2510
                    heads=out_channels // attention_head_dim,
                    dim_head=attention_head_dim,
2511
2512
                    rescale_output_factor=output_scale_factor,
                    eps=resnet_eps,
2513
                    norm_num_groups=resnet_groups if resnet_time_scale_shift != "spatial" else None,
YiYi Xu's avatar
YiYi Xu committed
2514
                    spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
2515
2516
2517
2518
                    residual_connection=True,
                    bias=True,
                    upcast_softmax=True,
                    _from_deprecated_attn_block=True,
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
                )
            )

        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

2530
2531
        self.resolution_idx = resolution_idx

2532
    def forward(self, hidden_states, temb=None, scale: float = 1.0):
2533
        for resnet, attn in zip(self.resnets, self.attentions):
2534
2535
2536
            hidden_states = resnet(hidden_states, temb=temb, scale=scale)
            cross_attention_kwargs = {"scale": scale}
            hidden_states = attn(hidden_states, temb=temb, **cross_attention_kwargs)
2537
2538
2539

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2540
                hidden_states = upsampler(hidden_states, scale=scale)
2541
2542
2543
2544

        return hidden_states


Patrick von Platen's avatar
Patrick von Platen committed
2545
class AttnSkipUpBlock2D(nn.Module):
2546
2547
2548
2549
2550
2551
    def __init__(
        self,
        in_channels: int,
        prev_output_channel: int,
        out_channels: int,
        temb_channels: int,
2552
        resolution_idx: int = None,
2553
2554
2555
2556
2557
2558
        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,
2559
        attention_head_dim=1,
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
        output_scale_factor=np.sqrt(2.0),
        add_upsample=True,
    ):
        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(
2572
                ResnetBlock2D(
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
                    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,
                )
            )

2587
2588
2589
2590
2591
2592
        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

2593
        self.attentions.append(
2594
            Attention(
2595
                out_channels,
2596
2597
                heads=out_channels // attention_head_dim,
                dim_head=attention_head_dim,
2598
2599
                rescale_output_factor=output_scale_factor,
                eps=resnet_eps,
2600
2601
2602
2603
2604
                norm_num_groups=32,
                residual_connection=True,
                bias=True,
                upcast_softmax=True,
                _from_deprecated_attn_block=True,
2605
2606
2607
2608
2609
            )
        )

        self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels)
        if add_upsample:
2610
            self.resnet_up = ResnetBlock2D(
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
                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,
2622
                use_in_shortcut=True,
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
                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

2637
2638
        self.resolution_idx = resolution_idx

2639
    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=None, scale: float = 1.0):
2640
2641
2642
2643
2644
2645
        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)

2646
            hidden_states = resnet(hidden_states, temb, scale=scale)
2647

2648
2649
        cross_attention_kwargs = {"scale": scale}
        hidden_states = self.attentions[0](hidden_states, **cross_attention_kwargs)
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662

        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

2663
            hidden_states = self.resnet_up(hidden_states, temb, scale=scale)
2664
2665
2666
2667

        return hidden_states, skip_sample


Patrick von Platen's avatar
Patrick von Platen committed
2668
class SkipUpBlock2D(nn.Module):
2669
2670
2671
2672
2673
2674
    def __init__(
        self,
        in_channels: int,
        prev_output_channel: int,
        out_channels: int,
        temb_channels: int,
2675
        resolution_idx: int = None,
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
        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,
        output_scale_factor=np.sqrt(2.0),
        add_upsample=True,
        upsample_padding=1,
    ):
        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(
2694
                ResnetBlock2D(
2695
2696
2697
2698
2699
2700
2701
2702
2703
2704
2705
2706
2707
2708
2709
2710
                    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:
2711
            self.resnet_up = ResnetBlock2D(
2712
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
                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,
2723
                use_in_shortcut=True,
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
                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

2738
2739
        self.resolution_idx = resolution_idx

2740
    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=None, scale: float = 1.0):
2741
2742
2743
2744
2745
2746
        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)

2747
            hidden_states = resnet(hidden_states, temb, scale=scale)
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758
2759
2760

        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

2761
            hidden_states = self.resnet_up(hidden_states, temb, scale=scale)
2762
2763

        return hidden_states, skip_sample
Will Berman's avatar
Will Berman committed
2764
2765
2766
2767
2768
2769
2770
2771
2772


class ResnetUpsampleBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        prev_output_channel: int,
        out_channels: int,
        temb_channels: int,
2773
        resolution_idx: int = None,
Will Berman's avatar
Will Berman committed
2774
2775
2776
2777
2778
2779
2780
2781
2782
        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,
        output_scale_factor=1.0,
        add_upsample=True,
2783
        skip_time_act=False,
Will Berman's avatar
Will Berman committed
2784
2785
2786
2787
2788
2789
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
2802
2803
    ):
        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,
2804
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
2816
2817
2818
2819
2820
2821
2822
2823
                )
            )

        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,
2824
                        skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
2825
2826
2827
2828
2829
2830
2831
2832
                        up=True,
                    )
                ]
            )
        else:
            self.upsamplers = None

        self.gradient_checkpointing = False
2833
        self.resolution_idx = resolution_idx
Will Berman's avatar
Will Berman committed
2834

2835
    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0):
Will Berman's avatar
Will Berman committed
2836
2837
2838
2839
2840
2841
2842
2843
2844
2845
2846
2847
2848
2849
        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

2850
2851
2852
2853
2854
2855
2856
2857
                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
2858
            else:
2859
                hidden_states = resnet(hidden_states, temb, scale=scale)
Will Berman's avatar
Will Berman committed
2860
2861
2862

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
2863
                hidden_states = upsampler(hidden_states, temb, scale=scale)
Will Berman's avatar
Will Berman committed
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874

        return hidden_states


class SimpleCrossAttnUpBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        prev_output_channel: int,
        temb_channels: int,
2875
        resolution_idx: int = None,
Will Berman's avatar
Will Berman committed
2876
2877
2878
2879
2880
2881
2882
        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,
2883
        attention_head_dim=1,
Will Berman's avatar
Will Berman committed
2884
2885
2886
        cross_attention_dim=1280,
        output_scale_factor=1.0,
        add_upsample=True,
2887
        skip_time_act=False,
2888
        only_cross_attention=False,
2889
        cross_attention_norm=None,
Will Berman's avatar
Will Berman committed
2890
2891
2892
2893
2894
2895
    ):
        super().__init__()
        resnets = []
        attentions = []

        self.has_cross_attention = True
2896
        self.attention_head_dim = attention_head_dim
Will Berman's avatar
Will Berman committed
2897

2898
        self.num_heads = out_channels // self.attention_head_dim
Will Berman's avatar
Will Berman committed
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915

        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,
2916
                    skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
2917
2918
                )
            )
2919
2920
2921
2922
2923

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

Will Berman's avatar
Will Berman committed
2924
            attentions.append(
Patrick von Platen's avatar
Patrick von Platen committed
2925
                Attention(
Will Berman's avatar
Will Berman committed
2926
2927
2928
                    query_dim=out_channels,
                    cross_attention_dim=out_channels,
                    heads=self.num_heads,
2929
                    dim_head=self.attention_head_dim,
Will Berman's avatar
Will Berman committed
2930
2931
2932
2933
                    added_kv_proj_dim=cross_attention_dim,
                    norm_num_groups=resnet_groups,
                    bias=True,
                    upcast_softmax=True,
2934
                    only_cross_attention=only_cross_attention,
2935
                    cross_attention_norm=cross_attention_norm,
2936
                    processor=processor,
Will Berman's avatar
Will Berman committed
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
                )
            )
        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,
2956
                        skip_time_act=skip_time_act,
Will Berman's avatar
Will Berman committed
2957
2958
2959
2960
2961
2962
2963
2964
                        up=True,
                    )
                ]
            )
        else:
            self.upsamplers = None

        self.gradient_checkpointing = False
2965
        self.resolution_idx = resolution_idx
Will Berman's avatar
Will Berman committed
2966
2967
2968

    def forward(
        self,
2969
2970
2971
2972
2973
2974
2975
2976
        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,
Will Berman's avatar
Will Berman committed
2977
    ):
2978
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
2979

2980
        lora_scale = cross_attention_kwargs.get("scale", 1.0)
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
        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
2992
2993
2994
2995
2996
2997
2998
        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)

2999
            if self.training and self.gradient_checkpointing:
Will Berman's avatar
Will Berman committed
3000

3001
3002
3003
3004
3005
3006
3007
3008
3009
3010
                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
3011
                hidden_states = attn(
3012
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
3013
3014
3015
3016
                    encoder_hidden_states=encoder_hidden_states,
                    attention_mask=mask,
                    **cross_attention_kwargs,
                )
3017
            else:
3018
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
3019
3020
3021
3022

                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
3023
                    attention_mask=mask,
3024
3025
                    **cross_attention_kwargs,
                )
Will Berman's avatar
Will Berman committed
3026
3027
3028

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

        return hidden_states
3032
3033
3034
3035
3036
3037
3038
3039


class KUpBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        temb_channels: int,
3040
        resolution_idx: int,
3041
3042
3043
3044
3045
3046
3047
3048
3049
3050
3051
3052
3053
3054
3055
3056
3057
3058
3059
3060
3061
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
3073
3074
3075
3076
3077
3078
3079
3080
3081
        dropout: float = 0.0,
        num_layers: int = 5,
        resnet_eps: float = 1e-5,
        resnet_act_fn: str = "gelu",
        resnet_group_size: Optional[int] = 32,
        add_upsample=True,
    ):
        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
3082
        self.resolution_idx = resolution_idx
3083

3084
    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0):
3085
3086
3087
3088
3089
3090
3091
3092
3093
3094
3095
3096
3097
        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

3098
3099
3100
3101
3102
3103
3104
3105
                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
                    )
3106
            else:
3107
                hidden_states = resnet(hidden_states, temb, scale=scale)
3108
3109
3110
3111
3112
3113
3114
3115
3116
3117
3118
3119
3120
3121

        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,
3122
        resolution_idx: int,
3123
3124
3125
3126
3127
        dropout: float = 0.0,
        num_layers: int = 4,
        resnet_eps: float = 1e-5,
        resnet_act_fn: str = "gelu",
        resnet_group_size: int = 32,
3128
        attention_head_dim=1,  # attention dim_head
3129
3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
3140
3141
        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
3142
        self.attention_head_dim = attention_head_dim
3143
3144
3145
3146
3147
3148
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

        # 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,
3178
                    k_out_channels // attention_head_dim
3179
                    if (i == num_layers - 1)
3180
3181
                    else out_channels // attention_head_dim,
                    attention_head_dim,
3182
3183
3184
3185
                    cross_attention_dim=cross_attention_dim,
                    temb_channels=temb_channels,
                    attention_bias=True,
                    add_self_attention=add_self_attention,
3186
                    cross_attention_norm="layer_norm",
3187
3188
3189
3190
3191
3192
3193
3194
3195
3196
3197
3198
3199
                    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
3200
        self.resolution_idx = resolution_idx
3201
3202
3203

    def forward(
        self,
3204
3205
3206
3207
3208
3209
3210
3211
        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,
3212
3213
3214
3215
3216
    ):
        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)

3217
        lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
3218
3219
3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
        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

3230
3231
3232
3233
3234
3235
3236
                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
3237
                hidden_states = attn(
3238
                    hidden_states,
ethansmith2000's avatar
ethansmith2000 committed
3239
3240
3241
3242
3243
3244
                    encoder_hidden_states=encoder_hidden_states,
                    emb=temb,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
                    encoder_attention_mask=encoder_attention_mask,
                )
3245
            else:
3246
                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
3247
3248
3249
3250
3251
3252
                hidden_states = attn(
                    hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    emb=temb,
                    attention_mask=attention_mask,
                    cross_attention_kwargs=cross_attention_kwargs,
3253
                    encoder_attention_mask=encoder_attention_mask,
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
3287
3288
3289
3290
3291
                )

        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.
        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
        num_embeds_ada_norm (:
            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
        attention_bias (:
            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
    """

    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,
3292
        cross_attention_norm: Optional[str] = None,
3293
3294
3295
3296
3297
3298
3299
3300
        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
3301
            self.attn1 = Attention(
3302
3303
3304
3305
3306
3307
                query_dim=dim,
                heads=num_attention_heads,
                dim_head=attention_head_dim,
                dropout=dropout,
                bias=attention_bias,
                cross_attention_dim=None,
3308
                cross_attention_norm=None,
3309
3310
3311
3312
            )

        # 2. Cross-Attn
        self.norm2 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size))
Patrick von Platen's avatar
Patrick von Platen committed
3313
        self.attn2 = Attention(
3314
3315
3316
3317
3318
3319
3320
3321
3322
3323
3324
3325
3326
3327
3328
3329
3330
3331
            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,
        )

    def _to_3d(self, hidden_states, height, weight):
        return hidden_states.permute(0, 2, 3, 1).reshape(hidden_states.shape[0], height * weight, -1)

    def _to_4d(self, hidden_states, height, weight):
        return hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, height, weight)

    def forward(
        self,
3332
3333
3334
3335
3336
3337
3338
3339
        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,
3340
3341
3342
3343
3344
3345
3346
3347
3348
3349
3350
3351
3352
    ):
        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,
3353
                attention_mask=attention_mask,
3354
3355
3356
3357
3358
3359
3360
3361
3362
3363
3364
3365
3366
3367
                **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,
3368
            attention_mask=attention_mask if encoder_hidden_states is None else encoder_attention_mask,
3369
3370
3371
3372
3373
3374
3375
            **cross_attention_kwargs,
        )
        attn_output = self._to_4d(attn_output, height, weight)

        hidden_states = attn_output + hidden_states

        return hidden_states