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(  UpperCamelCase_					)		->  Dict:







 '''simple docstring'''
 UpperCamelCase      =			ConvNextVaConfig()

 if "atto" in checkpoint_url:
  UpperCamelCase      =			[2, 2, 6, 2]
  UpperCamelCase      =			[40, 80, 160, 320]
 if "femto" in checkpoint_url:
  UpperCamelCase      =			[2, 2, 6, 2]
  UpperCamelCase      =			[48, 96, 192, 384]
 if "pico" in checkpoint_url:
  UpperCamelCase      =			[2, 2, 6, 2]
  UpperCamelCase      =			[64, 128, 256, 512]
 if "nano" in checkpoint_url:
  UpperCamelCase      =			[2, 2, 8, 2]
  UpperCamelCase      =			[80, 160, 320, 640]
 if "tiny" in checkpoint_url:
  UpperCamelCase      =			[3, 3, 9, 3]
  UpperCamelCase      =			[96, 192, 384, 768]
 if "base" in checkpoint_url:
  UpperCamelCase      =			[3, 3, 27, 3]
  UpperCamelCase      =			[128, 256, 512, 1024]
 if "large" in checkpoint_url:
  UpperCamelCase      =			[3, 3, 27, 3]
  UpperCamelCase      =			[192, 384, 768, 1536]
 if "huge" in checkpoint_url:
  UpperCamelCase      =			[3, 3, 27, 3]
  UpperCamelCase      =			[352, 704, 1408, 2816]

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

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

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

 return config, expected_shape





def  lowercase(  UpperCamelCase_					)		->  Any:







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

 return name





def  lowercase(  )		->  Any:







 '''simple docstring'''
 UpperCamelCase      =			"""http://images.cocodataset.org/val2017/000000039769.jpg"""
 UpperCamelCase      =			Image.open(requests.get(UpperCamelCase_       ,				stream=UpperCamelCase_					).raw					)
 return im





def  lowercase(  UpperCamelCase_					)		->  Any:







 '''simple docstring'''
 if "224" in checkpoint_url:
  UpperCamelCase      =			224
  UpperCamelCase      =			224 / 256
 elif "384" in checkpoint_url:
  UpperCamelCase      =			384
  UpperCamelCase      =			None
 else:
  UpperCamelCase      =			512
  UpperCamelCase      =			None

 return ConvNextImageProcessor(
     size=UpperCamelCase_       ,				crop_pct=UpperCamelCase_       ,				image_mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]       ,				image_std=[0.2_2_9, 0.2_2_4, 0.2_2_5]       ,				resample=PILImageResampling.BICUBIC       ,				)





