




'''simple docstring'''







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()
_lowerCamelCase							:						Optional[int]								=  logging.get_logger(__name__)






def 		_lowercase		(	lowercase_  ):



		'''simple docstring'''







		_snake_case	       =       ConvNextVaConfig()

		if "atto" in checkpoint_url:
				_snake_case	       =       [2, 2, 6, 2]
				_snake_case	       =       [40, 80, 160, 320]
		if "femto" in checkpoint_url:
				_snake_case	       =       [2, 2, 6, 2]
				_snake_case	       =       [48, 96, 192, 384]
		if "pico" in checkpoint_url:
				_snake_case	       =       [2, 2, 6, 2]
				_snake_case	       =       [64, 128, 256, 512]
		if "nano" in checkpoint_url:
				_snake_case	       =       [2, 2, 8, 2]
				_snake_case	       =       [80, 160, 320, 640]
		if "tiny" in checkpoint_url:
				_snake_case	       =       [3, 3, 9, 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 "huge" in checkpoint_url:
				_snake_case	       =       [3, 3, 27, 3]
				_snake_case	       =       [352, 704, 1408, 2816]

		_snake_case	       =       1000
		_snake_case	       =       "imagenet-1k-id2label.json"
		_snake_case	       =       (1, 1000)

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

		_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 		_lowercase		(	lowercase_  ):



		'''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 "gamma" in name:
				_snake_case	       =       name.replace("gamma"						,							"weight"  )
		if "beta" in name:
				_snake_case	       =       name.replace("beta"						,							"bias"  )
		if "stages" in name:
				_snake_case	       =       name.replace("stages"						,							"encoder.stages"  )
		if "norm" in name:
				_snake_case	       =       name.replace("norm"						,							"layernorm"  )
		if "head" in name:
				_snake_case	       =       name.replace("head"						,							"classifier"  )

		return name






def 		_lowercase		(	):



		'''simple docstring'''







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






def 		_lowercase		(	lowercase_  ):



		'''simple docstring'''







		if "224" in checkpoint_url:
				_snake_case	       =       224
				_snake_case	       =       224 / 256
		elif "384" in checkpoint_url:
				_snake_case	       =       384
				_snake_case	       =       None
		else:
				_snake_case	       =       512
				_snake_case	       =       None

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






@torch.no_grad()
def 		_lowercase		(	lowercase_						,							lowercase_						,							lowercase_						,							lowercase_  ):



		'''simple docstring'''







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

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

		# load HuggingFace model
		_snake_case	       =       ConvNextVaForImageClassification(lowercase_  )
		model.load_state_dict(lowercase_  )
		model.eval()

		# Check outputs on an image, prepared by ConvNextImageProcessor
		_snake_case	       =       convert_preprocessor(lowercase_  )
		_snake_case	       =       preprocessor(images=prepare_img()						,							return_tensors="pt"  )
		_snake_case	       =       model(**lowercase_  ).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":
				_snake_case	       =       torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
				_snake_case	       =       torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
				_snake_case	       =       torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
				_snake_case	       =       torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
				_snake_case	       =       torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
				_snake_case	       =       torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
				_snake_case	       =       torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
				_snake_case	       =       torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
				_snake_case	       =       torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
				_snake_case	       =       torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
				_snake_case	       =       torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
				_snake_case	       =       torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
				_snake_case	       =       torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
				_snake_case	       =       torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
				_snake_case	       =       torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
				_snake_case	       =       torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
				_snake_case	       =       torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309]  )
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
				_snake_case	       =       torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826]  )
		else:
				raise ValueError(f'''Unknown URL: {checkpoint_url}'''  )

		assert torch.allclose(logits[0, :5]						,							lowercase_						,							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(lowercase_  ):
						os.mkdir(lowercase_  )

				model.save_pretrained(lowercase_  )
				preprocessor.save_pretrained(lowercase_  )

		_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(lowercase_  )
				preprocessor.push_to_hub(lowercase_  )


if __name__ == "__main__":
						_lowerCamelCase							:						str								=  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')

						_lowerCamelCase							:						List[str]								=  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_torch_available, is_vision_available


UpperCamelCase__     :					Dict								=			{
    "configuration_donut_swin": ["DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP", "DonutSwinConfig"],
    "processing_donut": ["DonutProcessor"],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCamelCase__     :					Optional[int]								=			[
           "DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST",
           "DonutSwinModel",
           "DonutSwinPreTrainedModel",
       ]

try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCamelCase__     :					List[Any]								=			["DonutFeatureExtractor"]
       UpperCamelCase__     :					Dict								=			["DonutImageProcessor"]


if TYPE_CHECKING:
       from .configuration_donut_swin import DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, DonutSwinConfig
       from .processing_donut import DonutProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_donut_swin import (
                  DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
                  DonutSwinModel,
                  DonutSwinPreTrainedModel,
              )

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .feature_extraction_donut import DonutFeatureExtractor
              from .image_processing_donut import DonutImageProcessor

else:
       import sys

       UpperCamelCase__     :					Tuple								=			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


