import collections
import os
import re
from pathlib import Path


snake_case__				= 'src/transformers'


# Matches is_xxx_available()
snake_case__				= re.compile(R'is\_([a-z_]*)_available()')
# Catches a one-line _import_struct = {xxx}
snake_case__				= re.compile(R'^_import_structure\s+=\s+\{([^\}]+)\}')
# Catches a line with a key-values pattern: "bla": ["foo", "bar"]
snake_case__				= re.compile(R'\s+"\S*":\s+\[([^\]]*)\]')
# Catches a line if not is_foo_available
snake_case__				= re.compile(R'^\s*if\s+not\s+is\_[a-z_]*\_available\(\)')
# Catches a line _import_struct["bla"].append("foo")
snake_case__				= re.compile(R'^\s*_import_structure\["\S*"\]\.append\("(\S*)"\)')
# Catches a line _import_struct["bla"].extend(["foo", "bar"]) or _import_struct["bla"] = ["foo", "bar"]
snake_case__				= re.compile(R'^\s*_import_structure\[\S*\](?:\.extend\(|\s*=\s+)\[([^\]]*)\]')
# Catches a line with an object between quotes and a comma:     "MyModel",
snake_case__				= re.compile(R'^\s+"([^"]+)",')
# Catches a line with objects between brackets only:    ["foo", "bar"],
snake_case__				= re.compile(R'^\s+\[([^\]]+)\]')
# Catches a line with from foo import bar, bla, boo
snake_case__				= re.compile(R'\s+from\s+\S*\s+import\s+([^\(\s].*)\n')
# Catches a line with try:
snake_case__				= re.compile(R'^\s*try:')
# Catches a line with else:
snake_case__				= re.compile(R'^\s*else:')
def       __magic_name__(  __UpperCAmelCase		)						->    List[str]:



      '''simple docstring'''

      if _re_test_backend.search(__UpperCAmelCase		) is None:
            return None
      _lowerCamelCase			    =				[b[0] for b in _re_backend.findall(__UpperCAmelCase		)]
      backends.sort()
      return "_and_".join(__UpperCAmelCase		)
