


'''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()
__UpperCAmelCase     :List[Any]												=  logging.get_logger(__name__)




def       _a      (		_lowercase   :      List[Any]    ):





     '''simple docstring'''






     __UpperCAmelCase      :					str    				=     ConvNextVaConfig()

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

     __UpperCAmelCase      :					List[Any]    				=     1000
     __UpperCAmelCase      :					List[Any]    				=     '''imagenet-1k-id2label.json'''
     __UpperCAmelCase      :					Dict    				=     (1, 1000)

     __UpperCAmelCase      :					str    				=     '''huggingface/label-files'''
     __UpperCAmelCase      :					str    				=     num_labels
     __UpperCAmelCase      :					Any    				=     json.load(open(hf_hub_download(_lowercase     ,   _lowercase     ,   repo_type='''dataset'''    )     ,   '''r'''    )    )
     __UpperCAmelCase      :					int    				=     {int(_lowercase    ): v for k, v in idalabel.items()}

     __UpperCAmelCase      :					int    				=     idalabel
     __UpperCAmelCase      :					str    				=     {v: k for k, v in idalabel.items()}
     __UpperCAmelCase      :					Optional[Any]    				=     hidden_sizes
     __UpperCAmelCase      :					Optional[int]    				=     depths

     return config, expected_shape




def       _a      (		_lowercase   :      Union[str, Any]    ):





     '''simple docstring'''






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

     return name




def       _a      (		):





     '''simple docstring'''






     __UpperCAmelCase      :					Any    				=     '''http://images.cocodataset.org/val2017/000000039769.jpg'''
     __UpperCAmelCase      :					Tuple    				=     Image.open(requests.get(_lowercase     ,   stream=_lowercase    ).raw    )
     return im




