


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


UpperCAmelCase     =       logging.get_logger(__name__)







class lowerCAmelCase_		(					lowerCamelCase__			):

 '''simple docstring'''

 __snake_case				=    42
 __snake_case				=    42
 __snake_case				=    "PyTorch"



 @property
 def 		UpperCamelCase__	(      self   ):
       return torch.__version__



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   ):
       snake_case_			= self._prepare_inference_func(_UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   )
       return self._measure_speed(_inference   )



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   ):
       snake_case_			= self._prepare_inference_func(_UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   )
       return self._measure_memory(_inference   )



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   ):
       snake_case_			= self._prepare_train_func(_UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   )
       return self._measure_speed(_train   )



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   ):
       snake_case_			= self._prepare_train_func(_UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   )
       return self._measure_memory(_train   )



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   ):
       snake_case_			= self.config_dict[model_name]

       if self.args.torchscript:
             snake_case_			= True

       snake_case_			= (
           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:
                   snake_case_			= config.architectures[0]
                   snake_case_			= __import__('''transformers'''						, fromlist=[model_class]   )
                   snake_case_			= getattr(_UpperCAmelCase						, _UpperCAmelCase   )
                   snake_case_			= 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:
             snake_case_			= MODEL_MAPPING[config.__class__](_UpperCAmelCase   )

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

       # encoder-decoder has vocab size saved differently
       snake_case_			= config.vocab_size if hasattr(_UpperCAmelCase						, '''vocab_size'''   ) else config.encoder.vocab_size
       snake_case_			= torch.randint(_UpperCAmelCase						, (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():
                   snake_case_			= torch.jit.trace(_UpperCAmelCase						, _UpperCAmelCase   )
       else:
             snake_case_			= model

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

       def encoder_forward():
             with torch.no_grad():
                   snake_case_			= inference_model(_UpperCAmelCase   )
             return outputs

       snake_case_			= encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
       return _forward



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase						, _UpperCAmelCase						, _UpperCAmelCase   ):
       snake_case_			= self.config_dict[model_name]

       snake_case_			= (
           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:
                   snake_case_			= config.architectures[0]
                   snake_case_			= __import__('''transformers'''						, fromlist=[model_class]   )
                   snake_case_			= getattr(_UpperCAmelCase						, _UpperCAmelCase   )
                   snake_case_			= 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:
             snake_case_			= MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_UpperCAmelCase   )

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

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

       # encoder-decoder has vocab size saved differently
       snake_case_			= config.vocab_size if hasattr(_UpperCAmelCase						, '''vocab_size'''   ) else config.encoder.vocab_size
       snake_case_			= torch.randint(_UpperCAmelCase						, (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():
             snake_case_			= train_model(_UpperCAmelCase						, labels=_UpperCAmelCase   )[0]
             loss.backward()
             return loss

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

       snake_case_			= (
           compute_loss_and_backprob_encoder_decoder
           if config.is_encoder_decoder
           else compute_loss_and_backprob_encoder
       )
       return _train



 def 		UpperCamelCase__	(      self						, _UpperCAmelCase   ):
       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(
                       _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
             snake_case_			= timeit.repeat(
                 _UpperCAmelCase						, repeat=self.args.repeat						, number=10						, )

             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(_UpperCAmelCase   ) / 10.0
       except RuntimeError as e:
             self.print_fn(F'''Doesn\'t fit on GPU. {e}'''   )
             return "N/A"





 def 		UpperCamelCase__	(      self						, _UpperCAmelCase   ):
       try:
             if self.args.trace_memory_line_by_line:
                   snake_case_			= 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.'''   )
                         snake_case_			= '''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()
                         snake_case_			= nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx   )
                         snake_case_			= nvml.nvmlDeviceGetMemoryInfo(_UpperCAmelCase   )
                         snake_case_			= meminfo.used
                         snake_case_			= Memory(_UpperCAmelCase   )
                         # shutdown nvml
                         nvml.nvmlShutdown()
             else:
                   # cpu
                   snake_case_			= measure_peak_memory_cpu(_UpperCAmelCase   )
                   snake_case_			= Memory(_UpperCAmelCase   ) if isinstance(_UpperCAmelCase						, _UpperCAmelCase   ) else memory_bytes

             if self.args.trace_memory_line_by_line:
                   snake_case_			= stop_memory_tracing(_UpperCAmelCase   )
             else:
                   snake_case_			= None

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



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()
UpperCAmelCase     =       logging.get_logger(__name__)






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   )->    Tuple:




      """simple docstring"""

      snake_case_			= ConvNextConfig()

      if "tiny" in checkpoint_url:
            snake_case_			= [3, 3, 9, 3]
            snake_case_			= [96, 192, 384, 768]
      if "small" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [96, 192, 384, 768]
      if "base" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [128, 256, 512, 1024]
      if "large" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [192, 384, 768, 1536]
      if "xlarge" in checkpoint_url:
            snake_case_			= [3, 3, 27, 3]
            snake_case_			= [256, 512, 1024, 2048]

      if "1k" in checkpoint_url:
            snake_case_			= 1000
            snake_case_			= '''imagenet-1k-id2label.json'''
            snake_case_			= (1, 1000)
      else:
            snake_case_			= 2_1841
            snake_case_			= '''imagenet-22k-id2label.json'''
            snake_case_			= (1, 2_1841)

      snake_case_			= '''huggingface/label-files'''
      snake_case_			= num_labels
      snake_case_			= json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ,   repo_type='''dataset'''   )   ,   '''r'''   )   )
      snake_case_			= {int(SCREAMING_SNAKE_CASE   ): 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[1_5027]
      snake_case_			= idalabel
      snake_case_			= {v: k for k, v in idalabel.items()}
      snake_case_			= hidden_sizes
      snake_case_			= depths

      return config, expected_shape






def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   )->    Tuple:




      """simple docstring"""

      if "downsample_layers.0.0" in name:
            snake_case_			= name.replace('''downsample_layers.0.0'''   ,   '''embeddings.patch_embeddings'''   )
      if "downsample_layers.0.1" in name:
            snake_case_			= name.replace('''downsample_layers.0.1'''   ,   '''embeddings.norm'''   )  # we rename to layernorm later on
      if "downsample_layers.1.0" in name:
            snake_case_			= name.replace('''downsample_layers.1.0'''   ,   '''stages.1.downsampling_layer.0'''   )
      if "downsample_layers.1.1" in name:
            snake_case_			= name.replace('''downsample_layers.1.1'''   ,   '''stages.1.downsampling_layer.1'''   )
      if "downsample_layers.2.0" in name:
            snake_case_			= name.replace('''downsample_layers.2.0'''   ,   '''stages.2.downsampling_layer.0'''   )
      if "downsample_layers.2.1" in name:
            snake_case_			= name.replace('''downsample_layers.2.1'''   ,   '''stages.2.downsampling_layer.1'''   )
      if "downsample_layers.3.0" in name:
            snake_case_			= name.replace('''downsample_layers.3.0'''   ,   '''stages.3.downsampling_layer.0'''   )
      if "downsample_layers.3.1" in name:
            snake_case_			= 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.
            snake_case_			= name[: len('''stages.0'''   )] + '''.layers''' + name[len('''stages.0'''   ) :]
      if "stages" in name:
            snake_case_			= name.replace('''stages'''   ,   '''encoder.stages'''   )
      if "norm" in name:
            snake_case_			= name.replace('''norm'''   ,   '''layernorm'''   )
      if "gamma" in name:
            snake_case_			= name.replace('''gamma'''   ,   '''layer_scale_parameter'''   )
      if "head" in name:
            snake_case_			= name.replace('''head'''   ,   '''classifier'''   )

      return name