def       __magic_name__(  __UpperCAmelCase		)						->    str:



      '''simple docstring'''

      with open(__UpperCAmelCase ,						'''r''' ,						encoding='''utf-8''' ,						newline='''\n'''		) as f:
            _lowerCamelCase			    =				f.readlines()

      _lowerCamelCase			    =				0
      while line_index < len(__UpperCAmelCase		) and not lines[line_index].startswith('''_import_structure = {'''		):
            line_index += 1

      # If this is a traditional init, just return.
      if line_index >= len(__UpperCAmelCase		):
            return None

      # First grab the objects without a specific backend in _import_structure
      _lowerCamelCase			    =				[]
      while not lines[line_index].startswith('''if TYPE_CHECKING'''		) and find_backend(lines[line_index]		) is None:
            _lowerCamelCase			    =				lines[line_index]
            # If we have everything on a single line, let's deal with it.
            if _re_one_line_import_struct.search(__UpperCAmelCase		):
                  _lowerCamelCase			    =				_re_one_line_import_struct.search(__UpperCAmelCase		).groups()[0]
                  _lowerCamelCase			    =				re.findall(r'''\[([^\]]+)\]''' ,						__UpperCAmelCase		)
                  for imp in imports:
                        objects.extend([obj[1:-1] for obj in imp.split(''', '''		)]		)
                  line_index += 1
                  continue
            _lowerCamelCase			    =				_re_import_struct_key_value.search(__UpperCAmelCase		)
            if single_line_import_search is not None:
                  _lowerCamelCase			    =				[obj[1:-1] for obj in single_line_import_search.groups()[0].split(''', '''		) if len(__UpperCAmelCase		) > 0]
                  objects.extend(__UpperCAmelCase		)
            elif line.startswith(''' ''' * 8 + '''"'''		):
                  objects.append(line[9:-3]		)
            line_index += 1

      _lowerCamelCase			    =				{'''none''': objects}
      # Let's continue with backend-specific objects in _import_structure
      while not lines[line_index].startswith('''if TYPE_CHECKING'''		):
            # If the line is an if not is_backend_available, we grab all objects associated.
            _lowerCamelCase			    =				find_backend(lines[line_index]		)
            # Check if the backend declaration is inside a try block:
            if _re_try.search(lines[line_index - 1]		) is None:
                  _lowerCamelCase			    =				None

            if backend is not None:
                  line_index += 1

                  # Scroll until we hit the else block of try-except-else
                  while _re_else.search(lines[line_index]		) is None:
                        line_index += 1

                  line_index += 1

                  _lowerCamelCase			    =				[]
                  # Until we unindent, add backend objects to the list
                  while len(lines[line_index]		) <= 1 or lines[line_index].startswith(''' ''' * 4		):
                        _lowerCamelCase			    =				lines[line_index]
                        if _re_import_struct_add_one.search(__UpperCAmelCase		) is not None:
                              objects.append(_re_import_struct_add_one.search(__UpperCAmelCase		).groups()[0]		)
                        elif _re_import_struct_add_many.search(__UpperCAmelCase		) is not None:
                              _lowerCamelCase			    =				_re_import_struct_add_many.search(__UpperCAmelCase		).groups()[0].split(''', '''		)
                              _lowerCamelCase			    =				[obj[1:-1] for obj in imports if len(__UpperCAmelCase		) > 0]
                              objects.extend(__UpperCAmelCase		)
                        elif _re_between_brackets.search(__UpperCAmelCase		) is not None:
                              _lowerCamelCase			    =				_re_between_brackets.search(__UpperCAmelCase		).groups()[0].split(''', '''		)
                              _lowerCamelCase			    =				[obj[1:-1] for obj in imports if len(__UpperCAmelCase		) > 0]
                              objects.extend(__UpperCAmelCase		)
                        elif _re_quote_object.search(__UpperCAmelCase		) is not None:
                              objects.append(_re_quote_object.search(__UpperCAmelCase		).groups()[0]		)
                        elif line.startswith(''' ''' * 8 + '''"'''		):
                              objects.append(line[9:-3]		)
                        elif line.startswith(''' ''' * 12 + '''"'''		):
                              objects.append(line[13:-3]		)
                        line_index += 1

                  _lowerCamelCase			    =				objects
            else:
                  line_index += 1

    # At this stage we are in the TYPE_CHECKING part, first grab the objects without a specific backend
      _lowerCamelCase			    =				[]
      while (
          line_index < len(__UpperCAmelCase		)
          and find_backend(lines[line_index]		) is None
          and not lines[line_index].startswith('''else'''		)
      ):
            _lowerCamelCase			    =				lines[line_index]
            _lowerCamelCase			    =				_re_import.search(__UpperCAmelCase		)
            if single_line_import_search is not None:
                  objects.extend(single_line_import_search.groups()[0].split(''', '''		)		)
            elif line.startswith(''' ''' * 8		):
                  objects.append(line[8:-2]		)
            line_index += 1

      _lowerCamelCase			    =				{'''none''': objects}
      # Let's continue with backend-specific objects
      while line_index < len(__UpperCAmelCase		):
            # If the line is an if is_backend_available, we grab all objects associated.
            _lowerCamelCase			    =				find_backend(lines[line_index]		)
            # Check if the backend declaration is inside a try block:
            if _re_try.search(lines[line_index - 1]		) is None:
                  _lowerCamelCase			    =				None

            if backend is not None:
                  line_index += 1

                  # Scroll until we hit the else block of try-except-else
                  while _re_else.search(lines[line_index]		) is None:
                        line_index += 1

                  line_index += 1

                  _lowerCamelCase			    =				[]
                  # Until we unindent, add backend objects to the list
                  while len(lines[line_index]		) <= 1 or lines[line_index].startswith(''' ''' * 8		):
                        _lowerCamelCase			    =				lines[line_index]
                        _lowerCamelCase			    =				_re_import.search(__UpperCAmelCase		)
                        if single_line_import_search is not None:
                              objects.extend(single_line_import_search.groups()[0].split(''', '''		)		)
                        elif line.startswith(''' ''' * 12		):
                              objects.append(line[12:-2]		)
                        line_index += 1

                  _lowerCamelCase			    =				objects
            else:
                  line_index += 1

      return import_dict_objects, type_hint_objects
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    List[Any]:



      '''simple docstring'''

      def find_duplicates(__UpperCAmelCase		):
            return [k for k, v in collections.Counter(__UpperCAmelCase		).items() if v > 1]

      if list(import_dict_objects.keys()		) != list(type_hint_objects.keys()		):
            return ["Both sides of the init do not have the same backends!"]

      _lowerCamelCase			    =				[]
      for key in import_dict_objects.keys():
            _lowerCamelCase			    =				find_duplicates(import_dict_objects[key]		)
            if duplicate_imports:
                  errors.append(F'Duplicate _import_structure definitions for: {duplicate_imports}'		)
            _lowerCamelCase			    =				find_duplicates(type_hint_objects[key]		)
            if duplicate_type_hints:
                  errors.append(F'Duplicate TYPE_CHECKING objects for: {duplicate_type_hints}'		)

            if sorted(set(import_dict_objects[key]		)		) != sorted(set(type_hint_objects[key]		)		):
                  _lowerCamelCase			    =				'''base imports''' if key == '''none''' else F'{key} backend'
                  errors.append(F'Differences for {name}:'		)
                  for a in type_hint_objects[key]:
                        if a not in import_dict_objects[key]:
                              errors.append(F'  {a} in TYPE_HINT but not in _import_structure.'		)
                  for a in import_dict_objects[key]:
                        if a not in type_hint_objects[key]:
                              errors.append(F'  {a} in _import_structure but not in TYPE_HINT.'		)
      return errors
