

'''simple docstring'''

import argparse
import json
from pathlib import Path

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

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


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

def      __magic_name__							(			A				)      ->		List[str]:
	snake_case							=					ConvNextConfig()

	if "tiny" in checkpoint_url:
		snake_case							=					[3, 3, 9, 3]
		snake_case							=					[9_6, 1_9_2, 3_8_4, 7_6_8]
	if "small" in checkpoint_url:
		snake_case							=					[3, 3, 2_7, 3]
		snake_case							=					[9_6, 1_9_2, 3_8_4, 7_6_8]
	if "base" in checkpoint_url:
		snake_case							=					[3, 3, 2_7, 3]
		snake_case							=					[1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
	if "large" in checkpoint_url:
		snake_case							=					[3, 3, 2_7, 3]
		snake_case							=					[1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
	if "xlarge" in checkpoint_url:
		snake_case							=					[3, 3, 2_7, 3]
		snake_case							=					[2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]

	if "1k" in checkpoint_url:
		snake_case							=					1_0_0_0
		snake_case							=					'imagenet-1k-id2label.json'
		snake_case							=					(1, 1_0_0_0)
	else:
		snake_case							=					2_1_8_4_1
		snake_case							=					'imagenet-22k-id2label.json'
		snake_case							=					(1, 2_1_8_4_1)

	snake_case							=					'huggingface/label-files'
	snake_case							=					num_labels
	snake_case							=					json.load(open(hf_hub_download(A       , A       , repo_type='dataset'				)       , 'r'				)				)
	snake_case							=					{int(A				): v for k, v in idalabel.items()}
	if "1k" not in checkpoint_url:
		# this dataset contains 21843 labels but the model only has 21841
		# we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
		del idalabel[9_2_0_5]
		del idalabel[1_5_0_2_7]
	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      __magic_name__							(			A				)      ->		str:
	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 "stages" in name:
		snake_case							=					name.replace('stages'       , 'encoder.stages'				)
	if "norm" in name:
		snake_case							=					name.replace('norm'       , 'layernorm'				)
	if "gamma" in name:
		snake_case							=					name.replace('gamma'       , 'layer_scale_parameter'				)
	if "head" in name:
		snake_case							=					name.replace('head'       , 'classifier'				)

	return name

def      __magic_name__							(			)      ->		int:
	snake_case							=					'http://images.cocodataset.org/val2017/000000039769.jpg'
	snake_case							=					Image.open(requests.get(A       , stream=A				).raw				)
	return im

@torch.no_grad()
def      __magic_name__							(			A       , A				)      ->		Dict:
	snake_case			,			snake_case							=					get_convnext_config(A				)
	# load original state_dict from URL
	snake_case							=					torch.hub.load_state_dict_from_url(A				)['model']
	# rename keys
	for key in state_dict.copy().keys():
		snake_case							=					state_dict.pop(A				)
		snake_case							=					val
	# add prefix to all keys expect classifier head
	for key in state_dict.copy().keys():
		snake_case							=					state_dict.pop(A				)
		if not key.startswith('classifier'				):
			snake_case							=					'convnext.' + key
		snake_case							=					val

	# load HuggingFace model
	snake_case							=					ConvNextForImageClassification(A				)
	model.load_state_dict(A				)
	model.eval()

	# Check outputs on an image, prepared by ConvNextImageProcessor
	snake_case							=					2_2_4 if '224' in checkpoint_url else 3_8_4
	snake_case							=					ConvNextImageProcessor(size=A				)
	snake_case							=					image_processor(images=prepare_img()       , return_tensors='pt'				).pixel_values

	snake_case							=					model(A				).logits

	# note: the logits below were obtained without center cropping
	if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
		snake_case							=					torch.tensor([-0.1_210, -0.6_605, 0.1_918]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
		snake_case							=					torch.tensor([-0.4_473, -0.1_847, -0.6_365]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
		snake_case							=					torch.tensor([0.4_525, 0.7_539, 0.0_308]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
		snake_case							=					torch.tensor([0.3_561, 0.6_350, -0.0_384]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
		snake_case							=					torch.tensor([0.4_174, -0.0_989, 0.1_489]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
		snake_case							=					torch.tensor([0.2_513, -0.1_349, -0.1_613]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
		snake_case							=					torch.tensor([1.2_980, 0.3_631, -0.1_198]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
		snake_case							=					torch.tensor([1.2_963, 0.1_227, 0.1_723]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
		snake_case							=					torch.tensor([1.7_956, 0.8_390, 0.2_820]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
		snake_case							=					torch.tensor([-0.2_822, -0.0_502, -0.0_878]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
		snake_case							=					torch.tensor([-0.5_672, -0.0_730, -0.4_348]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
		snake_case							=					torch.tensor([0.2_681, 0.2_365, 0.6_246]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
		snake_case							=					torch.tensor([-0.2_642, 0.3_931, 0.5_116]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
		snake_case							=					torch.tensor([-0.6_677, -0.1_873, -0.8_379]				)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
		snake_case							=					torch.tensor([-0.7_749, -0.2_967, -0.6_444]				)
	else:
		raise ValueError(F'''Unknown URL: {checkpoint_url}'''				)

	assert torch.allclose(logits[0, :3]       , A       , atol=1E-3				)
	assert logits.shape == expected_shape

	Path(A				).mkdir(exist_ok=A				)
	print(F'''Saving model to {pytorch_dump_folder_path}'''				)
	model.save_pretrained(A				)
	print(F'''Saving image processor to {pytorch_dump_folder_path}'''				)
	image_processor.save_pretrained(A				)

	print('Pushing model to the hub...'				)
	snake_case							=					'convnext'
	if "tiny" in checkpoint_url:
		model_name += "-tiny"
	elif "small" in checkpoint_url:
		model_name += "-small"
	elif "base" in checkpoint_url:
		model_name += "-base"
	elif "xlarge" in checkpoint_url:
		model_name += "-xlarge"
	elif "large" in checkpoint_url:
		model_name += "-large"
	if "224" in checkpoint_url:
		model_name += "-224"
	elif "384" in checkpoint_url:
		model_name += "-384"
	if "22k" in checkpoint_url and "1k" not in checkpoint_url:
		model_name += "-22k"
	if "22k" in checkpoint_url and "1k" in checkpoint_url:
		model_name += "-22k-1k"

	model.push_to_hub(
	    repo_path_or_name=Path(A       , A				)       , organization='nielsr'       , commit_message='Add model'       , )


if __name__ == "__main__":
						lowerCAmelCase_     =      argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    "--checkpoint_url",
						    default="https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
						    type=str,
						    help="URL of the original ConvNeXT checkpoint you'd like to convert.",
						)
						parser.add_argument(
						    "--pytorch_dump_folder_path",
						    default=None,
						    type=str,
						    required=True,
						    help="Path to the output PyTorch model directory.",
						)

						lowerCAmelCase_     =      parser.parse_args()
						convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






"""simple docstring"""


import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


__A       			=      [nn.LayerNorm]

__A       			=      logging.get_logger(__name__)

__A       			=      version.parse(version.parse(torch.__version__).base_version)

__A       			=      parsed_torch_version_base >= version.parse("""2.0""")
__A       			=      parsed_torch_version_base >= version.parse("""1.12""")
__A       			=      parsed_torch_version_base >= version.parse("""1.11""")
__A       			=      parsed_torch_version_base < version.parse("""1.11""")

def 						UpperCamelCase				(	_lowerCAmelCase   :							Tuple				, _lowerCAmelCase   :							List[str]				, _lowerCAmelCase   :							Union[str, Any]				, _lowerCAmelCase   :							int				, self   :							Optional[Any]						):
    from torch import _softmax_backward_data

    if is_torch_less_than_1_11:
        return _softmax_backward_data(_lowerCAmelCase				, _lowerCAmelCase				, parent.dim				, self						)
    else:
        return _softmax_backward_data(_lowerCAmelCase				, _lowerCAmelCase				, parent.dim				, self.dtype						)

def 						UpperCamelCase				(	_lowerCAmelCase   :							nn.Linear				, _lowerCAmelCase   :							torch.LongTensor				, _lowerCAmelCase   :							int = 0						):
    __a			      =						index.to(layer.weight.device						)
    __a			      =						layer.weight.index_select(_lowerCAmelCase				, _lowerCAmelCase						).clone().detach()
    if layer.bias is not None:
        if dim == 1:
            __a			      =						layer.bias.clone().detach()
        else:
            __a			      =						layer.bias[index].clone().detach()
    __a			      =						list(layer.weight.size()						)
    __a			      =						len(_lowerCAmelCase						)
    __a			      =						nn.Linear(new_size[1]				, new_size[0]				, bias=layer.bias is not None						).to(layer.weight.device						)
    __a			      =						False
    new_layer.weight.copy_(W.contiguous()						)
    __a			      =						True
    if layer.bias is not None:
        __a			      =						False
        new_layer.bias.copy_(b.contiguous()						)
        __a			      =						True
    return new_layer






class 					a					(				nn.Module							):



  def __init__(		self   :    Optional[int]						,  lowerCamelCase_   :    Optional[Any]						,  lowerCamelCase_   :    List[Any]     )      -> List[str]:
      super().__init__()
      __a			      =						nf
      __a			      =						nn.Parameter(torch.empty(lowerCamelCase_						,  lowerCamelCase_     )     )
      __a			      =						nn.Parameter(torch.zeros(lowerCamelCase_     )     )
      nn.init.normal_(self.weight						,  std=0.02     )



  def        lowerCAmelCase_				(		self   :    int						,  lowerCamelCase_   :    Optional[Any]     )      -> Union[str, Any]:
      __a			      =						x.size()[:-1] + (self.nf,)
      __a			      =						torch.addmm(self.bias						,  x.view(-1						,  x.size(-1     )     )						,  self.weight     )
      __a			      =						x.view(lowerCamelCase_     )
      return x

def 						UpperCamelCase				(	_lowerCAmelCase   :							ConvaD				, _lowerCAmelCase   :							torch.LongTensor				, _lowerCAmelCase   :							int = 1						):
    __a			      =						index.to(layer.weight.device						)
    __a			      =						layer.weight.index_select(_lowerCAmelCase				, _lowerCAmelCase						).clone().detach()
    if dim == 0:
        __a			      =						layer.bias.clone().detach()
    else:
        __a			      =						layer.bias[index].clone().detach()
    __a			      =						list(layer.weight.size()						)
    __a			      =						len(_lowerCAmelCase						)
    __a			      =						ConvaD(new_size[1]				, new_size[0]						).to(layer.weight.device						)
    __a			      =						False
    new_layer.weight.copy_(W.contiguous()						)
    __a			      =						True
    __a			      =						False
    new_layer.bias.copy_(b.contiguous()						)
    __a			      =						True
    return new_layer

def 						UpperCamelCase				(	_lowerCAmelCase   :							Union[nn.Linear, ConvaD]				, _lowerCAmelCase   :							torch.LongTensor				, _lowerCAmelCase   :							Optional[int] = None						):
    if isinstance(_lowerCAmelCase				, nn.Linear						):
        return prune_linear_layer(_lowerCAmelCase				, _lowerCAmelCase				, dim=0 if dim is None else dim						)
    elif isinstance(_lowerCAmelCase				, _lowerCAmelCase						):
        return prune_convad_layer(_lowerCAmelCase				, _lowerCAmelCase				, dim=1 if dim is None else dim						)
    else:
        raise ValueError(f"""Can't prune layer of class {layer.__class__}"""						)

def 						UpperCamelCase				(	_lowerCAmelCase   :							Callable[..., torch.Tensor]				, _lowerCAmelCase   :							int				, _lowerCAmelCase   :							int				, *_lowerCAmelCase   :							List[Any]						):
    assert len(_lowerCAmelCase						) > 0, f"""{input_tensors} has to be a tuple/list of tensors"""

    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
    __a			      =						len(inspect.signature(_lowerCAmelCase						).parameters						)
    if num_args_in_forward_chunk_fn != len(_lowerCAmelCase						):
        raise ValueError(
            f"""forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(_lowerCAmelCase						)} input """
            """tensors are given"""						)

    if chunk_size > 0:
        __a			      =						input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    f"""All input tenors have to be of the same shape: {tensor_shape}, """
                    f"""found shape {input_tensor.shape[chunk_dim]}"""						)

        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                f"""The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk """
                f"""size {chunk_size}"""						)

        __a			      =						input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        __a			      =						tuple(input_tensor.chunk(_lowerCAmelCase				, dim=_lowerCAmelCase						) for input_tensor in input_tensors						)
        # apply forward fn to every tuple
        __a			      =						tuple(forward_fn(*_lowerCAmelCase						) for input_tensors_chunk in zip(*_lowerCAmelCase						)						)
        # concatenate output at same dimension
        return torch.cat(_lowerCAmelCase				, dim=_lowerCAmelCase						)

    return forward_fn(*_lowerCAmelCase						)

def 						UpperCamelCase				(	_lowerCAmelCase   :							List[int]				, _lowerCAmelCase   :							int				, _lowerCAmelCase   :							int				, _lowerCAmelCase   :							Set[int]						):
    __a			      =						torch.ones(_lowerCAmelCase				, _lowerCAmelCase						)
    __a			      =						set(_lowerCAmelCase						) - already_pruned_heads  # Convert to set and remove already pruned heads
    for head in heads:
        # Compute how many pruned heads are before the head and move the index accordingly
        __a			      =						head - sum(1 if h < head else 0 for h in already_pruned_heads						)
        __a			      =						0
    __a			      =						mask.view(-1						).contiguous().eq(1						)
    __a			      =						torch.arange(len(_lowerCAmelCase						)						)[mask].long()
    return heads, index

def 						UpperCamelCase				(	*_lowerCAmelCase   :							Union[torch.Tensor, List[torch.Tensor]]				, _lowerCAmelCase   :							Optional[str] = None						):
    return torch.meshgrid(*_lowerCAmelCase				, indexing=_lowerCAmelCase						)






def 						UpperCamelCase				(	_lowerCAmelCase   :							torch.Tensor						):
    return tensor.device, storage_ptr(_lowerCAmelCase						), storage_size(_lowerCAmelCase						)
