Spaces:
Sleeping
Sleeping
| # Copyright 2023 The HuggingFace Team. All rights reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| from typing import Any, Dict, Optional | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from diffusers.utils import maybe_allow_in_graph | |
| from .attention_processor import Attention | |
| from diffusers.models.embeddings import CombinedTimestepLabelEmbeddings | |
| # https://github.com/gligen/diffusers/blob/23a9a0fab1b48752c7b9bcc98f6fe3b1d8fa7990/src/diffusers/models/attention.py | |
| class GatedSelfAttentionDense(nn.Module): | |
| def __init__(self, query_dim, context_dim, n_heads, d_head): | |
| super().__init__() | |
| # we need a linear projection since we need cat visual feature and obj feature | |
| self.linear = nn.Linear(context_dim, query_dim) | |
| self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head) | |
| self.ff = FeedForward(query_dim, activation_fn="geglu") | |
| self.norm1 = nn.LayerNorm(query_dim) | |
| self.norm2 = nn.LayerNorm(query_dim) | |
| self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.))) | |
| self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.))) | |
| self.enabled = True | |
| def forward(self, x, objs, fuser_attn_kwargs={}): | |
| if not self.enabled: | |
| return x | |
| n_visual = x.shape[1] | |
| objs = self.linear(objs) | |
| x = x + self.alpha_attn.tanh() * self.attn(self.norm1(torch.cat([x, objs], dim=1)), **fuser_attn_kwargs)[:, :n_visual, :] | |
| x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x)) | |
| return x | |
| class BasicTransformerBlock(nn.Module): | |
| r""" | |
| A basic Transformer block. | |
| Parameters: | |
| dim (`int`): The number of channels in the input and output. | |
| num_attention_heads (`int`): The number of heads to use for multi-head attention. | |
| attention_head_dim (`int`): The number of channels in each head. | |
| dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. | |
| cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. | |
| only_cross_attention (`bool`, *optional*): | |
| Whether to use only cross-attention layers. In this case two cross attention layers are used. | |
| double_self_attention (`bool`, *optional*): | |
| Whether to use two self-attention layers. In this case no cross attention layers are used. | |
| activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. | |
| num_embeds_ada_norm (: | |
| obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. | |
| attention_bias (: | |
| obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. | |
| """ | |
| def __init__( | |
| self, | |
| dim: int, | |
| num_attention_heads: int, | |
| attention_head_dim: int, | |
| dropout=0.0, | |
| cross_attention_dim: Optional[int] = None, | |
| activation_fn: str = "geglu", | |
| num_embeds_ada_norm: Optional[int] = None, | |
| attention_bias: bool = False, | |
| only_cross_attention: bool = False, | |
| double_self_attention: bool = False, | |
| upcast_attention: bool = False, | |
| norm_elementwise_affine: bool = True, | |
| norm_type: str = "layer_norm", | |
| final_dropout: bool = False, | |
| use_gated_attention: bool = False, | |
| ): | |
| super().__init__() | |
| self.only_cross_attention = only_cross_attention | |
| self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" | |
| self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" | |
| if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: | |
| raise ValueError( | |
| f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" | |
| f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." | |
| ) | |
| # Define 3 blocks. Each block has its own normalization layer. | |
| # 1. Self-Attn | |
| if self.use_ada_layer_norm: | |
| self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) | |
| elif self.use_ada_layer_norm_zero: | |
| self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) | |
| else: | |
| self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) | |
| self.attn1 = Attention( | |
| query_dim=dim, | |
| heads=num_attention_heads, | |
| dim_head=attention_head_dim, | |
| dropout=dropout, | |
| bias=attention_bias, | |
| cross_attention_dim=cross_attention_dim if only_cross_attention else None, | |
| upcast_attention=upcast_attention, | |
| ) | |
| # 2. Cross-Attn | |
| if cross_attention_dim is not None or double_self_attention: | |
| # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. | |
| # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during | |
| # the second cross attention block. | |
| self.norm2 = ( | |
| AdaLayerNorm(dim, num_embeds_ada_norm) | |
| if self.use_ada_layer_norm | |
| else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) | |
| ) | |
| self.attn2 = Attention( | |
| query_dim=dim, | |
| cross_attention_dim=cross_attention_dim if not double_self_attention else None, | |
| heads=num_attention_heads, | |
| dim_head=attention_head_dim, | |
| dropout=dropout, | |
| bias=attention_bias, | |
| upcast_attention=upcast_attention, | |
| ) # is self-attn if encoder_hidden_states is none | |
| else: | |
| self.norm2 = None | |
| self.attn2 = None | |
| # 3. Feed-forward | |
| self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) | |
| self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) | |
| # 4. Fuser | |
| if use_gated_attention: | |
| self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) | |
| def forward( | |
| self, | |
| hidden_states: torch.FloatTensor, | |
| attention_mask: Optional[torch.FloatTensor] = None, | |
| encoder_hidden_states: Optional[torch.FloatTensor] = None, | |
| encoder_attention_mask: Optional[torch.FloatTensor] = None, | |
| timestep: Optional[torch.LongTensor] = None, | |
| cross_attention_kwargs: Dict[str, Any] = None, | |
| class_labels: Optional[torch.LongTensor] = None, | |
| return_cross_attention_probs: bool = None, | |
| ): | |
| # Notice that normalization is always applied before the real computation in the following blocks. | |
| # 0. Prepare GLIGEN inputs | |
| if 'gligen' in cross_attention_kwargs: | |
| cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} | |
| gligen_kwargs = cross_attention_kwargs.pop('gligen', None) | |
| else: | |
| cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} | |
| gligen_kwargs = None | |
| # 1. Self-Attention | |
| if self.use_ada_layer_norm: | |
| norm_hidden_states = self.norm1(hidden_states, timestep) | |
| elif self.use_ada_layer_norm_zero: | |
| norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( | |
| hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype | |
| ) | |
| else: | |
| norm_hidden_states = self.norm1(hidden_states) | |
| attn_output = self.attn1( | |
| norm_hidden_states, | |
| encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, | |
| attention_mask=attention_mask, | |
| **cross_attention_kwargs, | |
| ) | |
| if self.use_ada_layer_norm_zero: | |
| attn_output = gate_msa.unsqueeze(1) * attn_output | |
| hidden_states = attn_output + hidden_states | |
| # 1.5 GLIGEN Control | |
| if gligen_kwargs is not None: | |
| # print(gligen_kwargs) | |
| hidden_states = self.fuser(hidden_states, gligen_kwargs['objs'], fuser_attn_kwargs=gligen_kwargs.get("fuser_attn_kwargs", {})) | |
| # 1.5 ends | |
| # 2. Cross-Attention | |
| if self.attn2 is not None: | |
| norm_hidden_states = ( | |
| self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states) | |
| ) | |
| attn_output = self.attn2( | |
| norm_hidden_states, | |
| encoder_hidden_states=encoder_hidden_states, | |
| attention_mask=encoder_attention_mask, | |
| return_attntion_probs=return_cross_attention_probs, | |
| **cross_attention_kwargs, | |
| ) | |
| if return_cross_attention_probs: | |
| attn_output, cross_attention_probs = attn_output | |
| hidden_states = attn_output + hidden_states | |
| # 3. Feed-forward | |
| norm_hidden_states = self.norm3(hidden_states) | |
| if self.use_ada_layer_norm_zero: | |
| norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] | |
| ff_output = self.ff(norm_hidden_states) | |
| if self.use_ada_layer_norm_zero: | |
| ff_output = gate_mlp.unsqueeze(1) * ff_output | |
| hidden_states = ff_output + hidden_states | |
| if return_cross_attention_probs and self.attn2 is not None: | |
| return hidden_states, cross_attention_probs | |
| return hidden_states | |
| class FeedForward(nn.Module): | |
| r""" | |
| A feed-forward layer. | |
| Parameters: | |
| dim (`int`): The number of channels in the input. | |
| dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. | |
| mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. | |
| dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. | |
| activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. | |
| final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. | |
| """ | |
| def __init__( | |
| self, | |
| dim: int, | |
| dim_out: Optional[int] = None, | |
| mult: int = 4, | |
| dropout: float = 0.0, | |
| activation_fn: str = "geglu", | |
| final_dropout: bool = False, | |
| ): | |
| super().__init__() | |
| inner_dim = int(dim * mult) | |
| dim_out = dim_out if dim_out is not None else dim | |
| if activation_fn == "gelu": | |
| act_fn = GELU(dim, inner_dim) | |
| if activation_fn == "gelu-approximate": | |
| act_fn = GELU(dim, inner_dim, approximate="tanh") | |
| elif activation_fn == "geglu": | |
| act_fn = GEGLU(dim, inner_dim) | |
| elif activation_fn == "geglu-approximate": | |
| act_fn = ApproximateGELU(dim, inner_dim) | |
| self.net = nn.ModuleList([]) | |
| # project in | |
| self.net.append(act_fn) | |
| # project dropout | |
| self.net.append(nn.Dropout(dropout)) | |
| # project out | |
| self.net.append(nn.Linear(inner_dim, dim_out)) | |
| # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout | |
| if final_dropout: | |
| self.net.append(nn.Dropout(dropout)) | |
| def forward(self, hidden_states): | |
| for module in self.net: | |
| hidden_states = module(hidden_states) | |
| return hidden_states | |
| class GELU(nn.Module): | |
| r""" | |
| GELU activation function with tanh approximation support with `approximate="tanh"`. | |
| """ | |
| def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"): | |
| super().__init__() | |
| self.proj = nn.Linear(dim_in, dim_out) | |
| self.approximate = approximate | |
| def gelu(self, gate): | |
| if gate.device.type != "mps": | |
| return F.gelu(gate, approximate=self.approximate) | |
| # mps: gelu is not implemented for float16 | |
| return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(dtype=gate.dtype) | |
| def forward(self, hidden_states): | |
| hidden_states = self.proj(hidden_states) | |
| hidden_states = self.gelu(hidden_states) | |
| return hidden_states | |
| class GEGLU(nn.Module): | |
| r""" | |
| A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202. | |
| Parameters: | |
| dim_in (`int`): The number of channels in the input. | |
| dim_out (`int`): The number of channels in the output. | |
| """ | |
| def __init__(self, dim_in: int, dim_out: int): | |
| super().__init__() | |
| self.proj = nn.Linear(dim_in, dim_out * 2) | |
| def gelu(self, gate): | |
| if gate.device.type != "mps": | |
| return F.gelu(gate) | |
| # mps: gelu is not implemented for float16 | |
| return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype) | |
| def forward(self, hidden_states): | |
| hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1) | |
| return hidden_states * self.gelu(gate) | |
| class ApproximateGELU(nn.Module): | |
| """ | |
| The approximate form of Gaussian Error Linear Unit (GELU) | |
| For more details, see section 2: https://arxiv.org/abs/1606.08415 | |
| """ | |
| def __init__(self, dim_in: int, dim_out: int): | |
| super().__init__() | |
| self.proj = nn.Linear(dim_in, dim_out) | |
| def forward(self, x): | |
| x = self.proj(x) | |
| return x * torch.sigmoid(1.702 * x) | |
| class AdaLayerNorm(nn.Module): | |
| """ | |
| Norm layer modified to incorporate timestep embeddings. | |
| """ | |
| def __init__(self, embedding_dim, num_embeddings): | |
| super().__init__() | |
| self.emb = nn.Embedding(num_embeddings, embedding_dim) | |
| self.silu = nn.SiLU() | |
| self.linear = nn.Linear(embedding_dim, embedding_dim * 2) | |
| self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False) | |
| def forward(self, x, timestep): | |
| emb = self.linear(self.silu(self.emb(timestep))) | |
| scale, shift = torch.chunk(emb, 2) | |
| x = self.norm(x) * (1 + scale) + shift | |
| return x | |
| class AdaLayerNormZero(nn.Module): | |
| """ | |
| Norm layer adaptive layer norm zero (adaLN-Zero). | |
| """ | |
| def __init__(self, embedding_dim, num_embeddings): | |
| super().__init__() | |
| self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim) | |
| self.silu = nn.SiLU() | |
| self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) | |
| self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) | |
| def forward(self, x, timestep, class_labels, hidden_dtype=None): | |
| emb = self.linear(self.silu(self.emb(timestep, class_labels, hidden_dtype=hidden_dtype))) | |
| shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1) | |
| x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] | |
| return x, gate_msa, shift_mlp, scale_mlp, gate_mlp | |