def       __magic_name__(  )						->    int:



      '''simple docstring'''

      _lowerCamelCase			    =				[]
      for root, _, files in os.walk(__UpperCAmelCase		):
            if "__init__.py" in files:
                  _lowerCamelCase			    =				os.path.join(__UpperCAmelCase ,						'''__init__.py'''		)
                  _lowerCamelCase			    =				parse_init(__UpperCAmelCase		)
                  if objects is not None:
                        _lowerCamelCase			    =				analyze_results(*__UpperCAmelCase		)
                        if len(__UpperCAmelCase		) > 0:
                              _lowerCamelCase			    =				F'Problem in {fname}, both halves do not define the same objects.\n{errors[0]}'
                              failures.append('''\n'''.join(__UpperCAmelCase		)		)
      if len(__UpperCAmelCase		) > 0:
            raise ValueError('''\n\n'''.join(__UpperCAmelCase		)		)
def       __magic_name__(  )						->    str:



      '''simple docstring'''

      _lowerCamelCase			    =				[]
      for path, directories, files in os.walk(__UpperCAmelCase		):
            for folder in directories:
                  # Ignore private modules
                  if folder.startswith('''_'''		):
                        directories.remove(__UpperCAmelCase		)
                        continue
                  # Ignore leftovers from branches (empty folders apart from pycache)
                  if len(list((Path(__UpperCAmelCase		) / folder).glob('''*.py'''		)		)		) == 0:
                        continue
                  _lowerCamelCase			    =				str((Path(__UpperCAmelCase		) / folder).relative_to(__UpperCAmelCase		)		)
                  _lowerCamelCase			    =				short_path.replace(os.path.sep ,						'''.'''		)
                  submodules.append(__UpperCAmelCase		)
            for fname in files:
                  if fname == "__init__.py":
                        continue
                  _lowerCamelCase			    =				str((Path(__UpperCAmelCase		) / fname).relative_to(__UpperCAmelCase		)		)
                  _lowerCamelCase			    =				short_path.replace('''.py''' ,						''''''		).replace(os.path.sep ,						'''.'''		)
                  if len(submodule.split('''.'''		)		) == 1:
                        submodules.append(__UpperCAmelCase		)
      return submodules


snake_case__				= [
    'convert_pytorch_checkpoint_to_tf2',
    'modeling_flax_pytorch_utils',
    'models.esm.openfold_utils',
]
def       __magic_name__(  )						->    List[Any]:



      '''simple docstring'''

      from transformers.utils import direct_transformers_import

      _lowerCamelCase			    =				direct_transformers_import(__UpperCAmelCase		)

      _lowerCamelCase			    =				set(transformers._import_structure.keys()		)
      # This contains all the base keys of the _import_structure object defined in the init, but if the user is missing
      # some optional dependencies, they may not have all of them. Thus we read the init to read all additions and
      # (potentiall re-) add them.
      with open(os.path.join(__UpperCAmelCase ,						'''__init__.py'''		) ,						'''r'''		) as f:
            _lowerCamelCase			    =				f.read()
      import_structure_keys.update(set(re.findall(r'''import_structure\[\"([^\"]*)\"\]''' ,						__UpperCAmelCase		)		)		)

      _lowerCamelCase			    =				[
          module
          for module in get_transformers_submodules()
          if module not in IGNORE_SUBMODULES and module not in import_structure_keys
      ]

      if len(__UpperCAmelCase		) > 0:
            _lowerCamelCase			    =				'''\n'''.join(F'- {module}' for module in module_not_registered		)
            raise ValueError(
                '''The following submodules are not properly registed in the main init of Transformers:\n'''
                F'{list_of_modules}\n'
                '''Make sure they appear somewhere in the keys of `_import_structure` with an empty list as value.'''		)


