AnchoredAI/llm-grounded-diffusion
0
1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from dataclasses import dataclass15from typing import Any, Dict, Optional16 17import torch18import torch.nn.functional as F19from torch import nn20 21from diffusers.configuration_utils import ConfigMixin, register_to_config22from diffusers.models.embeddings import ImagePositionalEmbeddings23from diffusers.utils import BaseOutput, deprecate24from .attention import BasicTransformerBlock25from diffusers.models.embeddings import PatchEmbed26from diffusers.models.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 use_gated_attention: bool = False,102 ):103 super().__init__()104 self.use_linear_projection = use_linear_projection105 self.num_attention_heads = num_attention_heads106 self.attention_head_dim = attention_head_dim107 inner_dim = num_attention_heads * attention_head_dim108 109 # 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)`110 # Define whether input is continuous or discrete depending on configuration111 self.is_input_continuous = (in_channels is not None) and (patch_size is None)112 self.is_input_vectorized = num_vector_embeds is not None113 self.is_input_patches = in_channels is not None and patch_size is not None114 115 if norm_type == "layer_norm" and num_embeds_ada_norm is not None:116 deprecation_message = (117 f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"118 " incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config."119 " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"120 " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"121 " would be very nice if you could open a Pull request for the `transformer/config.json` file"122 )123 deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)124 norm_type = "ada_norm"125 126 if self.is_input_continuous and self.is_input_vectorized:127 raise ValueError(128 f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"129 " sure that either `in_channels` or `num_vector_embeds` is None."130 )131 elif self.is_input_vectorized and self.is_input_patches:132 raise ValueError(133 f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"134 " sure that either `num_vector_embeds` or `num_patches` is None."135 )136 elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:137 raise ValueError(138 f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"139 f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."140 )141 142 # 2. Define input layers143 if self.is_input_continuous:144 self.in_channels = in_channels145 146 self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)147 if use_linear_projection:148 self.proj_in = nn.Linear(in_channels, inner_dim)149 else:150 self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)151 elif self.is_input_vectorized:152 assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"153 assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"154 155 self.height = sample_size156 self.width = sample_size157 self.num_vector_embeds = num_vector_embeds158 self.num_latent_pixels = self.height * self.width159 160 self.latent_image_embedding = ImagePositionalEmbeddings(161 num_embed=num_vector_embeds, embed_dim=inner_dim, height=self.height, width=self.width162 )163 elif self.is_input_patches:164 assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"165 166 self.height = sample_size167 self.width = sample_size168 169 self.patch_size = patch_size170 self.pos_embed = PatchEmbed(171 height=sample_size,172 width=sample_size,173 patch_size=patch_size,174 in_channels=in_channels,175 embed_dim=inner_dim,176 )177 178 # 3. Define transformers blocks179 self.transformer_blocks = nn.ModuleList(180 [181 BasicTransformerBlock(182 inner_dim,183 num_attention_heads,184 attention_head_dim,185 dropout=dropout,186 cross_attention_dim=cross_attention_dim,187 activation_fn=activation_fn,188 num_embeds_ada_norm=num_embeds_ada_norm,189 attention_bias=attention_bias,190 only_cross_attention=only_cross_attention,191 upcast_attention=upcast_attention,192 norm_type=norm_type,193 norm_elementwise_affine=norm_elementwise_affine,194 use_gated_attention=use_gated_attention,195 )196 for d in range(num_layers)197 ]198 )199 200 # 4. Define output layers201 self.out_channels = in_channels if out_channels is None else out_channels202 if self.is_input_continuous:203 # TODO: should use out_channels for continuous projections204 if use_linear_projection:205 self.proj_out = nn.Linear(inner_dim, in_channels)206 else:207 self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)208 elif self.is_input_vectorized:209 self.norm_out = nn.LayerNorm(inner_dim)210 self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1)211 elif self.is_input_patches:212 self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)213 self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim)214 self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)215 216 def forward(217 self,218 hidden_states: torch.Tensor,219 encoder_hidden_states: Optional[torch.Tensor] = None,220 timestep: Optional[torch.LongTensor] = None,221 class_labels: Optional[torch.LongTensor] = None,222 cross_attention_kwargs: Dict[str, Any] = None,223 attention_mask: Optional[torch.Tensor] = None,224 encoder_attention_mask: Optional[torch.Tensor] = None,225 return_dict: bool = True,226 return_cross_attention_probs: bool = False,227 ):228 """229 Args:230 hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`.231 When continuous, `torch.FloatTensor` of shape `(batch size, channel, height, width)`): Input232 hidden_states233 encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):234 Conditional embeddings for cross attention layer. If not given, cross-attention defaults to235 self-attention.236 timestep ( `torch.LongTensor`, *optional*):237 Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step.238 class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):239 Optional class labels to be applied as an embedding in AdaLayerZeroNorm. Used to indicate class labels240 conditioning.241 encoder_attention_mask ( `torch.Tensor`, *optional* ).242 Cross-attention mask, applied to encoder_hidden_states. Two formats supported:243 Mask `(batch, sequence_length)` True = keep, False = discard. Bias `(batch, 1, sequence_length)` 0244 = keep, -10000 = discard.245 If ndim == 2: will be interpreted as a mask, then converted into a bias consistent with the format246 above. This bias will be added to the cross-attention scores.247 return_dict (`bool`, *optional*, defaults to `True`):248 Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.249 250 Returns:251 [`~models.transformer_2d.Transformer2DModelOutput`] or `tuple`:252 [`~models.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a `tuple`. When253 returning a tuple, the first element is the sample tensor.254 """255 # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.256 # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.257 # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.258 # expects mask of shape:259 # [batch, key_tokens]260 # adds singleton query_tokens dimension:261 # [batch, 1, key_tokens]262 # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:263 # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)264 # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)265 if attention_mask is not None and attention_mask.ndim == 2:266 # assume that mask is expressed as:267 # (1 = keep, 0 = discard)268 # convert mask into a bias that can be added to attention scores:269 # (keep = +0, discard = -10000.0)270 attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0271 attention_mask = attention_mask.unsqueeze(1)272 273 # convert encoder_attention_mask to a bias the same way we do for attention_mask274 if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:275 encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0276 encoder_attention_mask = encoder_attention_mask.unsqueeze(1)277 278 # 1. Input279 if self.is_input_continuous:280 batch, _, height, width = hidden_states.shape281 residual = hidden_states282 283 hidden_states = self.norm(hidden_states)284 if not self.use_linear_projection:285 hidden_states = self.proj_in(hidden_states)286 inner_dim = hidden_states.shape[1]287 hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)288 else:289 inner_dim = hidden_states.shape[1]290 hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)291 hidden_states = self.proj_in(hidden_states)292 elif self.is_input_vectorized:293 hidden_states = self.latent_image_embedding(hidden_states)294 elif self.is_input_patches:295 hidden_states = self.pos_embed(hidden_states)296 297 base_attn_key = cross_attention_kwargs["attn_key"]298 299 # 2. Blocks300 cross_attention_probs_all = []301 for block_ind, block in enumerate(self.transformer_blocks):302 cross_attention_kwargs["attn_key"] = base_attn_key + [block_ind]303 304 hidden_states = block(305 hidden_states,306 attention_mask=attention_mask,307 encoder_hidden_states=encoder_hidden_states,308 encoder_attention_mask=encoder_attention_mask,309 timestep=timestep,310 cross_attention_kwargs=cross_attention_kwargs,311 class_labels=class_labels,312 return_cross_attention_probs=return_cross_attention_probs,313 )314 if return_cross_attention_probs:315 hidden_states, cross_attention_probs = hidden_states316 cross_attention_probs_all.append(cross_attention_probs)317 318 # 3. Output319 if self.is_input_continuous:320 if not self.use_linear_projection:321 hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()322 hidden_states = self.proj_out(hidden_states)323 else:324 hidden_states = self.proj_out(hidden_states)325 hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()326 327 output = hidden_states + residual328 elif self.is_input_vectorized:329 hidden_states = self.norm_out(hidden_states)330 logits = self.out(hidden_states)331 # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)332 logits = logits.permute(0, 2, 1)333 334 # log(p(x_0))335 output = F.log_softmax(logits.double(), dim=1).float()336 elif self.is_input_patches:337 # TODO: cleanup!338 conditioning = self.transformer_blocks[0].norm1.emb(339 timestep, class_labels, hidden_dtype=hidden_states.dtype340 )341 shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)342 hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]343 hidden_states = self.proj_out_2(hidden_states)344 345 # unpatchify346 height = width = int(hidden_states.shape[1] ** 0.5)347 hidden_states = hidden_states.reshape(348 shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)349 )350 hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)351 output = hidden_states.reshape(352 shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)353 )354 355 if len(cross_attention_probs_all) == 1:356 # If we only have one transformer block in a Transformer2DModel, we do not create another nested level.357 cross_attention_probs_all = cross_attention_probs_all[0]358 359 if not return_dict:360 if return_cross_attention_probs:361 return (output, cross_attention_probs_all)362 return (output,)363 364 output = Transformer2DModelOutput(sample=output)365 if return_cross_attention_probs:366 return output, cross_attention_probs_all367 return output368