





'''simple docstring'''






import inspect
import unittest
from math import floor

from transformers import CvtConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import CvtForImageClassification, CvtModel
      from transformers.models.cvt.modeling_cvt import CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import AutoImageProcessor




class 					UpperCAmelCase	(     snake_case_  ):




      def 							lowercase__       (     self     :					Tuple    )							-> int:
            _lowerCAmelCase					=				self.config_class(**self.inputs_dict    )
            self.parent.assertTrue(hasattr(__snake_case     ,     """embed_dim"""    )    )
            self.parent.assertTrue(hasattr(__snake_case     ,     """num_heads"""    )    )




class 					UpperCAmelCase	:




      def __init__(     self     :					str     ,     __snake_case     :					List[str]     ,     __snake_case     :					Dict=13     ,     __snake_case     :					Dict=64     ,     __snake_case     :					List[str]=3     ,     __snake_case     :					str=[16, 48, 96]     ,     __snake_case     :					Tuple=[1, 3, 6]     ,     __snake_case     :					Optional[Any]=[1, 2, 10]     ,     __snake_case     :					str=[7, 3, 3]     ,     __snake_case     :					Optional[int]=[4, 2, 2]     ,     __snake_case     :					int=[2, 1, 1]     ,     __snake_case     :					str=[2, 2, 2]     ,     __snake_case     :					Any=[False, False, True]     ,     __snake_case     :					Optional[Any]=[0.0, 0.0, 0.0]     ,     __snake_case     :					List[Any]=0.02     ,     __snake_case     :					Optional[Any]=1E-1_2     ,     __snake_case     :					List[Any]=True     ,     __snake_case     :					List[str]=True     ,     __snake_case     :					Any=2     ,     )							-> Union[str, Any]:
            _lowerCAmelCase					=				parent
            _lowerCAmelCase					=				batch_size
            _lowerCAmelCase					=				image_size
            _lowerCAmelCase					=				patch_sizes
            _lowerCAmelCase					=				patch_stride
            _lowerCAmelCase					=				patch_padding
            _lowerCAmelCase					=				is_training
            _lowerCAmelCase					=				use_labels
            _lowerCAmelCase					=				num_labels
            _lowerCAmelCase					=				num_channels
            _lowerCAmelCase					=				embed_dim
            _lowerCAmelCase					=				num_heads
            _lowerCAmelCase					=				stride_kv
            _lowerCAmelCase					=				depth
            _lowerCAmelCase					=				cls_token
            _lowerCAmelCase					=				attention_drop_rate
            _lowerCAmelCase					=				initializer_range
            _lowerCAmelCase					=				layer_norm_eps




      def 							lowercase__       (     self     :					Dict    )							-> Optional[Any]:
            _lowerCAmelCase					=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

            _lowerCAmelCase					=				None
            if self.use_labels:
                  _lowerCAmelCase					=				ids_tensor([self.batch_size]     ,     self.num_labels    )

            _lowerCAmelCase					=				self.get_config()
            return config, pixel_values, labels




      def 							lowercase__       (     self     :					int    )							-> List[Any]:
            return CvtConfig(
                image_size=self.image_size     ,     num_labels=self.num_labels     ,     num_channels=self.num_channels     ,     embed_dim=self.embed_dim     ,     num_heads=self.num_heads     ,     patch_sizes=self.patch_sizes     ,     patch_padding=self.patch_padding     ,     patch_stride=self.patch_stride     ,     stride_kv=self.stride_kv     ,     depth=self.depth     ,     cls_token=self.cls_token     ,     attention_drop_rate=self.attention_drop_rate     ,     initializer_range=self.initializer_range     ,     )




      def 							lowercase__       (     self     :					List[str]     ,     __snake_case     :					Optional[int]     ,     __snake_case     :					List[str]     ,     __snake_case     :					int    )							-> int:
            _lowerCAmelCase					=				CvtModel(config=__snake_case    )
            model.to(__snake_case    )
            model.eval()
            _lowerCAmelCase					=				model(__snake_case    )
            _lowerCAmelCase					=				(self.image_size, self.image_size)
            _lowerCAmelCase	,				_lowerCAmelCase					=				image_size[0], image_size[1]
            for i in range(len(self.depth    )    ):
                  _lowerCAmelCase					=				floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1    )
                  _lowerCAmelCase					=				floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1    )
            self.parent.assertEqual(result.last_hidden_state.shape     ,     (self.batch_size, self.embed_dim[-1], height, width)    )




      def 							lowercase__       (     self     :					Tuple     ,     __snake_case     :					int     ,     __snake_case     :					List[str]     ,     __snake_case     :					Tuple    )							-> Tuple:
            _lowerCAmelCase					=				self.num_labels
            _lowerCAmelCase					=				CvtForImageClassification(__snake_case    )
            model.to(__snake_case    )
            model.eval()
            _lowerCAmelCase					=				model(__snake_case     ,     labels=__snake_case    )
            self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.num_labels)    )




      def 							lowercase__       (     self     :					Tuple    )							-> int:
            _lowerCAmelCase					=				self.prepare_config_and_inputs()
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				config_and_inputs
            _lowerCAmelCase					=				{"""pixel_values""": pixel_values}
            return config, inputs_dict




