




"""simple docstring"""

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 YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
a         =      logging.get_logger(__name__)






def 						_snake_case					( _snake_case							: int							)   -> Optional[int]:






      '''simple docstring'''


      _A        =       YolosConfig()

      # size of the architecture
      if "yolos_ti" in yolos_name:
            _A        =       1_92
            _A        =       7_68
            _A        =       12
            _A        =       3
            _A        =       [8_00, 13_33]
            _A        =       False
      elif yolos_name == "yolos_s_dWr":
            _A        =       3_30
            _A        =       14
            _A        =       6
            _A        =       13_20
      elif "yolos_s" in yolos_name:
            _A        =       3_84
            _A        =       15_36
            _A        =       12
            _A        =       6
      elif "yolos_b" in yolos_name:
            _A        =       [8_00, 13_44]

      _A        =       91
      _A        =       'huggingface/label-files'
      _A        =       'coco-detection-id2label.json'
      _A        =       json.load(open(hf_hub_download(lowerCamelCase__    ,  lowerCamelCase__    ,  repo_type='dataset'							)    ,  'r'							)							)
      _A        =       {int(lowerCamelCase__							): v for k, v in idalabel.items()}
      _A        =       idalabel
      _A        =       {v: k for k, v in idalabel.items()}

      return config






def 						_snake_case					( _snake_case							: int    ,  _snake_case							: str    ,  _snake_case							: int = False							)   -> Any:






      '''simple docstring'''


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






