




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




def    lowerCamelCase__			(		A__			:   Tuple      ):





       '''simple docstring'''



       __lowerCamelCase  =		ConvNextVaConfig()

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

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

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

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

       return config, expected_shape




def    lowerCamelCase__			(		A__			:   int      ):





       '''simple docstring'''



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

       return name




def    lowerCamelCase__			(		):





       '''simple docstring'''



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




def    lowerCamelCase__			(		A__			:   List[Any]      ):





       '''simple docstring'''



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

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




@torch.no_grad()
def    lowerCamelCase__			(		A__			:   List[Any]    ,     A__			:   Optional[Any]    ,     A__			:   str    ,     A__			:   Optional[int]      ):





       '''simple docstring'''



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

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

       # load HuggingFace model
       __lowerCamelCase  =		ConvNextVaForImageClassification(A__      )
       model.load_state_dict(A__      )
       model.eval()

       # Check outputs on an image, prepared by ConvNextImageProcessor
       __lowerCamelCase  =		convert_preprocessor(A__      )
       __lowerCamelCase  =		preprocessor(images=prepare_img()    ,     return_tensors="""pt"""      )
       __lowerCamelCase  =		model(**A__      ).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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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":
              __lowerCamelCase  =		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]    ,     A__    ,     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(A__      ):
                     os.mkdir(A__      )

              model.save_pretrained(A__      )
              preprocessor.save_pretrained(A__      )

       __lowerCamelCase  =		"""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(A__      )
              preprocessor.push_to_hub(A__      )


if __name__ == "__main__":
       UpperCAmelCase_   							=      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_   							=      parser.parse_args()
       convert_convnextva_checkpoint(
           args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
       )







from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor    ,     A__			:   bool      ):





       '''simple docstring'''



       if not inplace:
              __lowerCamelCase  =		ma + ma
       else:
              ma += ma

       return ma




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   List[int]      ):





       '''simple docstring'''



       __lowerCamelCase  =		-1 * len(A__      )
       __lowerCamelCase  =		list(range(len(tensor.shape[:zero_index]      )      )      )
       return tensor.permute(first_inds + [zero_index + i for i in inds]      )




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   int      ):





       '''simple docstring'''



       return t.reshape(t.shape[:-no_dims] + (-1,)      )




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor    ,     A__			:   int    ,     A__			:   float = 1E-4      ):





       '''simple docstring'''



       __lowerCamelCase  =		mask.expand(*value.shape      )
       return torch.sum(mask * value    ,     dim=A__      ) / (eps + torch.sum(A__    ,     dim=A__      ))




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.types.Number = 2.3_125    ,     A__			:   torch.types.Number = 21.6_875    ,     A__			:   int = 64      ):





       '''simple docstring'''



       __lowerCamelCase  =		torch.linspace(A__    ,     A__    ,     no_bins - 1    ,     device=pts.device      )
       __lowerCamelCase  =		torch.sqrt(torch.sum((pts.unsqueeze(-2      ) - pts.unsqueeze(-3      )) ** 2    ,     dim=-1      )      )
       return torch.bucketize(A__    ,     A__      )




def    lowerCamelCase__			(		A__			:   Callable[[list], Any]    ,     A__			:   List[dict]      ):





       '''simple docstring'''



       __lowerCamelCase  =		dicts[0]
       __lowerCamelCase  =		{}
       for k, v in first.items():
              __lowerCamelCase  =		[d[k] for d in dicts]
              if isinstance(A__    ,     A__      ):
                     __lowerCamelCase  =		dict_multimap(A__    ,     A__      )
              else:
                     __lowerCamelCase  =		fn(A__      )

       return new_dict




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor      ):





       '''simple docstring'''



       __lowerCamelCase  =		v_bins.view(((1,) * len(x.shape      )) + (len(A__      ),)      )
       __lowerCamelCase  =		x[..., None] - reshaped_bins
       __lowerCamelCase  =		torch.argmin(torch.abs(A__      )    ,     dim=-1      )
       return nn.functional.one_hot(A__    ,     num_classes=len(A__      )      ).float()




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor    ,     A__			:   int = 0    ,     A__			:   int = 0      ):





       '''simple docstring'''



       __lowerCamelCase  =		[]
       for i, s in enumerate(data.shape[:no_batch_dims]      ):
              __lowerCamelCase  =		torch.arange(A__      )
              __lowerCamelCase  =		r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape      ) - i - 1)))      )
              ranges.append(A__      )

       __lowerCamelCase  =		[slice(A__      ) for _ in range(len(data.shape      ) - no_batch_dims      )]
       __lowerCamelCase  =		inds
       ranges.extend(A__      )
       # Matt note: Editing this to get around the behaviour of using a list as an array index changing
       # in recent Numpy versions
       return data[tuple(A__      )]


UpperCAmelCase_   							=      TypeVar('T')




def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   Dict[Any, Union[dict, list, tuple, T]]    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       __lowerCamelCase  =		{}
       for k, v in dic.items():
              if isinstance(A__    ,     A__      ):
                     __lowerCamelCase  =		dict_map(A__    ,     A__    ,     A__      )
              else:
                     __lowerCamelCase  =		tree_map(A__    ,     A__    ,     A__      )

       return new_dict




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   T    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   dict    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   list    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   tuple    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




def    lowerCamelCase__			(		A__			:   Dict    ,     A__			:   List[str]    ,     A__			:   List[str]      ):





       '''simple docstring'''



       if isinstance(A__    ,     A__      ):
              return dict_map(A__    ,     A__    ,     A__      )
       elif isinstance(A__    ,     A__      ):
              return [tree_map(A__    ,     A__    ,     A__      ) for x in tree]
       elif isinstance(A__    ,     A__      ):
              return tuple(tree_map(A__    ,     A__    ,     A__      ) for x in tree      )
       elif isinstance(A__    ,     A__      ):
              return fn(A__      )
       else:
              print(type(A__      )      )
              raise ValueError("""Not supported"""      )




UpperCAmelCase_   							=      partial(tree_map, leaf_type=torch.Tensor)

