content: from dataclasses import dataclass
  from typing import Optional
  
  import torch
  import torch.nn.functional as F
  from torch import nn
  
  from StableFast3D.sf3d.models.utils import BaseModule
  
  
  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: 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 CrossAttention(nn.Module):
      def __init__(
          self,
          dim,
          kv_dim=None,
          num_heads=16,
          qkv_bias=False,
          attn_drop=0.0,
          proj_drop=0.0,
      ):
          super().__init__()
          self.num_heads = num_heads
          head_dim = dim // num_heads
          self.scale = head_dim**-0.5
          kv_dim = dim if not kv_dim else kv_dim
          self.wq = nn.Linear(dim, dim, bias=qkv_bias)
          self.wk = nn.Linear(kv_dim, dim, bias=qkv_bias)
          self.wv = nn.Linear(kv_dim, dim, bias=qkv_bias)
          self.attn_drop = attn_drop
          self.proj = nn.Linear(dim, dim)
          self.proj_drop = nn.Dropout(proj_drop)
  
      def forward(self, x_q, x_kv):
          B, N_q, C = x_q.shape
          B, N_kv, _ = x_kv.shape
          # [B, N_q, C] -> [B, N_q, H, C/H]
          q = self.wq(x_q).reshape(B, N_q, self.num_heads, C // self.num_heads)
          # [B, N_kv, C] -> [B, N_kv, H, C/H]
          k = self.wk(x_kv).reshape(B, N_kv, self.num_heads, C // self.num_heads)
          v = self.wv(x_kv).reshape(B, N_kv, self.num_heads, C // self.num_heads)
  
          #  attention
          x = torch.nn.functional.scaled_dot_product_attention(
              q.permute(0, 2, 1, 3),
              k.permute(0, 2, 1, 3),
              v.permute(0, 2, 1, 3),
              attn_mask=None,
              dropout_p=self.attn_drop,
              scale=self.scale,
          ).permute(0, 2, 1, 3)
  
          # [B, N_q, H, C/H] -> [B, N_q, C]
          x = x.reshape(B, N_q, C)
          x = self.proj(x)
          x = self.proj_drop(x)
          return x
  
  
  class FeedForward(nn.Module):
      def __init__(
          self,
          dim: int,
          dim_out: Optional[int] = None,
          mult: int = 4,
          dropout: float = 0.0,
      ):
          super().__init__()
          inner_dim = int(dim * mult)
          dim_out = dim_out if dim_out is not None else dim
          act_fn = GEGLU(dim, inner_dim)
          self.net = nn.ModuleList([])
          self.net.append(act_fn)
          self.net.append(nn.Dropout(dropout))
          self.net.append(nn.Linear(inner_dim, dim_out))
  
      def forward(self, x: torch.Tensor) -> torch.Tensor:
          for module in self.net:
              x = module(x)
          return x
  
  
  class BasicBlock(nn.Module):
      def __init__(
          self,
          dim: int,
          kv_dim: Optional[int] = None,
          num_heads: int = 16,
          qkv_bias: bool = False,
          attn_drop: float = 0.0,
          proj_drop: float = 0.0,
          ff_drop: float = 0.0,
      ):
          super().__init__()
          self.norm1 = nn.LayerNorm(dim)
          self.attn1 = CrossAttention(
              dim,
              kv_dim=dim,
              num_heads=num_heads,
              qkv_bias=qkv_bias,
              attn_drop=attn_drop,
              proj_drop=proj_drop,
          )
          self.norm2 = nn.LayerNorm(dim)
          self.attn2 = CrossAttention(
              dim,
              kv_dim=kv_dim,
              num_heads=num_heads,
              qkv_bias=qkv_bias,
              attn_drop=attn_drop,
              proj_drop=proj_drop,
          )
          self.norm3 = nn.LayerNorm(dim)
          self.ff = FeedForward(dim, dropout=ff_drop)
  
      def forward(self, z, x):
          z_norm = self.norm1(z)
          z = z + self.attn1(z_norm, z_norm)
          # TODO: do we need to have the second attention when x is None?
          z_norm = self.norm2(z)
          z = z + self.attn2(z_norm, x if x is not None else z_norm)
          z_norm = self.norm3(z)
          z = z + self.ff(z_norm)
          return z
  
  
  class SingleStreamTransformer(BaseModule):
      @dataclass
      class Config(BaseModule.Config):
          num_attention_heads: int = 16
          attention_head_dim: int = 88
          in_channels: Optional[int] = None
          out_channels: Optional[int] = None
          num_layers: int = 16
          dropout: float = 0.0
          norm_num_groups: int = 32
          cross_attention_dim: Optional[int] = None
          attention_bias: bool = False
  
      cfg: Config
  
      def configure(self) -> None:
          self.num_attention_heads = self.cfg.num_attention_heads
          self.attention_head_dim = self.cfg.attention_head_dim
          inner_dim = self.num_attention_heads * self.attention_head_dim
  
          # Define input layers
          self.norm = torch.nn.GroupNorm(
              num_groups=self.cfg.norm_num_groups,
              num_channels=self.cfg.in_channels,
              eps=1e-6,
              affine=True,
          )
          self.proj_in = nn.Linear(self.cfg.in_channels, inner_dim)
  
          # Define transformers blocks
          self.transformer_blocks = nn.ModuleList(
              [
                  BasicBlock(
                      inner_dim,
                      kv_dim=self.cfg.cross_attention_dim,
                      num_heads=self.num_attention_heads,
                      qkv_bias=self.cfg.attention_bias,
                      proj_drop=self.cfg.dropout,
                      ff_drop=self.cfg.dropout,
                  )
                  for d in range(self.cfg.num_layers)
              ]
          )
  
          # 4. Define output layers
          self.proj_out = nn.Linear(inner_dim, self.cfg.in_channels)
  
      def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
          residual = hidden_states
          hidden_states = self.norm(hidden_states)
          hidden_states = hidden_states.permute(0, 2, 1)
          hidden_states = self.proj_in(hidden_states)
          for block in self.transformer_blocks:
              hidden_states = block(hidden_states, encoder_hidden_states)
          hidden_states = self.proj_out(hidden_states).permute(0, 2, 1).contiguous()
          # TODO: do we really need to add the residual?
          hidden_states = hidden_states + residual
          return hidden_states
  
  
  class FuseBlock(nn.Module):
      """
      Fuse X in to Z with cross attention
      """
  
      def __init__(
          self,
          dim_z: int,
          dim_x: int,
          num_heads: int = 16,
          qkv_bias: bool = False,
          attn_drop: float = 0.0,
          proj_drop: float = 0.0,
          ff_drop: float = 0.0,
          norm_x_input: bool = True,
      ):
          super().__init__()
          self.norm_x_input = norm_x_input
          if self.norm_x_input:
              self.norm_x = nn.LayerNorm(dim_x)
          self.attn = CrossAttention(
              dim_z,
              kv_dim=dim_x,
              num_heads=num_heads,
              qkv_bias=qkv_bias,
              attn_drop=attn_drop,
              proj_drop=proj_drop,
          )
          self.norm_z1 = nn.LayerNorm(dim_z)
          self.norm_z2 = nn.LayerNorm(dim_z)
          self.ff = FeedForward(dim_z, dropout=ff_drop)
  
      def forward(self, z, x):
          # TODO: do we need to normalize x?
          z = z + self.attn(self.norm_z1(z), self.norm_x(x) if self.norm_x_input else x)
          z = z + self.ff(self.norm_z2(z))
          return z
  
  
  @torch.no_grad()
  def get_triplane_attention_mask(res):
      N = 3 * res * res
      attn_mask = torch.zeros(3, res, res, 3, res, res)
  
      i, j = torch.meshgrid(torch.arange(res), torch.arange(res))
  
      attn_mask[0, i, j, 1, i, :] = 1.0
      attn_mask[0, i, j, 2, j, :] = 1.0
      attn_mask[1, i, j, 0, i, :] = 1.0
      attn_mask[1, i, j, 2, :, j] = 1.0
      attn_mask[2, i, j, 0, :, i] = 1.0
      attn_mask[2, i, j, 1, :, j] = 1.0
      attn_mask = attn_mask.bool()
  
      attn_bias = torch.empty_like(attn_mask, dtype=torch.float)
      attn_bias.masked_fill_(attn_mask, 0.0)
      attn_bias.masked_fill_(~attn_mask, float("-inf"))
  
      return attn_bias.reshape(N, N)
  
  
  class TriplaneAttention(nn.Module):
      def __init__(
          self,
          dim: int,
          resolution: int,
          num_heads: int = 16,
          qkv_bias: bool = False,
          attn_drop: float = 0.0,
          proj_drop: float = 0.0,
          full_attention: bool = False,
      ):
          super().__init__()
          self.num_heads = num_heads
          head_dim = dim // num_heads
          self.scale = head_dim**-0.5
          self.wq = nn.Linear(dim, dim, bias=qkv_bias)
          self.wk = nn.Linear(dim, dim, bias=qkv_bias)
          self.wv = nn.Linear(dim, dim, bias=qkv_bias)
          self.attn_drop = attn_drop
          self.proj = nn.Linear(dim, dim)
          self.proj_drop = nn.Dropout(proj_drop)
  
          self.resolution = resolution
          self.full_attention = full_attention
          self.attn_mask = (
              get_triplane_attention_mask(resolution) if not full_attention else None
          )
  
      def forward(self, x):
          B, N, C = x.shape
          # [B, N, C] -> [B, N, H, C/H]
          q = self.wq(x).reshape(B, N, self.num_heads, C // self.num_heads)
          k = self.wk(x).reshape(B, N, self.num_heads, C // self.num_heads)
          v = self.wv(x).reshape(B, N, self.num_heads, C // self.num_heads)
  
          # detokenize the planes
          assert N == self.resolution**2 * 3
          attn_bias = (
              self.attn_mask.to(q)
              .unsqueeze(0)
              .unsqueeze(0)
              .expand(B, self.num_heads, -1, -1)
              if not self.full_attention
              else None
          )
  
          # full attention
          x = torch.nn.functional.scaled_dot_product_attention(
              q.permute(0, 2, 1, 3),
              k.permute(0, 2, 1, 3),
              v.permute(0, 2, 1, 3),
              attn_mask=attn_bias,
              dropout_p=self.attn_drop,
              scale=self.scale,
          ).permute(0, 2, 1, 3)
  
          # [B, N_q, H, C/H] -> [B, N_q, C]
          x = x.reshape(B, N, C)
          x = self.proj(x)
          x = self.proj_drop(x)
          return x
  
  
  class TwoStreamBlock(nn.Module):
      def __init__(
          self,
          dim_latent: int,
          dim_input: int,
          num_basic_blocks: int = 4,
          num_heads: int = 16,
          qkv_bias: bool = False,
          attn_drop: float = 0.0,
          proj_drop: float = 0.0,
          ff_drop: float = 0.0,
          norm_x_input: bool = True,
          dim_cross: Optional[int] = None,
      ):
          super().__init__()
  
          # Define the fuse block that fuse the input into the latent
          self.fuse_block_in = FuseBlock(
              dim_latent,
              dim_input,
              num_heads=num_heads,
              qkv_bias=qkv_bias,
              attn_drop=attn_drop,
              proj_drop=proj_drop,
              ff_drop=ff_drop,
              norm_x_input=norm_x_input,
          )
  
          # Define the transformer block that process the latent
          self.transformer_block = nn.ModuleList(
              [
                  BasicBlock(
                      dim_latent,
                      kv_dim=dim_cross,
                      num_heads=num_heads,
                      qkv_bias=qkv_bias,
                      proj_drop=proj_drop,
                      ff_drop=ff_drop,
                  )
                  for _ in range(num_basic_blocks)
              ]
          )
  
          # Define the fuse block that fuse the latent into the input
          self.fuse_block_out = FuseBlock(
              dim_input,
              dim_latent,
              num_heads=num_heads,
              qkv_bias=qkv_bias,
              attn_drop=attn_drop,
              proj_drop=proj_drop,
              ff_drop=ff_drop,
              norm_x_input=norm_x_input,
          )
  
      def forward(self, latent, input, cross_input):
          latent = self.fuse_block_in(latent, input)
          for block in self.transformer_block:
              latent = block(latent, cross_input)
          input = self.fuse_block_out(input, latent)
          return latent, input
  
  
  class TwoStreamInterleaveTransformer(BaseModule):
      @dataclass
      class Config(BaseModule.Config):
          num_attention_heads: int = 16
          attention_head_dim: int = 64
          raw_triplane_channels: int = 1024
          triplane_channels: int = 1024
          raw_image_channels: int = 1024
          num_latents: int = 1792
          num_blocks: int = 4
          num_basic_blocks: int = 3
          dropout: float = 0.0
          latent_init_std: float = 0.02
          norm_num_groups: int = 32
          attention_bias: bool = False
          norm_x_input: bool = False
          cross_attention_dim: int = 1024
          mix_latent: bool = True
  
      cfg: Config
  
      def configure(self) -> None:
          self.mix_latent = self.cfg.mix_latent
  
          # Define the dimensions
          self.num_attention_heads = self.cfg.num_attention_heads
          self.attention_head_dim = self.cfg.attention_head_dim
          self.num_latents = self.cfg.num_latents
          self.latent_dim = self.num_attention_heads * self.attention_head_dim
  
          # Define input layers
          if self.cfg.norm_num_groups > 0:
              self.norm_triplane = torch.nn.GroupNorm(
                  num_groups=self.cfg.norm_num_groups,
                  num_channels=self.cfg.raw_triplane_channels,
                  eps=1e-6,
                  affine=True,
              )
          else:
              self.norm_triplane = nn.LayerNorm(self.cfg.raw_triplane_channels)
          self.proj_triplane = nn.Linear(
              self.cfg.raw_triplane_channels, self.cfg.triplane_channels
          )
          if self.mix_latent:
              self.norm_image = nn.LayerNorm(self.cfg.raw_image_channels)
              self.proj_image = nn.Linear(self.cfg.raw_image_channels, self.latent_dim)
          self.norm_latent = nn.LayerNorm(self.latent_dim)
          self.proj_latent = nn.Linear(self.latent_dim, self.latent_dim)
  
          # Define the latents
          self.latent_init = nn.Parameter(
              torch.zeros(1, self.num_latents, self.latent_dim)
          )
          nn.init.normal_(self.latent_init, std=self.cfg.latent_init_std)
  
          # Define the transformer blocks
          self.main_blocks = nn.ModuleList(
              [
                  TwoStreamBlock(
                      self.latent_dim,
                      self.cfg.triplane_channels,
                      num_basic_blocks=self.cfg.num_basic_blocks,
                      num_heads=self.num_attention_heads,
                      qkv_bias=self.cfg.attention_bias,
                      proj_drop=self.cfg.dropout,
                      ff_drop=self.cfg.dropout,
                      norm_x_input=self.cfg.norm_x_input,
                      dim_cross=self.cfg.cross_attention_dim,
                  )
                  for _ in range(self.cfg.num_blocks)
              ]
          )
  
          # 4. Define output layers
          self.proj_out = nn.Linear(
              self.cfg.triplane_channels, self.cfg.raw_triplane_channels
          )
  
      def forward(self, hidden_states, encoder_hidden_states, **kwargs):
          # hidden_states: [B, triplane_dim, N_triplane] is triplane tokens
          # encoder_hidden_states: [B, N_image, image_dim] is the image tokens
          if isinstance(self.norm_triplane, nn.GroupNorm):
              triplane_tokens = self.norm_triplane(hidden_states)
              triplane_tokens = triplane_tokens.permute(
                  0, 2, 1
              )  # [B, N_triplane, triplane_dim]
          elif isinstance(self.norm_triplane, nn.LayerNorm):
              triplane_tokens = self.norm_triplane(hidden_states.permute(0, 2, 1))
          else:
              raise ValueError("Unknown normalization layer")
          triplane_tokens = self.proj_triplane(triplane_tokens)
          if self.mix_latent:
              image_tokens = self.norm_image(
                  encoder_hidden_states
              )  # [B, N_image, image_dim]
              image_tokens = self.proj_image(image_tokens)
          init_latents = self.latent_init.expand(
              hidden_states.shape[0], -1, -1
          )  # [B, N_latent_init, latent_dim]
          init_latents = self.norm_latent(init_latents)
          init_latents = self.proj_latent(init_latents)
          if self.mix_latent:
              latent_tokens = torch.cat(
                  [image_tokens, init_latents], dim=1
              )  # [B, N_latent, latent_dim]
          else:
              latent_tokens = init_latents
  
          # forward the main blocks
          for block in self.main_blocks:
              latent_tokens, triplane_tokens = block(
                  latent_tokens, triplane_tokens, encoder_hidden_states
              )
  
          # project the triplane tokens back to the original dimension
          triplane_tokens = self.proj_out(triplane_tokens).permute(0, 2, 1).contiguous()
          triplane_tokens = triplane_tokens + hidden_states
          return triplane_tokens
  