def       _a      (		_lowercase   :      str    ):





     '''simple docstring'''






     if "224" in checkpoint_url:
          __UpperCAmelCase      :					Optional[Any]    				=     224
          __UpperCAmelCase      :					int    				=     224 / 256
     elif "384" in checkpoint_url:
          __UpperCAmelCase      :					Optional[Any]    				=     384
          __UpperCAmelCase      :					List[Any]    				=     None
     else:
          __UpperCAmelCase      :					List[str]    				=     512
          __UpperCAmelCase      :					Union[str, Any]    				=     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       _a      (		_lowercase   :      Any     ,   _lowercase   :      Any     ,   _lowercase   :      Tuple     ,   _lowercase   :      Optional[int]    ):





     '''simple docstring'''






     print('''Downloading original model from checkpoint...'''    )
     # define ConvNeXTV2 configuration based on URL
     __UpperCAmelCase				,  __UpperCAmelCase      :					Dict    				=     get_convnextva_config(_lowercase    )
     # load original state_dict from URL
     __UpperCAmelCase      :					List[str]    				=     torch.hub.load_state_dict_from_url(_lowercase    )['''model''']

     print('''Converting model parameters...'''    )
     # rename keys
     for key in state_dict.copy().keys():
          __UpperCAmelCase      :					int    				=     state_dict.pop(_lowercase    )
          __UpperCAmelCase      :					str    				=     val
     # add prefix to all keys expect classifier head
     for key in state_dict.copy().keys():
          __UpperCAmelCase      :					List[Any]    				=     state_dict.pop(_lowercase    )
          if not key.startswith('''classifier'''    ):
               __UpperCAmelCase      :					Any    				=     '''convnextv2.''' + key
          __UpperCAmelCase      :					Tuple    				=     val

     # load HuggingFace model
     __UpperCAmelCase      :					str    				=     ConvNextVaForImageClassification(_lowercase    )
     model.load_state_dict(_lowercase    )
     model.eval()

     # Check outputs on an image, prepared by ConvNextImageProcessor
     __UpperCAmelCase      :					Union[str, Any]    				=     convert_preprocessor(_lowercase    )
     __UpperCAmelCase      :					int    				=     preprocessor(images=prepare_img()     ,   return_tensors='''pt'''    )
     __UpperCAmelCase      :					Dict    				=     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":
          __UpperCAmelCase      :					Optional[int]    				=     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":
          __UpperCAmelCase      :					Optional[Any]    				=     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":
          __UpperCAmelCase      :					Union[str, Any]    				=     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":
          __UpperCAmelCase      :					Any    				=     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":
          __UpperCAmelCase      :					List[str]    				=     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":
          __UpperCAmelCase      :					List[str]    				=     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":
          __UpperCAmelCase      :					Tuple    				=     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":
          __UpperCAmelCase      :					List[Any]    				=     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":
          __UpperCAmelCase      :					Optional[int]    				=     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":
          __UpperCAmelCase      :					List[Any]    				=     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":
          __UpperCAmelCase      :					List[Any]    				=     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":
          __UpperCAmelCase      :					Any    				=     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":
          __UpperCAmelCase      :					Optional[Any]    				=     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":
          __UpperCAmelCase      :					int    				=     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":
          __UpperCAmelCase      :					Optional[Any]    				=     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":
          __UpperCAmelCase      :					Optional[int]    				=     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":
          __UpperCAmelCase      :					Optional[Any]    				=     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":
          __UpperCAmelCase      :					Dict    				=     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    )

     __UpperCAmelCase      :					Optional[Any]    				=     '''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__":
      __UpperCAmelCase     :Tuple												=  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")

      __UpperCAmelCase     :Optional[int]												=  parser.parse_args()
      convert_convnextva_checkpoint(
          args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
      )






from typing import Any, Dict, List, Mapping

import numpy as np
import torch

from ...utils import is_cython_available, requires_backends


if is_cython_available():
				import pyximport

				pyximport.install(setup_args={"""include_dirs""": np.get_include()})
				from . import algos_graphormer  # noqa E402




def  A_(      A					,   A = 512 ):
	UpperCAmelCase_		    =  x.shape[1] if len(x.shape ) > 1 else 1
	UpperCAmelCase_		    =  1 + np.arange(0					,   feature_num * offset					,   A					,   dtype=np.intaa )
	UpperCAmelCase_		    =  x + feature_offset
	return x




def  A_(      A					,   A=True ):
	requires_backends(A					,   ["""cython"""] )

	if keep_features and "edge_attr" in item.keys():  # edge_attr
		UpperCAmelCase_		    =  np.asarray(item["""edge_attr"""]					,   dtype=np.intaa )
	else:
		UpperCAmelCase_		    =  np.ones((len(item["""edge_index"""][0] ), 1)					,   dtype=np.intaa )  # same embedding for all

	if keep_features and "node_feat" in item.keys():  # input_nodes
		UpperCAmelCase_		    =  np.asarray(item["""node_feat"""]					,   dtype=np.intaa )
	else:
		UpperCAmelCase_		    =  np.ones((item["""num_nodes"""], 1)					,   dtype=np.intaa )  # same embedding for all

	UpperCAmelCase_		    =  np.asarray(item["""edge_index"""]					,   dtype=np.intaa )

	UpperCAmelCase_		    =  convert_to_single_emb(A ) + 1
	UpperCAmelCase_		    =  item["""num_nodes"""]

	if len(edge_attr.shape ) == 1:
		UpperCAmelCase_		    =  edge_attr[:, None]
	UpperCAmelCase_		    =  np.zeros([num_nodes, num_nodes, edge_attr.shape[-1]]					,   dtype=np.intaa )
	UpperCAmelCase_		    =  convert_to_single_emb(A ) + 1

	# node adj matrix [num_nodes, num_nodes] bool
	UpperCAmelCase_		    =  np.zeros([num_nodes, num_nodes]					,   dtype=A )
	UpperCAmelCase_		    =  True

	UpperCAmelCase_    ,						UpperCAmelCase_		    =  algos_graphormer.floyd_warshall(A )
	UpperCAmelCase_		    =  np.amax(A )

	UpperCAmelCase_		    =  algos_graphormer.gen_edge_input(A					,   A					,   A )
	UpperCAmelCase_		    =  np.zeros([num_nodes + 1, num_nodes + 1]					,   dtype=np.single )  # with graph token

	# combine
	UpperCAmelCase_		    =  input_nodes + 1  # we shift all indices by one for padding
	UpperCAmelCase_		    =  attn_bias
	UpperCAmelCase_		    =  attn_edge_type
	UpperCAmelCase_		    =  shortest_path_result.astype(np.intaa ) + 1  # we shift all indices by one for padding
	UpperCAmelCase_		    =  np.sum(A					,   axis=1 ).reshape(-1 ) + 1  # we shift all indices by one for padding
	UpperCAmelCase_		    =  item["""in_degree"""]  # for undirected graph
	UpperCAmelCase_		    =  input_edges + 1  # we shift all indices by one for padding
	if "labels" not in item:
		UpperCAmelCase_		    =  item["""y"""]

	return item







class 	_UpperCamelCase	:







					'''simple docstring'''

					def __init__(    self     :	str			,       __lowercase     :	int=20			,       __lowercase     :	Union[str, Any]=False		):






						'''simple docstring'''







						if not is_cython_available():
							raise ImportError("""Graphormer preprocessing needs Cython (pyximport)"""		)

						UpperCAmelCase_		    =  spatial_pos_max
						UpperCAmelCase_		    =  on_the_fly_processing





					def __call__(    self     :	List[str]			,       __lowercase     :	List[dict]		):






						'''simple docstring'''







						if self.on_the_fly_processing:
							UpperCAmelCase_		    =  [preprocess_item(__lowercase		) for i in features]

						if not isinstance(features[0]			,       __lowercase		):
							UpperCAmelCase_		    =  [vars(__lowercase		) for f in features]
						UpperCAmelCase_		    =  {}

						UpperCAmelCase_		    =  max(len(i["""input_nodes"""]		) for i in features		)
						UpperCAmelCase_		    =  len(features[0]["""input_nodes"""][0]		)
						UpperCAmelCase_		    =  len(features[0]["""attn_edge_type"""][0][0]		)
						UpperCAmelCase_		    =  max(len(i["""input_edges"""][0][0]		) for i in features		)
						UpperCAmelCase_		    =  len(features[0]["""input_edges"""][0][0][0]		)
						UpperCAmelCase_		    =  len(__lowercase		)

						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       max_node_num + 1			,       max_node_num + 1			,       dtype=torch.float		)
						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       __lowercase			,       __lowercase			,       __lowercase			,       dtype=torch.long		)
						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       __lowercase			,       __lowercase			,       dtype=torch.long		)
						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       __lowercase			,       dtype=torch.long		)
						UpperCAmelCase_		    =  torch.zeros(__lowercase			,       __lowercase			,       __lowercase			,       dtype=torch.long		)
						UpperCAmelCase_		    =  torch.zeros(
						    __lowercase			,       __lowercase			,       __lowercase			,       __lowercase			,       __lowercase			,       dtype=torch.long		)

						for ix, f in enumerate(__lowercase		):
							for k in ["attn_bias", "attn_edge_type", "spatial_pos", "in_degree", "input_nodes", "input_edges"]:
								UpperCAmelCase_		    =  torch.tensor(f[k]		)

							if len(f["""attn_bias"""][1:, 1:][f["""spatial_pos"""] >= self.spatial_pos_max]		) > 0:
								UpperCAmelCase_		    =  float("""-inf"""		)

							UpperCAmelCase_		    =  f["""attn_bias"""]
							UpperCAmelCase_		    =  f[
							    """attn_edge_type"""
							]
							UpperCAmelCase_		    =  f["""spatial_pos"""]
							UpperCAmelCase_		    =  f["""in_degree"""]
							UpperCAmelCase_		    =  f["""input_nodes"""]
							UpperCAmelCase_		    =  f["""input_edges"""]

						UpperCAmelCase_		    =  batch["""in_degree"""]

						UpperCAmelCase_		    =  features[0]["""labels"""]
						if len(__lowercase		) == 1:  # one task
							if isinstance(sample[0]			,       __lowercase		):  # regression
								UpperCAmelCase_		    =  torch.from_numpy(np.concatenate([i["""labels"""] for i in features]		)		)
							else:  # binary classification
								UpperCAmelCase_		    =  torch.from_numpy(np.concatenate([i["""labels"""] for i in features]		)		)
						else:  # multi task classification, left to float to keep the NaNs
							UpperCAmelCase_		    =  torch.from_numpy(np.stack([i["""labels"""] for i in features]			,       axis=0		)		)

						return batch


