




from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
  import tensorflow as tf

  from transformers import TFBlipTextModel
  from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST


class     UpperCamelCase_   :

    '''simple docstring'''






    def __init__(      self				:     Optional[Any] ,			UpperCAmelCase__				:     int ,			UpperCAmelCase__				:     Optional[int]=12 ,			UpperCAmelCase__				:     Any=7 ,			UpperCAmelCase__				:     str=True ,			UpperCAmelCase__				:     Optional[int]=True ,			UpperCAmelCase__				:     Tuple=True ,			UpperCAmelCase__				:     int=99 ,			UpperCAmelCase__				:     List[str]=32 ,			UpperCAmelCase__				:     List[str]=32 ,			UpperCAmelCase__				:     Tuple=2 ,			UpperCAmelCase__				:     Any=4 ,			UpperCAmelCase__				:     Dict=37 ,			UpperCAmelCase__				:     Optional[int]=0.1 ,			UpperCAmelCase__				:     List[str]=0.1 ,			UpperCAmelCase__				:     Dict=512 ,			UpperCAmelCase__				:     Dict=0.02 ,			UpperCAmelCase__				:     Union[str, Any]=0 ,			UpperCAmelCase__				:     Any=None ,			)						->Union[str, Any]:





           '''simple docstring'''

           A__							=   parent
           A__							=   batch_size
           A__							=   seq_length
           A__							=   is_training
           A__							=   use_input_mask
           A__							=   use_labels
           A__							=   vocab_size
           A__							=   hidden_size
           A__							=   projection_dim
           A__							=   num_hidden_layers
           A__							=   num_attention_heads
           A__							=   intermediate_size
           A__							=   dropout
           A__							=   attention_dropout
           A__							=   max_position_embeddings
           A__							=   initializer_range
           A__							=   scope
           A__							=   bos_token_id


    def 	SCREAMING_SNAKE_CASE							(      self				:     List[str])						->Dict:





           '''simple docstring'''

           A__							=   ids_tensor([self.batch_size, self.seq_length] ,			self.vocab_size)

           A__							=   None
           if self.use_input_mask:
                  A__							=   random_attention_mask([self.batch_size, self.seq_length])

           if input_mask is not None:
                  A__							=   input_mask.numpy()
                  A__			,							A__							=   input_mask.shape
                  A__							=   np.random.randint(1 ,			seq_length - 1 ,			size=(batch_size,))
                  for batch_idx, start_index in enumerate(UpperCAmelCase__):
                         A__							=   1
                         A__							=   0

           A__							=   self.get_config()

           return config, input_ids, tf.convert_to_tensor(UpperCAmelCase__)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Any)						->Tuple:





           '''simple docstring'''

           return BlipTextConfig(
               vocab_size=self.vocab_size ,			hidden_size=self.hidden_size ,			projection_dim=self.projection_dim ,			num_hidden_layers=self.num_hidden_layers ,			num_attention_heads=self.num_attention_heads ,			intermediate_size=self.intermediate_size ,			dropout=self.dropout ,			attention_dropout=self.attention_dropout ,			max_position_embeddings=self.max_position_embeddings ,			initializer_range=self.initializer_range ,			bos_token_id=self.bos_token_id ,			)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[int] ,			UpperCAmelCase__				:     Tuple ,			UpperCAmelCase__				:     Optional[Any] ,			UpperCAmelCase__				:     Optional[Any])						->Dict:





           '''simple docstring'''

           A__							=   TFBlipTextModel(config=UpperCAmelCase__)
           A__							=   model(UpperCAmelCase__ ,			attention_mask=UpperCAmelCase__ ,			training=UpperCAmelCase__)
           A__							=   model(UpperCAmelCase__ ,			training=UpperCAmelCase__)
           self.parent.assertEqual(result.last_hidden_state.shape ,			(self.batch_size, self.seq_length, self.hidden_size))
           self.parent.assertEqual(result.pooler_output.shape ,			(self.batch_size, self.hidden_size))





    def 	SCREAMING_SNAKE_CASE							(      self				:     int)						->str:





           '''simple docstring'''

           A__							=   self.prepare_config_and_inputs()
           A__			,							A__			,							A__							=   config_and_inputs
           A__							=   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
           return config, inputs_dict