@require_torch
class 					UpperCAmelCase	(     snake_case_				,						snake_case_				,						unittest.TestCase  ):
      _lowercase:     int			=  (CvtModel, CvtForImageClassification) if is_torch_available() else ()
      _lowercase:     Tuple			=  (
          {'''feature-extraction''': CvtModel, '''image-classification''': CvtForImageClassification}
          if is_torch_available()
          else {}
      )

      _lowercase:     Tuple			=  False
      _lowercase:     Union[str, Any]			=  False
      _lowercase:     Tuple			=  False
      _lowercase:     List[str]			=  False
      _lowercase:     Optional[Any]			=  False




      def 							lowercase__       (     self     :					Dict    )							-> Any:
            _lowerCAmelCase					=				CvtModelTester(self    )
            _lowerCAmelCase					=				ConfigTester(self     ,     config_class=__snake_case     ,     has_text_modality=__snake_case     ,     hidden_size=37    )




      def 							lowercase__       (     self     :					Any    )							-> Dict:
            self.create_and_test_config_common_properties()
            self.config_tester.create_and_test_config_to_json_string()
            self.config_tester.create_and_test_config_to_json_file()
            self.config_tester.create_and_test_config_from_and_save_pretrained()
            self.config_tester.create_and_test_config_with_num_labels()
            self.config_tester.check_config_can_be_init_without_params()
            self.config_tester.check_config_arguments_init()




      def 							lowercase__       (     self     :					int    )							-> Any:
            return




      @unittest.skip(reason="""Cvt does not output attentions"""    )
      def 							lowercase__       (     self     :					Optional[Any]    )							-> Union[str, Any]:
            pass




      @unittest.skip(reason="""Cvt does not use inputs_embeds"""    )
      def 							lowercase__       (     self     :					Union[str, Any]    )							-> Optional[Any]:
            pass




      @unittest.skip(reason="""Cvt does not support input and output embeddings"""    )
      def 							lowercase__       (     self     :					List[Any]    )							-> Optional[int]:
            pass




      def 							lowercase__       (     self     :					int    )							-> Any:
            _lowerCAmelCase	,				_lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                  _lowerCAmelCase					=				model_class(__snake_case    )
                  _lowerCAmelCase					=				inspect.signature(model.forward    )
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  _lowerCAmelCase					=				[*signature.parameters.keys()]

                  _lowerCAmelCase					=				["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]     ,     __snake_case    )




      def 							lowercase__       (     self     :					str    )							-> Any:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*__snake_case    )




      def 							lowercase__       (     self     :					List[Any]    )							-> Optional[int]:
            def check_hidden_states_output(__snake_case     :					Union[str, Any]     ,     __snake_case     :					List[Any]     ,     __snake_case     :					Optional[Any]    ):
                  _lowerCAmelCase					=				model_class(__snake_case    )
                  model.to(__snake_case    )
                  model.eval()

                  with torch.no_grad():
                        _lowerCAmelCase					=				model(**self._prepare_for_class(__snake_case     ,     __snake_case    )    )

                  _lowerCAmelCase					=				outputs.hidden_states

                  _lowerCAmelCase					=				len(self.model_tester.depth    )
                  self.assertEqual(len(__snake_case    )     ,     __snake_case    )

                  # verify the first hidden states (first block)
                  self.assertListEqual(
                      list(hidden_states[0].shape[-3:]    )     ,     [
                          self.model_tester.embed_dim[0],
                          self.model_tester.image_size // 4,
                          self.model_tester.image_size // 4,
                      ]     ,     )

            _lowerCAmelCase	,				_lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                  _lowerCAmelCase					=				True
                  check_hidden_states_output(__snake_case     ,     __snake_case     ,     __snake_case    )

                  # check that output_hidden_states also work using config
                  del inputs_dict["output_hidden_states"]
                  _lowerCAmelCase					=				True

                  check_hidden_states_output(__snake_case     ,     __snake_case     ,     __snake_case    )




      def 							lowercase__       (     self     :					str    )							-> List[Any]:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*__snake_case    )




      @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""    )
      def 							lowercase__       (     self     :					Optional[int]    )							-> Dict:
            pass




      @slow
      def 							lowercase__       (     self     :					Union[str, Any]    )							-> Any:
            for model_name in CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  _lowerCAmelCase					=				CvtModel.from_pretrained(__snake_case    )
                  self.assertIsNotNone(__snake_case    )




