


'''simple docstring'''






import importlib
import json
import os
from collections import OrderedDict
from typing import Dict, Optional, Union

# Build the list of all image processors
from ...configuration_utils import PretrainedConfig
from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
from ...image_processing_utils import ImageProcessingMixin
from ...utils import CONFIG_NAME, IMAGE_PROCESSOR_NAME, get_file_from_repo, logging
from .auto_factory import _LazyAutoMapping
from .configuration_auto import (
    CONFIG_MAPPING_NAMES,
    AutoConfig,
    model_type_to_module_name,
    replace_list_option_in_docstrings,
)


__A :  Optional[int]    							=				logging.get_logger(__name__)

__A :  List[str]    							=				OrderedDict(
    [
        ("align", "EfficientNetImageProcessor"),
        ("beit", "BeitImageProcessor"),
        ("bit", "BitImageProcessor"),
        ("blip", "BlipImageProcessor"),
        ("blip-2", "BlipImageProcessor"),
        ("bridgetower", "BridgeTowerImageProcessor"),
        ("chinese_clip", "ChineseCLIPImageProcessor"),
        ("clip", "CLIPImageProcessor"),
        ("clipseg", "ViTImageProcessor"),
        ("conditional_detr", "ConditionalDetrImageProcessor"),
        ("convnext", "ConvNextImageProcessor"),
        ("convnextv2", "ConvNextImageProcessor"),
        ("cvt", "ConvNextImageProcessor"),
        ("data2vec-vision", "BeitImageProcessor"),
        ("deformable_detr", "DeformableDetrImageProcessor"),
        ("deit", "DeiTImageProcessor"),
        ("deta", "DetaImageProcessor"),
        ("detr", "DetrImageProcessor"),
        ("dinat", "ViTImageProcessor"),
        ("donut-swin", "DonutImageProcessor"),
        ("dpt", "DPTImageProcessor"),
        ("efficientformer", "EfficientFormerImageProcessor"),
        ("efficientnet", "EfficientNetImageProcessor"),
        ("flava", "FlavaImageProcessor"),
        ("focalnet", "BitImageProcessor"),
        ("git", "CLIPImageProcessor"),
        ("glpn", "GLPNImageProcessor"),
        ("groupvit", "CLIPImageProcessor"),
        ("imagegpt", "ImageGPTImageProcessor"),
        ("instructblip", "BlipImageProcessor"),
        ("layoutlmv2", "LayoutLMv2ImageProcessor"),
        ("layoutlmv3", "LayoutLMv3ImageProcessor"),
        ("levit", "LevitImageProcessor"),
        ("mask2former", "Mask2FormerImageProcessor"),
        ("maskformer", "MaskFormerImageProcessor"),
        ("mgp-str", "ViTImageProcessor"),
        ("mobilenet_v1", "MobileNetV1ImageProcessor"),
        ("mobilenet_v2", "MobileNetV2ImageProcessor"),
        ("mobilevit", "MobileViTImageProcessor"),
        ("mobilevit", "MobileViTImageProcessor"),
        ("mobilevitv2", "MobileViTImageProcessor"),
        ("nat", "ViTImageProcessor"),
        ("oneformer", "OneFormerImageProcessor"),
        ("owlvit", "OwlViTImageProcessor"),
        ("perceiver", "PerceiverImageProcessor"),
        ("pix2struct", "Pix2StructImageProcessor"),
        ("poolformer", "PoolFormerImageProcessor"),
        ("regnet", "ConvNextImageProcessor"),
        ("resnet", "ConvNextImageProcessor"),
        ("sam", "SamImageProcessor"),
        ("segformer", "SegformerImageProcessor"),
        ("swiftformer", "ViTImageProcessor"),
        ("swin", "ViTImageProcessor"),
        ("swin2sr", "Swin2SRImageProcessor"),
        ("swinv2", "ViTImageProcessor"),
        ("table-transformer", "DetrImageProcessor"),
        ("timesformer", "VideoMAEImageProcessor"),
        ("tvlt", "TvltImageProcessor"),
        ("upernet", "SegformerImageProcessor"),
        ("van", "ConvNextImageProcessor"),
        ("videomae", "VideoMAEImageProcessor"),
        ("vilt", "ViltImageProcessor"),
        ("vit", "ViTImageProcessor"),
        ("vit_hybrid", "ViTHybridImageProcessor"),
        ("vit_mae", "ViTImageProcessor"),
        ("vit_msn", "ViTImageProcessor"),
        ("xclip", "CLIPImageProcessor"),
        ("yolos", "YolosImageProcessor"),
    ]
)

