





'''simple docstring'''


import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor
def 							UpperCAmelCase_  (						__lowerCamelCase : str			):
   if "cls_token" in name:
      lowercase_      :List[Any]         =			name.replace("cls_token"     ,"vit.embeddings.cls_token"			)
   if "mask_token" in name:
      lowercase_      :str         =			name.replace("mask_token"     ,"decoder.mask_token"			)
   if "decoder_pos_embed" in name:
      lowercase_      :Optional[int]         =			name.replace("decoder_pos_embed"     ,"decoder.decoder_pos_embed"			)
   if "pos_embed" in name and "decoder" not in name:
      lowercase_      :List[Any]         =			name.replace("pos_embed"     ,"vit.embeddings.position_embeddings"			)
   if "patch_embed.proj" in name:
      lowercase_      :Dict         =			name.replace("patch_embed.proj"     ,"vit.embeddings.patch_embeddings.projection"			)
   if "patch_embed.norm" in name:
      lowercase_      :int         =			name.replace("patch_embed.norm"     ,"vit.embeddings.norm"			)
   if "decoder_blocks" in name:
      lowercase_      :Union[str, Any]         =			name.replace("decoder_blocks"     ,"decoder.decoder_layers"			)
   if "blocks" in name:
      lowercase_      :List[Any]         =			name.replace("blocks"     ,"vit.encoder.layer"			)
   if "attn.proj" in name:
      lowercase_      :List[Any]         =			name.replace("attn.proj"     ,"attention.output.dense"			)
   if "attn" in name:
      lowercase_      :Optional[int]         =			name.replace("attn"     ,"attention.self"			)
   if "norm1" in name:
      lowercase_      :Union[str, Any]         =			name.replace("norm1"     ,"layernorm_before"			)
   if "norm2" in name:
      lowercase_      :Union[str, Any]         =			name.replace("norm2"     ,"layernorm_after"			)
   if "mlp.fc1" in name:
      lowercase_      :Tuple         =			name.replace("mlp.fc1"     ,"intermediate.dense"			)
   if "mlp.fc2" in name:
      lowercase_      :Optional[int]         =			name.replace("mlp.fc2"     ,"output.dense"			)
   if "decoder_embed" in name:
      lowercase_      :Union[str, Any]         =			name.replace("decoder_embed"     ,"decoder.decoder_embed"			)
   if "decoder_norm" in name:
      lowercase_      :Any         =			name.replace("decoder_norm"     ,"decoder.decoder_norm"			)
   if "decoder_pred" in name:
      lowercase_      :Dict         =			name.replace("decoder_pred"     ,"decoder.decoder_pred"			)
   if "norm.weight" in name and "decoder" not in name:
      lowercase_      :Tuple         =			name.replace("norm.weight"     ,"vit.layernorm.weight"			)
   if "norm.bias" in name and "decoder" not in name:
      lowercase_      :Any         =			name.replace("norm.bias"     ,"vit.layernorm.bias"			)

   return name
def 							UpperCAmelCase_  (						__lowerCamelCase : List[str]     ,__lowerCamelCase : Dict			):
   for key in orig_state_dict.copy().keys():
      lowercase_      :Any         =			orig_state_dict.pop(__lowerCamelCase			)

      if "qkv" in key:
         lowercase_      :Optional[Any]         =			key.split("."			)
         lowercase_      :str         =			int(key_split[1]			)
         if "decoder_blocks" in key:
            lowercase_      :List[Any]         =			config.decoder_hidden_size
            lowercase_      :List[Any]         =			"decoder.decoder_layers."
            if "weight" in key:
               lowercase_      :str         =			val[:dim, :]
               lowercase_      :Any         =			val[dim : dim * 2, :]
               lowercase_      :List[str]         =			val[-dim:, :]
            elif "bias" in key:
               lowercase_      :Union[str, Any]         =			val[:dim]
               lowercase_      :Optional[Any]         =			val[dim : dim * 2]
               lowercase_      :Optional[int]         =			val[-dim:]
         else:
            lowercase_      :Tuple         =			config.hidden_size
            lowercase_      :Tuple         =			"vit.encoder.layer."
            if "weight" in key:
               lowercase_      :Tuple         =			val[:dim, :]
               lowercase_      :Optional[Any]         =			val[dim : dim * 2, :]
               lowercase_      :Optional[int]         =			val[-dim:, :]
            elif "bias" in key:
               lowercase_      :str         =			val[:dim]
               lowercase_      :Dict         =			val[dim : dim * 2]
               lowercase_      :str         =			val[-dim:]

      else:
         lowercase_      :Tuple         =			val

   return orig_state_dict