if __name__ == "__main__":
     check_all_inits()
     check_submodules()



import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
snake_case__				= logging.get_logger(__name__)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase=False		)						->    List[Any]:



      '''simple docstring'''

      _lowerCamelCase			    =				[]
      for i in range(config.num_hidden_layers		):
            # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
            rename_keys.append((F'blocks.{i}.norm1.weight', F'vit.encoder.layer.{i}.layernorm_before.weight')		)
            rename_keys.append((F'blocks.{i}.norm1.bias', F'vit.encoder.layer.{i}.layernorm_before.bias')		)
            rename_keys.append((F'blocks.{i}.attn.proj.weight', F'vit.encoder.layer.{i}.attention.output.dense.weight')		)
            rename_keys.append((F'blocks.{i}.attn.proj.bias', F'vit.encoder.layer.{i}.attention.output.dense.bias')		)
            rename_keys.append((F'blocks.{i}.norm2.weight', F'vit.encoder.layer.{i}.layernorm_after.weight')		)
            rename_keys.append((F'blocks.{i}.norm2.bias', F'vit.encoder.layer.{i}.layernorm_after.bias')		)
            rename_keys.append((F'blocks.{i}.mlp.fc1.weight', F'vit.encoder.layer.{i}.intermediate.dense.weight')		)
            rename_keys.append((F'blocks.{i}.mlp.fc1.bias', F'vit.encoder.layer.{i}.intermediate.dense.bias')		)
            rename_keys.append((F'blocks.{i}.mlp.fc2.weight', F'vit.encoder.layer.{i}.output.dense.weight')		)
            rename_keys.append((F'blocks.{i}.mlp.fc2.bias', F'vit.encoder.layer.{i}.output.dense.bias')		)

      # projection layer + position embeddings
      rename_keys.extend(
          [
              ('''cls_token''', '''vit.embeddings.cls_token'''),
              ('''patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
              ('''patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
              ('''pos_embed''', '''vit.embeddings.position_embeddings'''),
          ]		)

      if base_model:
            # layernorm + pooler
            rename_keys.extend(
                [
                    ('''norm.weight''', '''layernorm.weight'''),
                    ('''norm.bias''', '''layernorm.bias'''),
                ]		)

            # if just the base model, we should remove "vit" from all keys that start with "vit"
            _lowerCamelCase			    =				[(pair[0], pair[1][4:]) if pair[1].startswith('''vit'''		) else pair for pair in rename_keys]
      else:
            # layernorm + classification head
            rename_keys.extend(
                [
                    ('''norm.weight''', '''vit.layernorm.weight'''),
                    ('''norm.bias''', '''vit.layernorm.bias'''),
                    ('''head.weight''', '''classifier.weight'''),
                    ('''head.bias''', '''classifier.bias'''),
                ]		)

      return rename_keys
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase=False		)						->    str:



      '''simple docstring'''

      for i in range(config.num_hidden_layers		):
            if base_model:
                  _lowerCamelCase			    =				''''''
            else:
                  _lowerCamelCase			    =				'''vit.'''
            # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
            _lowerCamelCase			    =				state_dict.pop(F'blocks.{i}.attn.qkv.weight'		)
            _lowerCamelCase			    =				state_dict.pop(F'blocks.{i}.attn.qkv.bias'		)
            # next, add query, keys and values (in that order) to the state dict
            _lowerCamelCase			    =				in_proj_weight[
                : config.hidden_size, :
            ]
            _lowerCamelCase			    =				in_proj_bias[: config.hidden_size]
            _lowerCamelCase			    =				in_proj_weight[
                config.hidden_size : config.hidden_size * 2, :
            ]
            _lowerCamelCase			    =				in_proj_bias[
                config.hidden_size : config.hidden_size * 2
            ]
            _lowerCamelCase			    =				in_proj_weight[
                -config.hidden_size :, :
            ]
            _lowerCamelCase			    =				in_proj_bias[-config.hidden_size :]