__A :  Union[str, Any]    							=				_LazyAutoMapping(CONFIG_MAPPING_NAMES, IMAGE_PROCESSOR_MAPPING_NAMES)







def     UpperCamelCase_      (   A__ :       str ):
  '''simple docstring'''





  for module_name, extractors in IMAGE_PROCESSOR_MAPPING_NAMES.items():
    if class_name in extractors:
      lowerCAmelCase_      :				Optional[Any]      =       model_type_to_module_name(A__ )

      lowerCAmelCase_      :				str      =       importlib.import_module(f'.{module_name}'							,   """transformers.models""" )
      try:
        return getattr(A__							,   A__ )
      except AttributeError:
        continue

  for _, extractor in IMAGE_PROCESSOR_MAPPING._extra_content.items():
    if getattr(A__							,   """__name__"""							,   A__ ) == class_name:
      return extractor

    # We did not fine the class, but maybe it's because a dep is missing. In that case, the class will be in the main
    # init and we return the proper dummy to get an appropriate error message.
  lowerCAmelCase_      :				List[str]      =       importlib.import_module("""transformers""" )
  if hasattr(A__							,   A__ ):
    return getattr(A__							,   A__ )

  return None







def     UpperCamelCase_      (   A__ :       Union[str, os.PathLike]							,   A__ :       Optional[Union[str, os.PathLike]] = None							,   A__ :       bool = False							,   A__ :       bool = False							,   A__ :       Optional[Dict[str, str]] = None							,   A__ :       Optional[Union[bool, str]] = None							,   A__ :       Optional[str] = None							,   A__ :       bool = False							,   **A__ :       Tuple							,   ):
  '''simple docstring'''





  lowerCAmelCase_      :				Optional[Any]      =       get_file_from_repo(
      A__							,   A__							,   cache_dir=A__							,   force_download=A__							,   resume_download=A__							,   proxies=A__							,   use_auth_token=A__							,   revision=A__							,   local_files_only=A__							,   )
  if resolved_config_file is None:
    logger.info(
        """Could not locate the image processor configuration file, will try to use the model config instead.""" )
    return {}

  with open(A__							,   encoding="""utf-8""" ) as reader:
    return json.load(A__ )






