



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()
__snake_case		:Optional[int] 						=	logging.get_logger(__name__)





def 	__snake_case      (     _UpperCAmelCase							):
    __a					=							ConvNextVaConfig()

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

    __a					=							1000
    __a					=							'''imagenet-1k-id2label.json'''
    __a					=							(1, 1000)

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

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

    return config, expected_shape





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

    return name





def 	__snake_case      (     ):
    __a					=							'''http://images.cocodataset.org/val2017/000000039769.jpg'''
    __a					=							Image.open(requests.get(_UpperCAmelCase  ,					stream=_UpperCAmelCase							).raw							)
    return im





def 	__snake_case      (     _UpperCAmelCase							):
    if "224" in checkpoint_url:
        __a					=							224
        __a					=							224 / 256
    elif "384" in checkpoint_url:
        __a					=							384
        __a					=							None
    else:
        __a					=							512
        __a					=							None

    return ConvNextImageProcessor(
        size=_UpperCAmelCase  ,					crop_pct=_UpperCAmelCase  ,					image_mean=[0.4_85, 0.4_56, 0.4_06]  ,					image_std=[0.2_29, 0.2_24, 0.2_25]  ,					resample=PILImageResampling.BICUBIC  ,					)





@torch.no_grad()
def 	__snake_case      (     _UpperCAmelCase  ,					_UpperCAmelCase  ,					_UpperCAmelCase  ,					_UpperCAmelCase							):
    print('''Downloading original model from checkpoint...'''							)
    # define ConvNeXTV2 configuration based on URL
    __a			,  __a					=							get_convnextva_config(_UpperCAmelCase							)
    # load original state_dict from URL
    __a					=							torch.hub.load_state_dict_from_url(_UpperCAmelCase							)['''model''']

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

    # load HuggingFace model
    __a					=							ConvNextVaForImageClassification(_UpperCAmelCase							)
    model.load_state_dict(_UpperCAmelCase							)
    model.eval()

    # Check outputs on an image, prepared by ConvNextImageProcessor
    __a					=							convert_preprocessor(_UpperCAmelCase							)
    __a					=							preprocessor(images=prepare_img()  ,					return_tensors='''pt'''							)
    __a					=							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":
        __a					=							torch.tensor([-0.39_30, 0.17_47, -0.52_46, 0.41_77, 0.42_95]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
        __a					=							torch.tensor([-0.17_27, -0.53_41, -0.78_18, -0.47_45, -0.65_66]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
        __a					=							torch.tensor([-0.03_33, 0.15_63, -0.91_37, 0.10_54, 0.03_81]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
        __a					=							torch.tensor([-0.17_44, -0.15_55, -0.07_13, 0.09_50, -0.14_31]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
        __a					=							torch.tensor([0.99_96, 0.19_66, -0.43_86, -0.34_72, 0.66_61]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
        __a					=							torch.tensor([-0.25_53, -0.67_08, -0.13_59, 0.25_18, -0.24_88]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
        __a					=							torch.tensor([-0.06_73, -0.56_27, -0.37_53, -0.27_22, 0.01_78]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
        __a					=							torch.tensor([-0.63_77, -0.74_58, -0.21_50, 0.11_84, -0.05_97]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
        __a					=							torch.tensor([1.07_99, 0.23_22, -0.88_60, 1.02_19, 0.62_31]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
        __a					=							torch.tensor([0.37_66, 0.49_17, -1.14_26, 0.99_42, 0.60_24]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
        __a					=							torch.tensor([0.42_20, -0.69_19, -0.43_17, -0.28_81, -0.66_09]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
        __a					=							torch.tensor([0.10_82, -0.82_86, -0.50_95, 0.46_81, -0.80_85]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
        __a					=							torch.tensor([-0.24_19, -0.62_21, 0.21_76, -0.09_80, -0.75_27]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
        __a					=							torch.tensor([0.03_91, -0.43_71, 0.37_86, 0.12_51, -0.27_84]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
        __a					=							torch.tensor([-0.05_04, 0.56_36, -0.17_29, -0.65_07, -0.39_49]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
        __a					=							torch.tensor([0.35_60, 0.94_86, 0.31_49, -0.26_67, -0.51_38]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
        __a					=							torch.tensor([-0.24_69, -0.45_50, -0.58_53, -0.08_10, 0.03_09]							)
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
        __a					=							torch.tensor([-0.30_90, 0.08_02, -0.06_82, -0.19_79, -0.28_26]							)
    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							)

    __a					=							'''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__":
 __snake_case		:str 						=	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''')

 __snake_case		:int 						=	parser.parse_args()
 convert_convnextva_checkpoint(
     args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
 )










def 	__snake_case      (     _UpperCAmelCase							):
    __a					=							1
    __a					=							0
    __a					=							1
    __a					=							1
    __a					=							3
    __a					=							3
    __a					=							limit
    __a					=							0

    __a					=							''''''
    while counter != decimal + 1:
        if 4 * q + r - t < n * t:
            result += str(_UpperCAmelCase							)
            if counter == 0:
                result += "."

            if decimal == counter:
                break

            counter += 1
            __a					=							10 * (r - n * t)
            __a					=							((10 * (3 * q + r)) // t) - 10 * n
            q *= 10
            __a					=							nr
        else:
            __a					=							(2 * q + r) * l
            __a					=							(q * (7 * k) + 2 + (r * l)) // (t * l)
            q *= k
            t *= l
            l += 2
            k += 1
            __a					=							nn
            __a					=							nr
    return result





def 	__snake_case      (     ):
    print(f'{calculate_pi(50							) = }'							)
    import doctest

    doctest.testmod()


if __name__ == "__main__":
 main()





