





import argparse
import json
import os

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

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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



def   __lowercase  ( _SCREAMING_SNAKE_CASE      )					->  Union[str, Any]:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE						=		ConvNextVaConfig()

      if "atto" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		[2, 2, 6, 2]
            SCREAMING_SNAKE_CASE						=		[40, 80, 1_60, 3_20]
      if "femto" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		[2, 2, 6, 2]
            SCREAMING_SNAKE_CASE						=		[48, 96, 1_92, 3_84]
      if "pico" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		[2, 2, 6, 2]
            SCREAMING_SNAKE_CASE						=		[64, 1_28, 2_56, 5_12]
      if "nano" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		[2, 2, 8, 2]
            SCREAMING_SNAKE_CASE						=		[80, 1_60, 3_20, 6_40]
      if "tiny" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		[3, 3, 9, 3]
            SCREAMING_SNAKE_CASE						=		[96, 1_92, 3_84, 7_68]
      if "base" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		[3, 3, 27, 3]
            SCREAMING_SNAKE_CASE						=		[1_28, 2_56, 5_12, 10_24]
      if "large" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		[3, 3, 27, 3]
            SCREAMING_SNAKE_CASE						=		[1_92, 3_84, 7_68, 15_36]
      if "huge" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		[3, 3, 27, 3]
            SCREAMING_SNAKE_CASE						=		[3_52, 7_04, 14_08, 28_16]

      SCREAMING_SNAKE_CASE						=		10_00
      SCREAMING_SNAKE_CASE						=		"""imagenet-1k-id2label.json"""
      SCREAMING_SNAKE_CASE						=		(1, 10_00)

      SCREAMING_SNAKE_CASE						=		"""huggingface/label-files"""
      SCREAMING_SNAKE_CASE						=		num_labels
      SCREAMING_SNAKE_CASE						=		json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE					,		repo_type="""dataset"""      )					,		"""r"""      )      )
      SCREAMING_SNAKE_CASE						=		{int(_SCREAMING_SNAKE_CASE      ): v for k, v in idalabel.items()}

      SCREAMING_SNAKE_CASE						=		idalabel
      SCREAMING_SNAKE_CASE						=		{v: k for k, v in idalabel.items()}
      SCREAMING_SNAKE_CASE						=		hidden_sizes
      SCREAMING_SNAKE_CASE						=		depths

      return config, expected_shape



def   __lowercase  ( _SCREAMING_SNAKE_CASE      )					->  Optional[int]:


      '''simple docstring'''


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

      return name



def   __lowercase  ( )					->  Optional[int]:


      '''simple docstring'''


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



def   __lowercase  ( _SCREAMING_SNAKE_CASE      )					->  List[str]:


      '''simple docstring'''


      if "224" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		2_24
            SCREAMING_SNAKE_CASE						=		2_24 / 2_56
      elif "384" in checkpoint_url:
            SCREAMING_SNAKE_CASE						=		3_84
            SCREAMING_SNAKE_CASE						=		None
      else:
            SCREAMING_SNAKE_CASE						=		5_12
            SCREAMING_SNAKE_CASE						=		None

      return ConvNextImageProcessor(
          size=_SCREAMING_SNAKE_CASE					,		crop_pct=_SCREAMING_SNAKE_CASE					,		image_mean=[0.485, 0.456, 0.406]					,		image_std=[0.229, 0.224, 0.225]					,		resample=PILImageResampling.BICUBIC					,		)



