text: import torch.nn as nn
  from fairseq.models import FairseqEncoder, register_model, FairseqEncoderDecoderModel, register_model_architecture
  from fairseq.models.transformer import TransformerDecoder, Embedding, TransformerModel
  from fairseq.models.fairseq_encoder import EncoderOut
  from fairseq import utils
  
  # from timm.models.vision_transformer import HybridEmbed, PatchEmbed, Block
  from timm.models.layers import trunc_normal_
  import torch
  from torch.hub import load_state_dict_from_url
  
  from functools import partial
  import logging
  
  logger = logging.getLogger(__name__)
  
  DEFAULT_MAX_TARGET_POSITIONS = 1024
  
  @register_model('ViT_TR')
  class ViTTRModel(FairseqEncoderDecoderModel):
      
      @staticmethod
      def add_args(parser):
          TransformerModel.add_args(parser)
          # parser.add_argument('--decoder-embed-dim', type=int, metavar='N',
          #                     help='decoder embedding dimension')
          parser.add_argument(
              '--vit-img-size', type=int, metavar='N',
              help='the image size of h and w (h=w) of the ViT'
          )
          parser.add_argument(
              '--vit-patch-size', type=int, metavar='N',
              help='the patch size of h and w (h=w) of the ViT'
          )
          parser.add_argument(
              '--vit-dim', type=int, metavar='N',
              help='the hidden size of the ViT'
          )
          parser.add_argument(
              '--vit-depth', type=int, metavar='N',
              help='the layer num of the ViT'
          )
          parser.add_argument(
              '--vit-heads', type=int, metavar='N',
              help='the head num of the ViT'
          )
          parser.add_argument(
              '--vit-channels', type=int, metavar='N', default=3,
              help='the input image channels of the ViT'
          )
          parser.add_argument(
              '--vit-dropout', type=float, default=0.0,
              help='the dropout ratio of the ViT'
          )
          parser.add_argument(
              '--vit-atten-dropout', type=float, default=0.0,
              help='the input embedding dropout ratio of the ViT'
          )
          parser.add_argument(
              '--encoder-pretrained-url', type=str,
              help='the pretrained parameter url for the ViT encoder'
          )
  
  
      @classmethod
      def build_model(cls, args, task):
          encoder = ViTTREncoder(
              args = args,
              dictionary = task.source_dictionary
          )
          if args.encoder_pretrained_url:
              logger.info('load pretrianed encoder parameter from: {}'.format(args.encoder_pretrained_url))
              encoder_state_dict = load_state_dict_from_url(args.encoder_pretrained_url)
              encoder.load_state_dict(encoder_state_dict, strict=False)
  
          if getattr(args, "max_target_positions", None) is None:
              args.max_target_positions = DEFAULT_MAX_TARGET_POSITIONS
  
          decoder_embed_tokens = cls.build_embedding(
              args, task.target_dictionary, args.decoder_embed_dim, args.decoder_embed_path
          )
  
          decoder = TransformerDecoder(
              args = args,
              dictionary=task.target_dictionary,
              embed_tokens=decoder_embed_tokens,
              no_encoder_attn=False
          )
          model = cls(encoder, decoder)
          return model
  
      @classmethod
      def build_embedding(cls, args, dictionary, embed_dim, path=None):
          num_embeddings = len(dictionary)
          padding_idx = dictionary.pad()
  
          emb = Embedding(num_embeddings, embed_dim, padding_idx)
          # if provided, load from preloaded dictionaries
          if path:
              embed_dict = utils.parse_embedding(path)
              utils.load_embedding(embed_dict, dictionary, emb)
          return emb
  
      def forward(self, imgs, prev_output_tokens, **kwargs):
          encoder_out = self.encoder(imgs, **kwargs)
          decoder_out = self.decoder(
              prev_output_tokens, encoder_out=encoder_out, **kwargs
          )
          return decoder_out
  
  
  @register_model_architecture('ViT_TR', 'ViT_TR_base')
  def ViT_TR_base(args):
      # ViT Encoder  vit_base_patch16_224
      args.vit_img_size = getattr(args, "vit_img_size", 224)
      args.resize_img_size = args.vit_img_size
      args.vit_patch_size = getattr(args, "vit_patch_size", 16)
      args.vit_dim = getattr(args, "vit_dim", 768)
      args.vit_depth = getattr(args, "vit_depth", 12)
      args.vit_heads = getattr(args, "vit_heads", 12)
      args.encoder_pretrained_url = getattr(args, "encoder_pretrained_url", 
                      "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth")
  
      # Transformer Decoder
      args.encoder_embed_path = getattr(args, "encoder_embed_path", None)
      args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 768)
      args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 2048)
      args.encoder_layers = getattr(args, "encoder_layers", 6)
      args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 8)
      args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
      args.encoder_learned_pos = getattr(args, "encoder_learned_pos", False)
      args.decoder_embed_path = getattr(args, "decoder_embed_path", None)
      args.decoder_embed_dim = getattr(args, "decoder_embed_dim", args.encoder_embed_dim)
      args.decoder_ffn_embed_dim = getattr(
          args, "decoder_ffn_embed_dim", args.encoder_ffn_embed_dim
      )
      args.decoder_layers = getattr(args, "decoder_layers", 6)
      args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 8)
      args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
      args.decoder_learned_pos = getattr(args, "decoder_learned_pos", False)
      args.attention_dropout = getattr(args, "attention_dropout", 0.0)
      args.activation_dropout = getattr(args, "activation_dropout", 0.0)
      args.activation_fn = getattr(args, "activation_fn", "relu")
      args.dropout = getattr(args, "dropout", 0.1)
      args.adaptive_softmax_cutoff = getattr(args, "adaptive_softmax_cutoff", None)
      args.adaptive_softmax_dropout = getattr(args, "adaptive_softmax_dropout", 0)
      args.share_decoder_input_output_embed = getattr(
          args, "share_decoder_input_output_embed", False
      )
      args.share_all_embeddings = getattr(args, "share_all_embeddings", False)
      args.no_token_positional_embeddings = getattr(
          args, "no_token_positional_embeddings", False
      )
      args.adaptive_input = getattr(args, "adaptive_input", False)
      args.no_cross_attention = getattr(args, "no_cross_attention", False)
      args.cross_self_attention = getattr(args, "cross_self_attention", False)
  
      args.decoder_output_dim = getattr(
          args, "decoder_output_dim", args.decoder_embed_dim
      )
      args.decoder_input_dim = getattr(args, "decoder_input_dim", args.decoder_embed_dim)
  
      args.no_scale_embedding = getattr(args, "no_scale_embedding", False)
      args.layernorm_embedding = getattr(args, "layernorm_embedding", False)
      args.tie_adaptive_weights = getattr(args, "tie_adaptive_weights", False)
      args.checkpoint_activations = getattr(args, "checkpoint_activations", False)
      args.offload_activations = getattr(args, "offload_activations", False)
      if args.offload_activations:
          args.checkpoint_activations = True
      args.encoder_layers_to_keep = getattr(args, "encoder_layers_to_keep", None)
      args.decoder_layers_to_keep = getattr(args, "decoder_layers_to_keep", None)
      args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0)
      args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0)
      args.quant_noise_pq = getattr(args, "quant_noise_pq", 0)
      args.quant_noise_pq_block_size = getattr(args, "quant_noise_pq_block_size", 8)
      args.quant_noise_scalar = getattr(args, "quant_noise_scalar", 0)
  
  @register_model_architecture('ViT_TR', 'ViT_TR_large')
  def large_architecture(args):
      # ViT Encoder  vit_base_patch16_224
      args.vit_img_size = getattr(args, "vit_img_size", 384)
      args.resize_img_size = args.vit_img_size
      args.vit_patch_size = getattr(args, "vit_patch_size", 16)
      args.vit_dim = getattr(args, "vit_dim", 1024)
      args.vit_depth = getattr(args, "vit_depth", 24)
      args.vit_heads = getattr(args, "vit_heads", 16)
      args.encoder_pretrained_url = getattr(args, "encoder_pretrained_url", 
                      "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_384-b3be5167.pth")
  
      # Transformer Decoder
      args.encoder_embed_path = getattr(args, "encoder_embed_path", None)
      args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1024)
      args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 2048)
      args.encoder_layers = getattr(args, "encoder_layers", 6)
      args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 8)
      args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
      args.encoder_learned_pos = getattr(args, "encoder_learned_pos", False)
      args.decoder_embed_path = getattr(args, "decoder_embed_path", None)
      args.decoder_embed_dim = getattr(args, "decoder_embed_dim", args.encoder_embed_dim)
      args.decoder_ffn_embed_dim = getattr(
          args, "decoder_ffn_embed_dim", args.encoder_ffn_embed_dim
      )
      args.decoder_layers = getattr(args, "decoder_layers", 6)
      args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 8)
      args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
      args.decoder_learned_pos = getattr(args, "decoder_learned_pos", False)
      args.attention_dropout = getattr(args, "attention_dropout", 0.0)
      args.activation_dropout = getattr(args, "activation_dropout", 0.0)
      args.activation_fn = getattr(args, "activation_fn", "relu")
      args.dropout = getattr(args, "dropout", 0.1)
      args.adaptive_softmax_cutoff = getattr(args, "adaptive_softmax_cutoff", None)
      args.adaptive_softmax_dropout = getattr(args, "adaptive_softmax_dropout", 0)
      args.share_decoder_input_output_embed = getattr(
          args, "share_decoder_input_output_embed", False
      )
      args.share_all_embeddings = getattr(args, "share_all_embeddings", False)
      args.no_token_positional_embeddings = getattr(
          args, "no_token_positional_embeddings", False
      )
      args.adaptive_input = getattr(args, "adaptive_input", False)
      args.no_cross_attention = getattr(args, "no_cross_attention", False)
      args.cross_self_attention = getattr(args, "cross_self_attention", False)
  
      args.decoder_output_dim = getattr(
          args, "decoder_output_dim", args.decoder_embed_dim
      )
      args.decoder_input_dim = getattr(args, "decoder_input_dim", args.decoder_embed_dim)
  
      args.no_scale_embedding = getattr(args, "no_scale_embedding", False)
      args.layernorm_embedding = getattr(args, "layernorm_embedding", False)
      args.tie_adaptive_weights = getattr(args, "tie_adaptive_weights", False)
      args.checkpoint_activations = getattr(args, "checkpoint_activations", False)
      args.offload_activations = getattr(args, "offload_activations", False)
      if args.offload_activations:
          args.checkpoint_activations = True
      args.encoder_layers_to_keep = getattr(args, "encoder_layers_to_keep", None)
      args.decoder_layers_to_keep = getattr(args, "decoder_layers_to_keep", None)
      args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0)
      args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0)
      args.quant_noise_pq = getattr(args, "quant_noise_pq", 0)
      args.quant_noise_pq_block_size = getattr(args, "quant_noise_pq_block_size", 8)
      args.quant_noise_scalar = getattr(args, "quant_noise_scalar", 0)
  
  
  class ViTTREncoder(FairseqEncoder):
      def __init__(self, args, dictionary):
          super().__init__(dictionary)
  
          img_size = args.vit_img_size
          patch_size = args.vit_patch_size
          in_chans = args.vit_channels
          embed_dim = args.vit_dim
          depth = args.vit_depth
          num_heads = args.vit_heads
          mlp_ratio=4.
          qkv_bias=True
          qk_scale=None
          drop_rate = args.vit_dropout
          attn_drop_rate = args.vit_atten_dropout
          drop_path_rate=0.
          hybrid_backbone=None
          norm_layer=None
  
          self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
          norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
  
          if hybrid_backbone is not None:
              self.patch_embed = HybridEmbed(
                  hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim)
          else:
              self.patch_embed = PatchEmbed(
                  img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
          num_patches = self.patch_embed.num_patches
  
          self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
          self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
          self.pos_drop = nn.Dropout(p=drop_rate)
  
          dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
          self.blocks = nn.ModuleList([
              Block(
                  dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                  drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
              for i in range(depth)])
          self.norm = norm_layer(embed_dim)
  
          trunc_normal_(self.pos_embed, std=.02)
          trunc_normal_(self.cls_token, std=.02)
          self.apply(self._init_weights)
  
      def _init_weights(self, m):
          if isinstance(m, nn.Linear):
              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)
  
      def forward_features(self, x):
          B = x.shape[0]  # bs, num_patches, dim
          x = self.patch_embed(x)
  
          cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
          x = torch.cat((cls_tokens, x), dim=1)
          x = x + self.pos_embed
  
          encoder_embedding = x   # bs, n + 1, dim
          x = self.pos_drop(x)
  
          for blk in self.blocks:
              x = blk(x)
  
          x = self.norm(x)  # bs, n + 1, dim
          return x, encoder_embedding
  
      def forward(self, imgs):
          x, encoder_embedding = self.forward_features(imgs)  # bs, n + 1, dim
          x = x.transpose(0, 1) # n + 1, bs, dim
  
          encoder_padding_mask = torch.zeros(*x.shape[:2]).transpose(0, 1).to(imgs.device)
  
          return {
              "encoder_out": [x],  # T x B x C
              "encoder_padding_mask": [encoder_padding_mask],  # B x T
              "encoder_embedding": [encoder_embedding],  # B x T x C
              "encoder_states": [],  # List[T x B x C]
              "src_tokens": [],
              "src_lengths": [],
          }
  
      def reorder_encoder_out(self, encoder_out, new_order):
            """
            Reorder encoder output according to `new_order`.
  
            Args:
                encoder_out: output from the ``forward()`` method
                new_order (LongTensor): desired order
  
            Returns:
                `encoder_out` rearranged according to `new_order`
            """
            _encoder_out = encoder_out['encoder_out'][0]
            _encoder_padding_mask = encoder_out['encoder_padding_mask'][0]
            _encoder_embedding = encoder_out['encoder_embedding'][0]
            return {
                "encoder_out": [_encoder_out.index_select(1, new_order)],
                  "encoder_padding_mask": [_encoder_padding_mask.index_select(0, new_order)],  # B x T
                  "encoder_embedding": [_encoder_padding_mask.index_select(0, new_order)],  # B x T x C
                  "encoder_states": [], 
                  "src_tokens": [],
                  "src_lengths": [],
          }
  
      
  
  if __name__ == '__main__':
      pass