


import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Dict    ):
     if "cls_token" in name:
          __UpperCamelCase   							=name.replace('cls_token'      ,					'vit.embeddings.cls_token'    )
     if "mask_token" in name:
          __UpperCamelCase   							=name.replace('mask_token'      ,					'decoder.mask_token'    )
     if "decoder_pos_embed" in name:
          __UpperCamelCase   							=name.replace('decoder_pos_embed'      ,					'decoder.decoder_pos_embed'    )
     if "pos_embed" in name and "decoder" not in name:
          __UpperCamelCase   							=name.replace('pos_embed'      ,					'vit.embeddings.position_embeddings'    )
     if "patch_embed.proj" in name:
          __UpperCamelCase   							=name.replace('patch_embed.proj'      ,					'vit.embeddings.patch_embeddings.projection'    )
     if "patch_embed.norm" in name:
          __UpperCamelCase   							=name.replace('patch_embed.norm'      ,					'vit.embeddings.norm'    )
     if "decoder_blocks" in name:
          __UpperCamelCase   							=name.replace('decoder_blocks'      ,					'decoder.decoder_layers'    )
     if "blocks" in name:
          __UpperCamelCase   							=name.replace('blocks'      ,					'vit.encoder.layer'    )
     if "attn.proj" in name:
          __UpperCamelCase   							=name.replace('attn.proj'      ,					'attention.output.dense'    )
     if "attn" in name:
          __UpperCamelCase   							=name.replace('attn'      ,					'attention.self'    )
     if "norm1" in name:
          __UpperCamelCase   							=name.replace('norm1'      ,					'layernorm_before'    )
     if "norm2" in name:
          __UpperCamelCase   							=name.replace('norm2'      ,					'layernorm_after'    )
     if "mlp.fc1" in name:
          __UpperCamelCase   							=name.replace('mlp.fc1'      ,					'intermediate.dense'    )
     if "mlp.fc2" in name:
          __UpperCamelCase   							=name.replace('mlp.fc2'      ,					'output.dense'    )
     if "decoder_embed" in name:
          __UpperCamelCase   							=name.replace('decoder_embed'      ,					'decoder.decoder_embed'    )
     if "decoder_norm" in name:
          __UpperCamelCase   							=name.replace('decoder_norm'      ,					'decoder.decoder_norm'    )
     if "decoder_pred" in name:
          __UpperCamelCase   							=name.replace('decoder_pred'      ,					'decoder.decoder_pred'    )
     if "norm.weight" in name and "decoder" not in name:
          __UpperCamelCase   							=name.replace('norm.weight'      ,					'vit.layernorm.weight'    )
     if "norm.bias" in name and "decoder" not in name:
          __UpperCamelCase   							=name.replace('norm.bias'      ,					'vit.layernorm.bias'    )

     return name

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Any      ,					SCREAMING_SNAKE_CASE__      : str    ):
     for key in orig_state_dict.copy().keys():
          __UpperCamelCase   							=orig_state_dict.pop(SCREAMING_SNAKE_CASE__    )

          if "qkv" in key:
               __UpperCamelCase   							=key.split('.'    )
               __UpperCamelCase   							=int(key_split[1]    )
               if "decoder_blocks" in key:
                    __UpperCamelCase   							=config.decoder_hidden_size
                    __UpperCamelCase   							='decoder.decoder_layers.'
                    if "weight" in key:
                         __UpperCamelCase   							=val[:dim, :]
                         __UpperCamelCase   							=val[dim : dim * 2, :]
                         __UpperCamelCase   							=val[-dim:, :]
                    elif "bias" in key:
                         __UpperCamelCase   							=val[:dim]
                         __UpperCamelCase   							=val[dim : dim * 2]
                         __UpperCamelCase   							=val[-dim:]
               else:
                    __UpperCamelCase   							=config.hidden_size
                    __UpperCamelCase   							='vit.encoder.layer.'
                    if "weight" in key:
                         __UpperCamelCase   							=val[:dim, :]
                         __UpperCamelCase   							=val[dim : dim * 2, :]
                         __UpperCamelCase   							=val[-dim:, :]
                    elif "bias" in key:
                         __UpperCamelCase   							=val[:dim]
                         __UpperCamelCase   							=val[dim : dim * 2]
                         __UpperCamelCase   							=val[-dim:]

          else:
               __UpperCamelCase   							=val

     return orig_state_dict

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Tuple      ,					SCREAMING_SNAKE_CASE__      : Dict    ):
     __UpperCamelCase   							=ViTMAEConfig()
     if "large" in checkpoint_url:
          __UpperCamelCase   							=10_24
          __UpperCamelCase   							=40_96
          __UpperCamelCase   							=24
          __UpperCamelCase   							=16
     elif "huge" in checkpoint_url:
          __UpperCamelCase   							=14
          __UpperCamelCase   							=12_80
          __UpperCamelCase   							=51_20
          __UpperCamelCase   							=32
          __UpperCamelCase   							=16

     __UpperCamelCase   							=ViTMAEForPreTraining(SCREAMING_SNAKE_CASE__    )

     __UpperCamelCase   							=torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE__      ,					map_location='cpu'    )['model']

     __UpperCamelCase   							=ViTMAEImageProcessor(size=config.image_size    )

     __UpperCamelCase   							=convert_state_dict(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )

     model.load_state_dict(SCREAMING_SNAKE_CASE__    )
     model.eval()

     __UpperCamelCase   							='https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg'

     __UpperCamelCase   							=Image.open(requests.get(SCREAMING_SNAKE_CASE__      ,					stream=SCREAMING_SNAKE_CASE__    ).raw    )
     __UpperCamelCase   							=ViTMAEImageProcessor(size=config.image_size    )
     __UpperCamelCase   							=image_processor(images=SCREAMING_SNAKE_CASE__      ,					return_tensors='pt'    )

     # forward pass
     torch.manual_seed(2    )
     __UpperCamelCase   							=model(**SCREAMING_SNAKE_CASE__    )
     __UpperCamelCase   							=outputs.logits

     if "large" in checkpoint_url:
          __UpperCamelCase   							=torch.tensor(
              [[-0.7309, -0.7128, -1.0169], [-1.0161, -0.9058, -1.1878], [-1.0478, -0.9411, -1.1911]]    )
     elif "huge" in checkpoint_url:
          __UpperCamelCase   							=torch.tensor(
              [[-1.1599, -0.9199, -1.2221], [-1.1952, -0.9269, -1.2307], [-1.2143, -0.9337, -1.2262]]    )
     else:
          __UpperCamelCase   							=torch.tensor(
              [[-0.9192, -0.8481, -1.1259], [-1.1349, -1.0034, -1.2599], [-1.1757, -1.0429, -1.2726]]    )

     # verify logits
     assert torch.allclose(logits[0, :3, :3]      ,					SCREAMING_SNAKE_CASE__      ,					atol=1E-4    )

     print(F'Saving model to {pytorch_dump_folder_path}'    )
     model.save_pretrained(SCREAMING_SNAKE_CASE__    )

     print(F'Saving image processor to {pytorch_dump_folder_path}'    )
     image_processor.save_pretrained(SCREAMING_SNAKE_CASE__    )