def 			__lowerCAmelCase       ()->    Optional[Any]:




      """simple docstring"""

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






@torch.no_grad()
def 			__lowerCAmelCase       (SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )->    str:




      """simple docstring"""

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

      # load HuggingFace model
      snake_case_			= ConvNextForImageClassification(SCREAMING_SNAKE_CASE   )
      model.load_state_dict(SCREAMING_SNAKE_CASE   )
      model.eval()

      # Check outputs on an image, prepared by ConvNextImageProcessor
      snake_case_			= 224 if '''224''' in checkpoint_url else 384
      snake_case_			= ConvNextImageProcessor(size=SCREAMING_SNAKE_CASE   )
      snake_case_			= image_processor(images=prepare_img()   ,   return_tensors='''pt'''   ).pixel_values

      snake_case_			= model(SCREAMING_SNAKE_CASE   ).logits

      # note: the logits below were obtained without center cropping
      if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
            snake_case_			= torch.tensor([-0.1_2_1_0, -0.6_6_0_5, 0.1_9_1_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
            snake_case_			= torch.tensor([-0.4_4_7_3, -0.1_8_4_7, -0.6_3_6_5]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
            snake_case_			= torch.tensor([0.4_5_2_5, 0.7_5_3_9, 0.0_3_0_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
            snake_case_			= torch.tensor([0.3_5_6_1, 0.6_3_5_0, -0.0_3_8_4]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
            snake_case_			= torch.tensor([0.4_1_7_4, -0.0_9_8_9, 0.1_4_8_9]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
            snake_case_			= torch.tensor([0.2_5_1_3, -0.1_3_4_9, -0.1_6_1_3]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
            snake_case_			= torch.tensor([1.2_9_8_0, 0.3_6_3_1, -0.1_1_9_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
            snake_case_			= torch.tensor([1.2_9_6_3, 0.1_2_2_7, 0.1_7_2_3]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
            snake_case_			= torch.tensor([1.7_9_5_6, 0.8_3_9_0, 0.2_8_2_0]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
            snake_case_			= torch.tensor([-0.2_8_2_2, -0.0_5_0_2, -0.0_8_7_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
            snake_case_			= torch.tensor([-0.5_6_7_2, -0.0_7_3_0, -0.4_3_4_8]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
            snake_case_			= torch.tensor([0.2_6_8_1, 0.2_3_6_5, 0.6_2_4_6]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
            snake_case_			= torch.tensor([-0.2_6_4_2, 0.3_9_3_1, 0.5_1_1_6]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
            snake_case_			= torch.tensor([-0.6_6_7_7, -0.1_8_7_3, -0.8_3_7_9]   )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
            snake_case_			= torch.tensor([-0.7_7_4_9, -0.2_9_6_7, -0.6_4_4_4]   )
      else:
            raise ValueError(f'''Unknown URL: {checkpoint_url}'''   )

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

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

      print('''Pushing model to the hub...'''   )
      snake_case_			= '''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(SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   )   ,   organization='''nielsr'''   ,   commit_message='''Add model'''   ,   )


if __name__ == "__main__":
   UpperCAmelCase     =       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.""",
   )

   UpperCAmelCase     =       parser.parse_args()
   convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)