class        __snake_case	:







    """simple docstring"""



    def __init__(       self   :  Union[str, Any]		)       ->  str:
      raise EnvironmentError(
          """AutoImageProcessor is designed to be instantiated """
          """using the `AutoImageProcessor.from_pretrained(pretrained_model_name_or_path)` method."""		)

    @classmethod
    @replace_list_option_in_docstrings(lowerCamelCase		)
    def __lowercase       (       cls   :  Tuple   ,						lowerCamelCase   :  Dict   ,						**lowerCamelCase   :  Optional[Any]		)       ->  Dict:
      lowerCAmelCase_      :				List[str]      =       kwargs.pop("""config"""   ,						lowerCamelCase		)
      lowerCAmelCase_      :				Tuple      =       kwargs.pop("""trust_remote_code"""   ,						lowerCamelCase		)
      lowerCAmelCase_      :				List[str]      =       True

      lowerCAmelCase_,							lowerCAmelCase_      :				int      =       ImageProcessingMixin.get_image_processor_dict(lowerCamelCase   ,						**lowerCamelCase		)
      lowerCAmelCase_      :				List[Any]      =       config_dict.get("""image_processor_type"""   ,						lowerCamelCase		)
      lowerCAmelCase_      :				int      =       None
      if "AutoImageProcessor" in config_dict.get("""auto_map"""   ,						{}		):
        lowerCAmelCase_      :				str      =       config_dict["""auto_map"""]["""AutoImageProcessor"""]

      # If we still don't have the image processor class, check if we're loading from a previous feature extractor config
      # and if so, infer the image processor class from there.
      if image_processor_class is None and image_processor_auto_map is None:
        lowerCAmelCase_      :				List[Any]      =       config_dict.pop("""feature_extractor_type"""   ,						lowerCamelCase		)
        if feature_extractor_class is not None:
          logger.warning(
              """Could not find image processor class in the image processor config or the model config. Loading"""
              """ based on pattern matching with the model's feature extractor configuration."""		)
          lowerCAmelCase_      :				Optional[int]      =       feature_extractor_class.replace("""FeatureExtractor"""   ,						"""ImageProcessor"""		)
        if "AutoFeatureExtractor" in config_dict.get("""auto_map"""   ,						{}		):
          lowerCAmelCase_      :				Optional[Any]      =       config_dict["""auto_map"""]["""AutoFeatureExtractor"""]
          lowerCAmelCase_      :				Dict      =       feature_extractor_auto_map.replace("""FeatureExtractor"""   ,						"""ImageProcessor"""		)
          logger.warning(
              """Could not find image processor auto map in the image processor config or the model config."""
              """ Loading based on pattern matching with the model's feature extractor configuration."""		)

        # If we don't find the image processor class in the image processor config, let's try the model config.
      if image_processor_class is None and image_processor_auto_map is None:
        if not isinstance(lowerCamelCase   ,						lowerCamelCase		):
          lowerCAmelCase_      :				str      =       AutoConfig.from_pretrained(lowerCamelCase   ,						**lowerCamelCase		)
        # It could be in `config.image_processor_type``
        lowerCAmelCase_      :				Optional[Any]      =       getattr(lowerCamelCase   ,						"""image_processor_type"""   ,						lowerCamelCase		)
        if hasattr(lowerCamelCase   ,						"""auto_map"""		) and "AutoImageProcessor" in config.auto_map:
          lowerCAmelCase_      :				Dict      =       config.auto_map["""AutoImageProcessor"""]

      if image_processor_class is not None:
        lowerCAmelCase_      :				Optional[Any]      =       image_processor_class_from_name(lowerCamelCase		)

      lowerCAmelCase_      :				Optional[Any]      =       image_processor_auto_map is not None
      lowerCAmelCase_      :				Any      =       image_processor_class is not None or type(lowerCamelCase		) in IMAGE_PROCESSOR_MAPPING
      lowerCAmelCase_      :				Union[str, Any]      =       resolve_trust_remote_code(
          lowerCamelCase   ,						lowerCamelCase   ,						lowerCamelCase   ,						lowerCamelCase		)

      if has_remote_code and trust_remote_code:
        lowerCAmelCase_      :				Dict      =       get_class_from_dynamic_module(
            lowerCamelCase   ,						lowerCamelCase   ,						**lowerCamelCase		)
        lowerCAmelCase_      :				int      =       kwargs.pop("""code_revision"""   ,						lowerCamelCase		)
        if os.path.isdir(lowerCamelCase		):
          image_processor_class.register_for_auto_class()
        return image_processor_class.from_dict(lowerCamelCase   ,						**lowerCamelCase		)
      elif image_processor_class is not None:
        return image_processor_class.from_dict(lowerCamelCase   ,						**lowerCamelCase		)
      # Last try: we use the IMAGE_PROCESSOR_MAPPING.
      elif type(lowerCamelCase		) in IMAGE_PROCESSOR_MAPPING:
        lowerCAmelCase_      :				List[str]      =       IMAGE_PROCESSOR_MAPPING[type(lowerCamelCase		)]
        return image_processor_class.from_dict(lowerCamelCase   ,						**lowerCamelCase		)

      raise ValueError(
          F'Unrecognized image processor in {pretrained_model_name_or_path}. Should have a '
          F'`image_processor_type` key in its {IMAGE_PROCESSOR_NAME} of {CONFIG_NAME}, or one of the following '
          F'`model_type` keys in its {CONFIG_NAME}: {", ".join(c for c in IMAGE_PROCESSOR_MAPPING_NAMES.keys()		)}'		)





    @staticmethod
    def __lowercase       (       lowerCamelCase   :  Optional[Any]   ,						lowerCamelCase   :  Optional[Any]		)       ->  List[Any]:
      IMAGE_PROCESSOR_MAPPING.register(lowerCamelCase   ,						lowerCamelCase		)







