AnchoredAI/llm-grounded-diffusion
0
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 at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# 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 and13# limitations under the License.14from dataclasses import dataclass15from typing import Any, Dict, List, Optional, Tuple, Union16 17import torch18import torch.nn as nn19import torch.nn.functional as F20import torch.utils.checkpoint21 22from diffusers.configuration_utils import ConfigMixin, register_to_config23from diffusers.loaders import UNet2DConditionLoadersMixin24from diffusers.utils import BaseOutput, logging25from diffusers.models.embeddings import (26 GaussianFourierProjection,27 TextImageProjection,28 TextImageTimeEmbedding,29 TextTimeEmbedding,30 TimestepEmbedding,31 Timesteps,32)33from diffusers.models.modeling_utils import ModelMixin34from .unet_2d_blocks import (35 CrossAttnDownBlock2D,36 CrossAttnUpBlock2D,37 DownBlock2D,38 UNetMidBlock2DCrossAttn,39 UpBlock2D,40 get_down_block,41 get_up_block,42)43from .attention_processor import AttentionProcessor, AttnProcessor44 45 46logger = logging.get_logger(__name__) # pylint: disable=invalid-name47 48 49@dataclass50class UNet2DConditionOutput(BaseOutput):51 """52 Args:53 sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):54 Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.55 """56 57 sample: torch.FloatTensor58 cross_attention_probs_down: List[Any]59 cross_attention_probs_mid: List[Any]60 cross_attention_probs_up: List[Any]61 62 63class FourierEmbedder(nn.Module):64 def __init__(self, num_freqs=64, temperature=100):65 super().__init__()66 67 self.num_freqs = num_freqs68 self.temperature = temperature69 70 freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)71 freq_bands = freq_bands[None, None, None]72 self.register_buffer('freq_bands', freq_bands, persistent=False)73 74 def __call__(self, x):75 x = self.freq_bands * x.unsqueeze(-1)76 return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 1, 3, 4, 2).reshape(*x.shape[:2], -1)77 78 79class PositionNet(nn.Module):80 def __init__(self, positive_len, out_dim, fourier_freqs=8):81 super().__init__()82 self.positive_len = positive_len83 self.out_dim = out_dim 84 85 self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)86 self.position_dim = fourier_freqs * 2 * 4 # 2: sin/cos, 4: xyxy 87 88 self.linears = nn.Sequential(89 nn.Linear(self.positive_len + self.position_dim, 512),90 nn.SiLU(),91 nn.Linear(512, 512),92 nn.SiLU(),93 nn.Linear(512, out_dim),94 )95 96 self.null_positive_feature = torch.nn.Parameter(torch.zeros([self.positive_len]))97 self.null_position_feature = torch.nn.Parameter(torch.zeros([self.position_dim]))98 99 def forward(self, boxes, masks, positive_embeddings):100 masks = masks.unsqueeze(-1)101 102 # embedding position (it may includes padding as placeholder)103 xyxy_embedding = self.fourier_embedder(boxes) # B*N*4 -> B*N*C104 105 # learnable null embedding 106 positive_null = self.null_positive_feature.view(1, 1, -1)107 xyxy_null = self.null_position_feature.view(1, 1, -1)108 109 # replace padding with learnable null embedding 110 positive_embeddings = positive_embeddings * masks + (1 - masks) * positive_null111 xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null112 113 objs = self.linears(torch.cat([positive_embeddings, xyxy_embedding], dim=-1))114 return objs115 116 117 118class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):119 r"""120 UNet2DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a timestep121 and returns sample shaped output.122 123 This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library124 implements for all the models (such as downloading or saving, etc.)125 126 Parameters:127 sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):128 Height and width of input/output sample.129 in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.130 out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.131 center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.132 flip_sin_to_cos (`bool`, *optional*, defaults to `False`):133 Whether to flip the sin to cos in the time embedding.134 freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.135 down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):136 The tuple of downsample blocks to use.137 mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):138 The mid block type. Choose from `UNetMidBlock2DCrossAttn` or `UNetMidBlock2DSimpleCrossAttn`, will skip the139 mid block layer if `None`.140 up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D",)`):141 The tuple of upsample blocks to use.142 only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):143 Whether to include self-attention in the basic transformer blocks, see144 [`~models.attention.BasicTransformerBlock`].145 block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):146 The tuple of output channels for each block.147 layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.148 downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.149 mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.150 act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.151 norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.152 If `None`, it will skip the normalization and activation layers in post-processing153 norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.154 cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):155 The dimension of the cross attention features.156 encoder_hid_dim (`int`, *optional*, defaults to None):157 If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`158 dimension to `cross_attention_dim`.159 encoder_hid_dim_type (`str`, *optional*, defaults to None):160 If given, the `encoder_hidden_states` and potentially other embeddings will be down-projected to text161 embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.162 attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.163 resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config164 for resnet blocks, see [`~models.resnet.ResnetBlock2D`]. Choose from `default` or `scale_shift`.165 class_embed_type (`str`, *optional*, defaults to None):166 The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,167 `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.168 addition_embed_type (`str`, *optional*, defaults to None):169 Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or170 "text". "text" will use the `TextTimeEmbedding` layer.171 num_class_embeds (`int`, *optional*, defaults to None):172 Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing173 class conditioning with `class_embed_type` equal to `None`.174 time_embedding_type (`str`, *optional*, default to `positional`):175 The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.176 time_embedding_dim (`int`, *optional*, default to `None`):177 An optional override for the dimension of the projected time embedding.178 time_embedding_act_fn (`str`, *optional*, default to `None`):179 Optional activation function to use on the time embeddings only one time before they as passed to the rest180 of the unet. Choose from `silu`, `mish`, `gelu`, and `swish`.181 timestep_post_act (`str, *optional*, default to `None`):182 The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.183 time_cond_proj_dim (`int`, *optional*, default to `None`):184 The dimension of `cond_proj` layer in timestep embedding.185 conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer.186 conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer.187 projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when188 using the "projection" `class_embed_type`. Required when using the "projection" `class_embed_type`.189 class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time190 embeddings with the class embeddings.191 mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):192 Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If193 `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is None, the194 `only_cross_attention` value will be used as the value for `mid_block_only_cross_attention`. Else, it will195 default to `False`.196 """197 198 _supports_gradient_checkpointing = True199 200 @register_to_config201 def __init__(202 self,203 sample_size: Optional[int] = None,204 in_channels: int = 4,205 out_channels: int = 4,206 center_input_sample: bool = False,207 flip_sin_to_cos: bool = True,208 freq_shift: int = 0,209 down_block_types: Tuple[str] = (210 "CrossAttnDownBlock2D",211 "CrossAttnDownBlock2D",212 "CrossAttnDownBlock2D",213 "DownBlock2D",214 ),215 mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",216 up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),217 only_cross_attention: Union[bool, Tuple[bool]] = False,218 block_out_channels: Tuple[int] = (320, 640, 1280, 1280),219 layers_per_block: Union[int, Tuple[int]] = 2,220 downsample_padding: int = 1,221 mid_block_scale_factor: float = 1,222 act_fn: str = "silu",223 norm_num_groups: Optional[int] = 32,224 norm_eps: float = 1e-5,225 cross_attention_dim: Union[int, Tuple[int]] = 1280,226 encoder_hid_dim: Optional[int] = None,227 encoder_hid_dim_type: Optional[str] = None,228 attention_head_dim: Union[int, Tuple[int]] = 8,229 dual_cross_attention: bool = False,230 use_linear_projection: bool = False,231 class_embed_type: Optional[str] = None,232 addition_embed_type: Optional[str] = None,233 num_class_embeds: Optional[int] = None,234 upcast_attention: bool = False,235 resnet_time_scale_shift: str = "default",236 resnet_skip_time_act: bool = False,237 resnet_out_scale_factor: int = 1.0,238 time_embedding_type: str = "positional",239 time_embedding_dim: Optional[int] = None,240 time_embedding_act_fn: Optional[str] = None,241 timestep_post_act: Optional[str] = None,242 time_cond_proj_dim: Optional[int] = None,243 conv_in_kernel: int = 3,244 conv_out_kernel: int = 3,245 projection_class_embeddings_input_dim: Optional[int] = None,246 class_embeddings_concat: bool = False,247 mid_block_only_cross_attention: Optional[bool] = None,248 cross_attention_norm: Optional[str] = None,249 addition_embed_type_num_heads=64,250 use_gated_attention: bool = False,251 ):252 super().__init__()253 254 self.sample_size = sample_size255 256 # Check inputs257 if len(down_block_types) != len(up_block_types):258 raise ValueError(259 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}."260 )261 262 if len(block_out_channels) != len(down_block_types):263 raise ValueError(264 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}."265 )266 267 if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):268 raise ValueError(269 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}."270 )271 272 if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):273 raise ValueError(274 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}."275 )276 277 if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):278 raise ValueError(279 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}."280 )281 282 if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):283 raise ValueError(284 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}."285 )286 287 # input288 conv_in_padding = (conv_in_kernel - 1) // 2289 self.conv_in = nn.Conv2d(290 in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding291 )292 293 # time294 if time_embedding_type == "fourier":295 time_embed_dim = time_embedding_dim or block_out_channels[0] * 2296 if time_embed_dim % 2 != 0:297 raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.")298 self.time_proj = GaussianFourierProjection(299 time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos300 )301 timestep_input_dim = time_embed_dim302 elif time_embedding_type == "positional":303 time_embed_dim = time_embedding_dim or block_out_channels[0] * 4304 305 self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)306 timestep_input_dim = block_out_channels[0]307 else:308 raise ValueError(309 f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."310 )311 312 self.time_embedding = TimestepEmbedding(313 timestep_input_dim,314 time_embed_dim,315 act_fn=act_fn,316 post_act_fn=timestep_post_act,317 cond_proj_dim=time_cond_proj_dim,318 )319 320 if encoder_hid_dim_type is None and encoder_hid_dim is not None:321 encoder_hid_dim_type = "text_proj"322 logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")323 324 if encoder_hid_dim is None and encoder_hid_dim_type is not None:325 raise ValueError(326 f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."327 )328 329 if encoder_hid_dim_type == "text_proj":330 self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)331 elif encoder_hid_dim_type == "text_image_proj":332 # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much333 # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use334 # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`335 self.encoder_hid_proj = TextImageProjection(336 text_embed_dim=encoder_hid_dim,337 image_embed_dim=cross_attention_dim,338 cross_attention_dim=cross_attention_dim,339 )340 341 elif encoder_hid_dim_type is not None:342 raise ValueError(343 f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."344 )345 else:346 self.encoder_hid_proj = None347 348 # class embedding349 if class_embed_type is None and num_class_embeds is not None:350 self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)351 elif class_embed_type == "timestep":352 self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn)353 elif class_embed_type == "identity":354 self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)355 elif class_embed_type == "projection":356 if projection_class_embeddings_input_dim is None:357 raise ValueError(358 "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"359 )360 # The projection `class_embed_type` is the same as the timestep `class_embed_type` except361 # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings362 # 2. it projects from an arbitrary input dimension.363 #364 # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.365 # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.366 # As a result, `TimestepEmbedding` can be passed arbitrary vectors.367 self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)368 elif class_embed_type == "simple_projection":369 if projection_class_embeddings_input_dim is None:370 raise ValueError(371 "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"372 )373 self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim)374 else:375 self.class_embedding = None376 377 if addition_embed_type == "text":378 if encoder_hid_dim is not None:379 text_time_embedding_from_dim = encoder_hid_dim380 else:381 text_time_embedding_from_dim = cross_attention_dim382 383 self.add_embedding = TextTimeEmbedding(384 text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads385 )386 elif addition_embed_type == "text_image":387 # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much388 # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use389 # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`390 self.add_embedding = TextImageTimeEmbedding(391 text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim392 )393 elif addition_embed_type is not None:394 raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")395 396 if time_embedding_act_fn is None:397 self.time_embed_act = None398 elif time_embedding_act_fn == "swish":399 self.time_embed_act = lambda x: F.silu(x)400 elif time_embedding_act_fn == "mish":401 self.time_embed_act = nn.Mish()402 elif time_embedding_act_fn == "silu":403 self.time_embed_act = nn.SiLU()404 elif time_embedding_act_fn == "gelu":405 self.time_embed_act = nn.GELU()406 else:407 raise ValueError(f"Unsupported activation function: {time_embedding_act_fn}")408 409 self.down_blocks = nn.ModuleList([])410 self.up_blocks = nn.ModuleList([])411 412 if isinstance(only_cross_attention, bool):413 if mid_block_only_cross_attention is None:414 mid_block_only_cross_attention = only_cross_attention415 416 only_cross_attention = [only_cross_attention] * len(down_block_types)417 418 if mid_block_only_cross_attention is None:419 mid_block_only_cross_attention = False420 421 if isinstance(attention_head_dim, int):422 attention_head_dim = (attention_head_dim,) * len(down_block_types)423 424 if isinstance(cross_attention_dim, int):425 cross_attention_dim = (cross_attention_dim,) * len(down_block_types)426 else:427 assert not use_gated_attention, f"use_gated_attention is not supported with varying cross_attention_dim: {cross_attention_dim}"428 429 if isinstance(layers_per_block, int):430 layers_per_block = [layers_per_block] * len(down_block_types)431 432 if class_embeddings_concat:433 # The time embeddings are concatenated with the class embeddings. The dimension of the434 # time embeddings passed to the down, middle, and up blocks is twice the dimension of the435 # regular time embeddings436 blocks_time_embed_dim = time_embed_dim * 2437 else:438 blocks_time_embed_dim = time_embed_dim439 440 # down441 output_channel = block_out_channels[0]442 for i, down_block_type in enumerate(down_block_types):443 input_channel = output_channel444 output_channel = block_out_channels[i]445 is_final_block = i == len(block_out_channels) - 1446 447 down_block = get_down_block(448 down_block_type,449 num_layers=layers_per_block[i],450 in_channels=input_channel,451 out_channels=output_channel,452 temb_channels=blocks_time_embed_dim,453 add_downsample=not is_final_block,454 resnet_eps=norm_eps,455 resnet_act_fn=act_fn,456 resnet_groups=norm_num_groups,457 cross_attention_dim=cross_attention_dim[i],458 attn_num_head_channels=attention_head_dim[i],459 downsample_padding=downsample_padding,460 dual_cross_attention=dual_cross_attention,461 use_linear_projection=use_linear_projection,462 only_cross_attention=only_cross_attention[i],463 upcast_attention=upcast_attention,464 resnet_time_scale_shift=resnet_time_scale_shift,465 resnet_skip_time_act=resnet_skip_time_act,466 resnet_out_scale_factor=resnet_out_scale_factor,467 cross_attention_norm=cross_attention_norm,468 use_gated_attention=use_gated_attention,469 )470 self.down_blocks.append(down_block)471 472 # mid473 if mid_block_type == "UNetMidBlock2DCrossAttn":474 self.mid_block = UNetMidBlock2DCrossAttn(475 in_channels=block_out_channels[-1],476 temb_channels=blocks_time_embed_dim,477 resnet_eps=norm_eps,478 resnet_act_fn=act_fn,479 output_scale_factor=mid_block_scale_factor,480 resnet_time_scale_shift=resnet_time_scale_shift,481 cross_attention_dim=cross_attention_dim[-1],482 attn_num_head_channels=attention_head_dim[-1],483 resnet_groups=norm_num_groups,484 dual_cross_attention=dual_cross_attention,485 use_linear_projection=use_linear_projection,486 upcast_attention=upcast_attention,487 use_gated_attention=use_gated_attention,488 )489 elif mid_block_type is None:490 self.mid_block = None491 else:492 raise ValueError(f"unknown mid_block_type : {mid_block_type}")493 494 # count how many layers upsample the images495 self.num_upsamplers = 0496 497 # up498 reversed_block_out_channels = list(reversed(block_out_channels))499 reversed_attention_head_dim = list(reversed(attention_head_dim))500 reversed_layers_per_block = list(reversed(layers_per_block))501 reversed_cross_attention_dim = list(reversed(cross_attention_dim))502 only_cross_attention = list(reversed(only_cross_attention))503 504 output_channel = reversed_block_out_channels[0]505 for i, up_block_type in enumerate(up_block_types):506 is_final_block = i == len(block_out_channels) - 1507 508 prev_output_channel = output_channel509 output_channel = reversed_block_out_channels[i]510 input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]511 512 # add upsample block for all BUT final layer513 if not is_final_block:514 add_upsample = True515 self.num_upsamplers += 1516 else:517 add_upsample = False518 519 up_block = get_up_block(520 up_block_type,521 num_layers=reversed_layers_per_block[i] + 1,522 in_channels=input_channel,523 out_channels=output_channel,524 prev_output_channel=prev_output_channel,525 temb_channels=blocks_time_embed_dim,526 add_upsample=add_upsample,527 resnet_eps=norm_eps,528 resnet_act_fn=act_fn,529 resnet_groups=norm_num_groups,530 cross_attention_dim=reversed_cross_attention_dim[i],531 attn_num_head_channels=reversed_attention_head_dim[i],532 dual_cross_attention=dual_cross_attention,533 use_linear_projection=use_linear_projection,534 only_cross_attention=only_cross_attention[i],535 upcast_attention=upcast_attention,536 resnet_time_scale_shift=resnet_time_scale_shift,537 resnet_skip_time_act=resnet_skip_time_act,538 resnet_out_scale_factor=resnet_out_scale_factor,539 cross_attention_norm=cross_attention_norm,540 use_gated_attention=use_gated_attention,541 )542 self.up_blocks.append(up_block)543 prev_output_channel = output_channel544 545 # out546 if norm_num_groups is not None:547 self.conv_norm_out = nn.GroupNorm(548 num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps549 )550 551 if act_fn == "swish":552 self.conv_act = lambda x: F.silu(x)553 elif act_fn == "mish":554 self.conv_act = nn.Mish()555 elif act_fn == "silu":556 self.conv_act = nn.SiLU()557 elif act_fn == "gelu":558 self.conv_act = nn.GELU()559 else:560 raise ValueError(f"Unsupported activation function: {act_fn}")561 562 else:563 self.conv_norm_out = None564 self.conv_act = None565 566 conv_out_padding = (conv_out_kernel - 1) // 2567 self.conv_out = nn.Conv2d(568 block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding569 )570 571 if use_gated_attention:572 self.position_net = PositionNet(positive_len=768, out_dim=cross_attention_dim[-1])573 574 575 @property576 def attn_processors(self) -> Dict[str, AttentionProcessor]:577 r"""578 Returns:579 `dict` of attention processors: A dictionary containing all attention processors used in the model with580 indexed by its weight name.581 """582 # set recursively583 processors = {}584 585 def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):586 if hasattr(module, "set_processor"):587 processors[f"{name}.processor"] = module.processor588 589 for sub_name, child in module.named_children():590 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)591 592 return processors593 594 for name, module in self.named_children():595 fn_recursive_add_processors(name, module, processors)596 597 return processors598 599 def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):600 r"""601 Parameters:602 `processor (`dict` of `AttentionProcessor` or `AttentionProcessor`):603 The instantiated processor class or a dictionary of processor classes that will be set as the processor604 of **all** `Attention` layers.605 In case `processor` is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors.:606 607 """608 count = len(self.attn_processors.keys())609 610 if isinstance(processor, dict) and len(processor) != count:611 raise ValueError(612 f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"613 f" number of attention layers: {count}. Please make sure to pass {count} processor classes."614 )615 616 def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):617 if hasattr(module, "set_processor"):618 if not isinstance(processor, dict):619 module.set_processor(processor)620 else:621 module.set_processor(processor.pop(f"{name}.processor"))622 623 for sub_name, child in module.named_children():624 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)625 626 for name, module in self.named_children():627 fn_recursive_attn_processor(name, module, processor)628 629 def set_default_attn_processor(self):630 """631 Disables custom attention processors and sets the default attention implementation.632 """633 self.set_attn_processor(AttnProcessor())634 635 def set_attention_slice(self, slice_size):636 r"""637 Enable sliced attention computation.638 639 When this option is enabled, the attention module will split the input tensor in slices, to compute attention640 in several steps. This is useful to save some memory in exchange for a small speed decrease.641 642 Args:643 slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):644 When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If645 `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is646 provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`647 must be a multiple of `slice_size`.648 """649 sliceable_head_dims = []650 651 def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):652 if hasattr(module, "set_attention_slice"):653 sliceable_head_dims.append(module.sliceable_head_dim)654 655 for child in module.children():656 fn_recursive_retrieve_sliceable_dims(child)657 658 # retrieve number of attention layers659 for module in self.children():660 fn_recursive_retrieve_sliceable_dims(module)661 662 num_sliceable_layers = len(sliceable_head_dims)663 664 if slice_size == "auto":665 # half the attention head size is usually a good trade-off between666 # speed and memory667 slice_size = [dim // 2 for dim in sliceable_head_dims]668 elif slice_size == "max":669 # make smallest slice possible670 slice_size = num_sliceable_layers * [1]671 672 slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size673 674 if len(slice_size) != len(sliceable_head_dims):675 raise ValueError(676 f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"677 f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."678 )679 680 for i in range(len(slice_size)):681 size = slice_size[i]682 dim = sliceable_head_dims[i]683 if size is not None and size > dim:684 raise ValueError(f"size {size} has to be smaller or equal to {dim}.")685 686 # Recursively walk through all the children.687 # Any children which exposes the set_attention_slice method688 # gets the message689 def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):690 if hasattr(module, "set_attention_slice"):691 module.set_attention_slice(slice_size.pop())692 693 for child in module.children():694 fn_recursive_set_attention_slice(child, slice_size)695 696 reversed_slice_size = list(reversed(slice_size))697 for module in self.children():698 fn_recursive_set_attention_slice(module, reversed_slice_size)699 700 def _set_gradient_checkpointing(self, module, value=False):701 if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D, CrossAttnUpBlock2D, UpBlock2D)):702 module.gradient_checkpointing = value703 704 def forward(705 self,706 sample: torch.FloatTensor,707 timestep: Union[torch.Tensor, float, int],708 encoder_hidden_states: torch.Tensor,709 class_labels: Optional[torch.Tensor] = None,710 timestep_cond: Optional[torch.Tensor] = None,711 attention_mask: Optional[torch.Tensor] = None,712 cross_attention_kwargs: Optional[Dict[str, Any]] = None,713 added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,714 down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,715 mid_block_additional_residual: Optional[torch.Tensor] = None,716 encoder_attention_mask: Optional[torch.Tensor] = None,717 return_dict: bool = True,718 return_cross_attention_probs: bool = False719 ) -> Union[UNet2DConditionOutput, Tuple]:720 r"""721 Args:722 sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor723 timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps724 encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states725 encoder_attention_mask (`torch.Tensor`):726 (batch, sequence_length) cross-attention mask, applied to encoder_hidden_states. True = keep, False =727 discard. Mask will be converted into a bias, which adds large negative values to attention scores728 corresponding to "discard" tokens.729 return_dict (`bool`, *optional*, defaults to `True`):730 Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.731 cross_attention_kwargs (`dict`, *optional*):732 A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under733 `self.processor` in734 [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).735 added_cond_kwargs (`dict`, *optional*):736 A kwargs dictionary that if specified includes additonal conditions that can be used for additonal time737 embeddings or encoder hidden states projections. See the configurations `encoder_hid_dim_type` and738 `addition_embed_type` for more information.739 740 Returns:741 [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:742 [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When743 returning a tuple, the first element is the sample tensor.744 """745 # By default samples have to be AT least a multiple of the overall upsampling factor.746 # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).747 # However, the upsampling interpolation output size can be forced to fit any upsampling size748 # on the fly if necessary.749 default_overall_up_factor = 2**self.num_upsamplers750 751 # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`752 forward_upsample_size = False753 upsample_size = None754 755 if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):756 logger.info("Forward upsample size to force interpolation output size.")757 forward_upsample_size = True758 759 # ensure attention_mask is a bias, and give it a singleton query_tokens dimension760 # expects mask of shape:761 # [batch, key_tokens]762 # adds singleton query_tokens dimension:763 # [batch, 1, key_tokens]764 # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:765 # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)766 # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)767 if attention_mask is not None:768 # assume that mask is expressed as:769 # (1 = keep, 0 = discard)770 # convert mask into a bias that can be added to attention scores:771 # (keep = +0, discard = -10000.0)772 attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0773 attention_mask = attention_mask.unsqueeze(1)774 775 # convert encoder_attention_mask to a bias the same way we do for attention_mask776 if encoder_attention_mask is not None:777 encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0778 encoder_attention_mask = encoder_attention_mask.unsqueeze(1)779 780 # 0. center input if necessary781 if self.config.center_input_sample:782 sample = 2 * sample - 1.0783 784 # 1. time785 timesteps = timestep786 if not torch.is_tensor(timesteps):787 # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can788 # This would be a good case for the `match` statement (Python 3.10+)789 is_mps = sample.device.type == "mps"790 if isinstance(timestep, float):791 dtype = torch.float32 if is_mps else torch.float64792 else:793 dtype = torch.int32 if is_mps else torch.int64794 timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)795 elif len(timesteps.shape) == 0:796 timesteps = timesteps[None].to(sample.device)797 798 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML799 timesteps = timesteps.expand(sample.shape[0])800 801 t_emb = self.time_proj(timesteps)802 803 # `Timesteps` does not contain any weights and will always return f32 tensors804 # but time_embedding might actually be running in fp16. so we need to cast here.805 # there might be better ways to encapsulate this.806 t_emb = t_emb.to(dtype=sample.dtype)807 808 emb = self.time_embedding(t_emb, timestep_cond)809 810 if self.class_embedding is not None:811 if class_labels is None:812 raise ValueError("class_labels should be provided when num_class_embeds > 0")813 814 if self.config.class_embed_type == "timestep":815 class_labels = self.time_proj(class_labels)816 817 # `Timesteps` does not contain any weights and will always return f32 tensors818 # there might be better ways to encapsulate this.819 class_labels = class_labels.to(dtype=sample.dtype)820 821 class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)822 823 if self.config.class_embeddings_concat:824 emb = torch.cat([emb, class_emb], dim=-1)825 else:826 emb = emb + class_emb827 828 if self.config.addition_embed_type == "text":829 aug_emb = self.add_embedding(encoder_hidden_states)830 emb = emb + aug_emb831 elif self.config.addition_embed_type == "text_image":832 # Kadinsky 2.1 - style833 if "image_embeds" not in added_cond_kwargs:834 raise ValueError(835 f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"836 )837 838 image_embs = added_cond_kwargs.get("image_embeds")839 text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)840 841 aug_emb = self.add_embedding(text_embs, image_embs)842 emb = emb + aug_emb843 844 if self.time_embed_act is not None:845 emb = self.time_embed_act(emb)846 847 if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":848 encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)849 elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":850 # Kadinsky 2.1 - style851 if "image_embeds" not in added_cond_kwargs:852 raise ValueError(853 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`"854 )855 856 image_embeds = added_cond_kwargs.get("image_embeds")857 encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)858 859 # 2. pre-process860 sample = self.conv_in(sample)861 862 # 2.5 GLIGEN position net863 if cross_attention_kwargs is not None and cross_attention_kwargs.get('gligen', None) is not None:864 cross_attention_kwargs = cross_attention_kwargs.copy()865 cross_attention_kwargs['gligen'] = {866 'objs': self.position_net(867 boxes=cross_attention_kwargs['gligen']['boxes'],868 masks=cross_attention_kwargs['gligen']['masks'],869 positive_embeddings=cross_attention_kwargs['gligen']['positive_embeddings']870 ),871 'fuser_attn_kwargs': cross_attention_kwargs['gligen'].get('fuser_attn_kwargs', {})872 }873 874 # 3. down875 down_block_res_samples = (sample,)876 cross_attention_probs_down = []877 if cross_attention_kwargs is None:878 cross_attention_kwargs = {}879 880 for i, downsample_block in enumerate(self.down_blocks):881 cross_attention_kwargs["attn_key"] = ["down", i]882 883 if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:884 downsample_block_output = downsample_block(885 hidden_states=sample,886 temb=emb,887 encoder_hidden_states=encoder_hidden_states,888 attention_mask=attention_mask,889 cross_attention_kwargs=cross_attention_kwargs,890 encoder_attention_mask=encoder_attention_mask,891 return_cross_attention_probs=return_cross_attention_probs,892 )893 if return_cross_attention_probs:894 sample, res_samples, cross_attention_probs = downsample_block_output895 cross_attention_probs_down.append(cross_attention_probs)896 else:897 sample, res_samples = downsample_block_output898 else:899 sample, res_samples = downsample_block(hidden_states=sample, temb=emb)900 901 down_block_res_samples += res_samples902 903 if down_block_additional_residuals is not None:904 new_down_block_res_samples = ()905 906 for down_block_res_sample, down_block_additional_residual in zip(907 down_block_res_samples, down_block_additional_residuals908 ):909 down_block_res_sample = down_block_res_sample + down_block_additional_residual910 new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)911 912 down_block_res_samples = new_down_block_res_samples913 914 # 4. mid915 cross_attention_probs_mid = []916 if self.mid_block is not None:917 cross_attention_kwargs["attn_key"] = ["mid", 0]918 919 sample = self.mid_block(920 sample,921 emb,922 encoder_hidden_states=encoder_hidden_states,923 attention_mask=attention_mask,924 cross_attention_kwargs=cross_attention_kwargs,925 encoder_attention_mask=encoder_attention_mask,926 return_cross_attention_probs=return_cross_attention_probs,927 )928 if return_cross_attention_probs:929 sample, cross_attention_probs = sample930 cross_attention_probs_mid.append(cross_attention_probs)931 932 933 if mid_block_additional_residual is not None:934 sample = sample + mid_block_additional_residual935 936 cross_attention_probs_up = []937 # 5. up938 for i, upsample_block in enumerate(self.up_blocks):939 cross_attention_kwargs["attn_key"] = ["up", i]940 941 is_final_block = i == len(self.up_blocks) - 1942 943 res_samples = down_block_res_samples[-len(upsample_block.resnets) :]944 down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]945 946 # if we have not reached the final block and need to forward the947 # upsample size, we do it here948 if not is_final_block and forward_upsample_size:949 upsample_size = down_block_res_samples[-1].shape[2:]950 951 if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:952 sample = upsample_block(953 hidden_states=sample,954 temb=emb,955 res_hidden_states_tuple=res_samples,956 encoder_hidden_states=encoder_hidden_states,957 cross_attention_kwargs=cross_attention_kwargs,958 upsample_size=upsample_size,959 attention_mask=attention_mask,960 encoder_attention_mask=encoder_attention_mask,961 return_cross_attention_probs=return_cross_attention_probs,962 )963 if return_cross_attention_probs:964 sample, cross_attention_probs = sample965 cross_attention_probs_up.append(cross_attention_probs)966 else:967 sample = upsample_block(968 hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size969 )970 971 # 6. post-process972 if self.conv_norm_out:973 sample = self.conv_norm_out(sample)974 sample = self.conv_act(sample)975 sample = self.conv_out(sample)976 977 if not return_dict:978 return (sample,)979 980 return UNet2DConditionOutput(sample=sample, cross_attention_probs_down=cross_attention_probs_down, cross_attention_probs_mid=cross_attention_probs_mid, cross_attention_probs_up=cross_attention_probs_up)981