

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()
A_						:     List[str]			  =				logging.get_logger(__name__)


def  __a    (				SCREAMING_SNAKE_CASE ) ->       List[Any]:






    '''simple docstring'''
    __UpperCAmelCase			   =     ConvNextVaConfig()

    if "atto" in checkpoint_url:
        __UpperCAmelCase			   =     [2, 2, 6, 2]
        __UpperCAmelCase			   =     [4_0, 8_0, 1_6_0, 3_2_0]
    if "femto" in checkpoint_url:
        __UpperCAmelCase			   =     [2, 2, 6, 2]
        __UpperCAmelCase			   =     [4_8, 9_6, 1_9_2, 3_8_4]
    if "pico" in checkpoint_url:
        __UpperCAmelCase			   =     [2, 2, 6, 2]
        __UpperCAmelCase			   =     [6_4, 1_2_8, 2_5_6, 5_1_2]
    if "nano" in checkpoint_url:
        __UpperCAmelCase			   =     [2, 2, 8, 2]
        __UpperCAmelCase			   =     [8_0, 1_6_0, 3_2_0, 6_4_0]
    if "tiny" in checkpoint_url:
        __UpperCAmelCase			   =     [3, 3, 9, 3]
        __UpperCAmelCase			   =     [9_6, 1_9_2, 3_8_4, 7_6_8]
    if "base" in checkpoint_url:
        __UpperCAmelCase			   =     [3, 3, 2_7, 3]
        __UpperCAmelCase			   =     [1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
    if "large" in checkpoint_url:
        __UpperCAmelCase			   =     [3, 3, 2_7, 3]
        __UpperCAmelCase			   =     [1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
    if "huge" in checkpoint_url:
        __UpperCAmelCase			   =     [3, 3, 2_7, 3]
        __UpperCAmelCase			   =     [3_5_2, 7_0_4, 1_4_0_8, 2_8_1_6]

    __UpperCAmelCase			   =     1_0_0_0
    __UpperCAmelCase			   =     "imagenet-1k-id2label.json"
    __UpperCAmelCase			   =     (1, 1_0_0_0)

    __UpperCAmelCase			   =     "huggingface/label-files"
    __UpperCAmelCase			   =     num_labels
    __UpperCAmelCase			   =     json.load(open(hf_hub_download(_a    ,				_a    ,				repo_type='''dataset''' )    ,				'''r''' ) )
    __UpperCAmelCase			   =     {int(_a ): 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  __a    (				SCREAMING_SNAKE_CASE ) ->       List[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  __a    (				) ->       List[str]:






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


def  __a    (				SCREAMING_SNAKE_CASE ) ->       List[str]:






    '''simple docstring'''
    if "224" in checkpoint_url:
        __UpperCAmelCase			   =     2_2_4
        __UpperCAmelCase			   =     2_2_4 / 2_5_6
    elif "384" in checkpoint_url:
        __UpperCAmelCase			   =     3_8_4
        __UpperCAmelCase			   =     None
    else:
        __UpperCAmelCase			   =     5_1_2
        __UpperCAmelCase			   =     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  __a    (				SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE    ,				SCREAMING_SNAKE_CASE ) ->       Any:






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

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

    # load HuggingFace model
    __UpperCAmelCase			   =     ConvNextVaForImageClassification(_a )
    model.load_state_dict(_a )
    model.eval()

    # Check outputs on an image, prepared by ConvNextImageProcessor
    __UpperCAmelCase			   =     convert_preprocessor(_a )
    __UpperCAmelCase			   =     preprocessor(images=prepare_img()    ,				return_tensors='''pt''' )
    __UpperCAmelCase			   =     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":
        __UpperCAmelCase			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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			   =     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]    ,				_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 )

    __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(_a )
        preprocessor.push_to_hub(_a )


if __name__ == "__main__":
   A_						:     Optional[int]			  =				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')

   A_						:     Tuple			  =				parser.parse_args()
   convert_convnextva_checkpoint(
       args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
   )







# this script builds a small sample spm file tests/fixtures/test_sentencepiece_no_bos.model, with features needed by pegasus

# 1. pip install sentencepiece
#
# 2. wget https://raw.githubusercontent.com/google/sentencepiece/master/data/botchan.txt

# 3. build
import sentencepiece as spm

# pegasus:
# 1. no bos
# 2. eos_id is 1
# 3. unk_id is 2
# build a sample spm file accordingly
spm.SentencePieceTrainer.train(
    '--input=botchan.txt --model_prefix=test_sentencepiece_no_bos --bos_id=-1 --unk_id=2  --eos_id=1  --vocab_size=1000'
)

# 4. now update the fixture
# mv test_sentencepiece_no_bos.model ../../tests/fixtures/




