h1t commited on
Commit
ffedbd9
1 Parent(s): a7b0e90

Upload unet_2d_condition_woct.py

Browse files
Files changed (1) hide show
  1. oms_module/unet_2d_condition_woct.py +756 -0
oms_module/unet_2d_condition_woct.py ADDED
@@ -0,0 +1,756 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2023 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ from dataclasses import dataclass
15
+ from typing import Any, Dict, List, Optional, Tuple, Union
16
+
17
+ import torch
18
+ import torch.nn as nn
19
+ import torch.utils.checkpoint
20
+
21
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
22
+ from diffusers.loaders import UNet2DConditionLoadersMixin
23
+ from diffusers.utils import BaseOutput, logging
24
+ from diffusers.models.activations import get_activation
25
+ from diffusers.models.attention_processor import AttentionProcessor, AttnProcessor
26
+ from diffusers.models.embeddings import (
27
+ GaussianFourierProjection,
28
+ ImageHintTimeEmbedding,
29
+ ImageProjection,
30
+ ImageTimeEmbedding,
31
+ TextImageProjection,
32
+ TextImageTimeEmbedding,
33
+ TextTimeEmbedding,
34
+ TimestepEmbedding,
35
+ Timesteps,
36
+ )
37
+ from diffusers.models.modeling_utils import ModelMixin
38
+ from diffusers.models.unet_2d_blocks import (
39
+ CrossAttnDownBlock2D,
40
+ CrossAttnUpBlock2D,
41
+ DownBlock2D,
42
+ UNetMidBlock2DCrossAttn,
43
+ UNetMidBlock2DSimpleCrossAttn,
44
+ UpBlock2D,
45
+ get_down_block,
46
+ get_up_block,
47
+ )
48
+
49
+
50
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
51
+
52
+
53
+ @dataclass
54
+ class UNet2DConditionOutput(BaseOutput):
55
+ """
56
+ The output of [`UNet2DConditionModel`].
57
+
58
+ Args:
59
+ sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
60
+ The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
61
+ """
62
+
63
+ sample: torch.FloatTensor = None
64
+
65
+
66
+ class UNet2DConditionWoCTModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
67
+ r"""
68
+ A conditional 2D UNet model that takes a noisy sample, conditional state, but w/o a timestep and returns a sample
69
+ shaped output.
70
+
71
+ This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
72
+ for all models (such as downloading or saving).
73
+
74
+ Parameters:
75
+ sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
76
+ Height and width of input/output sample.
77
+ in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
78
+ out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
79
+ center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
80
+ down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
81
+ The tuple of downsample blocks to use.
82
+ mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
83
+ Block type for middle of UNet, it can be either `UNetMidBlock2DCrossAttn` or
84
+ `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.
85
+ up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
86
+ The tuple of upsample blocks to use.
87
+ only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
88
+ Whether to include self-attention in the basic transformer blocks, see
89
+ [`~models.attention.BasicTransformerBlock`].
90
+ block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
91
+ The tuple of output channels for each block.
92
+ layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
93
+ downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
94
+ mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
95
+ act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
96
+ norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
97
+ If `None`, normalization and activation layers is skipped in post-processing.
98
+ norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
99
+ cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
100
+ The dimension of the cross attention features.
101
+ transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
102
+ The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
103
+ [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
104
+ [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
105
+ encoder_hid_dim (`int`, *optional*, defaults to None):
106
+ If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
107
+ dimension to `cross_attention_dim`.
108
+ encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
109
+ If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
110
+ embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
111
+ attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
112
+ num_attention_heads (`int`, *optional*):
113
+ The number of attention heads. If not defined, defaults to `attention_head_dim`
114
+ conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer.
115
+ conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer.
116
+ mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):
117
+ Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If
118
+ `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the
119
+ `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False`
120
+ otherwise.
121
+ """
122
+
123
+ _supports_gradient_checkpointing = True
124
+
125
+ @register_to_config
126
+ def __init__(
127
+ self,
128
+ sample_size: Optional[int] = None,
129
+ in_channels: int = 4,
130
+ out_channels: int = 4,
131
+ center_input_sample: bool = False,
132
+ down_block_types: Tuple[str] = (
133
+ "CrossAttnDownBlock2D",
134
+ "CrossAttnDownBlock2D",
135
+ "CrossAttnDownBlock2D",
136
+ "DownBlock2D",
137
+ ),
138
+ mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
139
+ up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
140
+ only_cross_attention: Union[bool, Tuple[bool]] = False,
141
+ block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
142
+ layers_per_block: Union[int, Tuple[int]] = 2,
143
+ downsample_padding: int = 1,
144
+ mid_block_scale_factor: float = 1,
145
+ act_fn: str = "silu",
146
+ norm_num_groups: Optional[int] = 32,
147
+ norm_eps: float = 1e-5,
148
+ cross_attention_dim: Union[int, Tuple[int]] = 1280,
149
+ transformer_layers_per_block: Union[int, Tuple[int]] = 1,
150
+ encoder_hid_dim: Optional[int] = None,
151
+ encoder_hid_dim_type: Optional[str] = None,
152
+ attention_head_dim: Union[int, Tuple[int]] = 8,
153
+ num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
154
+ dual_cross_attention: bool = False,
155
+ use_linear_projection: bool = False,
156
+ upcast_attention: bool = False,
157
+ resnet_out_scale_factor: int = 1.0,
158
+ conv_in_kernel: int = 3,
159
+ conv_out_kernel: int = 3,
160
+ mid_block_only_cross_attention: Optional[bool] = None,
161
+ cross_attention_norm: Optional[str] = None,
162
+ ):
163
+ super().__init__()
164
+
165
+ self.sample_size = sample_size
166
+
167
+ if num_attention_heads is not None:
168
+ raise ValueError(
169
+ "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."
170
+ )
171
+
172
+ # If `num_attention_heads` is not defined (which is the case for most models)
173
+ # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
174
+ # The reason for this behavior is to correct for incorrectly named variables that were introduced
175
+ # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
176
+ # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
177
+ # which is why we correct for the naming here.
178
+ num_attention_heads = num_attention_heads or attention_head_dim
179
+
180
+ # Check inputs
181
+ if len(down_block_types) != len(up_block_types):
182
+ raise ValueError(
183
+ f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
184
+ )
185
+
186
+ if len(block_out_channels) != len(down_block_types):
187
+ raise ValueError(
188
+ f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
189
+ )
190
+
191
+ if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
192
+ raise ValueError(
193
+ f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
194
+ )
195
+
196
+ if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
197
+ raise ValueError(
198
+ f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
199
+ )
200
+
201
+ if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
202
+ raise ValueError(
203
+ f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
204
+ )
205
+
206
+ if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
207
+ raise ValueError(
208
+ f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
209
+ )
210
+
211
+ if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
212
+ raise ValueError(
213
+ f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
214
+ )
215
+
216
+ # input
217
+ conv_in_padding = (conv_in_kernel - 1) // 2
218
+ self.conv_in = nn.Conv2d(
219
+ in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
220
+ )
221
+
222
+ if encoder_hid_dim_type is None and encoder_hid_dim is not None:
223
+ encoder_hid_dim_type = "text_proj"
224
+ self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
225
+ logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
226
+
227
+ if encoder_hid_dim is None and encoder_hid_dim_type is not None:
228
+ raise ValueError(
229
+ f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
230
+ )
231
+
232
+ if encoder_hid_dim_type == "text_proj":
233
+ self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
234
+ elif encoder_hid_dim_type == "text_image_proj":
235
+ # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
236
+ # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
237
+ # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`
238
+ self.encoder_hid_proj = TextImageProjection(
239
+ text_embed_dim=encoder_hid_dim,
240
+ image_embed_dim=cross_attention_dim,
241
+ cross_attention_dim=cross_attention_dim,
242
+ )
243
+ elif encoder_hid_dim_type == "image_proj":
244
+ # Kandinsky 2.2
245
+ self.encoder_hid_proj = ImageProjection(
246
+ image_embed_dim=encoder_hid_dim,
247
+ cross_attention_dim=cross_attention_dim,
248
+ )
249
+ elif encoder_hid_dim_type is not None:
250
+ raise ValueError(
251
+ f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
252
+ )
253
+ else:
254
+ self.encoder_hid_proj = None
255
+
256
+ self.down_blocks = nn.ModuleList([])
257
+ self.up_blocks = nn.ModuleList([])
258
+
259
+ if isinstance(only_cross_attention, bool):
260
+ if mid_block_only_cross_attention is None:
261
+ mid_block_only_cross_attention = only_cross_attention
262
+
263
+ only_cross_attention = [only_cross_attention] * len(down_block_types)
264
+
265
+ if mid_block_only_cross_attention is None:
266
+ mid_block_only_cross_attention = False
267
+
268
+ if isinstance(num_attention_heads, int):
269
+ num_attention_heads = (num_attention_heads,) * len(down_block_types)
270
+
271
+ if isinstance(attention_head_dim, int):
272
+ attention_head_dim = (attention_head_dim,) * len(down_block_types)
273
+
274
+ if isinstance(cross_attention_dim, int):
275
+ cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
276
+
277
+ if isinstance(layers_per_block, int):
278
+ layers_per_block = [layers_per_block] * len(down_block_types)
279
+
280
+ if isinstance(transformer_layers_per_block, int):
281
+ transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
282
+
283
+ # disable time cond
284
+ time_embed_dim = None
285
+ blocks_time_embed_dim = time_embed_dim
286
+ resnet_time_scale_shift = None
287
+ resnet_skip_time_act = False
288
+
289
+ # down
290
+ output_channel = block_out_channels[0]
291
+ for i, down_block_type in enumerate(down_block_types):
292
+ input_channel = output_channel
293
+ output_channel = block_out_channels[i]
294
+ is_final_block = i == len(block_out_channels) - 1
295
+
296
+ down_block = get_down_block(
297
+ down_block_type,
298
+ num_layers=layers_per_block[i],
299
+ transformer_layers_per_block=transformer_layers_per_block[i],
300
+ in_channels=input_channel,
301
+ out_channels=output_channel,
302
+ temb_channels=blocks_time_embed_dim,
303
+ add_downsample=not is_final_block,
304
+ resnet_eps=norm_eps,
305
+ resnet_act_fn=act_fn,
306
+ resnet_groups=norm_num_groups,
307
+ cross_attention_dim=cross_attention_dim[i],
308
+ num_attention_heads=num_attention_heads[i],
309
+ downsample_padding=downsample_padding,
310
+ dual_cross_attention=dual_cross_attention,
311
+ use_linear_projection=use_linear_projection,
312
+ only_cross_attention=only_cross_attention[i],
313
+ upcast_attention=upcast_attention,
314
+ resnet_time_scale_shift=resnet_time_scale_shift,
315
+ resnet_skip_time_act=resnet_skip_time_act,
316
+ resnet_out_scale_factor=resnet_out_scale_factor,
317
+ cross_attention_norm=cross_attention_norm,
318
+ attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
319
+ )
320
+ self.down_blocks.append(down_block)
321
+
322
+ # mid
323
+ if mid_block_type == "UNetMidBlock2DCrossAttn":
324
+ self.mid_block = UNetMidBlock2DCrossAttn(
325
+ transformer_layers_per_block=transformer_layers_per_block[-1],
326
+ in_channels=block_out_channels[-1],
327
+ temb_channels=blocks_time_embed_dim,
328
+ resnet_eps=norm_eps,
329
+ resnet_act_fn=act_fn,
330
+ output_scale_factor=mid_block_scale_factor,
331
+ resnet_time_scale_shift=resnet_time_scale_shift,
332
+ cross_attention_dim=cross_attention_dim[-1],
333
+ num_attention_heads=num_attention_heads[-1],
334
+ resnet_groups=norm_num_groups,
335
+ dual_cross_attention=dual_cross_attention,
336
+ use_linear_projection=use_linear_projection,
337
+ upcast_attention=upcast_attention,
338
+ )
339
+ elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn":
340
+ self.mid_block = UNetMidBlock2DSimpleCrossAttn(
341
+ in_channels=block_out_channels[-1],
342
+ temb_channels=blocks_time_embed_dim,
343
+ resnet_eps=norm_eps,
344
+ resnet_act_fn=act_fn,
345
+ output_scale_factor=mid_block_scale_factor,
346
+ cross_attention_dim=cross_attention_dim[-1],
347
+ attention_head_dim=attention_head_dim[-1],
348
+ resnet_groups=norm_num_groups,
349
+ resnet_time_scale_shift=resnet_time_scale_shift,
350
+ skip_time_act=resnet_skip_time_act,
351
+ only_cross_attention=mid_block_only_cross_attention,
352
+ cross_attention_norm=cross_attention_norm,
353
+ )
354
+ elif mid_block_type is None:
355
+ self.mid_block = None
356
+ else:
357
+ raise ValueError(f"unknown mid_block_type : {mid_block_type}")
358
+
359
+ # count how many layers upsample the images
360
+ self.num_upsamplers = 0
361
+
362
+ # up
363
+ reversed_block_out_channels = list(reversed(block_out_channels))
364
+ reversed_num_attention_heads = list(reversed(num_attention_heads))
365
+ reversed_layers_per_block = list(reversed(layers_per_block))
366
+ reversed_cross_attention_dim = list(reversed(cross_attention_dim))
367
+ reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
368
+ only_cross_attention = list(reversed(only_cross_attention))
369
+
370
+ output_channel = reversed_block_out_channels[0]
371
+ for i, up_block_type in enumerate(up_block_types):
372
+ is_final_block = i == len(block_out_channels) - 1
373
+
374
+ prev_output_channel = output_channel
375
+ output_channel = reversed_block_out_channels[i]
376
+ input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
377
+
378
+ # add upsample block for all BUT final layer
379
+ if not is_final_block:
380
+ add_upsample = True
381
+ self.num_upsamplers += 1
382
+ else:
383
+ add_upsample = False
384
+
385
+ up_block = get_up_block(
386
+ up_block_type,
387
+ num_layers=reversed_layers_per_block[i] + 1,
388
+ transformer_layers_per_block=reversed_transformer_layers_per_block[i],
389
+ in_channels=input_channel,
390
+ out_channels=output_channel,
391
+ prev_output_channel=prev_output_channel,
392
+ temb_channels=blocks_time_embed_dim,
393
+ add_upsample=add_upsample,
394
+ resnet_eps=norm_eps,
395
+ resnet_act_fn=act_fn,
396
+ resnet_groups=norm_num_groups,
397
+ cross_attention_dim=reversed_cross_attention_dim[i],
398
+ num_attention_heads=reversed_num_attention_heads[i],
399
+ dual_cross_attention=dual_cross_attention,
400
+ use_linear_projection=use_linear_projection,
401
+ only_cross_attention=only_cross_attention[i],
402
+ upcast_attention=upcast_attention,
403
+ resnet_time_scale_shift=resnet_time_scale_shift,
404
+ resnet_skip_time_act=resnet_skip_time_act,
405
+ resnet_out_scale_factor=resnet_out_scale_factor,
406
+ cross_attention_norm=cross_attention_norm,
407
+ attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
408
+ )
409
+ self.up_blocks.append(up_block)
410
+ prev_output_channel = output_channel
411
+
412
+ # out
413
+ if norm_num_groups is not None:
414
+ self.conv_norm_out = nn.GroupNorm(
415
+ num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
416
+ )
417
+
418
+ self.conv_act = get_activation(act_fn)
419
+
420
+ else:
421
+ self.conv_norm_out = None
422
+ self.conv_act = None
423
+
424
+ conv_out_padding = (conv_out_kernel - 1) // 2
425
+ self.conv_out = nn.Conv2d(
426
+ block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding
427
+ )
428
+
429
+ @property
430
+ def attn_processors(self) -> Dict[str, AttentionProcessor]:
431
+ r"""
432
+ Returns:
433
+ `dict` of attention processors: A dictionary containing all attention processors used in the model with
434
+ indexed by its weight name.
435
+ """
436
+ # set recursively
437
+ processors = {}
438
+
439
+ def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
440
+ if hasattr(module, "set_processor"):
441
+ processors[f"{name}.processor"] = module.processor
442
+
443
+ for sub_name, child in module.named_children():
444
+ fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
445
+
446
+ return processors
447
+
448
+ for name, module in self.named_children():
449
+ fn_recursive_add_processors(name, module, processors)
450
+
451
+ return processors
452
+
453
+ def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
454
+ r"""
455
+ Sets the attention processor to use to compute attention.
456
+
457
+ Parameters:
458
+ processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
459
+ The instantiated processor class or a dictionary of processor classes that will be set as the processor
460
+ for **all** `Attention` layers.
461
+
462
+ If `processor` is a dict, the key needs to define the path to the corresponding cross attention
463
+ processor. This is strongly recommended when setting trainable attention processors.
464
+
465
+ """
466
+ count = len(self.attn_processors.keys())
467
+
468
+ if isinstance(processor, dict) and len(processor) != count:
469
+ raise ValueError(
470
+ f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
471
+ f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
472
+ )
473
+
474
+ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
475
+ if hasattr(module, "set_processor"):
476
+ if not isinstance(processor, dict):
477
+ module.set_processor(processor)
478
+ else:
479
+ module.set_processor(processor.pop(f"{name}.processor"))
480
+
481
+ for sub_name, child in module.named_children():
482
+ fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
483
+
484
+ for name, module in self.named_children():
485
+ fn_recursive_attn_processor(name, module, processor)
486
+
487
+ def set_default_attn_processor(self):
488
+ """
489
+ Disables custom attention processors and sets the default attention implementation.
490
+ """
491
+ self.set_attn_processor(AttnProcessor())
492
+
493
+ def set_attention_slice(self, slice_size):
494
+ r"""
495
+ Enable sliced attention computation.
496
+
497
+ When this option is enabled, the attention module splits the input tensor in slices to compute attention in
498
+ several steps. This is useful for saving some memory in exchange for a small decrease in speed.
499
+
500
+ Args:
501
+ slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
502
+ When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
503
+ `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
504
+ provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
505
+ must be a multiple of `slice_size`.
506
+ """
507
+ sliceable_head_dims = []
508
+
509
+ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
510
+ if hasattr(module, "set_attention_slice"):
511
+ sliceable_head_dims.append(module.sliceable_head_dim)
512
+
513
+ for child in module.children():
514
+ fn_recursive_retrieve_sliceable_dims(child)
515
+
516
+ # retrieve number of attention layers
517
+ for module in self.children():
518
+ fn_recursive_retrieve_sliceable_dims(module)
519
+
520
+ num_sliceable_layers = len(sliceable_head_dims)
521
+
522
+ if slice_size == "auto":
523
+ # half the attention head size is usually a good trade-off between
524
+ # speed and memory
525
+ slice_size = [dim // 2 for dim in sliceable_head_dims]
526
+ elif slice_size == "max":
527
+ # make smallest slice possible
528
+ slice_size = num_sliceable_layers * [1]
529
+
530
+ slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
531
+
532
+ if len(slice_size) != len(sliceable_head_dims):
533
+ raise ValueError(
534
+ f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
535
+ f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
536
+ )
537
+
538
+ for i in range(len(slice_size)):
539
+ size = slice_size[i]
540
+ dim = sliceable_head_dims[i]
541
+ if size is not None and size > dim:
542
+ raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
543
+
544
+ # Recursively walk through all the children.
545
+ # Any children which exposes the set_attention_slice method
546
+ # gets the message
547
+ def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
548
+ if hasattr(module, "set_attention_slice"):
549
+ module.set_attention_slice(slice_size.pop())
550
+
551
+ for child in module.children():
552
+ fn_recursive_set_attention_slice(child, slice_size)
553
+
554
+ reversed_slice_size = list(reversed(slice_size))
555
+ for module in self.children():
556
+ fn_recursive_set_attention_slice(module, reversed_slice_size)
557
+
558
+ def _set_gradient_checkpointing(self, module, value=False):
559
+ if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D, CrossAttnUpBlock2D, UpBlock2D)):
560
+ module.gradient_checkpointing = value
561
+
562
+ def forward(
563
+ self,
564
+ sample: torch.FloatTensor,
565
+ encoder_hidden_states: torch.Tensor,
566
+ attention_mask: Optional[torch.Tensor] = None,
567
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
568
+ added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
569
+ down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
570
+ mid_block_additional_residual: Optional[torch.Tensor] = None,
571
+ encoder_attention_mask: Optional[torch.Tensor] = None,
572
+ return_dict: bool = True,
573
+ ) -> Union[UNet2DConditionOutput, Tuple]:
574
+ r"""
575
+ The [`UNet2DConditionModel`] forward method.
576
+
577
+ Args:
578
+ sample (`torch.FloatTensor`):
579
+ The noisy input tensor with the following shape `(batch, channel, height, width)`.
580
+ encoder_hidden_states (`torch.FloatTensor`):
581
+ The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
582
+ encoder_attention_mask (`torch.Tensor`):
583
+ A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
584
+ `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
585
+ which adds large negative values to the attention scores corresponding to "discard" tokens.
586
+ return_dict (`bool`, *optional*, defaults to `True`):
587
+ Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
588
+ tuple.
589
+ cross_attention_kwargs (`dict`, *optional*):
590
+ A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].
591
+ added_cond_kwargs: (`dict`, *optional*):
592
+ A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that
593
+ are passed along to the UNet blocks.
594
+
595
+ Returns:
596
+ [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
597
+ If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise
598
+ a `tuple` is returned where the first element is the sample tensor.
599
+ """
600
+ # By default samples have to be AT least a multiple of the overall upsampling factor.
601
+ # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
602
+ # However, the upsampling interpolation output size can be forced to fit any upsampling size
603
+ # on the fly if necessary.
604
+ default_overall_up_factor = 2**self.num_upsamplers
605
+
606
+ # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
607
+ forward_upsample_size = False
608
+ upsample_size = None
609
+
610
+ if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
611
+ logger.info("Forward upsample size to force interpolation output size.")
612
+ forward_upsample_size = True
613
+
614
+ # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
615
+ # expects mask of shape:
616
+ # [batch, key_tokens]
617
+ # adds singleton query_tokens dimension:
618
+ # [batch, 1, key_tokens]
619
+ # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
620
+ # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
621
+ # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
622
+ if attention_mask is not None:
623
+ # assume that mask is expressed as:
624
+ # (1 = keep, 0 = discard)
625
+ # convert mask into a bias that can be added to attention scores:
626
+ # (keep = +0, discard = -10000.0)
627
+ attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
628
+ attention_mask = attention_mask.unsqueeze(1)
629
+
630
+ # convert encoder_attention_mask to a bias the same way we do for attention_mask
631
+ if encoder_attention_mask is not None:
632
+ encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
633
+ encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
634
+
635
+ # 0. center input if necessary
636
+ if self.config.center_input_sample:
637
+ sample = 2 * sample - 1.0
638
+
639
+ # 1. time (skip)
640
+ emb = None
641
+
642
+ if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
643
+ encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
644
+ elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
645
+ # Kadinsky 2.1 - style
646
+ if "image_embeds" not in added_cond_kwargs:
647
+ raise ValueError(
648
+ f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`"
649
+ )
650
+
651
+ image_embeds = added_cond_kwargs.get("image_embeds")
652
+ encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
653
+ elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
654
+ # Kandinsky 2.2 - style
655
+ if "image_embeds" not in added_cond_kwargs:
656
+ raise ValueError(
657
+ f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`"
658
+ )
659
+ image_embeds = added_cond_kwargs.get("image_embeds")
660
+ encoder_hidden_states = self.encoder_hid_proj(image_embeds)
661
+ # 2. pre-process
662
+ sample = self.conv_in(sample)
663
+
664
+ # 3. down
665
+
666
+ is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
667
+ is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None
668
+
669
+ down_block_res_samples = (sample,)
670
+ for downsample_block in self.down_blocks:
671
+ if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
672
+ # For t2i-adapter CrossAttnDownBlock2D
673
+ additional_residuals = {}
674
+ if is_adapter and len(down_block_additional_residuals) > 0:
675
+ additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0)
676
+
677
+ sample, res_samples = downsample_block(
678
+ hidden_states=sample,
679
+ temb=emb,
680
+ encoder_hidden_states=encoder_hidden_states,
681
+ attention_mask=attention_mask,
682
+ cross_attention_kwargs=cross_attention_kwargs,
683
+ encoder_attention_mask=encoder_attention_mask,
684
+ **additional_residuals,
685
+ )
686
+ else:
687
+ sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
688
+
689
+ if is_adapter and len(down_block_additional_residuals) > 0:
690
+ sample += down_block_additional_residuals.pop(0)
691
+
692
+ down_block_res_samples += res_samples
693
+
694
+ if is_controlnet:
695
+ new_down_block_res_samples = ()
696
+
697
+ for down_block_res_sample, down_block_additional_residual in zip(
698
+ down_block_res_samples, down_block_additional_residuals
699
+ ):
700
+ down_block_res_sample = down_block_res_sample + down_block_additional_residual
701
+ new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
702
+
703
+ down_block_res_samples = new_down_block_res_samples
704
+
705
+ # 4. mid
706
+ if self.mid_block is not None:
707
+ sample = self.mid_block(
708
+ sample,
709
+ emb,
710
+ encoder_hidden_states=encoder_hidden_states,
711
+ attention_mask=attention_mask,
712
+ cross_attention_kwargs=cross_attention_kwargs,
713
+ encoder_attention_mask=encoder_attention_mask,
714
+ )
715
+
716
+ if is_controlnet:
717
+ sample = sample + mid_block_additional_residual
718
+
719
+ # 5. up
720
+ for i, upsample_block in enumerate(self.up_blocks):
721
+ is_final_block = i == len(self.up_blocks) - 1
722
+
723
+ res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
724
+ down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
725
+
726
+ # if we have not reached the final block and need to forward the
727
+ # upsample size, we do it here
728
+ if not is_final_block and forward_upsample_size:
729
+ upsample_size = down_block_res_samples[-1].shape[2:]
730
+
731
+ if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
732
+ sample = upsample_block(
733
+ hidden_states=sample,
734
+ temb=emb,
735
+ res_hidden_states_tuple=res_samples,
736
+ encoder_hidden_states=encoder_hidden_states,
737
+ cross_attention_kwargs=cross_attention_kwargs,
738
+ upsample_size=upsample_size,
739
+ attention_mask=attention_mask,
740
+ encoder_attention_mask=encoder_attention_mask,
741
+ )
742
+ else:
743
+ sample = upsample_block(
744
+ hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
745
+ )
746
+
747
+ # 6. post-process
748
+ if self.conv_norm_out:
749
+ sample = self.conv_norm_out(sample)
750
+ sample = self.conv_act(sample)
751
+ sample = self.conv_out(sample)
752
+
753
+ if not return_dict:
754
+ return (sample,)
755
+
756
+ return UNet2DConditionOutput(sample=sample)