def 						_snake_case					( _snake_case							: Optional[Any]							)   -> Optional[Any]:






      '''simple docstring'''


      if "backbone" in name:
            _A        =       name.replace('backbone'    ,  'vit'							)
      if "cls_token" in name:
            _A        =       name.replace('cls_token'    ,  'embeddings.cls_token'							)
      if "det_token" in name:
            _A        =       name.replace('det_token'    ,  'embeddings.detection_tokens'							)
      if "mid_pos_embed" in name:
            _A        =       name.replace('mid_pos_embed'    ,  'encoder.mid_position_embeddings'							)
      if "pos_embed" in name:
            _A        =       name.replace('pos_embed'    ,  'embeddings.position_embeddings'							)
      if "patch_embed.proj" in name:
            _A        =       name.replace('patch_embed.proj'    ,  'embeddings.patch_embeddings.projection'							)
      if "blocks" in name:
            _A        =       name.replace('blocks'    ,  '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 "class_embed" in name:
            _A        =       name.replace('class_embed'    ,  'class_labels_classifier'							)
      if "bbox_embed" in name:
            _A        =       name.replace('bbox_embed'    ,  'bbox_predictor'							)
      if "vit.norm" in name:
            _A        =       name.replace('vit.norm'    ,  'vit.layernorm'							)

      return name






def 						_snake_case					( _snake_case							: List[Any]    ,  _snake_case							: List[str]							)   -> Dict:






      '''simple docstring'''


      for key in orig_state_dict.copy().keys():
            _A        =       orig_state_dict.pop(lowerCamelCase__							)

            if "qkv" in key:
                  _A        =       key.split('.'							)
                  _A        =       int(key_split[2]							)
                  _A        =       model.vit.encoder.layer[layer_num].attention.attention.all_head_size
                  if "weight" in key:
                        _A        =       val[:dim, :]
                        _A        =       val[
                            dim : dim * 2, :
                        ]
                        _A        =       val[-dim:, :]
                  else:
                        _A        =       val[:dim]
                        _A        =       val[dim : dim * 2]
                        _A        =       val[-dim:]
            else:
                  _A        =       val

      return orig_state_dict






def 						_snake_case					( )   -> str:






      '''simple docstring'''


      _A        =       'http://images.cocodataset.org/val2017/000000039769.jpg'
      _A        =       Image.open(requests.get(lowerCamelCase__    ,  stream=lowerCamelCase__							).raw							)
      return im






@torch.no_grad()
def 						_snake_case					( _snake_case							: Optional[Any]    ,  _snake_case							: Optional[int]    ,  _snake_case							: Dict    ,  _snake_case							: List[str] = False							)   -> Tuple:






      '''simple docstring'''


      _A        =       get_yolos_config(lowerCamelCase__							)

      # load original state_dict
      _A        =       torch.load(lowerCamelCase__    ,  map_location='cpu'							)['model']

      # load 🤗 model
      _A        =       YolosForObjectDetection(lowerCamelCase__							)
      model.eval()
      _A        =       convert_state_dict(lowerCamelCase__    ,  lowerCamelCase__							)
      model.load_state_dict(lowerCamelCase__							)

      # Check outputs on an image, prepared by YolosImageProcessor
      _A        =       8_00 if yolos_name != 'yolos_ti' else 5_12
      _A        =       YolosImageProcessor(format='coco_detection'    ,  size=lowerCamelCase__							)
      _A        =       image_processor(images=prepare_img()    ,  return_tensors='pt'							)
      _A        =       model(**lowerCamelCase__							)
      _A       ,   _A        =       outputs.logits, outputs.pred_boxes

      _A       ,   _A        =       None, None
      if yolos_name == "yolos_ti":
            _A        =       torch.tensor(
                [[-39.5022, -11.9820, -17.6888], [-29.9574, -9.9769, -17.7691], [-42.3281, -20.7200, -30.6294]]							)
            _A        =       torch.tensor(
                [[0.4021, 0.0836, 0.7979], [0.0184, 0.2609, 0.0364], [0.1781, 0.2004, 0.2095]]							)
      elif yolos_name == "yolos_s_200_pre":
            _A        =       torch.tensor(
                [[-24.0248, -10.3024, -14.8290], [-42.0392, -16.8200, -27.4334], [-27.2743, -11.8154, -18.7148]]							)
            _A        =       torch.tensor(
                [[0.2559, 0.5455, 0.4706], [0.2989, 0.7279, 0.1875], [0.7732, 0.4017, 0.4462]]							)
      elif yolos_name == "yolos_s_300_pre":
            _A        =       torch.tensor(
                [[-36.2220, -14.4385, -23.5457], [-35.6970, -14.7583, -21.3935], [-31.5939, -13.6042, -16.8049]]							)
            _A        =       torch.tensor(
                [[0.7614, 0.2316, 0.4728], [0.7168, 0.4495, 0.3855], [0.4996, 0.1466, 0.9996]]							)
      elif yolos_name == "yolos_s_dWr":
            _A        =       torch.tensor(
                [[-42.8668, -24.1049, -41.1690], [-34.7456, -14.1274, -24.9194], [-33.7898, -12.1946, -25.6495]]							)
            _A        =       torch.tensor(
                [[0.5587, 0.2773, 0.0605], [0.5004, 0.3014, 0.9994], [0.4999, 0.1548, 0.9994]]							)
      elif yolos_name == "yolos_base":
            _A        =       torch.tensor(
                [[-40.6064, -24.3084, -32.6447], [-55.1990, -30.7719, -35.5877], [-51.4311, -33.3507, -35.6462]]							)
            _A        =       torch.tensor(
                [[0.5555, 0.2794, 0.0655], [0.9049, 0.2664, 0.1894], [0.9183, 0.1984, 0.1635]]							)
      else:
            raise ValueError(F'''Unknown yolos_name: {yolos_name}'''							)

      assert torch.allclose(logits[0, :3, :3]    ,  lowerCamelCase__    ,  atol=1E-4							)
      assert torch.allclose(pred_boxes[0, :3, :3]    ,  lowerCamelCase__    ,  atol=1E-4							)

      Path(lowerCamelCase__							).mkdir(exist_ok=lowerCamelCase__							)
      print(F'''Saving model {yolos_name} 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 push_to_hub:
            _A        =       {
                'yolos_ti': 'yolos-tiny',
                'yolos_s_200_pre': 'yolos-small',
                'yolos_s_300_pre': 'yolos-small-300',
                'yolos_s_dWr': 'yolos-small-dwr',
                'yolos_base': 'yolos-base',
            }

            print('Pushing to the hub...'							)
            _A        =       model_mapping[yolos_name]
            image_processor.push_to_hub(lowerCamelCase__    ,  organization='hustvl'							)
            model.push_to_hub(lowerCamelCase__    ,  organization='hustvl'							)


if __name__ == "__main__":
   a         =      argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '''--yolos_name''',
       default='''yolos_s_200_pre''',
       type=str,
       help=(
           '''Name of the YOLOS model you\'d like to convert. Should be one of \'yolos_ti\', \'yolos_s_200_pre\','''
           ''' \'yolos_s_300_pre\', \'yolos_s_dWr\', \'yolos_base\'.'''
       ),
   )
   parser.add_argument(
       '''--checkpoint_path''', default=None, type=str, help='''Path to the original state dict (.pth file).'''
   )
   parser.add_argument(
       '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
   )
   parser.add_argument(
       '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
   )

   a         =      parser.parse_args()
   convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)







