multimodalart/EchoMimic-zero
8
1# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_condition.py2from dataclasses import dataclass3from typing import Any, Dict, List, Optional, Tuple, Union4 5import torch6import torch.nn as nn7import torch.utils.checkpoint8from diffusers.configuration_utils import ConfigMixin, register_to_config9from diffusers.loaders import UNet2DConditionLoadersMixin10from diffusers.models.activations import get_activation11from diffusers.models.attention_processor import (12 ADDED_KV_ATTENTION_PROCESSORS,13 CROSS_ATTENTION_PROCESSORS,14 AttentionProcessor,15 AttnAddedKVProcessor,16 AttnProcessor,17)18from diffusers.models.embeddings import (19 GaussianFourierProjection,20 ImageHintTimeEmbedding,21 ImageProjection,22 ImageTimeEmbedding,23 PositionNet,24 TextImageProjection,25 TextImageTimeEmbedding,26 TextTimeEmbedding,27 TimestepEmbedding,28 Timesteps,29)30from diffusers.models.modeling_utils import ModelMixin31from diffusers.utils import (32 USE_PEFT_BACKEND,33 BaseOutput,34 deprecate,35 logging,36 scale_lora_layers,37 unscale_lora_layers,38)39 40from .unet_2d_blocks import (41 UNetMidBlock2D,42 UNetMidBlock2DCrossAttn,43 get_down_block,44 get_up_block,45)46 47logger = logging.get_logger(__name__) # pylint: disable=invalid-name48 49 50@dataclass51class UNet2DConditionOutput(BaseOutput):52 """53 The output of [`UNet2DConditionModel`].54 55 Args:56 sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):57 The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.58 """59 60 sample: torch.FloatTensor = None61 ref_features: Tuple[torch.FloatTensor] = None62 63 64class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):65 r"""66 A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample67 shaped output.68 69 This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented70 for all models (such as downloading or saving).71 72 Parameters:73 sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):74 Height and width of input/output sample.75 in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.76 out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.77 center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.78 flip_sin_to_cos (`bool`, *optional*, defaults to `False`):79 Whether to flip the sin to cos in the time embedding.80 freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.81 down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):82 The tuple of downsample blocks to use.83 mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):84 Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or85 `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.86 up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):87 The tuple of upsample blocks to use.88 only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):89 Whether to include self-attention in the basic transformer blocks, see90 [`~models.attention.BasicTransformerBlock`].91 block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):92 The tuple of output channels for each block.93 layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.94 downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.95 mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.96 dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.97 act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.98 norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.99 If `None`, normalization and activation layers is skipped in post-processing.100 norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.101 cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):102 The dimension of the cross attention features.103 transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):104 The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for105 [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],106 [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].107 reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None):108 The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling109 blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for110 [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],111 [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].112 encoder_hid_dim (`int`, *optional*, defaults to None):113 If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`114 dimension to `cross_attention_dim`.115 encoder_hid_dim_type (`str`, *optional*, defaults to `None`):116 If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text117 embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.118 attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.119 num_attention_heads (`int`, *optional*):120 The number of attention heads. If not defined, defaults to `attention_head_dim`121 resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config122 for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.123 class_embed_type (`str`, *optional*, defaults to `None`):124 The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,125 `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.126 addition_embed_type (`str`, *optional*, defaults to `None`):127 Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or128 "text". "text" will use the `TextTimeEmbedding` layer.129 addition_time_embed_dim: (`int`, *optional*, defaults to `None`):130 Dimension for the timestep embeddings.131 num_class_embeds (`int`, *optional*, defaults to `None`):132 Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing133 class conditioning with `class_embed_type` equal to `None`.134 time_embedding_type (`str`, *optional*, defaults to `positional`):135 The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.136 time_embedding_dim (`int`, *optional*, defaults to `None`):137 An optional override for the dimension of the projected time embedding.138 time_embedding_act_fn (`str`, *optional*, defaults to `None`):139 Optional activation function to use only once on the time embeddings before they are passed to the rest of140 the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`.141 timestep_post_act (`str`, *optional*, defaults to `None`):142 The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.143 time_cond_proj_dim (`int`, *optional*, defaults to `None`):144 The dimension of `cond_proj` layer in the timestep embedding.145 conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. conv_out_kernel (`int`,146 *optional*, default to `3`): The kernel size of `conv_out` layer. projection_class_embeddings_input_dim (`int`,147 *optional*): The dimension of the `class_labels` input when148 `class_embed_type="projection"`. Required when `class_embed_type="projection"`.149 class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time150 embeddings with the class embeddings.151 mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):152 Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If153 `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the154 `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False`155 otherwise.156 """157 158 _supports_gradient_checkpointing = True159 160 @register_to_config161 def __init__(162 self,163 sample_size: Optional[int] = None,164 in_channels: int = 4,165 out_channels: int = 4,166 center_input_sample: bool = False,167 flip_sin_to_cos: bool = True,168 freq_shift: int = 0,169 down_block_types: Tuple[str] = (170 "CrossAttnDownBlock2D",171 "CrossAttnDownBlock2D",172 "CrossAttnDownBlock2D",173 "DownBlock2D",174 ),175 mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",176 up_block_types: Tuple[str] = (177 "UpBlock2D",178 "CrossAttnUpBlock2D",179 "CrossAttnUpBlock2D",180 "CrossAttnUpBlock2D",181 ),182 only_cross_attention: Union[bool, Tuple[bool]] = False,183 block_out_channels: Tuple[int] = (320, 640, 1280, 1280),184 layers_per_block: Union[int, Tuple[int]] = 2,185 downsample_padding: int = 1,186 mid_block_scale_factor: float = 1,187 dropout: float = 0.0,188 act_fn: str = "silu",189 norm_num_groups: Optional[int] = 32,190 norm_eps: float = 1e-5,191 cross_attention_dim: Union[int, Tuple[int]] = 1280,192 transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,193 reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,194 encoder_hid_dim: Optional[int] = None,195 encoder_hid_dim_type: Optional[str] = None,196 attention_head_dim: Union[int, Tuple[int]] = 8,197 num_attention_heads: Optional[Union[int, Tuple[int]]] = None,198 dual_cross_attention: bool = False,199 use_linear_projection: bool = False,200 class_embed_type: Optional[str] = None,201 addition_embed_type: Optional[str] = None,202 addition_time_embed_dim: Optional[int] = None,203 num_class_embeds: Optional[int] = None,204 upcast_attention: bool = False,205 resnet_time_scale_shift: str = "default",206 resnet_skip_time_act: bool = False,207 resnet_out_scale_factor: int = 1.0,208 time_embedding_type: str = "positional",209 time_embedding_dim: Optional[int] = None,210 time_embedding_act_fn: Optional[str] = None,211 timestep_post_act: Optional[str] = None,212 time_cond_proj_dim: Optional[int] = None,213 conv_in_kernel: int = 3,214 conv_out_kernel: int = 3,215 projection_class_embeddings_input_dim: Optional[int] = None,216 attention_type: str = "default",217 class_embeddings_concat: bool = False,218 mid_block_only_cross_attention: Optional[bool] = None,219 cross_attention_norm: Optional[str] = None,220 addition_embed_type_num_heads=64,221 ):222 super().__init__()223 224 self.sample_size = sample_size225 226 if num_attention_heads is not None:227 raise ValueError(228 "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."229 )230 231 # If `num_attention_heads` is not defined (which is the case for most models)232 # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.233 # The reason for this behavior is to correct for incorrectly named variables that were introduced234 # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131235 # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking236 # which is why we correct for the naming here.237 num_attention_heads = num_attention_heads or attention_head_dim238 239 # Check inputs240 if len(down_block_types) != len(up_block_types):241 raise ValueError(242 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}."243 )244 245 if len(block_out_channels) != len(down_block_types):246 raise ValueError(247 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}."248 )249 250 if not isinstance(only_cross_attention, bool) and len(251 only_cross_attention252 ) != len(down_block_types):253 raise ValueError(254 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}."255 )256 257 if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(258 down_block_types259 ):260 raise ValueError(261 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}."262 )263 264 if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(265 down_block_types266 ):267 raise ValueError(268 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}."269 )270 271 if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(272 down_block_types273 ):274 raise ValueError(275 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}."276 )277 278 if not isinstance(layers_per_block, int) and len(layers_per_block) != len(279 down_block_types280 ):281 raise ValueError(282 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}."283 )284 if (285 isinstance(transformer_layers_per_block, list)286 and reverse_transformer_layers_per_block is None287 ):288 for layer_number_per_block in transformer_layers_per_block:289 if isinstance(layer_number_per_block, list):290 raise ValueError(291 "Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet."292 )293 294 # input295 conv_in_padding = (conv_in_kernel - 1) // 2296 self.conv_in = nn.Conv2d(297 in_channels,298 block_out_channels[0],299 kernel_size=conv_in_kernel,300 padding=conv_in_padding,301 )302 303 # time304 if time_embedding_type == "fourier":305 time_embed_dim = time_embedding_dim or block_out_channels[0] * 2306 if time_embed_dim % 2 != 0:307 raise ValueError(308 f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}."309 )310 self.time_proj = GaussianFourierProjection(311 time_embed_dim // 2,312 set_W_to_weight=False,313 log=False,314 flip_sin_to_cos=flip_sin_to_cos,315 )316 timestep_input_dim = time_embed_dim317 elif time_embedding_type == "positional":318 time_embed_dim = time_embedding_dim or block_out_channels[0] * 4319 320 self.time_proj = Timesteps(321 block_out_channels[0], flip_sin_to_cos, freq_shift322 )323 timestep_input_dim = block_out_channels[0]324 else:325 raise ValueError(326 f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."327 )328 329 self.time_embedding = TimestepEmbedding(330 timestep_input_dim,331 time_embed_dim,332 act_fn=act_fn,333 post_act_fn=timestep_post_act,334 cond_proj_dim=time_cond_proj_dim,335 )336 337 if encoder_hid_dim_type is None and encoder_hid_dim is not None:338 encoder_hid_dim_type = "text_proj"339 self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)340 logger.info(341 "encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined."342 )343 344 if encoder_hid_dim is None and encoder_hid_dim_type is not None:345 raise ValueError(346 f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."347 )348 349 if encoder_hid_dim_type == "text_proj":350 self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)351 elif encoder_hid_dim_type == "text_image_proj":352 # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much353 # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use354 # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`355 self.encoder_hid_proj = TextImageProjection(356 text_embed_dim=encoder_hid_dim,357 image_embed_dim=cross_attention_dim,358 cross_attention_dim=cross_attention_dim,359 )360 elif encoder_hid_dim_type == "image_proj":361 # Kandinsky 2.2362 self.encoder_hid_proj = ImageProjection(363 image_embed_dim=encoder_hid_dim,364 cross_attention_dim=cross_attention_dim,365 )366 elif encoder_hid_dim_type is not None:367 raise ValueError(368 f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."369 )370 else:371 self.encoder_hid_proj = None372 373 # class embedding374 if class_embed_type is None and num_class_embeds is not None:375 self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)376 elif class_embed_type == "timestep":377 self.class_embedding = TimestepEmbedding(378 timestep_input_dim, time_embed_dim, act_fn=act_fn379 )380 elif class_embed_type == "identity":381 self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)382 elif class_embed_type == "projection":383 if projection_class_embeddings_input_dim is None:384 raise ValueError(385 "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"386 )387 # The projection `class_embed_type` is the same as the timestep `class_embed_type` except388 # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings389 # 2. it projects from an arbitrary input dimension.390 #391 # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.392 # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.393 # As a result, `TimestepEmbedding` can be passed arbitrary vectors.394 self.class_embedding = TimestepEmbedding(395 projection_class_embeddings_input_dim, time_embed_dim396 )397 elif class_embed_type == "simple_projection":398 if projection_class_embeddings_input_dim is None:399 raise ValueError(400 "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"401 )402 self.class_embedding = nn.Linear(403 projection_class_embeddings_input_dim, time_embed_dim404 )405 else:406 self.class_embedding = None407 408 if addition_embed_type == "text":409 if encoder_hid_dim is not None:410 text_time_embedding_from_dim = encoder_hid_dim411 else:412 text_time_embedding_from_dim = cross_attention_dim413 414 self.add_embedding = TextTimeEmbedding(415 text_time_embedding_from_dim,416 time_embed_dim,417 num_heads=addition_embed_type_num_heads,418 )419 elif addition_embed_type == "text_image":420 # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much421 # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use422 # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`423 self.add_embedding = TextImageTimeEmbedding(424 text_embed_dim=cross_attention_dim,425 image_embed_dim=cross_attention_dim,426 time_embed_dim=time_embed_dim,427 )428 elif addition_embed_type == "text_time":429 self.add_time_proj = Timesteps(430 addition_time_embed_dim, flip_sin_to_cos, freq_shift431 )432 self.add_embedding = TimestepEmbedding(433 projection_class_embeddings_input_dim, time_embed_dim434 )435 elif addition_embed_type == "image":436 # Kandinsky 2.2437 self.add_embedding = ImageTimeEmbedding(438 image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim439 )440 elif addition_embed_type == "image_hint":441 # Kandinsky 2.2 ControlNet442 self.add_embedding = ImageHintTimeEmbedding(443 image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim444 )445 elif addition_embed_type is not None:446 raise ValueError(447 f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'."448 )449 450 if time_embedding_act_fn is None:451 self.time_embed_act = None452 else:453 self.time_embed_act = get_activation(time_embedding_act_fn)454 455 self.down_blocks = nn.ModuleList([])456 self.up_blocks = nn.ModuleList([])457 458 if isinstance(only_cross_attention, bool):459 if mid_block_only_cross_attention is None:460 mid_block_only_cross_attention = only_cross_attention461 462 only_cross_attention = [only_cross_attention] * len(down_block_types)463 464 if mid_block_only_cross_attention is None:465 mid_block_only_cross_attention = False466 467 if isinstance(num_attention_heads, int):468 num_attention_heads = (num_attention_heads,) * len(down_block_types)469 470 if isinstance(attention_head_dim, int):471 attention_head_dim = (attention_head_dim,) * len(down_block_types)472 473 if isinstance(cross_attention_dim, int):474 cross_attention_dim = (cross_attention_dim,) * len(down_block_types)475 476 if isinstance(layers_per_block, int):477 layers_per_block = [layers_per_block] * len(down_block_types)478 479 if isinstance(transformer_layers_per_block, int):480 transformer_layers_per_block = [transformer_layers_per_block] * len(481 down_block_types482 )483 484 if class_embeddings_concat:485 # The time embeddings are concatenated with the class embeddings. The dimension of the486 # time embeddings passed to the down, middle, and up blocks is twice the dimension of the487 # regular time embeddings488 blocks_time_embed_dim = time_embed_dim * 2489 else:490 blocks_time_embed_dim = time_embed_dim491 492 # down493 output_channel = block_out_channels[0]494 for i, down_block_type in enumerate(down_block_types):495 input_channel = output_channel496 output_channel = block_out_channels[i]497 is_final_block = i == len(block_out_channels) - 1498 499 down_block = get_down_block(500 down_block_type,501 num_layers=layers_per_block[i],502 transformer_layers_per_block=transformer_layers_per_block[i],503 in_channels=input_channel,504 out_channels=output_channel,505 temb_channels=blocks_time_embed_dim,506 add_downsample=not is_final_block,507 resnet_eps=norm_eps,508 resnet_act_fn=act_fn,509 resnet_groups=norm_num_groups,510 cross_attention_dim=cross_attention_dim[i],511 num_attention_heads=num_attention_heads[i],512 downsample_padding=downsample_padding,513 dual_cross_attention=dual_cross_attention,514 use_linear_projection=use_linear_projection,515 only_cross_attention=only_cross_attention[i],516 upcast_attention=upcast_attention,517 resnet_time_scale_shift=resnet_time_scale_shift,518 attention_type=attention_type,519 resnet_skip_time_act=resnet_skip_time_act,520 resnet_out_scale_factor=resnet_out_scale_factor,521 cross_attention_norm=cross_attention_norm,522 attention_head_dim=attention_head_dim[i]523 if attention_head_dim[i] is not None524 else output_channel,525 dropout=dropout,526 )527 self.down_blocks.append(down_block)528 529 # mid530 if mid_block_type == "UNetMidBlock2DCrossAttn":531 self.mid_block = UNetMidBlock2DCrossAttn(532 transformer_layers_per_block=transformer_layers_per_block[-1],533 in_channels=block_out_channels[-1],534 temb_channels=blocks_time_embed_dim,535 dropout=dropout,536 resnet_eps=norm_eps,537 resnet_act_fn=act_fn,538 output_scale_factor=mid_block_scale_factor,539 resnet_time_scale_shift=resnet_time_scale_shift,540 cross_attention_dim=cross_attention_dim[-1],541 num_attention_heads=num_attention_heads[-1],542 resnet_groups=norm_num_groups,543 dual_cross_attention=dual_cross_attention,544 use_linear_projection=use_linear_projection,545 upcast_attention=upcast_attention,546 attention_type=attention_type,547 )548 elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn":549 raise NotImplementedError(f"Unsupport mid_block_type: {mid_block_type}")550 elif mid_block_type == "UNetMidBlock2D":551 self.mid_block = UNetMidBlock2D(552 in_channels=block_out_channels[-1],553 temb_channels=blocks_time_embed_dim,554 dropout=dropout,555 num_layers=0,556 resnet_eps=norm_eps,557 resnet_act_fn=act_fn,558 output_scale_factor=mid_block_scale_factor,559 resnet_groups=norm_num_groups,560 resnet_time_scale_shift=resnet_time_scale_shift,561 add_attention=False,562 )563 elif mid_block_type is None:564 self.mid_block = None565 else:566 raise ValueError(f"unknown mid_block_type : {mid_block_type}")567 568 # count how many layers upsample the images569 self.num_upsamplers = 0570 571 # up572 reversed_block_out_channels = list(reversed(block_out_channels))573 reversed_num_attention_heads = list(reversed(num_attention_heads))574 reversed_layers_per_block = list(reversed(layers_per_block))575 reversed_cross_attention_dim = list(reversed(cross_attention_dim))576 reversed_transformer_layers_per_block = (577 list(reversed(transformer_layers_per_block))578 if reverse_transformer_layers_per_block is None579 else reverse_transformer_layers_per_block580 )581 only_cross_attention = list(reversed(only_cross_attention))582 583 output_channel = reversed_block_out_channels[0]584 for i, up_block_type in enumerate(up_block_types):585 is_final_block = i == len(block_out_channels) - 1586 587 prev_output_channel = output_channel588 output_channel = reversed_block_out_channels[i]589 input_channel = reversed_block_out_channels[590 min(i + 1, len(block_out_channels) - 1)591 ]592 593 # add upsample block for all BUT final layer594 if not is_final_block:595 add_upsample = True596 self.num_upsamplers += 1597 else:598 add_upsample = False599 600 up_block = get_up_block(601 up_block_type,602 num_layers=reversed_layers_per_block[i] + 1,603 transformer_layers_per_block=reversed_transformer_layers_per_block[i],604 in_channels=input_channel,605 out_channels=output_channel,606 prev_output_channel=prev_output_channel,607 temb_channels=blocks_time_embed_dim,608 add_upsample=add_upsample,609 resnet_eps=norm_eps,610 resnet_act_fn=act_fn,611 resolution_idx=i,612 resnet_groups=norm_num_groups,613 cross_attention_dim=reversed_cross_attention_dim[i],614 num_attention_heads=reversed_num_attention_heads[i],615 dual_cross_attention=dual_cross_attention,616 use_linear_projection=use_linear_projection,617 only_cross_attention=only_cross_attention[i],618 upcast_attention=upcast_attention,619 resnet_time_scale_shift=resnet_time_scale_shift,620 attention_type=attention_type,621 resnet_skip_time_act=resnet_skip_time_act,622 resnet_out_scale_factor=resnet_out_scale_factor,623 cross_attention_norm=cross_attention_norm,624 attention_head_dim=attention_head_dim[i]625 if attention_head_dim[i] is not None626 else output_channel,627 dropout=dropout,628 )629 self.up_blocks.append(up_block)630 prev_output_channel = output_channel631 632 # out633 if norm_num_groups is not None:634 self.conv_norm_out = nn.GroupNorm(635 num_channels=block_out_channels[0],636 num_groups=norm_num_groups,637 eps=norm_eps,638 )639 640 self.conv_act = get_activation(act_fn)641 642 else:643 self.conv_norm_out = None644 self.conv_act = None645 self.conv_norm_out = None646 647 conv_out_padding = (conv_out_kernel - 1) // 2648 # self.conv_out = nn.Conv2d(649 # block_out_channels[0],650 # out_channels,651 # kernel_size=conv_out_kernel,652 # padding=conv_out_padding,653 # )654 655 if attention_type in ["gated", "gated-text-image"]:656 positive_len = 768657 if isinstance(cross_attention_dim, int):658 positive_len = cross_attention_dim659 elif isinstance(cross_attention_dim, tuple) or isinstance(660 cross_attention_dim, list661 ):662 positive_len = cross_attention_dim[0]663 664 feature_type = "text-only" if attention_type == "gated" else "text-image"665 self.position_net = PositionNet(666 positive_len=positive_len,667 out_dim=cross_attention_dim,668 feature_type=feature_type,669 )670 671 @property672 def attn_processors(self) -> Dict[str, AttentionProcessor]:673 r"""674 Returns:675 `dict` of attention processors: A dictionary containing all attention processors used in the model with676 indexed by its weight name.677 """678 # set recursively679 processors = {}680 681 def fn_recursive_add_processors(682 name: str,683 module: torch.nn.Module,684 processors: Dict[str, AttentionProcessor],685 ):686 if hasattr(module, "get_processor"):687 processors[f"{name}.processor"] = module.get_processor(688 return_deprecated_lora=True689 )690 691 for sub_name, child in module.named_children():692 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)693 694 return processors695 696 for name, module in self.named_children():697 fn_recursive_add_processors(name, module, processors)698 699 return processors700 701 def set_attn_processor(702 self,703 processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]],704 _remove_lora=False,705 ):706 r"""707 Sets the attention processor to use to compute attention.708 709 Parameters:710 processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):711 The instantiated processor class or a dictionary of processor classes that will be set as the processor712 for **all** `Attention` layers.713 714 If `processor` is a dict, the key needs to define the path to the corresponding cross attention715 processor. This is strongly recommended when setting trainable attention processors.716 717 """718 count = len(self.attn_processors.keys())719 720 if isinstance(processor, dict) and len(processor) != count:721 raise ValueError(722 f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"723 f" number of attention layers: {count}. Please make sure to pass {count} processor classes."724 )725 726 def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):727 if hasattr(module, "set_processor"):728 if not isinstance(processor, dict):729 module.set_processor(processor, _remove_lora=_remove_lora)730 else:731 module.set_processor(732 processor.pop(f"{name}.processor"), _remove_lora=_remove_lora733 )734 735 for sub_name, child in module.named_children():736 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)737 738 for name, module in self.named_children():739 fn_recursive_attn_processor(name, module, processor)740 741 def set_default_attn_processor(self):742 """743 Disables custom attention processors and sets the default attention implementation.744 """745 if all(746 proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS747 for proc in self.attn_processors.values()748 ):749 processor = AttnAddedKVProcessor()750 elif all(751 proc.__class__ in CROSS_ATTENTION_PROCESSORS752 for proc in self.attn_processors.values()753 ):754 processor = AttnProcessor()755 else:756 raise ValueError(757 f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"758 )759 760 self.set_attn_processor(processor, _remove_lora=True)761 762 def set_attention_slice(self, slice_size):763 r"""764 Enable sliced attention computation.765 766 When this option is enabled, the attention module splits the input tensor in slices to compute attention in767 several steps. This is useful for saving some memory in exchange for a small decrease in speed.768 769 Args:770 slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):771 When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If772 `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is773 provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`774 must be a multiple of `slice_size`.775 """776 sliceable_head_dims = []777 778 def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):779 if hasattr(module, "set_attention_slice"):780 sliceable_head_dims.append(module.sliceable_head_dim)781 782 for child in module.children():783 fn_recursive_retrieve_sliceable_dims(child)784 785 # retrieve number of attention layers786 for module in self.children():787 fn_recursive_retrieve_sliceable_dims(module)788 789 num_sliceable_layers = len(sliceable_head_dims)790 791 if slice_size == "auto":792 # half the attention head size is usually a good trade-off between793 # speed and memory794 slice_size = [dim // 2 for dim in sliceable_head_dims]795 elif slice_size == "max":796 # make smallest slice possible797 slice_size = num_sliceable_layers * [1]798 799 slice_size = (800 num_sliceable_layers * [slice_size]801 if not isinstance(slice_size, list)802 else slice_size803 )804 805 if len(slice_size) != len(sliceable_head_dims):806 raise ValueError(807 f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"808 f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."809 )810 811 for i in range(len(slice_size)):812 size = slice_size[i]813 dim = sliceable_head_dims[i]814 if size is not None and size > dim:815 raise ValueError(f"size {size} has to be smaller or equal to {dim}.")816 817 # Recursively walk through all the children.818 # Any children which exposes the set_attention_slice method819 # gets the message820 def fn_recursive_set_attention_slice(821 module: torch.nn.Module, slice_size: List[int]822 ):823 if hasattr(module, "set_attention_slice"):824 module.set_attention_slice(slice_size.pop())825 826 for child in module.children():827 fn_recursive_set_attention_slice(child, slice_size)828 829 reversed_slice_size = list(reversed(slice_size))830 for module in self.children():831 fn_recursive_set_attention_slice(module, reversed_slice_size)832 833 def _set_gradient_checkpointing(self, module, value=False):834 if hasattr(module, "gradient_checkpointing"):835 module.gradient_checkpointing = value836 837 def enable_freeu(self, s1, s2, b1, b2):838 r"""Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497.839 840 The suffixes after the scaling factors represent the stage blocks where they are being applied.841 842 Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that843 are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.844 845 Args:846 s1 (`float`):847 Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to848 mitigate the "oversmoothing effect" in the enhanced denoising process.849 s2 (`float`):850 Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to851 mitigate the "oversmoothing effect" in the enhanced denoising process.852 b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.853 b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.854 """855 for i, upsample_block in enumerate(self.up_blocks):856 setattr(upsample_block, "s1", s1)857 setattr(upsample_block, "s2", s2)858 setattr(upsample_block, "b1", b1)859 setattr(upsample_block, "b2", b2)860 861 def disable_freeu(self):862 """Disables the FreeU mechanism."""863 freeu_keys = {"s1", "s2", "b1", "b2"}864 for i, upsample_block in enumerate(self.up_blocks):865 for k in freeu_keys:866 if (867 hasattr(upsample_block, k)868 or getattr(upsample_block, k, None) is not None869 ):870 setattr(upsample_block, k, None)871 872 def forward(873 self,874 sample: torch.FloatTensor,875 timestep: Union[torch.Tensor, float, int],876 encoder_hidden_states: torch.Tensor,877 class_labels: Optional[torch.Tensor] = None,878 timestep_cond: Optional[torch.Tensor] = None,879 attention_mask: Optional[torch.Tensor] = None,880 cross_attention_kwargs: Optional[Dict[str, Any]] = None,881 added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,882 down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,883 mid_block_additional_residual: Optional[torch.Tensor] = None,884 down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,885 encoder_attention_mask: Optional[torch.Tensor] = None,886 return_dict: bool = True,887 ) -> Union[UNet2DConditionOutput, Tuple]:888 r"""889 The [`UNet2DConditionModel`] forward method.890 891 Args:892 sample (`torch.FloatTensor`):893 The noisy input tensor with the following shape `(batch, channel, height, width)`.894 timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.895 encoder_hidden_states (`torch.FloatTensor`):896 The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.897 class_labels (`torch.Tensor`, *optional*, defaults to `None`):898 Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.899 timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):900 Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed901 through the `self.time_embedding` layer to obtain the timestep embeddings.902 attention_mask (`torch.Tensor`, *optional*, defaults to `None`):903 An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask904 is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large905 negative values to the attention scores corresponding to "discard" tokens.906 cross_attention_kwargs (`dict`, *optional*):907 A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under908 `self.processor` in909 [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).910 added_cond_kwargs: (`dict`, *optional*):911 A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that912 are passed along to the UNet blocks.913 down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):914 A tuple of tensors that if specified are added to the residuals of down unet blocks.915 mid_block_additional_residual: (`torch.Tensor`, *optional*):916 A tensor that if specified is added to the residual of the middle unet block.917 encoder_attention_mask (`torch.Tensor`):918 A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If919 `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,920 which adds large negative values to the attention scores corresponding to "discard" tokens.921 return_dict (`bool`, *optional*, defaults to `True`):922 Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain923 tuple.924 cross_attention_kwargs (`dict`, *optional*):925 A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].926 added_cond_kwargs: (`dict`, *optional*):927 A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that928 are passed along to the UNet blocks.929 down_block_additional_residuals (`tuple` of `torch.Tensor`, *optional*):930 additional residuals to be added to UNet long skip connections from down blocks to up blocks for931 example from ControlNet side model(s)932 mid_block_additional_residual (`torch.Tensor`, *optional*):933 additional residual to be added to UNet mid block output, for example from ControlNet side model934 down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*):935 additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)936 937 Returns:938 [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:939 If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise940 a `tuple` is returned where the first element is the sample tensor.941 """942 # By default samples have to be AT least a multiple of the overall upsampling factor.943 # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).944 # However, the upsampling interpolation output size can be forced to fit any upsampling size945 # on the fly if necessary.946 default_overall_up_factor = 2**self.num_upsamplers947 948 # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`949 forward_upsample_size = False950 upsample_size = None951 952 for dim in sample.shape[-2:]:953 if dim % default_overall_up_factor != 0:954 # Forward upsample size to force interpolation output size.955 forward_upsample_size = True956 break957 958 # ensure attention_mask is a bias, and give it a singleton query_tokens dimension959 # expects mask of shape:960 # [batch, key_tokens]961 # adds singleton query_tokens dimension:962 # [batch, 1, key_tokens]963 # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:964 # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)965 # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)966 if attention_mask is not None:967 # assume that mask is expressed as:968 # (1 = keep, 0 = discard)969 # convert mask into a bias that can be added to attention scores:970 # (keep = +0, discard = -10000.0)971 attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0972 attention_mask = attention_mask.unsqueeze(1)973 974 # convert encoder_attention_mask to a bias the same way we do for attention_mask975 if encoder_attention_mask is not None:976 encoder_attention_mask = (977 1 - encoder_attention_mask.to(sample.dtype)978 ) * -10000.0979 encoder_attention_mask = encoder_attention_mask.unsqueeze(1)980 981 # 0. center input if necessary982 if self.config.center_input_sample:983 sample = 2 * sample - 1.0984 985 # 1. time986 timesteps = timestep987 if not torch.is_tensor(timesteps):988 # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can989 # This would be a good case for the `match` statement (Python 3.10+)990 is_mps = sample.device.type == "mps"991 if isinstance(timestep, float):992 dtype = torch.float32 if is_mps else torch.float64993 else:994 dtype = torch.int32 if is_mps else torch.int64995 timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)996 elif len(timesteps.shape) == 0:997 timesteps = timesteps[None].to(sample.device)998 999 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML1000 timesteps = timesteps.expand(sample.shape[0])1001 1002 t_emb = self.time_proj(timesteps)1003 1004 # `Timesteps` does not contain any weights and will always return f32 tensors1005 # but time_embedding might actually be running in fp16. so we need to cast here.1006 # there might be better ways to encapsulate this.1007 t_emb = t_emb.to(dtype=sample.dtype)1008 1009 emb = self.time_embedding(t_emb, timestep_cond)1010 aug_emb = None1011 1012 if self.class_embedding is not None:1013 if class_labels is None:1014 raise ValueError(1015 "class_labels should be provided when num_class_embeds > 0"1016 )1017 1018 if self.config.class_embed_type == "timestep":1019 class_labels = self.time_proj(class_labels)1020 1021 # `Timesteps` does not contain any weights and will always return f32 tensors1022 # there might be better ways to encapsulate this.1023 class_labels = class_labels.to(dtype=sample.dtype)1024 1025 class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)1026 1027 if self.config.class_embeddings_concat:1028 emb = torch.cat([emb, class_emb], dim=-1)1029 else:1030 emb = emb + class_emb1031 1032 if self.config.addition_embed_type == "text":1033 aug_emb = self.add_embedding(encoder_hidden_states)1034 elif self.config.addition_embed_type == "text_image":1035 # Kandinsky 2.1 - style1036 if "image_embeds" not in added_cond_kwargs:1037 raise ValueError(1038 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`"1039 )1040 1041 image_embs = added_cond_kwargs.get("image_embeds")1042 text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)1043 aug_emb = self.add_embedding(text_embs, image_embs)1044 elif self.config.addition_embed_type == "text_time":1045 # SDXL - style1046 if "text_embeds" not in added_cond_kwargs:1047 raise ValueError(1048 f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"1049 )1050 text_embeds = added_cond_kwargs.get("text_embeds")1051 if "time_ids" not in added_cond_kwargs:1052 raise ValueError(1053 f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"1054 )1055 time_ids = added_cond_kwargs.get("time_ids")1056 time_embeds = self.add_time_proj(time_ids.flatten())1057 time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))1058 add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)1059 add_embeds = add_embeds.to(emb.dtype)1060 aug_emb = self.add_embedding(add_embeds)1061 elif self.config.addition_embed_type == "image":1062 # Kandinsky 2.2 - style1063 if "image_embeds" not in added_cond_kwargs:1064 raise ValueError(1065 f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"1066 )1067 image_embs = added_cond_kwargs.get("image_embeds")1068 aug_emb = self.add_embedding(image_embs)1069 elif self.config.addition_embed_type == "image_hint":1070 # Kandinsky 2.2 - style1071 if (1072 "image_embeds" not in added_cond_kwargs1073 or "hint" not in added_cond_kwargs1074 ):1075 raise ValueError(1076 f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"1077 )1078 image_embs = added_cond_kwargs.get("image_embeds")1079 hint = added_cond_kwargs.get("hint")1080 aug_emb, hint = self.add_embedding(image_embs, hint)1081 sample = torch.cat([sample, hint], dim=1)1082 1083 emb = emb + aug_emb if aug_emb is not None else emb1084 1085 if self.time_embed_act is not None:1086 emb = self.time_embed_act(emb)1087 1088 if (1089 self.encoder_hid_proj is not None1090 and self.config.encoder_hid_dim_type == "text_proj"1091 ):1092 encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)1093 elif (1094 self.encoder_hid_proj is not None1095 and self.config.encoder_hid_dim_type == "text_image_proj"1096 ):1097 # Kadinsky 2.1 - style1098 if "image_embeds" not in added_cond_kwargs:1099 raise ValueError(1100 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`"1101 )1102 1103 image_embeds = added_cond_kwargs.get("image_embeds")1104 encoder_hidden_states = self.encoder_hid_proj(1105 encoder_hidden_states, image_embeds1106 )1107 elif (1108 self.encoder_hid_proj is not None1109 and self.config.encoder_hid_dim_type == "image_proj"1110 ):1111 # Kandinsky 2.2 - style1112 if "image_embeds" not in added_cond_kwargs:1113 raise ValueError(1114 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`"1115 )1116 image_embeds = added_cond_kwargs.get("image_embeds")1117 encoder_hidden_states = self.encoder_hid_proj(image_embeds)1118 elif (1119 self.encoder_hid_proj is not None1120 and self.config.encoder_hid_dim_type == "ip_image_proj"1121 ):1122 if "image_embeds" not in added_cond_kwargs:1123 raise ValueError(1124 f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`"1125 )1126 image_embeds = added_cond_kwargs.get("image_embeds")1127 image_embeds = self.encoder_hid_proj(image_embeds).to(1128 encoder_hidden_states.dtype1129 )1130 encoder_hidden_states = torch.cat(1131 [encoder_hidden_states, image_embeds], dim=11132 )1133 1134 # 2. pre-process1135 sample = self.conv_in(sample)1136 1137 # 2.5 GLIGEN position net1138 if (1139 cross_attention_kwargs is not None1140 and cross_attention_kwargs.get("gligen", None) is not None1141 ):1142 cross_attention_kwargs = cross_attention_kwargs.copy()1143 gligen_args = cross_attention_kwargs.pop("gligen")1144 cross_attention_kwargs["gligen"] = {1145 "objs": self.position_net(**gligen_args)1146 }1147 1148 # 3. down1149 lora_scale = (1150 cross_attention_kwargs.get("scale", 1.0)1151 if cross_attention_kwargs is not None1152 else 1.01153 )1154 if USE_PEFT_BACKEND:1155 # weight the lora layers by setting `lora_scale` for each PEFT layer1156 scale_lora_layers(self, lora_scale)1157 1158 is_controlnet = (1159 mid_block_additional_residual is not None1160 and down_block_additional_residuals is not None1161 )1162 # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets1163 is_adapter = down_intrablock_additional_residuals is not None1164 # maintain backward compatibility for legacy usage, where1165 # T2I-Adapter and ControlNet both use down_block_additional_residuals arg1166 # but can only use one or the other1167 if (1168 not is_adapter1169 and mid_block_additional_residual is None1170 and down_block_additional_residuals is not None1171 ):1172 deprecate(1173 "T2I should not use down_block_additional_residuals",1174 "1.3.0",1175 "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \1176 and will be removed in diffusers 1.3.0. `down_block_additional_residuals` should only be used \1177 for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",1178 standard_warn=False,1179 )1180 down_intrablock_additional_residuals = down_block_additional_residuals1181 is_adapter = True1182 1183 down_block_res_samples = (sample,)1184 tot_referece_features = ()1185 for downsample_block in self.down_blocks:1186 if (1187 hasattr(downsample_block, "has_cross_attention")1188 and downsample_block.has_cross_attention1189 ):1190 # For t2i-adapter CrossAttnDownBlock2D1191 additional_residuals = {}1192 if is_adapter and len(down_intrablock_additional_residuals) > 0:1193 additional_residuals[1194 "additional_residuals"1195 ] = down_intrablock_additional_residuals.pop(0)1196 1197 sample, res_samples = downsample_block(1198 hidden_states=sample,1199 temb=emb,1200 encoder_hidden_states=encoder_hidden_states,