'''simple docstring'''


import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
 import torch

 from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
 import pyanvml.pyanvml as nvml


snake_case__     :			str      			=  logging.get_logger(__name__)






class      __SCREAMING_SNAKE_CASE						(       lowerCamelCase_       ):







  '''simple docstring'''






  lowerCamelCase_		:PyTorchBenchmarkArguments
  lowerCamelCase_		:PretrainedConfig
  lowerCamelCase_		:str						    =		"PyTorch"


  @property
  def _UpperCamelCase     ( self						):




       '''simple docstring'''

       return torch.__version__


  def _UpperCamelCase     ( self      ,		snake_case_      ,		snake_case_      ,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:   Any   =      self._prepare_inference_func(snake_case_      ,		snake_case_      ,		snake_case_						)
       return self._measure_speed(_inference						)


  def _UpperCamelCase     ( self      ,		snake_case_      ,		snake_case_      ,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:   str   =      self._prepare_inference_func(snake_case_      ,		snake_case_      ,		snake_case_						)
       return self._measure_memory(_inference						)


  def _UpperCamelCase     ( self      ,		snake_case_      ,		snake_case_      ,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:   Tuple   =      self._prepare_train_func(snake_case_      ,		snake_case_      ,		snake_case_						)
       return self._measure_speed(_train						)


  def _UpperCamelCase     ( self      ,		snake_case_      ,		snake_case_      ,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:   Union[str, Any]   =      self._prepare_train_func(snake_case_      ,		snake_case_      ,		snake_case_						)
       return self._measure_memory(_train						)


  def _UpperCamelCase     ( self      ,		snake_case_      ,		snake_case_      ,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:   Optional[int]   =      self.config_dict[model_name]

       if self.args.torchscript:
            UpperCAmelCase_							:   Dict   =      True

       UpperCAmelCase_							:   Dict   =      (
           hasattr(snake_case_      ,		'architectures'						)
           and isinstance(config.architectures      ,		snake_case_						)
           and len(config.architectures						) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                 UpperCAmelCase_							:   List[str]   =      config.architectures[0]
                 UpperCAmelCase_							:   int   =      __import__('transformers'      ,		fromlist=[model_class]						)
                 UpperCAmelCase_							:   Dict   =      getattr(snake_case_      ,		snake_case_						)
                 UpperCAmelCase_							:   Union[str, Any]   =      model_cls(snake_case_						)
            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:
            UpperCAmelCase_							:   Optional[Any]   =      MODEL_MAPPING[config.__class__](snake_case_						)

       model.eval()
       model.to(self.args.device						)

       # encoder-decoder has vocab size saved differently
       UpperCAmelCase_							:   Tuple   =      config.vocab_size if hasattr(snake_case_      ,		'vocab_size'						) else config.encoder.vocab_size
       UpperCAmelCase_							:   List[str]   =      torch.randint(snake_case_      ,		(batch_size, sequence_length)      ,		dtype=torch.long      ,		device=self.args.device						)

       if self.args.fpaa:
            logger.info('Running training in Mixed Precision...'						)
            if not self.args.is_gpu:
                 raise ValueError('Mixed precision is possible only for GPU.'						)
            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

       if self.args.torchscript:
            with torch.no_grad():
                 UpperCAmelCase_							:   int   =      torch.jit.trace(snake_case_      ,		snake_case_						)
       else:
            UpperCAmelCase_							:   Optional[int]   =      model

       def encoder_decoder_forward():
            with torch.no_grad():
                 UpperCAmelCase_							:   Tuple   =      inference_model(snake_case_      ,		decoder_input_ids=snake_case_						)
            return outputs

       def encoder_forward():
            with torch.no_grad():
                 UpperCAmelCase_							:   List[str]   =      inference_model(snake_case_						)
            return outputs

       UpperCAmelCase_							:   List[str]   =      encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
       return _forward


  def _UpperCamelCase     ( self      ,		snake_case_      ,		snake_case_      ,		snake_case_						):




       '''simple docstring'''

       UpperCAmelCase_							:   List[Any]   =      self.config_dict[model_name]

       UpperCAmelCase_							:   str   =      (
           hasattr(snake_case_      ,		'architectures'						)
           and isinstance(config.architectures      ,		snake_case_						)
           and len(config.architectures						) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
                 UpperCAmelCase_							:   Tuple   =      config.architectures[0]
                 UpperCAmelCase_							:   int   =      __import__('transformers'      ,		fromlist=[model_class]						)
                 UpperCAmelCase_							:   Optional[int]   =      getattr(snake_case_      ,		snake_case_						)
                 UpperCAmelCase_							:   List[str]   =      model_cls(snake_case_						)
            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:
            UpperCAmelCase_							:   Optional[Any]   =      MODEL_WITH_LM_HEAD_MAPPING[config.__class__](snake_case_						)

       if self.args.torchscript:
            raise NotImplementedError('Training for torchscript is currently not implemented'						)
       else:
            UpperCAmelCase_							:   str   =      model

       model.train()
       model.to(self.args.device						)

       # encoder-decoder has vocab size saved differently
       UpperCAmelCase_							:   Tuple   =      config.vocab_size if hasattr(snake_case_      ,		'vocab_size'						) else config.encoder.vocab_size
       UpperCAmelCase_							:   Optional[Any]   =      torch.randint(snake_case_      ,		(batch_size, sequence_length)      ,		dtype=torch.long      ,		device=self.args.device						)

       if self.args.fpaa:
            logger.info('Running training in Mixed Precision...'						)
            if not self.args.is_gpu:
                 raise ValueError('Mixed precision is possible only for GPU.'						)

            # amp seems to have memory leaks so that memory usage
            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
            model.half()

       def compute_loss_and_backprob_encoder():
            UpperCAmelCase_							:   Optional[Any]   =      train_model(snake_case_      ,		labels=snake_case_						)[0]
            loss.backward()
            return loss

       def compute_loss_and_backprob_encoder_decoder():
            UpperCAmelCase_							:   Optional[int]   =      train_model(snake_case_      ,		decoder_input_ids=snake_case_      ,		labels=snake_case_						)[0]
            loss.backward()
            return loss

       UpperCAmelCase_							:   List[str]   =      (
           compute_loss_and_backprob_encoder_decoder
           if config.is_encoder_decoder
           else compute_loss_and_backprob_encoder
       )
       return _train


  def _UpperCamelCase     ( self      ,		snake_case_						):




       '''simple docstring'''

       try:
            if self.args.is_tpu or self.args.torchscript:
                 # run additional 10 times to stabilize compilation for tpu and torchscript
                 logger.info('Do inference on TPU or torchscript. Running model 5 times to stabilize compilation'						)
                 timeit.repeat(
                     snake_case_      ,		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
            UpperCAmelCase_							:   Tuple   =      timeit.repeat(
                snake_case_      ,		repeat=self.args.repeat      ,		number=1_0      ,		)

            if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
                 import torch_xla.debug.metrics as met

                 self.print_fn(met.metrics_report()						)

            return min(snake_case_						) / 10.0
       except RuntimeError as e:
            self.print_fn(F'''Doesn\'t fit on GPU. {e}'''						)
            return "N/A"






  def _UpperCamelCase     ( self      ,		snake_case_						):




       '''simple docstring'''

       try:
            if self.args.trace_memory_line_by_line:
                 UpperCAmelCase_							:   int   =      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'
                     ' `--no-memory` or `args.memory=False`'						)
            elif self.args.is_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.'						)
                      UpperCAmelCase_							:   Tuple   =      '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()
                      UpperCAmelCase_							:   Tuple   =      nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx						)
                      UpperCAmelCase_							:   Dict   =      nvml.nvmlDeviceGetMemoryInfo(snake_case_						)
                      UpperCAmelCase_							:   List[str]   =      meminfo.used
                      UpperCAmelCase_							:   int   =      Memory(snake_case_						)
                      # shutdown nvml
                      nvml.nvmlShutdown()
            else:
                 # cpu
                 UpperCAmelCase_							:   Optional[Any]   =      measure_peak_memory_cpu(snake_case_						)
                 UpperCAmelCase_							:   int   =      Memory(snake_case_						) if isinstance(snake_case_      ,		snake_case_						) else memory_bytes

            if self.args.trace_memory_line_by_line:
                 UpperCAmelCase_							:   Dict   =      stop_memory_tracing(snake_case_						)
            else:
                 UpperCAmelCase_							:   Tuple   =      None

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







'''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 ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
snake_case__     :			Optional[Any]      			=  logging.get_logger(__name__)


def 					_lowerCamelCase    (      lowerCamelCase_					:							List[str]   ):



     """simple docstring"""





     UpperCAmelCase_							:   Dict   =      ConvNextConfig()

     if "tiny" in checkpoint_url:
          UpperCAmelCase_							:   Any   =      [3, 3, 9, 3]
          UpperCAmelCase_							:   Any   =      [96, 192, 384, 768]
     if "small" in checkpoint_url:
          UpperCAmelCase_							:   int   =      [3, 3, 27, 3]
          UpperCAmelCase_							:   Optional[int]   =      [96, 192, 384, 768]
     if "base" in checkpoint_url:
          UpperCAmelCase_							:   str   =      [3, 3, 27, 3]
          UpperCAmelCase_							:   Dict   =      [128, 256, 512, 1024]
     if "large" in checkpoint_url:
          UpperCAmelCase_							:   Optional[Any]   =      [3, 3, 27, 3]
          UpperCAmelCase_							:   List[Any]   =      [192, 384, 768, 1536]
     if "xlarge" in checkpoint_url:
          UpperCAmelCase_							:   str   =      [3, 3, 27, 3]
          UpperCAmelCase_							:   Any   =      [256, 512, 1024, 2048]

     if "1k" in checkpoint_url:
          UpperCAmelCase_							:   int   =      1000
          UpperCAmelCase_							:   Any   =      'imagenet-1k-id2label.json'
          UpperCAmelCase_							:   str   =      (1, 1000)
     else:
          UpperCAmelCase_							:   Dict   =      21841
          UpperCAmelCase_							:   Optional[Any]   =      'imagenet-22k-id2label.json'
          UpperCAmelCase_							:   Optional[Any]   =      (1, 21841)

     UpperCAmelCase_							:   int   =      'huggingface/label-files'
     UpperCAmelCase_							:   Any   =      num_labels
     UpperCAmelCase_							:   Dict   =      json.load(open(hf_hub_download(lowerCamelCase_	,     lowerCamelCase_	,     repo_type='dataset'   )	,     'r'   )   )
     UpperCAmelCase_							:   List[Any]   =      {int(lowerCamelCase_   ): v for k, v in idalabel.items()}
     if "1k" not in checkpoint_url:
          # this dataset contains 21843 labels but the model only has 21841
          # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
          del idalabel[9205]
          del idalabel[15027]
     UpperCAmelCase_							:   Optional[Any]   =      idalabel
     UpperCAmelCase_							:   Dict   =      {v: k for k, v in idalabel.items()}
     UpperCAmelCase_							:   Tuple   =      hidden_sizes
     UpperCAmelCase_							:   List[str]   =      depths

     return config, expected_shape


def 					_lowerCamelCase    (      lowerCamelCase_					:							int   ):



     """simple docstring"""





     if "downsample_layers.0.0" in name:
          UpperCAmelCase_							:   List[Any]   =      name.replace('downsample_layers.0.0'	,     'embeddings.patch_embeddings'   )
     if "downsample_layers.0.1" in name:
          UpperCAmelCase_							:   List[Any]   =      name.replace('downsample_layers.0.1'	,     'embeddings.norm'   )  # we rename to layernorm later on
     if "downsample_layers.1.0" in name:
          UpperCAmelCase_							:   int   =      name.replace('downsample_layers.1.0'	,     'stages.1.downsampling_layer.0'   )
     if "downsample_layers.1.1" in name:
          UpperCAmelCase_							:   List[str]   =      name.replace('downsample_layers.1.1'	,     'stages.1.downsampling_layer.1'   )
     if "downsample_layers.2.0" in name:
          UpperCAmelCase_							:   str   =      name.replace('downsample_layers.2.0'	,     'stages.2.downsampling_layer.0'   )
     if "downsample_layers.2.1" in name:
          UpperCAmelCase_							:   List[str]   =      name.replace('downsample_layers.2.1'	,     'stages.2.downsampling_layer.1'   )
     if "downsample_layers.3.0" in name:
          UpperCAmelCase_							:   Optional[int]   =      name.replace('downsample_layers.3.0'	,     'stages.3.downsampling_layer.0'   )
     if "downsample_layers.3.1" in name:
          UpperCAmelCase_							:   Dict   =      name.replace('downsample_layers.3.1'	,     'stages.3.downsampling_layer.1'   )
     if "stages" in name and "downsampling_layer" not in name:
          # stages.0.0. for instance should be renamed to stages.0.layers.0.
          UpperCAmelCase_							:   str   =      name[: len('stages.0'   )] + '.layers' + name[len('stages.0'   ) :]
     if "stages" in name:
          UpperCAmelCase_							:   int   =      name.replace('stages'	,     'encoder.stages'   )
     if "norm" in name:
          UpperCAmelCase_							:   List[str]   =      name.replace('norm'	,     'layernorm'   )
     if "gamma" in name:
          UpperCAmelCase_							:   Any   =      name.replace('gamma'	,     'layer_scale_parameter'   )
     if "head" in name:
          UpperCAmelCase_							:   List[str]   =      name.replace('head'	,     'classifier'   )

     return name


def 					_lowerCamelCase    (      ):



     """simple docstring"""





     UpperCAmelCase_							:   List[str]   =      'http://images.cocodataset.org/val2017/000000039769.jpg'
     UpperCAmelCase_							:   Optional[int]   =      Image.open(requests.get(lowerCamelCase_	,     stream=lowerCamelCase_   ).raw   )
     return im


@torch.no_grad()
def 					_lowerCamelCase    (      lowerCamelCase_					:							int	,     lowerCamelCase_					:							Optional[int]   ):



     """simple docstring"""





     UpperCAmelCase_   ,			UpperCAmelCase_							:   Optional[Any]   =      get_convnext_config(lowerCamelCase_   )
     # load original state_dict from URL
     UpperCAmelCase_							:   Dict   =      torch.hub.load_state_dict_from_url(lowerCamelCase_   )['model']
     # rename keys
     for key in state_dict.copy().keys():
          UpperCAmelCase_							:   Optional[int]   =      state_dict.pop(lowerCamelCase_   )
          UpperCAmelCase_							:   Tuple   =      val
     # add prefix to all keys expect classifier head
     for key in state_dict.copy().keys():
          UpperCAmelCase_							:   int   =      state_dict.pop(lowerCamelCase_   )
          if not key.startswith('classifier'   ):
               UpperCAmelCase_							:   int   =      'convnext.' + key
          UpperCAmelCase_							:   Tuple   =      val

     # load HuggingFace model
     UpperCAmelCase_							:   Tuple   =      ConvNextForImageClassification(lowerCamelCase_   )
     model.load_state_dict(lowerCamelCase_   )
     model.eval()

     # Check outputs on an image, prepared by ConvNextImageProcessor
     UpperCAmelCase_							:   Tuple   =      224 if '224' in checkpoint_url else 384
     UpperCAmelCase_							:   Optional[Any]   =      ConvNextImageProcessor(size=lowerCamelCase_   )
     UpperCAmelCase_							:   Tuple   =      image_processor(images=prepare_img()	,     return_tensors='pt'   ).pixel_values

     UpperCAmelCase_							:   Dict   =      model(lowerCamelCase_   ).logits

     # note: the logits below were obtained without center cropping
     if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
          UpperCAmelCase_							:   Tuple   =      torch.tensor([-0.1210, -0.6605, 0.1918]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
          UpperCAmelCase_							:   str   =      torch.tensor([-0.4473, -0.1847, -0.6365]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
          UpperCAmelCase_							:   Tuple   =      torch.tensor([0.4525, 0.7539, 0.0308]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
          UpperCAmelCase_							:   List[Any]   =      torch.tensor([0.3561, 0.6350, -0.0384]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
          UpperCAmelCase_							:   List[str]   =      torch.tensor([0.4174, -0.0989, 0.1489]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
          UpperCAmelCase_							:   Any   =      torch.tensor([0.2513, -0.1349, -0.1613]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
          UpperCAmelCase_							:   Dict   =      torch.tensor([1.2980, 0.3631, -0.1198]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
          UpperCAmelCase_							:   Union[str, Any]   =      torch.tensor([1.2963, 0.1227, 0.1723]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
          UpperCAmelCase_							:   Optional[Any]   =      torch.tensor([1.7956, 0.8390, 0.2820]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
          UpperCAmelCase_							:   Tuple   =      torch.tensor([-0.2822, -0.0502, -0.0878]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
          UpperCAmelCase_							:   Union[str, Any]   =      torch.tensor([-0.5672, -0.0730, -0.4348]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
          UpperCAmelCase_							:   Optional[Any]   =      torch.tensor([0.2681, 0.2365, 0.6246]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
          UpperCAmelCase_							:   Optional[int]   =      torch.tensor([-0.2642, 0.3931, 0.5116]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
          UpperCAmelCase_							:   Union[str, Any]   =      torch.tensor([-0.6677, -0.1873, -0.8379]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
          UpperCAmelCase_							:   str   =      torch.tensor([-0.7749, -0.2967, -0.6444]   )
     else:
          raise ValueError(F'''Unknown URL: {checkpoint_url}'''   )

     assert torch.allclose(logits[0, :3]	,     lowerCamelCase_	,     atol=1e-3   )
     assert logits.shape == expected_shape

     Path(lowerCamelCase_   ).mkdir(exist_ok=lowerCamelCase_   )
     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_   )

     print('Pushing model to the hub...'   )
     UpperCAmelCase_							:   Union[str, Any]   =      'convnext'
     if "tiny" in checkpoint_url:
          model_name += "-tiny"
     elif "small" in checkpoint_url:
          model_name += "-small"
     elif "base" in checkpoint_url:
          model_name += "-base"
     elif "xlarge" in checkpoint_url:
          model_name += "-xlarge"
     elif "large" in checkpoint_url:
          model_name += "-large"
     if "224" in checkpoint_url:
          model_name += "-224"
     elif "384" in checkpoint_url:
          model_name += "-384"
     if "22k" in checkpoint_url and "1k" not in checkpoint_url:
          model_name += "-22k"
     if "22k" in checkpoint_url and "1k" in checkpoint_url:
          model_name += "-22k-1k"

     model.push_to_hub(
         repo_path_or_name=Path(lowerCamelCase_	,     lowerCamelCase_   )	,     organization='nielsr'	,     commit_message='Add model'	,     )


if __name__ == "__main__":
 snake_case__     :			List[Any]      			=  argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--checkpoint_url''',
     default='''https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth''',
     type=str,
     help='''URL of the original ConvNeXT checkpoint you\'d like to convert.''',
 )
 parser.add_argument(
     '''--pytorch_dump_folder_path''',
     default=None,
     type=str,
     required=True,
     help='''Path to the output PyTorch model directory.''',
 )

 snake_case__     :			int      			=  parser.parse_args()
 convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