def 							UpperCAmelCase_  (						__lowerCamelCase : str     ,__lowerCamelCase : Union[str, Any]			):
   lowercase_      :Union[str, Any]         =			ViTMAEConfig()
   if "large" in checkpoint_url:
      lowercase_      :List[str]         =			10_24
      lowercase_      :Optional[int]         =			40_96
      lowercase_      :List[Any]         =			24
      lowercase_      :Union[str, Any]         =			16
   elif "huge" in checkpoint_url:
      lowercase_      :str         =			14
      lowercase_      :List[Any]         =			12_80
      lowercase_      :int         =			51_20
      lowercase_      :int         =			32
      lowercase_      :Optional[int]         =			16

   lowercase_      :Union[str, Any]         =			ViTMAEForPreTraining(__lowerCamelCase			)

   lowercase_      :Optional[Any]         =			torch.hub.load_state_dict_from_url(__lowerCamelCase     ,map_location="cpu"			)["model"]

   lowercase_      :Optional[Any]         =			ViTMAEImageProcessor(size=config.image_size			)

   lowercase_      :int         =			convert_state_dict(__lowerCamelCase     ,__lowerCamelCase			)

   model.load_state_dict(__lowerCamelCase			)
   model.eval()

   lowercase_      :Optional[int]         =			"https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg"

   lowercase_      :Union[str, Any]         =			Image.open(requests.get(__lowerCamelCase     ,stream=__lowerCamelCase			).raw			)
   lowercase_      :int         =			ViTMAEImageProcessor(size=config.image_size			)
   lowercase_      :Dict         =			image_processor(images=__lowerCamelCase     ,return_tensors="pt"			)

   # forward pass
   torch.manual_seed(2			)
   lowercase_      :List[Any]         =			model(**__lowerCamelCase			)
   lowercase_      :Optional[int]         =			outputs.logits

   if "large" in checkpoint_url:
      lowercase_      :List[Any]         =			torch.tensor(
          [[-0.7_309, -0.7_128, -1.0_169], [-1.0_161, -0.9_058, -1.1_878], [-1.0_478, -0.9_411, -1.1_911]]			)
   elif "huge" in checkpoint_url:
      lowercase_      :Tuple         =			torch.tensor(
          [[-1.1_599, -0.9_199, -1.2_221], [-1.1_952, -0.9_269, -1.2_307], [-1.2_143, -0.9_337, -1.2_262]]			)
   else:
      lowercase_      :Optional[int]         =			torch.tensor(
          [[-0.9_192, -0.8_481, -1.1_259], [-1.1_349, -1.0_034, -1.2_599], [-1.1_757, -1.0_429, -1.2_726]]			)

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

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

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


