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


def       UpperCAmelCase_					(  __lowerCAmelCase		)	->  Dict:
    __lowercase						:     Optional[int]		      =    ConvNextVaConfig()

    if "atto" in checkpoint_url:
        __lowercase						:     Optional[int]		      =    [2, 2, 6, 2]
        __lowercase						:     Optional[int]		      =    [40, 80, 160, 320]
    if "femto" in checkpoint_url:
        __lowercase						:     Tuple		      =    [2, 2, 6, 2]
        __lowercase						:     Tuple		      =    [48, 96, 192, 384]
    if "pico" in checkpoint_url:
        __lowercase						:     str		      =    [2, 2, 6, 2]
        __lowercase						:     Optional[int]		      =    [64, 128, 256, 512]
    if "nano" in checkpoint_url:
        __lowercase						:     List[str]		      =    [2, 2, 8, 2]
        __lowercase						:     Optional[Any]		      =    [80, 160, 320, 640]
    if "tiny" in checkpoint_url:
        __lowercase						:     List[Any]		      =    [3, 3, 9, 3]
        __lowercase						:     Optional[Any]		      =    [96, 192, 384, 768]
    if "base" in checkpoint_url:
        __lowercase						:     Optional[Any]		      =    [3, 3, 27, 3]
        __lowercase						:     Optional[Any]		      =    [128, 256, 512, 1_024]
    if "large" in checkpoint_url:
        __lowercase						:     int		      =    [3, 3, 27, 3]
        __lowercase						:     List[str]		      =    [192, 384, 768, 1_536]
    if "huge" in checkpoint_url:
        __lowercase						:     int		      =    [3, 3, 27, 3]
        __lowercase						:     Optional[Any]		      =    [352, 704, 1_408, 2_816]

    __lowercase						:     Tuple		      =    1_000
    __lowercase						:     Dict		      =    '''imagenet-1k-id2label.json'''
    __lowercase						:     Dict		      =    (1, 1_000)

    __lowercase						:     List[Any]		      =    '''huggingface/label-files'''
    __lowercase						:     int		      =    num_labels
    __lowercase						:     Union[str, Any]		      =    json.load(open(hf_hub_download(__lowerCAmelCase		,   __lowerCAmelCase		,   repo_type='''dataset'''		)		,   '''r'''		)		)
    __lowercase						:     Tuple		      =    {int(__lowerCAmelCase		): v for k, v in idalabel.items()}

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

    return config, expected_shape


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

    return name


def       UpperCAmelCase_					(  )	->  Union[str, Any]:
    __lowercase						:     Optional[int]		      =    '''http://images.cocodataset.org/val2017/000000039769.jpg'''
    __lowercase						:     List[str]		      =    Image.open(requests.get(__lowerCAmelCase		,   stream=__lowerCAmelCase		).raw		)
    return im


def       UpperCAmelCase_					(  __lowerCAmelCase		)	->  Union[str, Any]:
    if "224" in checkpoint_url:
        __lowercase						:     List[str]		      =    224
        __lowercase						:     List[Any]		      =    224 / 256
    elif "384" in checkpoint_url:
        __lowercase						:     Optional[int]		      =    384
        __lowercase						:     Union[str, Any]		      =    None
    else:
        __lowercase						:     Tuple		      =    512
        __lowercase						:     Any		      =    None

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


@torch.no_grad()
def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		,   __lowerCAmelCase		,   __lowerCAmelCase		)	->  Any:
    print('''Downloading original model from checkpoint...'''		)
    # define ConvNeXTV2 configuration based on URL
    __lowercase      , __lowercase						:     List[Any]		      =    get_convnextva_config(__lowerCAmelCase		)
    # load original state_dict from URL
    __lowercase						:     Optional[Any]		      =    torch.hub.load_state_dict_from_url(__lowerCAmelCase		)['''model''']

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

    # load HuggingFace model
    __lowercase						:     Tuple		      =    ConvNextVaForImageClassification(__lowerCAmelCase		)
    model.load_state_dict(__lowerCAmelCase		)
    model.eval()

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

        model.save_pretrained(__lowerCAmelCase		)
        preprocessor.save_pretrained(__lowerCAmelCase		)

    __lowercase						:     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(__lowerCAmelCase		)
        preprocessor.push_to_hub(__lowerCAmelCase		)


if __name__ == "__main__":
 __lowerCAmelCase   :     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")

 __lowerCAmelCase   :     str      =						parser.parse_args()
 convert_convnextva_checkpoint(
     args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
 )




from __future__ import annotations

from math import pi, pow


