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, Optional16 17import torch18import torch.nn.functional as F19from torch import nn20 21from ..configuration_utils import ConfigMixin, register_to_config22from ..models.embeddings import ImagePositionalEmbeddings23from ..utils import BaseOutput, deprecate24from .attention import BasicTransformerBlock25from .embeddings import PatchEmbed26from .modeling_utils import ModelMixin27 28 29@dataclass30class Transformer2DModelOutput(BaseOutput):31 """32 Args:33 sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):34 Hidden states conditioned on `encoder_hidden_states` input. If discrete, returns probability distributions35 for the unnoised latent pixels.36 """37 38 sample: torch.FloatTensor39 40 41class Transformer2DModel(ModelMixin, ConfigMixin):42 """43 Transformer model for image-like data. Takes either discrete (classes of vector embeddings) or continuous (actual44 embeddings) inputs.45 46 When input is continuous: First, project the input (aka embedding) and reshape to b, t, d. Then apply standard47 transformer action. Finally, reshape to image.48 49 When input is discrete: First, input (classes of latent pixels) is converted to embeddings and has positional50 embeddings applied, see `ImagePositionalEmbeddings`. Then apply standard transformer action. Finally, predict51 classes of unnoised image.52 53 Note that it is assumed one of the input classes is the masked latent pixel. The predicted classes of the unnoised54 image do not contain a prediction for the masked pixel as the unnoised image cannot be masked.55 56 Parameters:57 num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.58 attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.59 in_channels (`int`, *optional*):60 Pass if the input is continuous. The number of channels in the input and output.61 num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.62 dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.63 cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use.64 sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images.65 Note that this is fixed at training time as it is used for learning a number of position embeddings. See66 `ImagePositionalEmbeddings`.67 num_vector_embeds (`int`, *optional*):68 Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels.69 Includes the class for the masked latent pixel.70 activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.71 num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`.72 The number of diffusion steps used during training. Note that this is fixed at training time as it is used73 to learn a number of embeddings that are added to the hidden states. During inference, you can denoise for74 up to but not more than steps than `num_embeds_ada_norm`.75 attention_bias (`bool`, *optional*):76 Configure if the TransformerBlocks' attention should contain a bias parameter.77 """78 79 @register_to_config80 def __init__(81 self,82 num_attention_heads: int = 16,83 attention_head_dim: int = 88,84 in_channels: Optional[int] = None,85 out_channels: Optional[int] = None,86 num_layers: int = 1,87 dropout: float = 0.0,88 norm_num_groups: int = 32,89 cross_attention_dim: Optional[int] = None,90 attention_bias: bool = False,91 sample_size: Optional[int] = None,92 num_vector_embeds: Optional[int] = None,93 patch_size: Optional[int] = None,94 activation_fn: str = "geglu",95 num_embeds_ada_norm: Optional[int] = None,96 use_linear_projection: bool = False,97 only_cross_attention: bool = False,98 upcast_attention: bool = False,99 norm_type: str = "layer_norm",100 norm_elementwise_affine: bool = True,101 ):102 super().__init__()103 self.use_linear_projection = use_linear_projection104 self.num_attention_heads = num_attention_heads105 self.attention_head_dim = attention_head_dim106 inner_dim = num_attention_heads * attention_head_dim107 108 # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`109 # Define whether input is continuous or discrete depending on configuration110 self.is_input_continuous = (in_channels is not None) and (patch_size is None)111 self.is_input_vectorized = num_vector_embeds is not None112 self.is_input_patches = in_channels is not None and patch_size is not None113 114 if norm_type == "layer_norm" and num_embeds_ada_norm is not None:115 deprecation_message = (116 f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"117 " incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config."118 " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"119 " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"120 " would be very nice if you could open a Pull request for the `transformer/config.json` file"121 )122 deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)123 norm_type = "ada_norm"124 125 if self.is_input_continuous and self.is_input_vectorized:126 raise ValueError(127 f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"128 " sure that either `in_channels` or `num_vector_embeds` is None."129 )130 elif self.is_input_vectorized and self.is_input_patches:131 raise ValueError(132 f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"133 " sure that either `num_vector_embeds` or `num_patches` is None."134 )135 elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:136 raise ValueError(137 f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"138 f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."139 )140 141 # 2. Define input layers142 if self.is_input_continuous:143 self.in_channels = in_channels144 145 self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)146 if use_linear_projection:147 self.proj_in = nn.Linear(in_channels, inner_dim)148 else:149 self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)150 elif self.is_input_vectorized:151 assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"152 assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"153 154 self.height = sample_size155 self.width = sample_size156 self.num_vector_embeds = num_vector_embeds157 self.num_latent_pixels = self.height * self.width158 159 self.latent_image_embedding = ImagePositionalEmbeddings(160 num_embed=num_vector_embeds, embed_dim=inner_dim, height=self.height, width=self.width161 )162 elif self.is_input_patches:163 assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"164 165 self.height = sample_size166 self.width = sample_size167 168 self.patch_size = patch_size169 self.pos_embed = PatchEmbed(170 height=sample_size,171 width=sample_size,172 patch_size=patch_size,173 in_channels=in_channels,174 embed_dim=inner_dim,175 )176 177 # 3. Define transformers blocks178 self.transformer_blocks = nn.ModuleList(179 [180 BasicTransformerBlock(181 inner_dim,182 num_attention_heads,183 attention_head_dim,184 dropout=dropout,185 cross_attention_dim=cross_attention_dim,186 activation_fn=activation_fn,187 num_embeds_ada_norm=num_embeds_ada_norm,188 attention_bias=attention_bias,189 only_cross_attention=only_cross_attention,190 upcast_attention=upcast_attention,191 norm_type=norm_type,192 norm_elementwise_affine=norm_elementwise_affine,193 )194 for d in range(num_layers)195 ]196 )197 198 # 4. Define output layers199 self.out_channels = in_channels if out_channels is None else out_channels200 if self.is_input_continuous:201 # TODO: should use out_channels for continuous projections202 if use_linear_projection:203 self.proj_out = nn.Linear(inner_dim, in_channels)204 else:205 self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)206 elif self.is_input_vectorized:207 self.norm_out = nn.LayerNorm(inner_dim)208 self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1)209 elif self.is_input_patches:210 self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)211 self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim)212 self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)213 214 def forward(215 self,216 hidden_states: torch.Tensor,217 encoder_hidden_states: Optional[torch.Tensor] = None,218 timestep: Optional[torch.LongTensor] = None,219 class_labels: Optional[torch.LongTensor] = None,220 cross_attention_kwargs: Dict[str, Any] = None,221 attention_mask: Optional[torch.Tensor] = None,222 encoder_attention_mask: Optional[torch.Tensor] = None,223 return_dict: bool = True,224 ):225 """226 Args:227 hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`.228 When continuous, `torch.FloatTensor` of shape `(batch size, channel, height, width)`): Input229 hidden_states230 encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*):231 Conditional embeddings for cross attention layer. If not given, cross-attention defaults to232 self-attention.233 timestep ( `torch.LongTensor`, *optional*):234 Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step.235 class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):236 Optional class labels to be applied as an embedding in AdaLayerZeroNorm. Used to indicate class labels237 conditioning.238 attention_mask ( `torch.Tensor` of shape (batch size, num latent pixels), *optional* ).239 Bias to add to attention scores.240 encoder_attention_mask ( `torch.Tensor` of shape (batch size, num encoder tokens), *optional* ).241 Bias to add to cross-attention scores.242 return_dict (`bool`, *optional*, defaults to `True`):243 Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.244 245 Returns:246 [`~models.transformer_2d.Transformer2DModelOutput`] or `tuple`:247 [`~models.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a `tuple`. When248 returning a tuple, the first element is the sample tensor.249 """250 # 1. Input251 if self.is_input_continuous:252 batch, _, height, width = hidden_states.shape253 residual = hidden_states254 255 hidden_states = self.norm(hidden_states)256 if not self.use_linear_projection:257 hidden_states = self.proj_in(hidden_states)258 inner_dim = hidden_states.shape[1]259 hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)260 else:261 inner_dim = hidden_states.shape[1]262 hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)263 hidden_states = self.proj_in(hidden_states)264 elif self.is_input_vectorized:265 hidden_states = self.latent_image_embedding(hidden_states)266 elif self.is_input_patches:267 hidden_states = self.pos_embed(hidden_states)268 269 # 2. Blocks270 for block in self.transformer_blocks:271 hidden_states = block(272 hidden_states,273 attention_mask=attention_mask,274 encoder_hidden_states=encoder_hidden_states,275 encoder_attention_mask=encoder_attention_mask,276 timestep=timestep,277 cross_attention_kwargs=cross_attention_kwargs,278 class_labels=class_labels,279 )280 281 # 3. Output282 if self.is_input_continuous:283 if not self.use_linear_projection:284 hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()285 hidden_states = self.proj_out(hidden_states)286 else:287 hidden_states = self.proj_out(hidden_states)288 hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()289 290 output = hidden_states + residual291 elif self.is_input_vectorized:292 hidden_states = self.norm_out(hidden_states)293 logits = self.out(hidden_states)294 # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)295 logits = logits.permute(0, 2, 1)296 297 # log(p(x_0))298 output = F.log_softmax(logits.double(), dim=1).float()299 elif self.is_input_patches:300 # TODO: cleanup!301 conditioning = self.transformer_blocks[0].norm1.emb(302 timestep, class_labels, hidden_dtype=hidden_states.dtype303 )304 shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)305 hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]306 hidden_states = self.proj_out_2(hidden_states)307 308 # unpatchify309 height = width = int(hidden_states.shape[1] ** 0.5)310 hidden_states = hidden_states.reshape(311 shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)312 )313 hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)314 output = hidden_states.reshape(315 shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)316 )317 318 if not return_dict:319 return (output,)320 321 return Transformer2DModelOutput(sample=output)322