if __name__ == "__main__":
 lowerCAmelCase :					Tuple 							=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.'''
 )

 lowerCAmelCase :					Tuple 							=parser.parse_args()
 convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









'''simple docstring'''


from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers import (
     TFDeiTForImageClassification,
     TFDeiTForImageClassificationWithTeacher,
     TFDeiTForMaskedImageModeling,
     TFDeiTModel,
 )
 from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import DeiTImageProcessor
class 				a_       :





     def __init__(			self		: int     ,				lowercase		: Optional[int]     ,				lowercase		: List[Any]=13     ,				lowercase		: Dict=30     ,				lowercase		: Dict=2     ,				lowercase		: Tuple=3     ,				lowercase		: Dict=True     ,				lowercase		: Dict=True     ,				lowercase		: Tuple=32     ,				lowercase		: List[Any]=2     ,				lowercase		: List[str]=4     ,				lowercase		: List[Any]=37     ,				lowercase		: Union[str, Any]="gelu"     ,				lowercase		: List[str]=0.1     ,				lowercase		: List[Any]=0.1     ,				lowercase		: List[str]=10     ,				lowercase		: Any=0.02     ,				lowercase		: Union[str, Any]=3     ,				lowercase		: Tuple=None     ,				lowercase		: List[str]=2     ,				):





        """simple docstring"""




        lowercase_      :Union[str, Any]         =			parent
        lowercase_      :Optional[int]         =			batch_size
        lowercase_      :Tuple         =			image_size
        lowercase_      :Any         =			patch_size
        lowercase_      :List[Any]         =			num_channels
        lowercase_      :Optional[Any]         =			is_training
        lowercase_      :str         =			use_labels
        lowercase_      :Any         =			hidden_size
        lowercase_      :Optional[int]         =			num_hidden_layers
        lowercase_      :List[Any]         =			num_attention_heads
        lowercase_      :str         =			intermediate_size
        lowercase_      :Optional[int]         =			hidden_act
        lowercase_      :List[Any]         =			hidden_dropout_prob
        lowercase_      :Optional[Any]         =			attention_probs_dropout_prob
        lowercase_      :Any         =			type_sequence_label_size
        lowercase_      :Dict         =			initializer_range
        lowercase_      :List[Any]         =			scope
        lowercase_      :List[str]         =			encoder_stride

        # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
        lowercase_      :str         =			(image_size // patch_size) ** 2
        lowercase_      :Any         =			num_patches + 2





     def 							lowercase__				(			self		: List[Any]						):





        """simple docstring"""




        lowercase_      :List[str]         =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

        lowercase_      :Optional[Any]         =			None
        if self.use_labels:
           lowercase_      :Optional[Any]         =			ids_tensor([self.batch_size]     ,				self.type_sequence_label_size						)

        lowercase_      :Any         =			self.get_config()

        return config, pixel_values, labels





     def 							lowercase__				(			self		: List[Any]						):





        """simple docstring"""




        return DeiTConfig(
            image_size=self.image_size     ,				patch_size=self.patch_size     ,				num_channels=self.num_channels     ,				hidden_size=self.hidden_size     ,				num_hidden_layers=self.num_hidden_layers     ,				num_attention_heads=self.num_attention_heads     ,				intermediate_size=self.intermediate_size     ,				hidden_act=self.hidden_act     ,				hidden_dropout_prob=self.hidden_dropout_prob     ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,				is_decoder=lowercase     ,				initializer_range=self.initializer_range     ,				encoder_stride=self.encoder_stride     ,				)





     def 							lowercase__				(			self		: Tuple     ,				lowercase		: int     ,				lowercase		: str     ,				lowercase		: Tuple						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			TFDeiTModel(config=lowercase						)
        lowercase_      :str         =			model(lowercase						)
        self.parent.assertEqual(result.last_hidden_state.shape     ,				(self.batch_size, self.seq_length, self.hidden_size)						)





     def 							lowercase__				(			self		: str     ,				lowercase		: str     ,				lowercase		: Optional[Any]     ,				lowercase		: str						):





        """simple docstring"""




        lowercase_      :List[str]         =			TFDeiTForMaskedImageModeling(config=lowercase						)
        lowercase_      :Optional[int]         =			model(lowercase						)
        self.parent.assertEqual(
            result.reconstruction.shape     ,				(self.batch_size, self.num_channels, self.image_size, self.image_size)						)

        # test greyscale images
        lowercase_      :List[str]         =			1
        lowercase_      :List[str]         =			TFDeiTForMaskedImageModeling(lowercase						)

        lowercase_      :Union[str, Any]         =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
        lowercase_      :List[str]         =			model(lowercase						)
        self.parent.assertEqual(result.reconstruction.shape     ,				(self.batch_size, 1, self.image_size, self.image_size)						)





     def 							lowercase__				(			self		: List[str]     ,				lowercase		: Any     ,				lowercase		: Union[str, Any]     ,				lowercase		: Union[str, Any]						):





        """simple docstring"""




        lowercase_      :Tuple         =			self.type_sequence_label_size
        lowercase_      :str         =			TFDeiTForImageClassification(lowercase						)
        lowercase_      :str         =			model(lowercase     ,				labels=lowercase						)
        self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.type_sequence_label_size)						)

        # test greyscale images
        lowercase_      :Optional[int]         =			1
        lowercase_      :Optional[int]         =			TFDeiTForImageClassification(lowercase						)

        lowercase_      :List[Any]         =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
        lowercase_      :int         =			model(lowercase     ,				labels=lowercase						)
        self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.type_sequence_label_size)						)







     def 							lowercase__				(			self		: Dict						):





        """simple docstring"""




        lowercase_      :Tuple         =			self.prepare_config_and_inputs()
        lowercase_      ,							lowercase_      ,							lowercase_      :Tuple         =			config_and_inputs
        lowercase_      :Tuple         =			{"pixel_values": pixel_values}
        return config, inputs_dict







@require_tf
class 				a_       (    _lowerCAmelCase     ,						_lowerCAmelCase     ,						unittest.TestCase							):
     __A								=      (
         (
             TFDeiTModel,
             TFDeiTForImageClassification,
             TFDeiTForImageClassificationWithTeacher,
             TFDeiTForMaskedImageModeling,
         )
         if is_tf_available()
         else ()
     )
     __A								=      (
         {
             "feature-extraction": TFDeiTModel,
             "image-classification": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
         }
         if is_tf_available()
         else {}
     )

     __A								=      False
     __A								=      False
     __A								=      False
     __A								=      False





     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        lowercase_      :List[Any]         =			TFDeiTModelTester(self						)
        lowercase_      :List[Any]         =			ConfigTester(self     ,				config_class=lowercase     ,				has_text_modality=lowercase     ,				hidden_size=37						)





     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




        self.config_tester.run_common_tests()





     @unittest.skip(reason="DeiT does not use inputs_embeds"						)
     def 							lowercase__				(			self		: Optional[int]						):





        """simple docstring"""




        pass





     def 							lowercase__				(			self		: Tuple						):





        """simple docstring"""




        lowercase_      ,							lowercase_      :List[str]         =			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_      :Optional[Any]         =			model_class(lowercase						)
           self.assertIsInstance(model.get_input_embeddings()     ,				(tf.keras.layers.Layer)						)
           lowercase_      :Tuple         =			model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(lowercase     ,				tf.keras.layers.Dense						)						)





     def 							lowercase__				(			self		: Any						):





        """simple docstring"""




        lowercase_      ,							lowercase_      :List[str]         =			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_      :Tuple         =			model_class(lowercase						)
           lowercase_      :Optional[int]         =			inspect.signature(model.call						)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           lowercase_      :Tuple         =			[*signature.parameters.keys()]

           lowercase_      :Any         =			["pixel_values"]
           self.assertListEqual(arg_names[:1]     ,				lowercase						)





     def 							lowercase__				(			self		: Optional[Any]						):





        """simple docstring"""




        lowercase_      :Optional[Any]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase						)





     def 							lowercase__				(			self		: str						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_image_modeling(*lowercase						)





     def 							lowercase__				(			self		: List[str]						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*lowercase						)





     def 							lowercase__				(			self		: List[Any]     ,				lowercase		: Optional[Any]     ,				lowercase		: List[str]     ,				lowercase		: Any=False						):





        """simple docstring"""




        lowercase_      :Dict         =			super()._prepare_for_class(lowercase     ,				lowercase     ,				return_labels=lowercase						)

        if return_labels:
           if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call						).parameters:
              del inputs_dict["labels"]

        return inputs_dict







     @slow
     def 							lowercase__				(			self		: List[Any]						):





        """simple docstring"""




        for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase_      :List[Any]         =			TFDeiTModel.from_pretrained(lowercase						)
           self.assertIsNotNone(lowercase						)







def 							UpperCAmelCase_  (						):
   lowercase_      :List[str]         =			Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			)
   return image


@require_tf
@require_vision
class 				a_       (    unittest.TestCase							):





     @cached_property
     def 							lowercase__				(			self		: Dict						):





        """simple docstring"""




        return (
            DeiTImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224"						)
            if is_vision_available()
            else None
        )







     @slow
     def 							lowercase__				(			self		: Optional[Any]						):





        """simple docstring"""




        lowercase_      :Optional[int]         =			TFDeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"						)

        lowercase_      :Optional[int]         =			self.default_image_processor
        lowercase_      :Union[str, Any]         =			prepare_img()
        lowercase_      :Dict         =			image_processor(images=lowercase     ,				return_tensors="tf"						)

        # forward pass
        lowercase_      :Union[str, Any]         =			model(**lowercase						)

        # verify the logits
        lowercase_      :Union[str, Any]         =			tf.TensorShape((1, 1_000)						)
        self.assertEqual(outputs.logits.shape     ,				lowercase						)

        lowercase_      :Dict         =			tf.constant([-1.02_66, 0.19_12, -1.28_61]						)

        self.assertTrue(np.allclose(outputs.logits[0, :3]     ,				lowercase     ,				atol=1e-4						)						)