'''simple docstring'''






import argparse
import json
from pathlib import Path

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

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


logging.set_verbosity_info()
__A :  List[str]    							=				logging.get_logger(__name__)







def     UpperCamelCase_      (   A__ :       Union[str, Any]							,   A__ :       Tuple=False ):
  '''simple docstring'''





  lowerCAmelCase_      :				Tuple      =       []
  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"""),
            ("""pre_logits.fc.weight""", """pooler.dense.weight"""),
            ("""pre_logits.fc.bias""", """pooler.dense.bias"""),
        ] )

    # if just the base model, we should remove "vit" from all keys that start with "vit"
    lowerCAmelCase_      :				List[str]      =       [(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     UpperCamelCase_      (   A__ :       Any							,   A__ :       Any							,   A__ :       Tuple=False ):
  '''simple docstring'''





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







def     UpperCamelCase_      (   A__ :       str ):
  '''simple docstring'''





  lowerCAmelCase_      :				Dict      =       ["""head.weight""", """head.bias"""]
  for k in ignore_keys:
    state_dict.pop(A__							,   A__ )







def     UpperCamelCase_      (   A__ :       List[Any]							,   A__ :       Optional[Any]							,   A__ :       Dict ):
  '''simple docstring'''





  lowerCAmelCase_      :				Tuple      =       dct.pop(A__ )
  lowerCAmelCase_      :				Tuple      =       val







def     UpperCamelCase_      (   ):
  '''simple docstring'''





  lowerCAmelCase_      :				str      =       """http://images.cocodataset.org/val2017/000000039769.jpg"""
  lowerCAmelCase_      :				Optional[int]      =       Image.open(requests.get(A__							,   stream=A__ ).raw )
  return im