@torch.no_grad()
def   __lowercase  ( _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE      )					->  Tuple:


      '''simple docstring'''


      print("""Downloading original model from checkpoint..."""      )
      # define ConvNeXTV2 configuration based on URL
      SCREAMING_SNAKE_CASE,   SCREAMING_SNAKE_CASE						=		get_convnextva_config(_SCREAMING_SNAKE_CASE      )
      # load original state_dict from URL
      SCREAMING_SNAKE_CASE						=		torch.hub.load_state_dict_from_url(_SCREAMING_SNAKE_CASE      )["""model"""]

      print("""Converting model parameters..."""      )
      # rename keys
      for key in state_dict.copy().keys():
            SCREAMING_SNAKE_CASE						=		state_dict.pop(_SCREAMING_SNAKE_CASE      )
            SCREAMING_SNAKE_CASE						=		val
      # add prefix to all keys expect classifier head
      for key in state_dict.copy().keys():
            SCREAMING_SNAKE_CASE						=		state_dict.pop(_SCREAMING_SNAKE_CASE      )
            if not key.startswith("""classifier"""      ):
                  SCREAMING_SNAKE_CASE						=		"""convnextv2.""" + key
            SCREAMING_SNAKE_CASE						=		val

      # load HuggingFace model
      SCREAMING_SNAKE_CASE						=		ConvNextVaForImageClassification(_SCREAMING_SNAKE_CASE      )
      model.load_state_dict(_SCREAMING_SNAKE_CASE      )
      model.eval()

      # Check outputs on an image, prepared by ConvNextImageProcessor
      SCREAMING_SNAKE_CASE						=		convert_preprocessor(_SCREAMING_SNAKE_CASE      )
      SCREAMING_SNAKE_CASE						=		preprocessor(images=prepare_img()					,		return_tensors="""pt"""      )
      SCREAMING_SNAKE_CASE						=		model(**_SCREAMING_SNAKE_CASE      ).logits

      # note: the logits below were obtained without center cropping
      if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.3_930, 0.1_747, -0.5_246, 0.4_177, 0.4_295]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.1_727, -0.5_341, -0.7_818, -0.4_745, -0.6_566]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.0_333, 0.1_563, -0.9_137, 0.1_054, 0.0_381]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.1_744, -0.1_555, -0.0_713, 0.0_950, -0.1_431]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([0.9_996, 0.1_966, -0.4_386, -0.3_472, 0.6_661]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.2_553, -0.6_708, -0.1_359, 0.2_518, -0.2_488]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.0_673, -0.5_627, -0.3_753, -0.2_722, 0.0_178]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.6_377, -0.7_458, -0.2_150, 0.1_184, -0.0_597]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([1.0_799, 0.2_322, -0.8_860, 1.0_219, 0.6_231]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([0.3_766, 0.4_917, -1.1_426, 0.9_942, 0.6_024]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([0.4_220, -0.6_919, -0.4_317, -0.2_881, -0.6_609]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([0.1_082, -0.8_286, -0.5_095, 0.4_681, -0.8_085]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.2_419, -0.6_221, 0.2_176, -0.0_980, -0.7_527]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([0.0_391, -0.4_371, 0.3_786, 0.1_251, -0.2_784]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.0_504, 0.5_636, -0.1_729, -0.6_507, -0.3_949]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([0.3_560, 0.9_486, 0.3_149, -0.2_667, -0.5_138]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.2_469, -0.4_550, -0.5_853, -0.0_810, 0.0_309]      )
      elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
            SCREAMING_SNAKE_CASE						=		torch.tensor([-0.3_090, 0.0_802, -0.0_682, -0.1_979, -0.2_826]      )
      else:
            raise ValueError(F"""Unknown URL: {checkpoint_url}"""      )

      assert torch.allclose(logits[0, :5]					,		_SCREAMING_SNAKE_CASE					,		atol=1E-3      )
      assert logits.shape == expected_shape
      print("""Model outputs match the original results!"""      )

      if save_model:
            print("""Saving model to local..."""      )
            # Create folder to save model
            if not os.path.isdir(_SCREAMING_SNAKE_CASE      ):
                  os.mkdir(_SCREAMING_SNAKE_CASE      )

            model.save_pretrained(_SCREAMING_SNAKE_CASE      )
            preprocessor.save_pretrained(_SCREAMING_SNAKE_CASE      )

      SCREAMING_SNAKE_CASE						=		"""convnextv2"""
      if "atto" in checkpoint_url:
            model_name += "-atto"
      if "femto" in checkpoint_url:
            model_name += "-femto"
      if "pico" in checkpoint_url:
            model_name += "-pico"
      if "nano" in checkpoint_url:
            model_name += "-nano"
      elif "tiny" in checkpoint_url:
            model_name += "-tiny"
      elif "base" in checkpoint_url:
            model_name += "-base"
      elif "large" in checkpoint_url:
            model_name += "-large"
      elif "huge" in checkpoint_url:
            model_name += "-huge"
      if "22k" in checkpoint_url and "1k" not in checkpoint_url:
            model_name += "-22k"
      elif "22k" in checkpoint_url and "1k" in checkpoint_url:
            model_name += "-22k-1k"
      elif "1k" in checkpoint_url:
            model_name += "-1k"
      if "224" in checkpoint_url:
            model_name += "-224"
      elif "384" in checkpoint_url:
            model_name += "-384"
      elif "512" in checkpoint_url:
            model_name += "-512"

      if push_to_hub:
            print(F"""Pushing {model_name} to the hub..."""      )
            model.push_to_hub(_SCREAMING_SNAKE_CASE      )
            preprocessor.push_to_hub(_SCREAMING_SNAKE_CASE      )


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE_		  =			argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--checkpoint_url""",
          default="""https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt""",
          type=str,
          help="""URL of the original ConvNeXTV2 checkpoint you'd like to convert.""",
      )
      parser.add_argument(
          """--pytorch_dump_folder_path""",
          default="""model""",
          type=str,
          help="""Path to the output PyTorch model directory.""",
      )
      parser.add_argument("""--save_model""", action="""store_true""", help="""Save model to local""")
      parser.add_argument("""--push_to_hub""", action="""store_true""", help="""Push model and image preprocessor to the hub""")

      SCREAMING_SNAKE_CASE_		  =			parser.parse_args()
      convert_convnextva_checkpoint(
          args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
      )









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


SCREAMING_SNAKE_CASE_		  =			{
    """configuration_markuplm""": ["""MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MarkupLMConfig"""],
    """feature_extraction_markuplm""": ["""MarkupLMFeatureExtractor"""],
    """processing_markuplm""": ["""MarkupLMProcessor"""],
    """tokenization_markuplm""": ["""MarkupLMTokenizer"""],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_		  =			["""MarkupLMTokenizerFast"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_		  =			[
          """MARKUPLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """MarkupLMForQuestionAnswering""",
          """MarkupLMForSequenceClassification""",
          """MarkupLMForTokenClassification""",
          """MarkupLMModel""",
          """MarkupLMPreTrainedModel""",
      ]


if TYPE_CHECKING:
      from .configuration_markuplm import MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP, MarkupLMConfig
      from .feature_extraction_markuplm import MarkupLMFeatureExtractor
      from .processing_markuplm import MarkupLMProcessor
      from .tokenization_markuplm import MarkupLMTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_markuplm_fast import MarkupLMTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_markuplm import (
                MARKUPLM_PRETRAINED_MODEL_ARCHIVE_LIST,
                MarkupLMForQuestionAnswering,
                MarkupLMForSequenceClassification,
                MarkupLMForTokenClassification,
                MarkupLMModel,
                MarkupLMPreTrainedModel,
            )


else:
      import sys

      SCREAMING_SNAKE_CASE_		  =			_LazyModule(__name__, globals()["""__file__"""], _import_structure)


