content: # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py
  
  from dataclasses import dataclass
  from typing import Optional
  
  import torch
  import torch.nn.functional as F
  from torch import nn
  
  from diffusers.configuration_utils import ConfigMixin, register_to_config
  from diffusers import ModelMixin
  from diffusers.utils import BaseOutput
  from diffusers.utils.import_utils import is_xformers_available
  from diffusers.models.attention import CrossAttention, FeedForward, AdaLayerNorm
  
  from einops import rearrange, repeat
  
  
  @dataclass
  class Transformer3DModelOutput(BaseOutput):
      sample: torch.FloatTensor
  
  
  if is_xformers_available():
      import xformers
      import xformers.ops
  else:
      xformers = None
  
  
  class Transformer3DModel(ModelMixin, ConfigMixin):
      @register_to_config
      def __init__(
          self,
          num_attention_heads: int = 16,
          attention_head_dim: int = 88,
          in_channels: Optional[int] = None,
          num_layers: int = 1,
          dropout: float = 0.0,
          norm_num_groups: int = 32,
          cross_attention_dim: Optional[int] = None,
          attention_bias: bool = False,
          activation_fn: str = "geglu",
          num_embeds_ada_norm: Optional[int] = None,
          use_linear_projection: bool = False,
          only_cross_attention: bool = False,
          upcast_attention: bool = False,
          use_attn_temp: bool = False,
      ):
          super().__init__()
          self.use_linear_projection = use_linear_projection
          self.num_attention_heads = num_attention_heads
          self.attention_head_dim = attention_head_dim
          inner_dim = num_attention_heads * attention_head_dim
  
          # Define input layers
          self.in_channels = in_channels
  
          self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
          if use_linear_projection:
              self.proj_in = nn.Linear(in_channels, inner_dim)
          else:
              self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
  
          # Define transformers blocks
          self.transformer_blocks = nn.ModuleList(
              [
                  BasicTransformerBlock(
                      inner_dim,
                      num_attention_heads,
                      attention_head_dim,
                      dropout=dropout,
                      cross_attention_dim=cross_attention_dim,
                      activation_fn=activation_fn,
                      num_embeds_ada_norm=num_embeds_ada_norm,
                      attention_bias=attention_bias,
                      only_cross_attention=only_cross_attention,
                      upcast_attention=upcast_attention,
                      use_attn_temp = use_attn_temp,
                  )
                  for d in range(num_layers)
              ]
          )
  
          # 4. Define output layers
          if use_linear_projection:
              self.proj_out = nn.Linear(in_channels, inner_dim)
          else:
              self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
  
      def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, return_dict: bool = True):
          # Input
          assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
          video_length = hidden_states.shape[2]
          hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
          encoder_hidden_states = repeat(encoder_hidden_states, 'b n c -> (b f) n c', f=video_length)
  
          batch, channel, height, weight = hidden_states.shape
          residual = hidden_states
  
          hidden_states = self.norm(hidden_states)
          if not self.use_linear_projection:
              hidden_states = self.proj_in(hidden_states)
              inner_dim = hidden_states.shape[1]
              hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
          else:
              inner_dim = hidden_states.shape[1]
              hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
              hidden_states = self.proj_in(hidden_states)
  
          # Blocks
          for block in self.transformer_blocks:
              hidden_states = block(
                  hidden_states,
                  encoder_hidden_states=encoder_hidden_states,
                  timestep=timestep,
                  video_length=video_length
              )
  
          # Output
          if not self.use_linear_projection:
              hidden_states = (
                  hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
              )
              hidden_states = self.proj_out(hidden_states)
          else:
              hidden_states = self.proj_out(hidden_states)
              hidden_states = (
                  hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
              )
  
          output = hidden_states + residual
  
          output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
          if not return_dict:
              return (output,)
  
          return Transformer3DModelOutput(sample=output)
  
  
  class BasicTransformerBlock(nn.Module):
      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,
          upcast_attention: bool = False,
          use_attn_temp: bool = False
      ):
          super().__init__()
          self.only_cross_attention = only_cross_attention
          self.use_ada_layer_norm = num_embeds_ada_norm is not None
          self.use_attn_temp = use_attn_temp
          # SC-Attn
          self.attn1 = SparseCausalAttention(
              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,
          )
          self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
  
          # Cross-Attn
          if cross_attention_dim is not None:
              self.attn2 = CrossAttention(
                  query_dim=dim,
                  cross_attention_dim=cross_attention_dim,
                  heads=num_attention_heads,
                  dim_head=attention_head_dim,
                  dropout=dropout,
                  bias=attention_bias,
                  upcast_attention=upcast_attention,
              )
          else:
              self.attn2 = None
  
          if cross_attention_dim is not None:
              self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
          else:
              self.norm2 = None
  
          # Feed-forward
          self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
          self.norm3 = nn.LayerNorm(dim)
  
          # Temp-Attn
          if self.use_attn_temp:
              self.attn_temp = CrossAttention(
                  query_dim=dim,
                  heads=num_attention_heads,
                  dim_head=attention_head_dim,
                  dropout=dropout,
                  bias=attention_bias,
                  upcast_attention=upcast_attention,
              )
              nn.init.zeros_(self.attn_temp.to_out[0].weight.data)
              self.norm_temp = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
  
      def set_use_memory_efficient_attention_xformers(self, use_memory_efficient_attention_xformers: bool):
          if not is_xformers_available():
              print("Here is how to install it")
              raise ModuleNotFoundError(
                  "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
                  " xformers",
                  name="xformers",
              )
          elif not torch.cuda.is_available():
              raise ValueError(
                  "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is only"
                  " available for GPU "
              )
          else:
              try:
                  # Make sure we can run the memory efficient attention
                  _ = xformers.ops.memory_efficient_attention(
                      torch.randn((1, 2, 40), device="cuda"),
                      torch.randn((1, 2, 40), device="cuda"),
                      torch.randn((1, 2, 40), device="cuda"),
                  )
              except Exception as e:
                  raise e
              self.attn1._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
              if self.attn2 is not None:
                  self.attn2._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
              #self.attn_temp._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
  
      def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, attention_mask=None, video_length=None):
          # SparseCausal-Attention
          norm_hidden_states = (
              self.norm1(hidden_states, timestep) if self.use_ada_layer_norm else self.norm1(hidden_states)
          )
  
          if self.only_cross_attention:
              hidden_states = (
                  self.attn1(norm_hidden_states, encoder_hidden_states, attention_mask=attention_mask) + hidden_states
              )
          else:
              hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, video_length=video_length) + hidden_states
  
          if self.attn2 is not None:
              # Cross-Attention
              norm_hidden_states = (
                  self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
              )
              hidden_states = (
                  self.attn2(
                      norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
                  )
                  + hidden_states
              )
  
          # Feed-forward
          hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
  
          # Temporal-Attention
          if self.use_attn_temp:
              d = hidden_states.shape[1]
              hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=video_length)
              norm_hidden_states = (
                  self.norm_temp(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_temp(hidden_states)
              )
              hidden_states = self.attn_temp(norm_hidden_states) + hidden_states
              hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
  
          return hidden_states
  
  
  class SparseCausalAttention(CrossAttention):
      def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None, use_full_attn=True):
          batch_size, sequence_length, _ = hidden_states.shape
  
          encoder_hidden_states = encoder_hidden_states
  
          if self.group_norm is not None:
              hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
  
          query = self.to_q(hidden_states)
          # query = rearrange(query, "(b f) d c -> b (f d) c", f=video_length)
          dim = query.shape[-1]
          query = self.reshape_heads_to_batch_dim(query)
  
          if self.added_kv_proj_dim is not None:
              raise NotImplementedError
  
          encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
          key = self.to_k(encoder_hidden_states)
          value = self.to_v(encoder_hidden_states)
  
          former_frame_index = torch.arange(video_length) - 1
          former_frame_index[0] = 0
  
          key = rearrange(key, "(b f) d c -> b f d c", f=video_length)
          if not use_full_attn:
              key = torch.cat([key[:, [0] * video_length], key[:, former_frame_index]], dim=2)
          else:
              # key = torch.cat([key[:, [0] * video_length], key[:, [1] * video_length], key[:, [2] * video_length], key[:, [3] * video_length]], dim=2)
              key_video_length = [key[:, [i] * video_length] for i in range(video_length)]
              key = torch.cat(key_video_length, dim=2)
          key = rearrange(key, "b f d c -> (b f) d c")
  
          value = rearrange(value, "(b f) d c -> b f d c", f=video_length)
          if not use_full_attn:
              value = torch.cat([value[:, [0] * video_length], value[:, former_frame_index]], dim=2)
          else:
              # value = torch.cat([value[:, [0] * video_length], value[:, [1] * video_length], value[:, [2] * video_length], value[:, [3] * video_length]], dim=2)
              value_video_length = [value[:, [i] * video_length] for i in range(video_length)]
              value = torch.cat(value_video_length, dim=2)
          value = rearrange(value, "b f d c -> (b f) d c")
  
          key = self.reshape_heads_to_batch_dim(key)
          value = self.reshape_heads_to_batch_dim(value)
  
          if attention_mask is not None:
              if attention_mask.shape[-1] != query.shape[1]:
                  target_length = query.shape[1]
                  attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
                  attention_mask = attention_mask.repeat_interleave(self.heads, dim=0)
  
          # attention, what we cannot get enough of
          if self._use_memory_efficient_attention_xformers:
              hidden_states = self._memory_efficient_attention_xformers(query, key, value, attention_mask)
              # Some versions of xformers return output in fp32, cast it back to the dtype of the input
              hidden_states = hidden_states.to(query.dtype)
          else:
              if self._slice_size is None or query.shape[0] // self._slice_size == 1:
                  hidden_states = self._attention(query, key, value, attention_mask)
              else:
                  hidden_states = self._sliced_attention(query, key, value, sequence_length, dim, attention_mask)
  
          # linear proj
          hidden_states = self.to_out[0](hidden_states)
  
          # dropout
          hidden_states = self.to_out[1](hidden_states)
          return hidden_states
  