@torch.no_grad()
def     UpperCamelCase_      (   A__ :       Union[str, Any]							,   A__ :       List[Any] ):
  '''simple docstring'''





  lowerCAmelCase_      :				Optional[Any]      =       ViTConfig()
  lowerCAmelCase_      :				Any      =       False
  # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
  if vit_name[-5:] == "in21k":
    lowerCAmelCase_      :				int      =       True
    lowerCAmelCase_      :				Tuple      =       int(vit_name[-12:-10] )
    lowerCAmelCase_      :				Optional[int]      =       int(vit_name[-9:-6] )
  else:
    lowerCAmelCase_      :				Optional[int]      =       10_00
    lowerCAmelCase_      :				Tuple      =       """huggingface/label-files"""
    lowerCAmelCase_      :				Any      =       """imagenet-1k-id2label.json"""
    lowerCAmelCase_      :				Dict      =       json.load(open(hf_hub_download(A__							,   A__							,   repo_type="""dataset""" )							,   """r""" ) )
    lowerCAmelCase_      :				Union[str, Any]      =       {int(A__ ): v for k, v in idalabel.items()}
    lowerCAmelCase_      :				Union[str, Any]      =       idalabel
    lowerCAmelCase_      :				Union[str, Any]      =       {v: k for k, v in idalabel.items()}
    lowerCAmelCase_      :				Optional[int]      =       int(vit_name[-6:-4] )
    lowerCAmelCase_      :				Dict      =       int(vit_name[-3:] )
  # size of the architecture
  if "deit" in vit_name:
    if vit_name[9:].startswith("""tiny""" ):
      lowerCAmelCase_      :				int      =       1_92
      lowerCAmelCase_      :				List[str]      =       7_68
      lowerCAmelCase_      :				List[str]      =       12
      lowerCAmelCase_      :				int      =       3
    elif vit_name[9:].startswith("""small""" ):
      lowerCAmelCase_      :				Optional[Any]      =       3_84
      lowerCAmelCase_      :				Optional[int]      =       15_36
      lowerCAmelCase_      :				Dict      =       12
      lowerCAmelCase_      :				str      =       6
    else:
      pass
  else:
    if vit_name[4:].startswith("""small""" ):
      lowerCAmelCase_      :				Tuple      =       7_68
      lowerCAmelCase_      :				Any      =       23_04
      lowerCAmelCase_      :				List[str]      =       8
      lowerCAmelCase_      :				List[str]      =       8
    elif vit_name[4:].startswith("""base""" ):
      pass
    elif vit_name[4:].startswith("""large""" ):
      lowerCAmelCase_      :				Dict      =       10_24
      lowerCAmelCase_      :				List[Any]      =       40_96
      lowerCAmelCase_      :				Any      =       24
      lowerCAmelCase_      :				List[str]      =       16
    elif vit_name[4:].startswith("""huge""" ):
      lowerCAmelCase_      :				Optional[int]      =       12_80
      lowerCAmelCase_      :				Dict      =       51_20
      lowerCAmelCase_      :				Union[str, Any]      =       32
      lowerCAmelCase_      :				Optional[int]      =       16

    # load original model from timm
  lowerCAmelCase_      :				Union[str, Any]      =       timm.create_model(A__							,   pretrained=A__ )
  timm_model.eval()

  # load state_dict of original model, remove and rename some keys
  lowerCAmelCase_      :				int      =       timm_model.state_dict()
  if base_model:
    remove_classification_head_(A__ )
  lowerCAmelCase_      :				str      =       create_rename_keys(A__							,   A__ )
  for src, dest in rename_keys:
    rename_key(A__							,   A__							,   A__ )
  read_in_q_k_v(A__							,   A__							,   A__ )

  # load HuggingFace model
  if vit_name[-5:] == "in21k":
    lowerCAmelCase_      :				int      =       ViTModel(A__ ).eval()
  else:
    lowerCAmelCase_      :				Optional[int]      =       ViTForImageClassification(A__ ).eval()
  model.load_state_dict(A__ )

  # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
  if "deit" in vit_name:
    lowerCAmelCase_      :				Any      =       DeiTImageProcessor(size=config.image_size )
  else:
    lowerCAmelCase_      :				Any      =       ViTImageProcessor(size=config.image_size )
  lowerCAmelCase_      :				Tuple      =       image_processor(images=prepare_img()							,   return_tensors="""pt""" )
  lowerCAmelCase_      :				int      =       encoding["""pixel_values"""]
  lowerCAmelCase_      :				int      =       model(A__ )

  if base_model:
    lowerCAmelCase_      :				Union[str, Any]      =       timm_model.forward_features(A__ )
    assert timm_pooled_output.shape == outputs.pooler_output.shape
    assert torch.allclose(A__							,   outputs.pooler_output							,   atol=1E-3 )
  else:
    lowerCAmelCase_      :				Union[str, Any]      =       timm_model(A__ )
    assert timm_logits.shape == outputs.logits.shape
    assert torch.allclose(A__							,   outputs.logits							,   atol=1E-3 )

  Path(A__ ).mkdir(exist_ok=A__ )
  print(f'Saving model {vit_name} to {pytorch_dump_folder_path}' )
  model.save_pretrained(A__ )
  print(f'Saving image processor to {pytorch_dump_folder_path}' )
  image_processor.save_pretrained(A__ )


if __name__ == "__main__":
  __A :  Optional[int]    							=				argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--vit_name",
      default="vit_base_patch16_224",
      type=str,
      help="Name of the ViT timm model 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 :  Union[str, Any]    							=				parser.parse_args()
  convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)



