content: # 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.
  #
  # --------
  #
  # Modified 2024 by the Tripo AI and Stability AI Team.
  #
  # Copyright (c) 2024 Tripo AI & Stability AI
  #
  # Permission is hereby granted, free of charge, to any person obtaining a copy
  # of this software and associated documentation files (the "Software"), to deal
  # in the Software without restriction, including without limitation the rights
  # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  # copies of the Software, and to permit persons to whom the Software is
  # furnished to do so, subject to the following conditions:
  #
  # The above copyright notice and this permission notice shall be included in all
  # copies or substantial portions of the Software.
  #
  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  # SOFTWARE.
  
  from typing import Optional
  
  import torch
  import torch.nn.functional as F
  from torch import nn
  
  from .attention import Attention
  
  
  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.
          activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
          attention_bias (:
              obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
          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.
          upcast_attention (`bool`, *optional*):
              Whether to upcast the attention computation to float32. This is useful for mixed precision training.
          norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
              Whether to use learnable elementwise affine parameters for normalization.
          norm_type (`str`, *optional*, defaults to `"layer_norm"`):
              The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
          final_dropout (`bool` *optional*, defaults to False):
              Whether to apply a final dropout after the last feed-forward layer.
      """
  
      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",
          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,
      ):
          super().__init__()
          self.only_cross_attention = only_cross_attention
  
          assert norm_type == "layer_norm"
  
          # Define 3 blocks. Each block has its own normalization layer.
          # 1. Self-Attn
          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 = 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,
          )
  
          # let chunk size default to None
          self._chunk_size = None
          self._chunk_dim = 0
  
      def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
          # Sets chunk feed-forward
          self._chunk_size = chunk_size
          self._chunk_dim = 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,
      ) -> torch.FloatTensor:
          # Notice that normalization is always applied before the real computation in the following blocks.
          # 0. Self-Attention
          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,
          )
  
          hidden_states = attn_output + hidden_states
  
          # 3. Cross-Attention
          if self.attn2 is not None:
              norm_hidden_states = self.norm2(hidden_states)
  
              attn_output = self.attn2(
                  norm_hidden_states,
                  encoder_hidden_states=encoder_hidden_states,
                  attention_mask=encoder_attention_mask,
              )
              hidden_states = attn_output + hidden_states
  
          # 4. Feed-forward
          norm_hidden_states = self.norm3(hidden_states)
  
          if self._chunk_size is not None:
              # "feed_forward_chunk_size" can be used to save memory
              if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
                  raise ValueError(
                      f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
                  )
  
              num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
              ff_output = torch.cat(
                  [
                      self.ff(hid_slice)
                      for hid_slice in norm_hidden_states.chunk(
                          num_chunks, dim=self._chunk_dim
                      )
                  ],
                  dim=self._chunk_dim,
              )
          else:
              ff_output = self.ff(norm_hidden_states)
  
          hidden_states = ff_output + hidden_states
  
          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
          linear_cls = nn.Linear
  
          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(linear_cls(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: torch.Tensor) -> torch.Tensor:
          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"`.
  
      Parameters:
          dim_in (`int`): The number of channels in the input.
          dim_out (`int`): The number of channels in the output.
          approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
      """
  
      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: torch.Tensor) -> torch.Tensor:
          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__()
          linear_cls = nn.Linear
  
          self.proj = linear_cls(dim_in, dim_out * 2)
  
      def gelu(self, gate: torch.Tensor) -> torch.Tensor:
          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, scale: float = 1.0):
          args = ()
          hidden_states, gate = self.proj(hidden_states, *args).chunk(2, dim=-1)
          return hidden_states * self.gelu(gate)
  
  
  class ApproximateGELU(nn.Module):
      r"""
      The approximate form of Gaussian Error Linear Unit (GELU). For more details, see section 2:
      https://arxiv.org/abs/1606.08415.
  
      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)
  
      def forward(self, x: torch.Tensor) -> torch.Tensor:
          x = self.proj(x)
          return x * torch.sigmoid(1.702 * x)
  