if __name__ == "__main__":
   _A			=						argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--checkpoint_url',
       default='https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth',
       type=str,
       help='URL of the checkpoint you\'d like to convert.',
   )
   parser.add_argument(
       '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
   )

   _A			=						parser.parse_args()
   convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









import warnings
from contextlib import contextmanager

from ....processing_utils import ProcessorMixin






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    UpperCAmelCase__						:		Optional[int]              =  "MCTCTFeatureExtractor"
    UpperCAmelCase__						:		str              =  "AutoTokenizer"






    def __init__(      self ,       A_ ,       A_   )  ->	Dict:
         super().__init__(A_ ,       A_   )
         __UpperCamelCase   							=self.feature_extractor
         __UpperCamelCase   							=False






    def __call__(      self ,       *A_ ,       **A_   )  ->	Optional[Any]:
         # For backward compatibility
         if self._in_target_context_manager:
              return self.current_processor(*A_ ,       **A_   )

         if "raw_speech" in kwargs:
              warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'   )
              __UpperCamelCase   							=kwargs.pop('raw_speech'   )
         else:
              __UpperCamelCase   							=kwargs.pop('audio' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('sampling_rate' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('text' ,       A_   )
         if len(A_   ) > 0:
              __UpperCamelCase   							=args[0]
              __UpperCamelCase   							=args[1:]

         if audio is None and text is None:
              raise ValueError('You need to specify either an `audio` or `text` input to process.'   )

         if audio is not None:
              __UpperCamelCase   							=self.feature_extractor(A_ ,       *A_ ,       sampling_rate=A_ ,       **A_   )
         if text is not None:
              __UpperCamelCase   							=self.tokenizer(A_ ,       **A_   )

         if text is None:
              return inputs
         elif audio is None:
              return encodings
         else:
              __UpperCamelCase   							=encodings['input_ids']
              return inputs






    def _a (      self ,       *A_ ,       **A_   )  ->	str:
         return self.tokenizer.batch_decode(*A_ ,       **A_   )






    def _a (      self ,       *A_ ,       **A_   )  ->	List[Any]:
         # For backward compatibility
         if self._in_target_context_manager:
              return self.current_processor.pad(*A_ ,       **A_   )

         __UpperCamelCase   							=kwargs.pop('input_features' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('labels' ,       A_   )
         if len(A_   ) > 0:
              __UpperCamelCase   							=args[0]
              __UpperCamelCase   							=args[1:]

         if input_features is not None:
              __UpperCamelCase   							=self.feature_extractor.pad(A_ ,       *A_ ,       **A_   )
         if labels is not None:
              __UpperCamelCase   							=self.tokenizer.pad(A_ ,       **A_   )

         if labels is None:
              return input_features
         elif input_features is None:
              return labels
         else:
              __UpperCamelCase   							=labels['input_ids']
              return input_features






    def _a (      self ,       *A_ ,       **A_   )  ->	Optional[int]:
         return self.tokenizer.decode(*A_ ,       **A_   )






    @contextmanager
    def _a (      self   )  ->	str:
         warnings.warn(
             '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
             'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
             'your audio inputs, or in a separate call.'   )
         __UpperCamelCase   							=True
         __UpperCamelCase   							=self.tokenizer
         yield
         __UpperCamelCase   							=self.feature_extractor
         __UpperCamelCase   							=False