@require_tf
class     UpperCamelCase_   (			UpperCAmelCase__		, unittest.TestCase					):

    '''simple docstring'''






    UpperCAmelCase__							  =				(TFBlipTextModel,) if is_tf_available() else ()
    UpperCAmelCase__							  =				False
    UpperCAmelCase__							  =				False
    UpperCAmelCase__							  =				False


    def 	SCREAMING_SNAKE_CASE							(      self				:     Tuple)						->Optional[Any]:





           '''simple docstring'''

           A__							=   BlipTextModelTester(self)
           A__							=   ConfigTester(self ,			config_class=UpperCAmelCase__ ,			hidden_size=37)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Tuple)						->Tuple:





           '''simple docstring'''

           self.config_tester.run_common_tests()


    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[int])						->Tuple:





           '''simple docstring'''

           A__							=   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCAmelCase__)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Tuple)						->int:





           '''simple docstring'''

           pass


    def 	SCREAMING_SNAKE_CASE							(      self				:     Union[str, Any])						->int:





           '''simple docstring'''

           pass


    @unittest.skip(reason='''Blip does not use inputs_embeds''')
    def 	SCREAMING_SNAKE_CASE							(      self				:     int)						->Optional[Any]:





           '''simple docstring'''

           pass


    @unittest.skip(reason='''BlipTextModel has no base class and is not available in MODEL_MAPPING''')
    def 	SCREAMING_SNAKE_CASE							(      self				:     Dict)						->Tuple:





           '''simple docstring'''

           pass


    @unittest.skip(reason='''BlipTextModel has no base class and is not available in MODEL_MAPPING''')
    def 	SCREAMING_SNAKE_CASE							(      self				:     Any)						->Union[str, Any]:





           '''simple docstring'''

           pass


    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     str)						->Optional[Any]:





           '''simple docstring'''

           for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A__							=   TFBlipTextModel.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)





    def 	SCREAMING_SNAKE_CASE							(      self				:     Union[str, Any] ,			UpperCAmelCase__				:     str=True)						->str:





           '''simple docstring'''

           super().test_pt_tf_model_equivalence(allow_missing_keys=UpperCAmelCase__)








import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	str:



       """simple docstring"""

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

       return name







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	str:



       """simple docstring"""

       for key in orig_state_dict.copy().keys():
              A__							=   orig_state_dict.pop(lowercase_						)

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

              else:
                     A__							=   val

       return orig_state_dict







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	Tuple:



       """simple docstring"""

       A__							=   ViTMAEConfig()
       if "large" in checkpoint_url:
              A__							=   1_024
              A__							=   4_096
              A__							=   24
              A__							=   16
       elif "huge" in checkpoint_url:
              A__							=   14
              A__							=   1_280
              A__							=   5_120
              A__							=   32
              A__							=   16

       A__							=   ViTMAEForPreTraining(lowercase_						)

       A__							=   torch.hub.load_state_dict_from_url(lowercase_			,    map_location='''cpu'''						)['''model''']

       A__							=   ViTMAEImageProcessor(size=config.image_size						)

       A__							=   convert_state_dict(lowercase_			,    lowercase_						)

       model.load_state_dict(lowercase_						)
       model.eval()

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

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

       # forward pass
       torch.manual_seed(2						)
       A__							=   model(**lowercase_						)
       A__							=   outputs.logits

       if "large" in checkpoint_url:
              A__							=   torch.tensor(
                  [[-0.73_09, -0.71_28, -1.01_69], [-1.01_61, -0.90_58, -1.18_78], [-1.04_78, -0.94_11, -1.19_11]]						)
       elif "huge" in checkpoint_url:
              A__							=   torch.tensor(
                  [[-1.15_99, -0.91_99, -1.22_21], [-1.19_52, -0.92_69, -1.23_07], [-1.21_43, -0.93_37, -1.22_62]]						)
       else:
              A__							=   torch.tensor(
                  [[-0.91_92, -0.84_81, -1.12_59], [-1.13_49, -1.00_34, -1.25_99], [-1.17_57, -1.04_29, -1.27_26]]						)

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

       print(f"""Saving model to {pytorch_dump_folder_path}"""						)
       model.save_pretrained(lowercase_						)

       print(f"""Saving image processor to {pytorch_dump_folder_path}"""						)
       image_processor.save_pretrained(lowercase_						)


if __name__ == "__main__":
  _lowerCamelCase     :		Any								=       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     :		Optional[int]								=       parser.parse_args()
  convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