"""simple docstring"""

import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
       import tensorflow as tf
       from tensorflow.python.framework.errors_impl import ResourceExhaustedError

       from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
       import pyanvml.pyanvml as nvml

a         =      logging.get_logger(__name__)






def 						_snake_case					( _snake_case							: bool    ,  _snake_case							: bool							)   -> Tuple:






      '''simple docstring'''


      def run_func(_snake_case							: Any							):
            @wraps(_snake_case							)
            def run_in_eager_mode(*_snake_case							: List[str]    ,  **_snake_case							: Tuple							):
                  return func(*_snake_case    ,  **_snake_case							)

            @wraps(_snake_case							)
            @tf.function(experimental_compile=_snake_case							)
            def run_in_graph_mode(*_snake_case							: Dict    ,  **_snake_case							: Tuple							):
                  return func(*_snake_case    ,  **_snake_case							)

            if do_eager_mode is True:
                  if use_xla is not False:
                        raise ValueError(
                            'Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.'							)
                  return run_in_eager_mode
            else:
                  return run_in_graph_mode

      return run_func






def 						_snake_case					( _snake_case							: int    ,  _snake_case							: int    ,  _snake_case							: int							)   -> ["tf.Tensor"]:






      '''simple docstring'''


      _A        =       random.Random()
      _A        =       [rng.randint(0    ,  vocab_size - 1							) for i in range(batch_size * sequence_length							)]
      return tf.constant(_snake_case    ,  shape=(batch_size, sequence_length)    ,  dtype=tf.intaa							)







