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.14import math15from typing import Any, Callable, Dict, Optional16 17import torch18import torch.nn.functional as F19from torch import nn20 21from ..utils.import_utils import is_xformers_available22from .attention_processor import Attention23from .embeddings import CombinedTimestepLabelEmbeddings24 25 26if is_xformers_available():27 import xformers28 import xformers.ops29else:30 xformers = None31 32 33class AttentionBlock(nn.Module):34 """35 An attention block that allows spatial positions to attend to each other. Originally ported from here, but adapted36 to the N-d case.37 https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.38 Uses three q, k, v linear layers to compute attention.39 40 Parameters:41 channels (`int`): The number of channels in the input and output.42 num_head_channels (`int`, *optional*):43 The number of channels in each head. If None, then `num_heads` = 1.44 norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for group norm.45 rescale_output_factor (`float`, *optional*, defaults to 1.0): The factor to rescale the output by.46 eps (`float`, *optional*, defaults to 1e-5): The epsilon value to use for group norm.47 """48 49 # IMPORTANT;TODO(Patrick, William) - this class will be deprecated soon. Do not use it anymore50 51 def __init__(52 self,53 channels: int,54 num_head_channels: Optional[int] = None,55 norm_num_groups: int = 32,56 rescale_output_factor: float = 1.0,57 eps: float = 1e-5,58 ):59 super().__init__()60 self.channels = channels61 62 self.num_heads = channels // num_head_channels if num_head_channels is not None else 163 self.num_head_size = num_head_channels64 self.group_norm = nn.GroupNorm(num_channels=channels, num_groups=norm_num_groups, eps=eps, affine=True)65 66 # define q,k,v as linear layers67 self.query = nn.Linear(channels, channels)68 self.key = nn.Linear(channels, channels)69 self.value = nn.Linear(channels, channels)70 71 self.rescale_output_factor = rescale_output_factor72 self.proj_attn = nn.Linear(channels, channels, bias=True)73 74 self._use_memory_efficient_attention_xformers = False75 self._attention_op = None76 77 def reshape_heads_to_batch_dim(self, tensor):78 batch_size, seq_len, dim = tensor.shape79 head_size = self.num_heads80 tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)81 tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size * head_size, seq_len, dim // head_size)82 return tensor83 84 def reshape_batch_dim_to_heads(self, tensor):85 batch_size, seq_len, dim = tensor.shape86 head_size = self.num_heads87 tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)88 tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)89 return tensor90 91 def set_use_memory_efficient_attention_xformers(92 self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None93 ):94 if use_memory_efficient_attention_xformers:95 if not is_xformers_available():96 raise ModuleNotFoundError(97 (98 "Refer to https://github.com/facebookresearch/xformers for more information on how to install"99 " xformers"100 ),101 name="xformers",102 )103 elif not torch.cuda.is_available():104 raise ValueError(105 "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"106 " only available for GPU "107 )108 else:109 try:110 # Make sure we can run the memory efficient attention111 _ = xformers.ops.memory_efficient_attention(112 torch.randn((1, 2, 40), device="cuda"),113 torch.randn((1, 2, 40), device="cuda"),114 torch.randn((1, 2, 40), device="cuda"),115 )116 except Exception as e:117 raise e118 self._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers119 self._attention_op = attention_op120 121 def forward(self, hidden_states):122 residual = hidden_states123 batch, channel, height, width = hidden_states.shape124 125 # norm126 hidden_states = self.group_norm(hidden_states)127 128 hidden_states = hidden_states.view(batch, channel, height * width).transpose(1, 2)129 130 # proj to q, k, v131 query_proj = self.query(hidden_states)132 key_proj = self.key(hidden_states)133 value_proj = self.value(hidden_states)134 135 scale = 1 / math.sqrt(self.channels / self.num_heads)136 137 query_proj = self.reshape_heads_to_batch_dim(query_proj)138 key_proj = self.reshape_heads_to_batch_dim(key_proj)139 value_proj = self.reshape_heads_to_batch_dim(value_proj)140 141 if self._use_memory_efficient_attention_xformers:142 # Memory efficient attention143 hidden_states = xformers.ops.memory_efficient_attention(144 query_proj, key_proj, value_proj, attn_bias=None, op=self._attention_op145 )146 hidden_states = hidden_states.to(query_proj.dtype)147 else:148 attention_scores = torch.baddbmm(149 torch.empty(150 query_proj.shape[0],151 query_proj.shape[1],152 key_proj.shape[1],153 dtype=query_proj.dtype,154 device=query_proj.device,155 ),156 query_proj,157 key_proj.transpose(-1, -2),158 beta=0,159 alpha=scale,160 )161 attention_probs = torch.softmax(attention_scores.float(), dim=-1).type(attention_scores.dtype)162 hidden_states = torch.bmm(attention_probs, value_proj)163 164 # reshape hidden_states165 hidden_states = self.reshape_batch_dim_to_heads(hidden_states)166 167 # compute next hidden_states168 hidden_states = self.proj_attn(hidden_states)169 170 hidden_states = hidden_states.transpose(-1, -2).reshape(batch, channel, height, width)171 172 # res connect and rescale173 hidden_states = (hidden_states + residual) / self.rescale_output_factor174 return hidden_states175 176 177class BasicTransformerBlock(nn.Module):178 r"""179 A basic Transformer block.180 181 Parameters:182 dim (`int`): The number of channels in the input and output.183 num_attention_heads (`int`): The number of heads to use for multi-head attention.184 attention_head_dim (`int`): The number of channels in each head.185 dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.186 cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.187 only_cross_attention (`bool`, *optional*):188 Whether to use only cross-attention layers. In this case two cross attention layers are used.189 double_self_attention (`bool`, *optional*):190 Whether to use two self-attention layers. In this case no cross attention layers are used.191 activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.192 num_embeds_ada_norm (:193 obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.194 attention_bias (:195 obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.196 """197 198 def __init__(199 self,200 dim: int,201 num_attention_heads: int,202 attention_head_dim: int,203 dropout=0.0,204 cross_attention_dim: Optional[int] = None,205 activation_fn: str = "geglu",206 num_embeds_ada_norm: Optional[int] = None,207 attention_bias: bool = False,208 only_cross_attention: bool = False,209 double_self_attention: bool = False,210 upcast_attention: bool = False,211 norm_elementwise_affine: bool = True,212 norm_type: str = "layer_norm",213 final_dropout: bool = False,214 ):215 super().__init__()216 self.only_cross_attention = only_cross_attention217 218 self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"219 self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"220 221 if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:222 raise ValueError(223 f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"224 f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."225 )226 227 # 1. Self-Attn228 self.attn1 = Attention(229 query_dim=dim,230 heads=num_attention_heads,231 dim_head=attention_head_dim,232 dropout=dropout,233 bias=attention_bias,234 cross_attention_dim=cross_attention_dim if only_cross_attention else None,235 upcast_attention=upcast_attention,236 )237 238 self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout)239 240 # 2. Cross-Attn241 if cross_attention_dim is not None or double_self_attention:242 self.attn2 = Attention(243 query_dim=dim,244 cross_attention_dim=cross_attention_dim if not double_self_attention else None,245 heads=num_attention_heads,246 dim_head=attention_head_dim,247 dropout=dropout,248 bias=attention_bias,249 upcast_attention=upcast_attention,250 ) # is self-attn if encoder_hidden_states is none251 else:252 self.attn2 = None253 254 if self.use_ada_layer_norm:255 self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)256 elif self.use_ada_layer_norm_zero:257 self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)258 else:259 self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)260 261 if cross_attention_dim is not None or double_self_attention:262 # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.263 # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during264 # the second cross attention block.265 self.norm2 = (266 AdaLayerNorm(dim, num_embeds_ada_norm)267 if self.use_ada_layer_norm268 else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)269 )270 else:271 self.norm2 = None272 273 # 3. Feed-forward274 self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)275 276 def forward(277 self,278 hidden_states: torch.FloatTensor,279 attention_mask: Optional[torch.FloatTensor] = None,280 encoder_hidden_states: Optional[torch.FloatTensor] = None,281 encoder_attention_mask: Optional[torch.FloatTensor] = None,282 timestep: Optional[torch.LongTensor] = None,283 cross_attention_kwargs: Dict[str, Any] = None,284 class_labels: Optional[torch.LongTensor] = None,285 ):286 if self.use_ada_layer_norm:287 norm_hidden_states = self.norm1(hidden_states, timestep)288 elif self.use_ada_layer_norm_zero:289 norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(290 hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype291 )292 else:293 norm_hidden_states = self.norm1(hidden_states)294 295 # 1. Self-Attention296 cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}297 attn_output = self.attn1(298 norm_hidden_states,299 encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,300 attention_mask=attention_mask,301 **cross_attention_kwargs,302 )303 if self.use_ada_layer_norm_zero:304 attn_output = gate_msa.unsqueeze(1) * attn_output305 hidden_states = attn_output + hidden_states306 307 if self.attn2 is not None:308 norm_hidden_states = (309 self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)310 )311 312 # 2. Cross-Attention313 attn_output = self.attn2(314 norm_hidden_states,315 encoder_hidden_states=encoder_hidden_states,316 attention_mask=encoder_attention_mask,317 **cross_attention_kwargs,318 )319 hidden_states = attn_output + hidden_states320 321 # 3. Feed-forward322 norm_hidden_states = self.norm3(hidden_states)323 324 if self.use_ada_layer_norm_zero:325 norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]326 327 ff_output = self.ff(norm_hidden_states)328 329 if self.use_ada_layer_norm_zero:330 ff_output = gate_mlp.unsqueeze(1) * ff_output331 332 hidden_states = ff_output + hidden_states333 334 return hidden_states335 336 337class FeedForward(nn.Module):338 r"""339 A feed-forward layer.340 341 Parameters:342 dim (`int`): The number of channels in the input.343 dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.344 mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.345 dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.346 activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.347 final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.348 """349 350 def __init__(351 self,352 dim: int,353 dim_out: Optional[int] = None,354 mult: int = 4,355 dropout: float = 0.0,356 activation_fn: str = "geglu",357 final_dropout: bool = False,358 ):359 super().__init__()360 inner_dim = int(dim * mult)361 dim_out = dim_out if dim_out is not None else dim362 363 if activation_fn == "gelu":364 act_fn = GELU(dim, inner_dim)365 if activation_fn == "gelu-approximate":366 act_fn = GELU(dim, inner_dim, approximate="tanh")367 elif activation_fn == "geglu":368 act_fn = GEGLU(dim, inner_dim)369 elif activation_fn == "geglu-approximate":370 act_fn = ApproximateGELU(dim, inner_dim)371 372 self.net = nn.ModuleList([])373 # project in374 self.net.append(act_fn)375 # project dropout376 self.net.append(nn.Dropout(dropout))377 # project out378 self.net.append(nn.Linear(inner_dim, dim_out))379 # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout380 if final_dropout:381 self.net.append(nn.Dropout(dropout))382 383 def forward(self, hidden_states):384 for module in self.net:385 hidden_states = module(hidden_states)386 return hidden_states387 388 389class GELU(nn.Module):390 r"""391 GELU activation function with tanh approximation support with `approximate="tanh"`.392 """393 394 def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"):395 super().__init__()396 self.proj = nn.Linear(dim_in, dim_out)397 self.approximate = approximate398 399 def gelu(self, gate):400 if gate.device.type != "mps":401 return F.gelu(gate, approximate=self.approximate)402 # mps: gelu is not implemented for float16403 return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(dtype=gate.dtype)404 405 def forward(self, hidden_states):406 hidden_states = self.proj(hidden_states)407 hidden_states = self.gelu(hidden_states)408 return hidden_states409 410 411class GEGLU(nn.Module):412 r"""413 A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.414 415 Parameters:416 dim_in (`int`): The number of channels in the input.417 dim_out (`int`): The number of channels in the output.418 """419 420 def __init__(self, dim_in: int, dim_out: int):421 super().__init__()422 self.proj = nn.Linear(dim_in, dim_out * 2)423 424 def gelu(self, gate):425 if gate.device.type != "mps":426 return F.gelu(gate)427 # mps: gelu is not implemented for float16428 return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)429 430 def forward(self, hidden_states):431 hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1)432 return hidden_states * self.gelu(gate)433 434 435class ApproximateGELU(nn.Module):436 """437 The approximate form of Gaussian Error Linear Unit (GELU)438 439 For more details, see section 2: https://arxiv.org/abs/1606.08415440 """441 442 def __init__(self, dim_in: int, dim_out: int):443 super().__init__()444 self.proj = nn.Linear(dim_in, dim_out)445 446 def forward(self, x):447 x = self.proj(x)448 return x * torch.sigmoid(1.702 * x)449 450 451class AdaLayerNorm(nn.Module):452 """453 Norm layer modified to incorporate timestep embeddings.454 """455 456 def __init__(self, embedding_dim, num_embeddings):457 super().__init__()458 self.emb = nn.Embedding(num_embeddings, embedding_dim)459 self.silu = nn.SiLU()460 self.linear = nn.Linear(embedding_dim, embedding_dim * 2)461 self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False)462 463 def forward(self, x, timestep):464 emb = self.linear(self.silu(self.emb(timestep)))465 scale, shift = torch.chunk(emb, 2)466 x = self.norm(x) * (1 + scale) + shift467 return x468 469 470class AdaLayerNormZero(nn.Module):471 """472 Norm layer adaptive layer norm zero (adaLN-Zero).473 """474 475 def __init__(self, embedding_dim, num_embeddings):476 super().__init__()477 478 self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim)479 480 self.silu = nn.SiLU()481 self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True)482 self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)483 484 def forward(self, x, timestep, class_labels, hidden_dtype=None):485 emb = self.linear(self.silu(self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)))486 shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1)487 x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]488 return x, gate_msa, shift_mlp, scale_mlp, gate_mlp489 490 491class AdaGroupNorm(nn.Module):492 """493 GroupNorm layer modified to incorporate timestep embeddings.494 """495 496 def __init__(497 self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: Optional[str] = None, eps: float = 1e-5498 ):499 super().__init__()500 self.num_groups = num_groups501 self.eps = eps502 self.act = None503 if act_fn == "swish":504 self.act = lambda x: F.silu(x)505 elif act_fn == "mish":506 self.act = nn.Mish()507 elif act_fn == "silu":508 self.act = nn.SiLU()509 elif act_fn == "gelu":510 self.act = nn.GELU()511 512 self.linear = nn.Linear(embedding_dim, out_dim * 2)513 514 def forward(self, x, emb):515 if self.act:516 emb = self.act(emb)517 emb = self.linear(emb)518 emb = emb[:, :, None, None]519 scale, shift = emb.chunk(2, dim=1)520 521 x = F.group_norm(x, self.num_groups, eps=self.eps)522 x = x * (1 + scale) + shift523 return x524