@torch.no_grad()
def  lowercase(  UpperCamelCase_       ,				UpperCamelCase_       ,				UpperCamelCase_       ,				UpperCamelCase_					)		->  Union[str, Any]:







 '''simple docstring'''
 print("""Downloading original model from checkpoint..."""					)
 # define ConvNeXTV2 configuration based on URL
 UpperCamelCase     ,    UpperCamelCase      =			get_convnextva_config(UpperCamelCase_					)
 # load original state_dict from URL
 UpperCamelCase      =			torch.hub.load_state_dict_from_url(UpperCamelCase_					)["""model"""]

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

 # load HuggingFace model
 UpperCamelCase      =			ConvNextVaForImageClassification(UpperCamelCase_					)
 model.load_state_dict(UpperCamelCase_					)
 model.eval()

 # Check outputs on an image, prepared by ConvNextImageProcessor
 UpperCamelCase      =			convert_preprocessor(UpperCamelCase_					)
 UpperCamelCase      =			preprocessor(images=prepare_img()       ,				return_tensors="""pt"""					)
 UpperCamelCase      =			model(**UpperCamelCase_					).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":
  UpperCamelCase      =			torch.tensor([-0.3_9_3_0, 0.1_7_4_7, -0.5_2_4_6, 0.4_1_7_7, 0.4_2_9_5]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.1_7_2_7, -0.5_3_4_1, -0.7_8_1_8, -0.4_7_4_5, -0.6_5_6_6]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.0_3_3_3, 0.1_5_6_3, -0.9_1_3_7, 0.1_0_5_4, 0.0_3_8_1]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.1_7_4_4, -0.1_5_5_5, -0.0_7_1_3, 0.0_9_5_0, -0.1_4_3_1]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([0.9_9_9_6, 0.1_9_6_6, -0.4_3_8_6, -0.3_4_7_2, 0.6_6_6_1]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.2_5_5_3, -0.6_7_0_8, -0.1_3_5_9, 0.2_5_1_8, -0.2_4_8_8]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.0_6_7_3, -0.5_6_2_7, -0.3_7_5_3, -0.2_7_2_2, 0.0_1_7_8]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.6_3_7_7, -0.7_4_5_8, -0.2_1_5_0, 0.1_1_8_4, -0.0_5_9_7]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([1.0_7_9_9, 0.2_3_2_2, -0.8_8_6_0, 1.0_2_1_9, 0.6_2_3_1]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
  UpperCamelCase      =			torch.tensor([0.3_7_6_6, 0.4_9_1_7, -1.1_4_2_6, 0.9_9_4_2, 0.6_0_2_4]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([0.4_2_2_0, -0.6_9_1_9, -0.4_3_1_7, -0.2_8_8_1, -0.6_6_0_9]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
  UpperCamelCase      =			torch.tensor([0.1_0_8_2, -0.8_2_8_6, -0.5_0_9_5, 0.4_6_8_1, -0.8_0_8_5]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.2_4_1_9, -0.6_2_2_1, 0.2_1_7_6, -0.0_9_8_0, -0.7_5_2_7]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
  UpperCamelCase      =			torch.tensor([0.0_3_9_1, -0.4_3_7_1, 0.3_7_8_6, 0.1_2_5_1, -0.2_7_8_4]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.0_5_0_4, 0.5_6_3_6, -0.1_7_2_9, -0.6_5_0_7, -0.3_9_4_9]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
  UpperCamelCase      =			torch.tensor([0.3_5_6_0, 0.9_4_8_6, 0.3_1_4_9, -0.2_6_6_7, -0.5_1_3_8]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.2_4_6_9, -0.4_5_5_0, -0.5_8_5_3, -0.0_8_1_0, 0.0_3_0_9]					)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
  UpperCamelCase      =			torch.tensor([-0.3_0_9_0, 0.0_8_0_2, -0.0_6_8_2, -0.1_9_7_9, -0.2_8_2_6]					)
 else:
  raise ValueError(f"""Unknown URL: {checkpoint_url}"""					)

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

  model.save_pretrained(UpperCamelCase_					)
  preprocessor.save_pretrained(UpperCamelCase_					)

 UpperCamelCase      =			"""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(UpperCamelCase_					)
  preprocessor.push_to_hub(UpperCamelCase_					)


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
     )

import unittest

from datasets import load_dataset

from transformers import load_tool

from .test_tools_common import ToolTesterMixin
class SCREAMING_SNAKE_CASE_		(			unittest.TestCase     ,    __lowerCAmelCase		):







   def 	lowerCamelCase_  (    self	:	Optional[Any]							):





    """simple docstring"""


    UpperCamelCase      =			load_tool("""document-question-answering"""							)
    self.tool.setup()
    UpperCamelCase      =			load_tool("""document-question-answering"""      ,      remote=lowerCamelCase_							)







   def 	lowerCamelCase_  (    self	:	Optional[Any]							):





    """simple docstring"""


    UpperCamelCase      =			load_dataset("""hf-internal-testing/example-documents"""      ,      split="""test"""							)
    UpperCamelCase      =			dataset[0]["""image"""]

    UpperCamelCase      =			self.tool(lowerCamelCase_      ,      """When is the coffee break?"""							)
    self.assertEqual(lowerCamelCase_      ,      """11-14 to 11:39 a.m."""							)







   def 	lowerCamelCase_  (    self	:	Union[str, Any]							):





    """simple docstring"""


    UpperCamelCase      =			load_dataset("""hf-internal-testing/example-documents"""      ,      split="""test"""							)
    UpperCamelCase      =			dataset[0]["""image"""]

    UpperCamelCase      =			self.remote_tool(lowerCamelCase_      ,      """When is the coffee break?"""							)
    self.assertEqual(lowerCamelCase_      ,      """11-14 to 11:39 a.m."""							)







   def 	lowerCamelCase_  (    self	:	int							):





    """simple docstring"""


    UpperCamelCase      =			load_dataset("""hf-internal-testing/example-documents"""      ,      split="""test"""							)
    UpperCamelCase      =			dataset[0]["""image"""]

    self.tool(document=lowerCamelCase_      ,      question="""When is the coffee break?"""							)







   def 	lowerCamelCase_  (    self	:	Any							):





    """simple docstring"""


    UpperCamelCase      =			load_dataset("""hf-internal-testing/example-documents"""      ,      split="""test"""							)
    UpperCamelCase      =			dataset[0]["""image"""]

    UpperCamelCase      =			self.remote_tool(document=lowerCamelCase_      ,      question="""When is the coffee break?"""							)
    self.assertEqual(lowerCamelCase_      ,      """11-14 to 11:39 a.m."""							)