class     lowercase_	(					__lowerCAmelCase							):



 '''simple docstring'''
 UpperCAmelCase   :					TensorFlowBenchmarkArguments
 UpperCAmelCase   :					PretrainedConfig
 UpperCAmelCase   :					str				  =     "TensorFlow"

 @property
 def       lowerCAmelCase_							(			self	:	str		):
       return tf.__version__

 def       lowerCAmelCase_							(			self	:	List[str]   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	int		):
       # initialize GPU on separate process
       _A        =       self.args.strategy
       if strategy is None:
             raise ValueError('A device strategy has to be initialized before using TensorFlow.'		)
       _A        =       self._prepare_inference_func(_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase		)
       return self._measure_speed(_inference		)

 def       lowerCAmelCase_							(			self	:	Optional[Any]   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	int		):
       _A        =       self.args.strategy
       if strategy is None:
             raise ValueError('A device strategy has to be initialized before using TensorFlow.'		)
       _A        =       self._prepare_train_func(_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase		)
       return self._measure_speed(_train		)

 def       lowerCAmelCase_							(			self	:	Dict   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	int		):
       # initialize GPU on separate process
       if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]   ,							_UpperCAmelCase		)
       _A        =       self.args.strategy
       if strategy is None:
             raise ValueError('A device strategy has to be initialized before using TensorFlow.'		)
       _A        =       self._prepare_inference_func(_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase		)
       return self._measure_memory(_inference		)

 def       lowerCAmelCase_							(			self	:	int   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	int		):
       if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]   ,							_UpperCAmelCase		)
       _A        =       self.args.strategy
       if strategy is None:
             raise ValueError('A device strategy has to be initialized before using TensorFlow.'		)

       _A        =       self._prepare_train_func(_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase		)
       return self._measure_memory(_train		)

 def       lowerCAmelCase_							(			self	:	List[Any]   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	int		):
       _A        =       self.config_dict[model_name]

       if self.args.fpaa:
             raise NotImplementedError('Mixed precision is currently not supported.'		)

       _A        =       (
           hasattr(_UpperCAmelCase   ,							'architectures'		)
           and isinstance(config.architectures   ,							_UpperCAmelCase		)
           and len(config.architectures		) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                   _A        =       'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
                   _A        =       __import__('transformers'   ,							fromlist=[model_class]		)
                   _A        =       getattr(_UpperCAmelCase   ,							_UpperCAmelCase		)
                   _A        =       model_cls(_UpperCAmelCase		)
             except ImportError:
                   raise ImportError(
                       F'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                       ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'		)
       else:
             _A        =       TF_MODEL_MAPPING[config.__class__](_UpperCAmelCase		)

       # encoder-decoder has vocab size saved differently
       _A        =       config.vocab_size if hasattr(_UpperCAmelCase   ,							'vocab_size'		) else config.encoder.vocab_size
       _A        =       random_input_ids(_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase		)

       @run_with_tf_optimizations(self.args.eager_mode   ,							self.args.use_xla		)
       def encoder_decoder_forward():
             return model(_UpperCAmelCase   ,							decoder_input_ids=_UpperCAmelCase   ,							training=_UpperCAmelCase		)

       @run_with_tf_optimizations(self.args.eager_mode   ,							self.args.use_xla		)
       def encoder_forward():
             return model(_UpperCAmelCase   ,							training=_UpperCAmelCase		)

       _A        =       encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

       return _inference

 def       lowerCAmelCase_							(			self	:	str   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	int		):
       _A        =       self.config_dict[model_name]

       if self.args.eager_mode is not False:
             raise ValueError('Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.'		)

       if self.args.fpaa:
             raise NotImplementedError('Mixed precision is currently not supported.'		)

       _A        =       (
           hasattr(_UpperCAmelCase   ,							'architectures'		)
           and isinstance(config.architectures   ,							_UpperCAmelCase		)
           and len(config.architectures		) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                   _A        =       'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
                   _A        =       __import__('transformers'   ,							fromlist=[model_class]		)
                   _A        =       getattr(_UpperCAmelCase   ,							_UpperCAmelCase		)
                   _A        =       model_cls(_UpperCAmelCase		)
             except ImportError:
                   raise ImportError(
                       F'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                       ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'		)
       else:
             _A        =       TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_UpperCAmelCase		)

       # encoder-decoder has vocab size saved differently
       _A        =       config.vocab_size if hasattr(_UpperCAmelCase   ,							'vocab_size'		) else config.encoder.vocab_size
       _A        =       random_input_ids(_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase		)

       @run_with_tf_optimizations(self.args.eager_mode   ,							self.args.use_xla		)
       def encoder_decoder_train():
             _A        =       model(_UpperCAmelCase   ,							decoder_input_ids=_UpperCAmelCase   ,							labels=_UpperCAmelCase   ,							training=_UpperCAmelCase		)[0]
             _A        =       tf.gradients(_UpperCAmelCase   ,							model.trainable_variables		)
             return gradients

       @run_with_tf_optimizations(self.args.eager_mode   ,							self.args.use_xla		)
       def encoder_train():
             _A        =       model(_UpperCAmelCase   ,							labels=_UpperCAmelCase   ,							training=_UpperCAmelCase		)[0]
             _A        =       tf.gradients(_UpperCAmelCase   ,							model.trainable_variables		)
             return gradients

       _A        =       encoder_decoder_train if config.is_encoder_decoder else encoder_train

       return _train

 def       lowerCAmelCase_							(			self	:	Tuple   ,							_UpperCAmelCase	:	int		):
       with self.args.strategy.scope():
             try:
                   if self.args.is_tpu or self.args.use_xla:
                         # run additional 10 times to stabilize compilation for tpu
                         logger.info('Do inference on TPU. Running model 5 times to stabilize compilation'		)
                         timeit.repeat(_UpperCAmelCase   ,							repeat=1   ,							number=5		)

                   # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
                   _A        =       timeit.repeat(
                       _UpperCAmelCase   ,							repeat=self.args.repeat   ,							number=10   ,							)

                   return min(_UpperCAmelCase		) / 10.0
             except ResourceExhaustedError as e:
                   self.print_fn(F'''Doesn\'t fit on GPU. {e}'''		)



 def       lowerCAmelCase_							(			self	:	Optional[int]   ,							_UpperCAmelCase	:	Callable[[], None]		):
       logger.info(
           'Note that TensorFlow allocates more memory than '
           'it might need to speed up computation. '
           'The memory reported here corresponds to the memory '
           'reported by `nvidia-smi`, which can vary depending '
           'on total available memory on the GPU that is used.'		)
       with self.args.strategy.scope():
             try:
                   if self.args.trace_memory_line_by_line:
                         if not self.args.eager_mode:
                               raise ValueError(
                                   '`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory'
                                   ' consumption line by line.'		)
                         _A        =       start_memory_tracing('transformers'		)

                   if self.args.is_tpu:
                         # tpu
                         raise NotImplementedError(
                             'Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking'
                             ' with `args.memory=False`'		)
                   elif self.args.is_gpu:
                         # gpu
                         if not is_pyanvml_available():
                               logger.warning(
                                   'py3nvml not installed, we won\'t log GPU memory usage. '
                                   'Install py3nvml (pip install py3nvml) to log information about GPU.'		)
                               _A        =       'N/A'
                         else:
                               logger.info(
                                   'Measuring total GPU usage on GPU device. Make sure to not have additional processes'
                                   ' running on the same GPU.'		)
                               # init nvml
                               nvml.nvmlInit()
                               func()
                               _A        =       nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx		)
                               _A        =       nvml.nvmlDeviceGetMemoryInfo(_UpperCAmelCase		)
                               _A        =       meminfo.used
                               _A        =       Memory(_UpperCAmelCase		)
                               # shutdown nvml
                               nvml.nvmlShutdown()
                   else:
                         # cpu
                         if self.args.trace_memory_line_by_line:
                               logger.info(
                                   'When enabling line by line tracing, the max peak memory for CPU is inaccurate in'
                                   ' TensorFlow.'		)
                               _A        =       None
                         else:
                               _A        =       measure_peak_memory_cpu(_UpperCAmelCase		)
                               _A        =       Memory(_UpperCAmelCase		) if isinstance(_UpperCAmelCase   ,							_UpperCAmelCase		) else memory_bytes
                   if self.args.trace_memory_line_by_line:
                         _A        =       stop_memory_tracing(_UpperCAmelCase		)
                         if memory is None:
                               _A        =       summary.total
                   else:
                         _A        =       None

                   return memory, summary
             except ResourceExhaustedError as e:
                   self.print_fn(F'''Doesn\'t fit on GPU. {e}'''		)
                   return "N/A", None

