text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from vit_pytorch.face_losses import CosFace
  
  # Pre-defined CCT Models
  __all__ = ['cct_2', 'cct_4', 'cct_6', 'cct_7', 'cct_8', 'cct_14', 'cct_16']
  
  
  def cct_2(*args, **kwargs):
      return _cct(num_layers=2, num_heads=2, mlp_ratio=1, embedding_dim=128,
                  *args, **kwargs)
  
  
  def cct_4(*args, **kwargs):
      return _cct(num_layers=4, num_heads=2, mlp_ratio=1, embedding_dim=128,
                  *args, **kwargs)
  
  
  def cct_6(*args, **kwargs):
      return _cct(num_layers=6, num_heads=4, mlp_ratio=2, embedding_dim=256,
                  *args, **kwargs)
  
  
  def cct_7(*args, **kwargs):
      return _cct(num_layers=7, num_heads=4, mlp_ratio=2, embedding_dim=256,
                  *args, **kwargs)
  
  
  def cct_8(*args, **kwargs):
      return _cct(num_layers=8, num_heads=4, mlp_ratio=2, embedding_dim=256,
                  *args, **kwargs)
  
  
  def cct_14(*args, **kwargs):
      return _cct(num_layers=14, num_heads=6, mlp_ratio=3, embedding_dim=384,
                  *args, **kwargs)
  
  
  def cct_16(*args, **kwargs):
      return _cct(num_layers=16, num_heads=6, mlp_ratio=3, embedding_dim=384,
                  *args, **kwargs)
  
  
  def _cct(num_layers, num_heads, mlp_ratio, embedding_dim,
           kernel_size=3, stride=None, padding=None,
           *args, **kwargs):
      stride = stride if stride is not None else max(1, (kernel_size // 2) - 1)
      padding = padding if padding is not None else max(1, (kernel_size // 2))
      return CCT(num_layers=num_layers,
                 num_heads=num_heads,
                 mlp_ratio=mlp_ratio,
                 embedding_dim=embedding_dim,
                 kernel_size=kernel_size,
                 stride=stride,
                 padding=padding,
                 *args, **kwargs)
  
  
  class Residual(nn.Module):
      def __init__(self, fn):
          super().__init__()
          self.fn = fn
  
      def forward(self, x, **kwargs):
          return self.fn(x, **kwargs) + x
  
  
  class PreNorm(nn.Module):
      def __init__(self, dim, fn):
          super().__init__()
          self.norm = nn.LayerNorm(dim)
          self.fn = fn
  
      def forward(self, x, **kwargs):
          return self.fn(self.norm(x), **kwargs)
  
  
  class FeedForward(nn.Module):
      def __init__(self, dim, hidden_dim, dropout=0.):
          super().__init__()
          self.net = nn.Sequential(
              nn.Linear(dim, hidden_dim),
              nn.GELU(),
              nn.Dropout(dropout),
              nn.Linear(hidden_dim, dim),
              nn.Dropout(dropout)
          )
  
      def forward(self, x):
          return self.net(x)
  
  
  class Attention(nn.Module):
      def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
          super().__init__()
          inner_dim = dim_head * heads
          self.heads = heads
          self.scale = dim ** -0.5
  
          self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
          self.to_out = nn.Sequential(
              nn.Linear(inner_dim, dim),
              nn.Dropout(dropout)
          )
  
      def forward(self, x, mask=None):
          b, n, _, h = *x.shape, self.heads
          qkv = self.to_qkv(x).chunk(3, dim=-1)
  
          q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv)
          dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale
          mask_value = -torch.finfo(dots.dtype).max
          # embed()
          if mask is not None:
              mask = F.pad(mask.flatten(1), (1, 0), value=True)
              assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'  # Something is changing here
              mask = mask[:, None, :] * mask[:, :, None]
              dots.masked_fill_(~mask, mask_value)
              del mask
  
          attn = dots.softmax(dim=-1)
  
          out = torch.einsum('bhij,bhjd->bhid', attn, v)
          out = rearrange(out, 'b h n d -> b n (h d)')
          out = self.to_out(out)
  
          return out
  
  
  def drop_path(x, drop_prob: float = 0., training: bool = False):
      """
      Obtained from: github.com:rwightman/pytorch-image-models
      Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
      This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
      the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
      See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
      changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
      'survival rate' as the argument.
      """
      if drop_prob == 0. or not training:
          return x
      keep_prob = 1 - drop_prob
      shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
      random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
      random_tensor.floor_()  # binarize
      output = x.div(keep_prob) * random_tensor
      return output
  
  
  class Transformer(nn.Module):
      def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout):
          super().__init__()
          self.layers = nn.ModuleList([])
          for _ in range(depth):
              self.layers.append(nn.ModuleList([
                  Residual(PreNorm(dim, Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout))),
                  Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout)))
              ]))
  
      def forward(self, x, mask=None):
          for attn, ff in self.layers:
              x = attn(x, mask=mask)
              # embed()
              x = ff(x)
          return x
  
  
  class DropPath(nn.Module):
      """
      Obtained from: github.com:rwightman/pytorch-image-models
      Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
      """
  
      def __init__(self, drop_prob=None):
          super(DropPath, self).__init__()
          self.drop_prob = drop_prob
  
      def forward(self, x):
          return drop_path(x, self.drop_prob, self.training)
  
  
  class Tokenizer(nn.Module):
      def __init__(self,
                   kernel_size, stride, padding,
                   pooling_kernel_size=3, pooling_stride=2, pooling_padding=1,
                   n_conv_layers=1,
                   n_input_channels=3,
                   n_output_channels=64,
                   in_planes=64,
                   activation=None,
                   max_pool=True,
                   conv_bias=False):
          super(Tokenizer, self).__init__()
  
          n_filter_list = [n_input_channels] + \
                          [in_planes for _ in range(n_conv_layers - 1)] + \
                          [n_output_channels]
  
          self.conv_layers = nn.Sequential(
              *[nn.Sequential(
                  nn.Conv2d(n_filter_list[i], n_filter_list[i + 1],
                            kernel_size=(kernel_size, kernel_size),
                            stride=(stride, stride),
                            padding=(padding, padding), bias=conv_bias),
                  nn.Identity() if activation is None else activation(),
                  nn.MaxPool2d(kernel_size=pooling_kernel_size,
                               stride=pooling_stride,
                               padding=pooling_padding) if max_pool else nn.Identity()
              )
                  for i in range(n_conv_layers)
              ])
  
          self.flattener = nn.Flatten(2, 3)
          self.apply(self.init_weight)
  
      def sequence_length(self, n_channels=3, height=224, width=224):
          return self.forward(torch.zeros((1, n_channels, height, width))).shape[1]
  
      def forward(self, x):
          return self.flattener(self.conv_layers(x)).transpose(-2, -1)
  
      @staticmethod
      def init_weight(m):
          if isinstance(m, nn.Conv2d):
              nn.init.kaiming_normal_(m.weight)
  
  
  class TransformerClassifier(nn.Module):
      def __init__(self,
                   loss_type,
                   GPU_ID,
                   seq_pool=True,
                   embedding_dim=768,
                   num_layers=12,
                   num_heads=12,
                   mlp_ratio=4.0,
                   num_classes=1000,
                   dropout_rate=0.1,
                   attention_dropout=0.1,
                   stochastic_depth_rate=0.1,
                   positional_embedding='sine',
                   sequence_length=None,
                   *args, **kwargs):
          super().__init__()
          positional_embedding = positional_embedding if \
              positional_embedding in ['sine', 'learnable', 'none'] else 'sine'
  
          dim_feedforward = int(embedding_dim * mlp_ratio)
          self.embedding_dim = embedding_dim
          self.sequence_length = sequence_length
          self.seq_pool = seq_pool
  
          assert sequence_length is not None or positional_embedding == 'none', \
              f"Positional embedding is set to {positional_embedding} and" \
              f" the sequence length was not specified."
  
          if not seq_pool:
              sequence_length += 1
              self.class_emb = nn.Parameter(torch.zeros(1, 1, self.embedding_dim),
                                            requires_grad=True)
          else:
              self.attention_pool = nn.Linear(self.embedding_dim, 1)
  
          if positional_embedding != 'none':
              if positional_embedding == 'learnable':
                  self.positional_emb = nn.Parameter(torch.zeros(1, sequence_length, embedding_dim),
                                                     requires_grad=True)
                  nn.init.trunc_normal_(self.positional_emb, std=0.2)
              else:
                  self.positional_emb = nn.Parameter(self.sinusoidal_embedding(sequence_length, embedding_dim),
                                                     requires_grad=False)
          else:
              self.positional_emb = None
  
          self.dropout = nn.Dropout(p=dropout_rate)
          dpr = [x.item() for x in torch.linspace(0, stochastic_depth_rate, num_layers)]
          self.blocks = nn.ModuleList([
              Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
              for i in range(num_layers)])
  
          self.norm = nn.Sequential = nn.Sequential(
              nn.LayerNorm(embedding_dim),
          )
          self.loss_type = loss_type
          self.GPU_ID = GPU_ID
          if self.loss_type == 'None':
              print("no loss for vit_face")
          else:
              if self.loss_type == 'Softmax':
                  self.loss = Softmax(in_features=embedding_dim, out_features=num_classes, device_id=self.GPU_ID)
              elif self.loss_type == 'CosFace':
                  self.loss = CosFace(in_features=embedding_dim, out_features=num_classes, device_id=self.GPU_ID)
              elif self.loss_type == 'ArcFace':
                  self.loss = ArcFace(in_features=embedding_dim, out_features=num_classes, device_id=self.GPU_ID)
              elif self.loss_type == 'SFace':
                  self.loss = SFaceLoss(in_features=embedding_dim, out_features=num_classes, device_id=self.GPU_ID)
          self.apply(self.init_weight)
  
      def forward(self, x, label=None, mask=None):
          if self.positional_emb is None and x.size(1) < self.sequence_length:
              x = F.pad(x, (0, 0, 0, self.n_channels - x.size(1)), mode='constant', value=0)
  
          if not self.seq_pool:
              cls_token = self.class_emb.expand(x.shape[0], -1, -1)
              x = torch.cat((cls_token, x), dim=1)
  
          if self.positional_emb is not None:
              x += self.positional_emb
  
          x = self.dropout(x)
  
          for blk in self.blocks:
              x = blk(x, mask)
  
          x = self.norm(x)
  
          if self.seq_pool:
              x = torch.matmul(F.softmax(self.attention_pool(x), dim=1).transpose(-1, -2), x).squeeze(-2)
          else:
              x = x[:, 0]
  
          emb = self.norm(x)
          if label is not None:
              x = self.loss(x, label)
              return x, emb
          else:
              return emb
  
      @staticmethod
      def init_weight(m):
          if isinstance(m, nn.Linear):
              nn.init.trunc_normal_(m.weight, std=.02)
              if isinstance(m, nn.Linear) and m.bias is not None:
                  nn.init.constant_(m.bias, 0)
          elif isinstance(m, nn.LayerNorm):
              nn.init.constant_(m.bias, 0)
              nn.init.constant_(m.weight, 1.0)
  
      @staticmethod
      def sinusoidal_embedding(n_channels, dim):
          pe = torch.FloatTensor([[p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
                                  for p in range(n_channels)])
          pe[:, 0::2] = torch.sin(pe[:, 0::2])
          pe[:, 1::2] = torch.cos(pe[:, 1::2])
          return pe.unsqueeze(0)
  
  
  # CCT Main model
  class CCT(nn.Module):
      def __init__(self,
                   img_size=224,
                   embedding_dim=768,
                   n_input_channels=3,
                   n_conv_layers=1,
                   kernel_size=7,
                   stride=2,
                   padding=3,
                   pooling_kernel_size=3,
                   pooling_stride=2,
                   pooling_padding=1,
                   *args, **kwargs):
          super(CCT, self).__init__()
  
          self.tokenizer = Tokenizer(n_input_channels=n_input_channels,
                                     n_output_channels=embedding_dim,
                                     kernel_size=kernel_size,
                                     stride=stride,
                                     padding=padding,
                                     pooling_kernel_size=pooling_kernel_size,
                                     pooling_stride=pooling_stride,
                                     pooling_padding=pooling_padding,
                                     max_pool=True,
                                     activation=nn.ReLU,
                                     n_conv_layers=n_conv_layers,
                                     conv_bias=False)
  
          self.classifier = TransformerClassifier(
              sequence_length=self.tokenizer.sequence_length(n_channels=n_input_channels,
                                                             height=img_size,
                                                             width=img_size),
              embedding_dim=embedding_dim,
              seq_pool=True,
              dropout_rate=0.,
              attention_dropout=0.1,
              stochastic_depth=0.1,
              *args, **kwargs)
  
      def forward(self, x, mask=None):
          x = self.tokenizer(x)
          return self.classifier(x)
  