declare-lab/tango2
92
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.utils.checkpoint20 21from ..configuration_utils import ConfigMixin, register_to_config22from ..loaders import UNet2DConditionLoadersMixin23from ..utils import BaseOutput, logging24from .attention_processor import AttentionProcessor, AttnProcessor25from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps26from .modeling_utils import ModelMixin27from .unet_2d_blocks import (28 CrossAttnDownBlock2D,29 CrossAttnUpBlock2D,30 DownBlock2D,31 UNetMidBlock2DCrossAttn,32 UNetMidBlock2DSimpleCrossAttn,33 UpBlock2D,34 get_down_block,35 get_up_block,36)37 38 39logger = logging.get_logger(__name__) # pylint: disable=invalid-name40 41 42@dataclass43class UNet2DConditionOutput(BaseOutput):44 """45 Args:46 sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):47 Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.48 """49 50 sample: torch.FloatTensor51 52 53class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):54 r"""55 UNet2DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a timestep56 and returns sample shaped output.57 58 This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library59 implements for all the models (such as downloading or saving, etc.)60 61 Parameters:62 sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):63 Height and width of input/output sample.64 in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.65 out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.66 center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.67 flip_sin_to_cos (`bool`, *optional*, defaults to `False`):68 Whether to flip the sin to cos in the time embedding.69 freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.70 down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):71 The tuple of downsample blocks to use.72 mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):73 The mid block type. Choose from `UNetMidBlock2DCrossAttn` or `UNetMidBlock2DSimpleCrossAttn`, will skip the74 mid block layer if `None`.75 up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D",)`):76 The tuple of upsample blocks to use.77 only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):78 Whether to include self-attention in the basic transformer blocks, see79 [`~models.attention.BasicTransformerBlock`].80 block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):81 The tuple of output channels for each block.82 layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.83 downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.84 mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.85 act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.86 norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.87 If `None`, it will skip the normalization and activation layers in post-processing88 norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.89 cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):90 The dimension of the cross attention features.91 attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.92 resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config93 for resnet blocks, see [`~models.resnet.ResnetBlock2D`]. Choose from `default` or `scale_shift`.94 class_embed_type (`str`, *optional*, defaults to None):95 The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,96 `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.97 num_class_embeds (`int`, *optional*, defaults to None):98 Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing99 class conditioning with `class_embed_type` equal to `None`.100 time_embedding_type (`str`, *optional*, default to `positional`):101 The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.102 timestep_post_act (`str, *optional*, default to `None`):103 The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.104 time_cond_proj_dim (`int`, *optional*, default to `None`):105 The dimension of `cond_proj` layer in timestep embedding.106 conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer.107 conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer.108 projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when109 using the "projection" `class_embed_type`. Required when using the "projection" `class_embed_type`.110 class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time111 embeddings with the class embeddings.112 """113 114 _supports_gradient_checkpointing = True115 116 @register_to_config117 def __init__(118 self,119 sample_size: Optional[int] = None,120 in_channels: int = 4,121 out_channels: int = 4,122 center_input_sample: bool = False,123 flip_sin_to_cos: bool = True,124 freq_shift: int = 0,125 down_block_types: Tuple[str] = (126 "CrossAttnDownBlock2D",127 "CrossAttnDownBlock2D",128 "CrossAttnDownBlock2D",129 "DownBlock2D",130 ),131 mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",132 up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),133 only_cross_attention: Union[bool, Tuple[bool]] = False,134 block_out_channels: Tuple[int] = (320, 640, 1280, 1280),135 layers_per_block: int = 2,136 downsample_padding: int = 1,137 mid_block_scale_factor: float = 1,138 act_fn: str = "silu",139 norm_num_groups: Optional[int] = 32,140 norm_eps: float = 1e-5,141 cross_attention_dim: Union[int, Tuple[int]] = 1280,142 attention_head_dim: Union[int, Tuple[int]] = 8,143 dual_cross_attention: bool = False,144 use_linear_projection: bool = False,145 class_embed_type: Optional[str] = None,146 num_class_embeds: Optional[int] = None,147 upcast_attention: bool = False,148 resnet_time_scale_shift: str = "default",149 time_embedding_type: str = "positional",150 timestep_post_act: Optional[str] = None,151 time_cond_proj_dim: Optional[int] = None,152 conv_in_kernel: int = 3,153 conv_out_kernel: int = 3,154 projection_class_embeddings_input_dim: Optional[int] = None,155 class_embeddings_concat: bool = False,156 ):157 super().__init__()158 159 self.sample_size = sample_size160 161 # Check inputs162 if len(down_block_types) != len(up_block_types):163 raise ValueError(164 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}."165 )166 167 if len(block_out_channels) != len(down_block_types):168 raise ValueError(169 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}."170 )171 172 if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):173 raise ValueError(174 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}."175 )176 177 if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):178 raise ValueError(179 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}."180 )181 182 if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):183 raise ValueError(184 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}."185 )186 187 # input188 conv_in_padding = (conv_in_kernel - 1) // 2189 self.conv_in = nn.Conv2d(190 in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding191 )192 193 # time194 if time_embedding_type == "fourier":195 time_embed_dim = block_out_channels[0] * 2196 if time_embed_dim % 2 != 0:197 raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.")198 self.time_proj = GaussianFourierProjection(199 time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos200 )201 timestep_input_dim = time_embed_dim202 elif time_embedding_type == "positional":203 time_embed_dim = block_out_channels[0] * 4204 205 self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)206 timestep_input_dim = block_out_channels[0]207 else:208 raise ValueError(209 f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."210 )211 212 self.time_embedding = TimestepEmbedding(213 timestep_input_dim,214 time_embed_dim,215 act_fn=act_fn,216 post_act_fn=timestep_post_act,217 cond_proj_dim=time_cond_proj_dim,218 )219 220 # class embedding221 if class_embed_type is None and num_class_embeds is not None:222 self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)223 elif class_embed_type == "timestep":224 self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)225 elif class_embed_type == "identity":226 self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)227 elif class_embed_type == "projection":228 if projection_class_embeddings_input_dim is None:229 raise ValueError(230 "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"231 )232 # The projection `class_embed_type` is the same as the timestep `class_embed_type` except233 # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings234 # 2. it projects from an arbitrary input dimension.235 #236 # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.237 # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.238 # As a result, `TimestepEmbedding` can be passed arbitrary vectors.239 self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)240 elif class_embed_type == "simple_projection":241 if projection_class_embeddings_input_dim is None:242 raise ValueError(243 "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"244 )245 self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim)246 else:247 self.class_embedding = None248 249 self.down_blocks = nn.ModuleList([])250 self.up_blocks = nn.ModuleList([])251 252 if isinstance(only_cross_attention, bool):253 only_cross_attention = [only_cross_attention] * len(down_block_types)254 255 if isinstance(attention_head_dim, int):256 attention_head_dim = (attention_head_dim,) * len(down_block_types)257 258 if isinstance(cross_attention_dim, int):259 cross_attention_dim = (cross_attention_dim,) * len(down_block_types)260 261 if class_embeddings_concat:262 # The time embeddings are concatenated with the class embeddings. The dimension of the263 # time embeddings passed to the down, middle, and up blocks is twice the dimension of the264 # regular time embeddings265 blocks_time_embed_dim = time_embed_dim * 2266 else:267 blocks_time_embed_dim = time_embed_dim268 269 # down270 output_channel = block_out_channels[0]271 for i, down_block_type in enumerate(down_block_types):272 input_channel = output_channel273 output_channel = block_out_channels[i]274 is_final_block = i == len(block_out_channels) - 1275 276 down_block = get_down_block(277 down_block_type,278 num_layers=layers_per_block,279 in_channels=input_channel,280 out_channels=output_channel,281 temb_channels=blocks_time_embed_dim,282 add_downsample=not is_final_block,283 resnet_eps=norm_eps,284 resnet_act_fn=act_fn,285 resnet_groups=norm_num_groups,286 cross_attention_dim=cross_attention_dim[i],287 attn_num_head_channels=attention_head_dim[i],288 downsample_padding=downsample_padding,289 dual_cross_attention=dual_cross_attention,290 use_linear_projection=use_linear_projection,291 only_cross_attention=only_cross_attention[i],292 upcast_attention=upcast_attention,293 resnet_time_scale_shift=resnet_time_scale_shift,294 )295 self.down_blocks.append(down_block)296 297 # mid298 if mid_block_type == "UNetMidBlock2DCrossAttn":299 self.mid_block = UNetMidBlock2DCrossAttn(300 in_channels=block_out_channels[-1],301 temb_channels=blocks_time_embed_dim,302 resnet_eps=norm_eps,303 resnet_act_fn=act_fn,304 output_scale_factor=mid_block_scale_factor,305 resnet_time_scale_shift=resnet_time_scale_shift,306 cross_attention_dim=cross_attention_dim[-1],307 attn_num_head_channels=attention_head_dim[-1],308 resnet_groups=norm_num_groups,309 dual_cross_attention=dual_cross_attention,310 use_linear_projection=use_linear_projection,311 upcast_attention=upcast_attention,312 )313 elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn":314 self.mid_block = UNetMidBlock2DSimpleCrossAttn(315 in_channels=block_out_channels[-1],316 temb_channels=blocks_time_embed_dim,317 resnet_eps=norm_eps,318 resnet_act_fn=act_fn,319 output_scale_factor=mid_block_scale_factor,320 cross_attention_dim=cross_attention_dim[-1],321 attn_num_head_channels=attention_head_dim[-1],322 resnet_groups=norm_num_groups,323 resnet_time_scale_shift=resnet_time_scale_shift,324 )325 elif mid_block_type is None:326 self.mid_block = None327 else:328 raise ValueError(f"unknown mid_block_type : {mid_block_type}")329 330 # count how many layers upsample the images331 self.num_upsamplers = 0332 333 # up334 reversed_block_out_channels = list(reversed(block_out_channels))335 reversed_attention_head_dim = list(reversed(attention_head_dim))336 reversed_cross_attention_dim = list(reversed(cross_attention_dim))337 only_cross_attention = list(reversed(only_cross_attention))338 339 output_channel = reversed_block_out_channels[0]340 for i, up_block_type in enumerate(up_block_types):341 is_final_block = i == len(block_out_channels) - 1342 343 prev_output_channel = output_channel344 output_channel = reversed_block_out_channels[i]345 input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]346 347 # add upsample block for all BUT final layer348 if not is_final_block:349 add_upsample = True350 self.num_upsamplers += 1351 else:352 add_upsample = False353 354 up_block = get_up_block(355 up_block_type,356 num_layers=layers_per_block + 1,357 in_channels=input_channel,358 out_channels=output_channel,359 prev_output_channel=prev_output_channel,360 temb_channels=blocks_time_embed_dim,361 add_upsample=add_upsample,362 resnet_eps=norm_eps,363 resnet_act_fn=act_fn,364 resnet_groups=norm_num_groups,365 cross_attention_dim=reversed_cross_attention_dim[i],366 attn_num_head_channels=reversed_attention_head_dim[i],367 dual_cross_attention=dual_cross_attention,368 use_linear_projection=use_linear_projection,369 only_cross_attention=only_cross_attention[i],370 upcast_attention=upcast_attention,371 resnet_time_scale_shift=resnet_time_scale_shift,372 )373 self.up_blocks.append(up_block)374 prev_output_channel = output_channel375 376 # out377 if norm_num_groups is not None:378 self.conv_norm_out = nn.GroupNorm(379 num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps380 )381 self.conv_act = nn.SiLU()382 else:383 self.conv_norm_out = None384 self.conv_act = None385 386 conv_out_padding = (conv_out_kernel - 1) // 2387 self.conv_out = nn.Conv2d(388 block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding389 )390 391 @property392 def attn_processors(self) -> Dict[str, AttentionProcessor]:393 r"""394 Returns:395 `dict` of attention processors: A dictionary containing all attention processors used in the model with396 indexed by its weight name.397 """398 # set recursively399 processors = {}400 401 def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):402 if hasattr(module, "set_processor"):403 processors[f"{name}.processor"] = module.processor404 405 for sub_name, child in module.named_children():406 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)407 408 return processors409 410 for name, module in self.named_children():411 fn_recursive_add_processors(name, module, processors)412 413 return processors414 415 def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):416 r"""417 Parameters:418 `processor (`dict` of `AttentionProcessor` or `AttentionProcessor`):419 The instantiated processor class or a dictionary of processor classes that will be set as the processor420 of **all** `Attention` layers.421 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.:422 423 """424 count = len(self.attn_processors.keys())425 426 if isinstance(processor, dict) and len(processor) != count:427 raise ValueError(428 f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"429 f" number of attention layers: {count}. Please make sure to pass {count} processor classes."430 )431 432 def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):433 if hasattr(module, "set_processor"):434 if not isinstance(processor, dict):435 module.set_processor(processor)436 else:437 module.set_processor(processor.pop(f"{name}.processor"))438 439 for sub_name, child in module.named_children():440 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)441 442 for name, module in self.named_children():443 fn_recursive_attn_processor(name, module, processor)444 445 def set_default_attn_processor(self):446 """447 Disables custom attention processors and sets the default attention implementation.448 """449 self.set_attn_processor(AttnProcessor())450 451 def set_attention_slice(self, slice_size):452 r"""453 Enable sliced attention computation.454 455 When this option is enabled, the attention module will split the input tensor in slices, to compute attention456 in several steps. This is useful to save some memory in exchange for a small speed decrease.457 458 Args:459 slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):460 When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If461 `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is462 provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`463 must be a multiple of `slice_size`.464 """465 sliceable_head_dims = []466 467 def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):468 if hasattr(module, "set_attention_slice"):469 sliceable_head_dims.append(module.sliceable_head_dim)470 471 for child in module.children():472 fn_recursive_retrieve_sliceable_dims(child)473 474 # retrieve number of attention layers475 for module in self.children():476 fn_recursive_retrieve_sliceable_dims(module)477 478 num_sliceable_layers = len(sliceable_head_dims)479 480 if slice_size == "auto":481 # half the attention head size is usually a good trade-off between482 # speed and memory483 slice_size = [dim // 2 for dim in sliceable_head_dims]484 elif slice_size == "max":485 # make smallest slice possible486 slice_size = num_sliceable_layers * [1]487 488 slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size489 490 if len(slice_size) != len(sliceable_head_dims):491 raise ValueError(492 f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"493 f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."494 )495 496 for i in range(len(slice_size)):497 size = slice_size[i]498 dim = sliceable_head_dims[i]499 if size is not None and size > dim:500 raise ValueError(f"size {size} has to be smaller or equal to {dim}.")501 502 # Recursively walk through all the children.503 # Any children which exposes the set_attention_slice method504 # gets the message505 def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):506 if hasattr(module, "set_attention_slice"):507 module.set_attention_slice(slice_size.pop())508 509 for child in module.children():510 fn_recursive_set_attention_slice(child, slice_size)511 512 reversed_slice_size = list(reversed(slice_size))513 for module in self.children():514 fn_recursive_set_attention_slice(module, reversed_slice_size)515 516 def _set_gradient_checkpointing(self, module, value=False):517 if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D, CrossAttnUpBlock2D, UpBlock2D)):518 module.gradient_checkpointing = value519 520 def forward(521 self,522 sample: torch.FloatTensor,523 timestep: Union[torch.Tensor, float, int],524 encoder_hidden_states: torch.Tensor,525 class_labels: Optional[torch.Tensor] = None,526 timestep_cond: Optional[torch.Tensor] = None,527 attention_mask: Optional[torch.Tensor] = None,528 cross_attention_kwargs: Optional[Dict[str, Any]] = None,529 down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,530 mid_block_additional_residual: Optional[torch.Tensor] = None,531 encoder_attention_mask: Optional[torch.Tensor] = None,532 return_dict: bool = True,533 ) -> Union[UNet2DConditionOutput, Tuple]:534 r"""535 Args:536 sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor537 timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps538 encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states539 encoder_attention_mask (`torch.Tensor`):540 (batch, sequence_length) cross-attention mask (or bias), applied to encoder_hidden_states. If a541 BoolTensor is provided, it will be turned into a bias, by adding a large negative value. False = hide542 token. Other tensor types will be used as-is as bias values.543 return_dict (`bool`, *optional*, defaults to `True`):544 Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.545 cross_attention_kwargs (`dict`, *optional*):546 A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under547 `self.processor` in548 [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).549 550 Returns:551 [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:552 [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When553 returning a tuple, the first element is the sample tensor.554 """555 # By default samples have to be AT least a multiple of the overall upsampling factor.556 # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).557 # However, the upsampling interpolation output size can be forced to fit any upsampling size558 # on the fly if necessary.559 default_overall_up_factor = 2**self.num_upsamplers560 561 # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`562 forward_upsample_size = False563 upsample_size = None564 565 if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):566 logger.info("Forward upsample size to force interpolation output size.")567 forward_upsample_size = True568 569 # prepare attention_mask570 if attention_mask is not None:571 attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0572 attention_mask = attention_mask.unsqueeze(1)573 574 # ensure encoder_attention_mask is a bias, and make it broadcastable over multi-head-attention channels575 if encoder_attention_mask is not None:576 # if it's a mask: turn it into a bias. otherwise: assume it's already a bias577 if encoder_attention_mask.dtype is torch.bool:578 encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0579 encoder_attention_mask = encoder_attention_mask.unsqueeze(1)580 581 # 0. center input if necessary582 if self.config.center_input_sample:583 sample = 2 * sample - 1.0584 585 # 1. time586 timesteps = timestep587 if not torch.is_tensor(timesteps):588 # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can589 # This would be a good case for the `match` statement (Python 3.10+)590 is_mps = sample.device.type == "mps"591 if isinstance(timestep, float):592 dtype = torch.float32 if is_mps else torch.float64593 else:594 dtype = torch.int32 if is_mps else torch.int64595 timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)596 elif len(timesteps.shape) == 0:597 timesteps = timesteps[None].to(sample.device)598 599 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML600 timesteps = timesteps.expand(sample.shape[0])601 602 t_emb = self.time_proj(timesteps)603 604 # timesteps does not contain any weights and will always return f32 tensors605 # but time_embedding might actually be running in fp16. so we need to cast here.606 # there might be better ways to encapsulate this.607 t_emb = t_emb.to(dtype=self.dtype)608 609 emb = self.time_embedding(t_emb, timestep_cond)610 611 if self.class_embedding is not None:612 if class_labels is None:613 raise ValueError("class_labels should be provided when num_class_embeds > 0")614 615 if self.config.class_embed_type == "timestep":616 class_labels = self.time_proj(class_labels)617 618 class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)619 620 if self.config.class_embeddings_concat:621 emb = torch.cat([emb, class_emb], dim=-1)622 else:623 emb = emb + class_emb624 625 # 2. pre-process626 sample = self.conv_in(sample)627 628 # 3. down629 down_block_res_samples = (sample,)630 for downsample_block in self.down_blocks:631 if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:632 sample, res_samples = downsample_block(633 hidden_states=sample,634 temb=emb,635 encoder_hidden_states=encoder_hidden_states,636 attention_mask=attention_mask,637 cross_attention_kwargs=cross_attention_kwargs,638 encoder_attention_mask=encoder_attention_mask,639 )640 else:641 sample, res_samples = downsample_block(hidden_states=sample, temb=emb)642 643 down_block_res_samples += res_samples644 645 if down_block_additional_residuals is not None:646 new_down_block_res_samples = ()647 648 for down_block_res_sample, down_block_additional_residual in zip(649 down_block_res_samples, down_block_additional_residuals650 ):651 down_block_res_sample = down_block_res_sample + down_block_additional_residual652 new_down_block_res_samples += (down_block_res_sample,)653 654 down_block_res_samples = new_down_block_res_samples655 656 # 4. mid657 if self.mid_block is not None:658 sample = self.mid_block(659 sample,660 emb,661 encoder_hidden_states=encoder_hidden_states,662 attention_mask=attention_mask,663 cross_attention_kwargs=cross_attention_kwargs,664 encoder_attention_mask=encoder_attention_mask,665 )666 667 if mid_block_additional_residual is not None:668 sample = sample + mid_block_additional_residual669 670 # 5. up671 for i, upsample_block in enumerate(self.up_blocks):672 is_final_block = i == len(self.up_blocks) - 1673 674 res_samples = down_block_res_samples[-len(upsample_block.resnets) :]675 down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]676 677 # if we have not reached the final block and need to forward the678 # upsample size, we do it here679 if not is_final_block and forward_upsample_size:680 upsample_size = down_block_res_samples[-1].shape[2:]681 682 if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:683 sample = upsample_block(684 hidden_states=sample,685 temb=emb,686 res_hidden_states_tuple=res_samples,687 encoder_hidden_states=encoder_hidden_states,688 cross_attention_kwargs=cross_attention_kwargs,689 upsample_size=upsample_size,690 attention_mask=attention_mask,691 encoder_attention_mask=encoder_attention_mask,692 )693 else:694 sample = upsample_block(695 hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size696 )697 698 # 6. post-process699 if self.conv_norm_out:700 sample = self.conv_norm_out(sample)701 sample = self.conv_act(sample)702 sample = self.conv_out(sample)703 704 if not return_dict:705 return (sample,)706 707 return UNet2DConditionOutput(sample=sample)708