def       __magic_name__(  __UpperCAmelCase		)						->    Dict:



      '''simple docstring'''

      _lowerCamelCase			    =				['''head.weight''', '''head.bias''']
      for k in ignore_keys:
            state_dict.pop(__UpperCAmelCase ,						__UpperCAmelCase		)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    List[Any]:



      '''simple docstring'''

      _lowerCamelCase			    =				dct.pop(__UpperCAmelCase		)
      _lowerCamelCase			    =				val
def       __magic_name__(  )						->    List[str]:



      '''simple docstring'''

      _lowerCamelCase			    =				'''http://images.cocodataset.org/val2017/000000039769.jpg'''
      _lowerCamelCase			    =				Image.open(requests.get(__UpperCAmelCase ,						stream=__UpperCAmelCase		).raw		)
      return im
@torch.no_grad()
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase=True		)						->    str:



      '''simple docstring'''

      _lowerCamelCase			    =				ViTConfig()
      # patch_size
      if model_name[-1] == "8":
            _lowerCamelCase			    =				8
      # set labels if required
      if not base_model:
            _lowerCamelCase			    =				1000
            _lowerCamelCase			    =				'''huggingface/label-files'''
            _lowerCamelCase			    =				'''imagenet-1k-id2label.json'''
            _lowerCamelCase			    =				json.load(open(hf_hub_download(__UpperCAmelCase ,						__UpperCAmelCase ,						repo_type='''dataset'''		) ,						'''r'''		)		)
            _lowerCamelCase			    =				{int(__UpperCAmelCase		): v for k, v in idalabel.items()}
            _lowerCamelCase			    =				idalabel
            _lowerCamelCase			    =				{v: k for k, v in idalabel.items()}
      # size of the architecture
      if model_name in ["dino_vits8", "dino_vits16"]:
            _lowerCamelCase			    =				384
            _lowerCamelCase			    =				1536
            _lowerCamelCase			    =				12
            _lowerCamelCase			    =				6

      # load original model from torch hub
      _lowerCamelCase			    =				torch.hub.load('''facebookresearch/dino:main''' ,						__UpperCAmelCase		)
      original_model.eval()

      # load state_dict of original model, remove and rename some keys
      _lowerCamelCase			    =				original_model.state_dict()
      if base_model:
            remove_classification_head_(__UpperCAmelCase		)
      _lowerCamelCase			    =				create_rename_keys(__UpperCAmelCase ,						base_model=__UpperCAmelCase		)
      for src, dest in rename_keys:
            rename_key(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)
      read_in_q_k_v(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)

      # load HuggingFace model
      if base_model:
            _lowerCamelCase			    =				ViTModel(__UpperCAmelCase ,						add_pooling_layer=__UpperCAmelCase		).eval()
      else:
            _lowerCamelCase			    =				ViTForImageClassification(__UpperCAmelCase		).eval()
      model.load_state_dict(__UpperCAmelCase		)

      # Check outputs on an image, prepared by ViTImageProcessor
      _lowerCamelCase			    =				ViTImageProcessor()
      _lowerCamelCase			    =				image_processor(images=prepare_img() ,						return_tensors='''pt'''		)
      _lowerCamelCase			    =				encoding['''pixel_values''']
      _lowerCamelCase			    =				model(__UpperCAmelCase		)

      if base_model:
            _lowerCamelCase			    =				original_model(__UpperCAmelCase		)
            assert torch.allclose(__UpperCAmelCase ,						outputs.last_hidden_state[:, 0, :] ,						atol=1E-1		)
      else:
            _lowerCamelCase			    =				original_model(__UpperCAmelCase		)
            assert logits.shape == outputs.logits.shape
            assert torch.allclose(__UpperCAmelCase ,						outputs.logits ,						atol=1E-3		)

      Path(__UpperCAmelCase		).mkdir(exist_ok=__UpperCAmelCase		)
      print(F'Saving model {model_name} to {pytorch_dump_folder_path}'		)
      model.save_pretrained(__UpperCAmelCase		)
      print(F'Saving image processor to {pytorch_dump_folder_path}'		)
      image_processor.save_pretrained(__UpperCAmelCase		)


if __name__ == "__main__":
     snake_case__				= argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--model_name',
         default='dino_vitb16',
         type=str,
         help='Name of the model trained with DINO you\'d like to convert.',
     )
     parser.add_argument(
         '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
     )
     parser.add_argument(
         '--base_model',
         action='store_true',
         help='Whether to only convert the base model (no projection head weights).',
     )

     parser.set_defaults(base_model=True)
     snake_case__				= parser.parse_args()
     convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)