def       UpperCAmelCase_					(  __lowerCAmelCase		)	->  float:
    if side_length < 0:
        raise ValueError('''vol_cube() only accepts non-negative values'''		)
    return pow(__lowerCAmelCase		,   3		)


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    if height < 0 or radius < 0:
        raise ValueError('''vol_spherical_cap() only accepts non-negative values'''		)
    # Volume is 1/3 pi * height squared * (3 * radius - height)
    return 1 / 3 * pi * pow(__lowerCAmelCase		,   2		) * (3 * radius - height)


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    if radius_a < 0 or radius_a < 0 or centers_distance < 0:
        raise ValueError('''vol_spheres_intersect() only accepts non-negative values'''		)
    if centers_distance == 0:
        return vol_sphere(min(__lowerCAmelCase		,   __lowerCAmelCase		)		)

    __lowercase						:     Dict		      =    (
        (radius_a - radius_a + centers_distance)
        * (radius_a + radius_a - centers_distance)
        / (2 * centers_distance)
    )
    __lowercase						:     Optional[int]		      =    (
        (radius_a - radius_a + centers_distance)
        * (radius_a + radius_a - centers_distance)
        / (2 * centers_distance)
    )

    return vol_spherical_cap(__lowerCAmelCase		,   __lowerCAmelCase		) + vol_spherical_cap(__lowerCAmelCase		,   __lowerCAmelCase		)


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:

    if radius_a <= 0 or radius_a <= 0 or centers_distance < 0:
        raise ValueError(
            '''vol_spheres_union() only accepts non-negative values, non-zero radius'''		)

    if centers_distance == 0:
        return vol_sphere(max(__lowerCAmelCase		,   __lowerCAmelCase		)		)

    return (
        vol_sphere(__lowerCAmelCase		)
        + vol_sphere(__lowerCAmelCase		)
        - vol_spheres_intersect(__lowerCAmelCase		,   __lowerCAmelCase		,   __lowerCAmelCase		)
    )


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    if width < 0 or height < 0 or length < 0:
        raise ValueError('''vol_cuboid() only accepts non-negative values'''		)
    return float(width * height * length		)


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    if height < 0 or area_of_base < 0:
        raise ValueError('''vol_cone() only accepts non-negative values'''		)
    return area_of_base * height / 3.0


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    if height < 0 or radius < 0:
        raise ValueError('''vol_right_circ_cone() only accepts non-negative values'''		)
    return pi * pow(__lowerCAmelCase		,   2		) * height / 3.0


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    if height < 0 or area_of_base < 0:
        raise ValueError('''vol_prism() only accepts non-negative values'''		)
    return float(area_of_base * height		)


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    if height < 0 or area_of_base < 0:
        raise ValueError('''vol_pyramid() only accepts non-negative values'''		)
    return area_of_base * height / 3.0


def       UpperCAmelCase_					(  __lowerCAmelCase		)	->  float:
    if radius < 0:
        raise ValueError('''vol_sphere() only accepts non-negative values'''		)
    # Volume is 4/3 * pi * radius cubed
    return 4 / 3 * pi * pow(__lowerCAmelCase		,   3		)


def       UpperCAmelCase_					(  __lowerCAmelCase		)	->  float:
    if radius < 0:
        raise ValueError('''vol_hemisphere() only accepts non-negative values'''		)
    # Volume is radius cubed * pi * 2/3
    return pow(__lowerCAmelCase		,   3		) * pi * 2 / 3


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    if height < 0 or radius < 0:
        raise ValueError('''vol_circular_cylinder() only accepts non-negative values'''		)
    # Volume is radius squared * height * pi
    return pow(__lowerCAmelCase		,   2		) * height * pi


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    # Volume - (outer_radius squared - inner_radius squared) * pi * height
    if inner_radius < 0 or outer_radius < 0 or height < 0:
        raise ValueError(
            '''vol_hollow_circular_cylinder() only accepts non-negative values'''		)
    if outer_radius <= inner_radius:
        raise ValueError('''outer_radius must be greater than inner_radius'''		)
    return pi * (pow(__lowerCAmelCase		,   2		) - pow(__lowerCAmelCase		,   2		)) * height


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    # Volume is 1/3 * pi * height *
    #           (radius_1 squared + radius_2 squared + radius_1 * radius_2)
    if radius_a < 0 or radius_a < 0 or height < 0:
        raise ValueError('''vol_conical_frustum() only accepts non-negative values'''		)
    return (
        1
        / 3
        * pi
        * height
        * (pow(__lowerCAmelCase		,   2		) + pow(__lowerCAmelCase		,   2		) + radius_a * radius_a)
    )


def       UpperCAmelCase_					(  __lowerCAmelCase		,   __lowerCAmelCase		)	->  float:
    if torus_radius < 0 or tube_radius < 0:
        raise ValueError('''vol_torus() only accepts non-negative values'''		)
    return 2 * pow(__lowerCAmelCase		,   2		) * torus_radius * pow(__lowerCAmelCase		,   2		)


def       UpperCAmelCase_					(  )	->  Optional[Any]:
    print('''Volumes:'''		)
    print(F'Cube: {vol_cube(2		) = }'		)  # = 8
    print(F'Cuboid: {vol_cuboid(2		,   2		,   2		) = }'		)  # = 8
    print(F'Cone: {vol_cone(2		,   2		) = }'		)  # ~= 1.33
    print(F'Right Circular Cone: {vol_right_circ_cone(2		,   2		) = }'		)  # ~= 8.38
    print(F'Prism: {vol_prism(2		,   2		) = }'		)  # = 4
    print(F'Pyramid: {vol_pyramid(2		,   2		) = }'		)  # ~= 1.33
    print(F'Sphere: {vol_sphere(2		) = }'		)  # ~= 33.5
    print(F'Hemisphere: {vol_hemisphere(2		) = }'		)  # ~= 16.75
    print(F'Circular Cylinder: {vol_circular_cylinder(2		,   2		) = }'		)  # ~= 25.1
    print(F'Torus: {vol_torus(2		,   2		) = }'		)  # ~= 157.9
    print(F'Conical Frustum: {vol_conical_frustum(2		,   2		,   4		) = }'		)  # ~= 58.6
    print(F'Spherical cap: {vol_spherical_cap(1		,   2		) = }'		)  # ~= 5.24
    print(F'Spheres intersetion: {vol_spheres_intersect(2		,   2		,   1		) = }'		)  # ~= 21.21
    print(F'Spheres union: {vol_spheres_union(2		,   2		,   1		) = }'		)  # ~= 45.81
    print(
        F'Hollow Circular Cylinder: {vol_hollow_circular_cylinder(1		,   2		,   3		) = }'		)  # ~= 28.3


if __name__ == "__main__":
 main()