def   UpperCamelCase__ (						):





      """simple docstring"""


      _lowerCAmelCase					=				Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
      return image





@require_torch
@require_vision
class 					UpperCAmelCase	(     unittest.TestCase  ):




      @cached_property
      def 							lowercase__       (     self     :					List[str]    )							-> Any:
            return AutoImageProcessor.from_pretrained(CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]    )




      @slow
      def 							lowercase__       (     self     :					List[Any]    )							-> Tuple:
            _lowerCAmelCase					=				CvtForImageClassification.from_pretrained(CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]    ).to(__snake_case    )

            _lowerCAmelCase					=				self.default_image_processor
            _lowerCAmelCase					=				prepare_img()
            _lowerCAmelCase					=				image_processor(images=__snake_case     ,     return_tensors="""pt"""    ).to(__snake_case    )

            # forward pass
            with torch.no_grad():
                  _lowerCAmelCase					=				model(**__snake_case    )

            # verify the logits
            _lowerCAmelCase					=				torch.Size((1, 10_00)    )
            self.assertEqual(outputs.logits.shape     ,     __snake_case    )

            _lowerCAmelCase					=				torch.tensor([0.92_85, 0.90_15, -0.31_50]    ).to(__snake_case    )

            self.assertTrue(torch.allclose(outputs.logits[0, :3]     ,     __snake_case     ,     atol=1E-4    )    )













'''simple docstring'''






import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor
def   UpperCamelCase__ (						lowerCAmelCase    ):





      """simple docstring"""


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

      return name
def   UpperCamelCase__ (						lowerCAmelCase				,      lowerCAmelCase    ):





      """simple docstring"""


      for key in orig_state_dict.copy().keys():
            _lowerCAmelCase					=				orig_state_dict.pop(lowerCAmelCase    )

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

            else:
                  _lowerCAmelCase					=				val

      return orig_state_dict
def   UpperCamelCase__ (						lowerCAmelCase				,      lowerCAmelCase    ):





      """simple docstring"""


      _lowerCAmelCase					=				ViTMAEConfig()
      if "large" in checkpoint_url:
            _lowerCAmelCase					=				10_24
            _lowerCAmelCase					=				40_96
            _lowerCAmelCase					=				24
            _lowerCAmelCase					=				16
      elif "huge" in checkpoint_url:
            _lowerCAmelCase					=				14
            _lowerCAmelCase					=				12_80
            _lowerCAmelCase					=				51_20
            _lowerCAmelCase					=				32
            _lowerCAmelCase					=				16

      _lowerCAmelCase					=				ViTMAEForPreTraining(lowerCAmelCase    )

      _lowerCAmelCase					=				torch.hub.load_state_dict_from_url(lowerCAmelCase				,      map_location="""cpu"""    )["""model"""]

      _lowerCAmelCase					=				ViTMAEImageProcessor(size=config.image_size    )

      _lowerCAmelCase					=				convert_state_dict(lowerCAmelCase				,      lowerCAmelCase    )

      model.load_state_dict(lowerCAmelCase    )
      model.eval()

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

      _lowerCAmelCase					=				Image.open(requests.get(lowerCAmelCase				,      stream=lowerCAmelCase    ).raw    )
      _lowerCAmelCase					=				ViTMAEImageProcessor(size=config.image_size    )
      _lowerCAmelCase					=				image_processor(images=lowerCAmelCase				,      return_tensors="""pt"""    )

      # forward pass
      torch.manual_seed(2    )
      _lowerCAmelCase					=				model(**lowerCAmelCase    )
      _lowerCAmelCase					=				outputs.logits

      if "large" in checkpoint_url:
            _lowerCAmelCase					=				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:
            _lowerCAmelCase					=				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:
            _lowerCAmelCase					=				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__":
      A__							:  Dict    =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__							:  List[Any]    =parser.parse_args()
      convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






