





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




def        lowerCAmelCase_       (							UpperCamelCase_     )   ->						Optional[int]:
       UpperCamelCase_      				=     ConvNextConfig()

       if "tiny" in checkpoint_url:
              UpperCamelCase_      				=     [3, 3, 9, 3]
              UpperCamelCase_      				=     [96, 192, 384, 768]
       if "small" in checkpoint_url:
              UpperCamelCase_      				=     [3, 3, 27, 3]
              UpperCamelCase_      				=     [96, 192, 384, 768]
       if "base" in checkpoint_url:
              UpperCamelCase_      				=     [3, 3, 27, 3]
              UpperCamelCase_      				=     [128, 256, 512, 1024]
       if "large" in checkpoint_url:
              UpperCamelCase_      				=     [3, 3, 27, 3]
              UpperCamelCase_      				=     [192, 384, 768, 1536]
       if "xlarge" in checkpoint_url:
              UpperCamelCase_      				=     [3, 3, 27, 3]
              UpperCamelCase_      				=     [256, 512, 1024, 2048]

       if "1k" in checkpoint_url:
              UpperCamelCase_      				=     1000
              UpperCamelCase_      				=     "imagenet-1k-id2label.json"
              UpperCamelCase_      				=     (1, 1000)
       else:
              UpperCamelCase_      				=     21841
              UpperCamelCase_      				=     "imagenet-22k-id2label.json"
              UpperCamelCase_      				=     (1, 21841)

       UpperCamelCase_      				=     "huggingface/label-files"
       UpperCamelCase_      				=     num_labels
       UpperCamelCase_      				=     json.load(open(hf_hub_download(UpperCamelCase_    ,						UpperCamelCase_    ,						repo_type="dataset"     )    ,						"r"     )     )
       UpperCamelCase_      				=     {int(UpperCamelCase_     ): 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[9205]
              del idalabel[15027]
       UpperCamelCase_      				=     idalabel
       UpperCamelCase_      				=     {v: k for k, v in idalabel.items()}
       UpperCamelCase_      				=     hidden_sizes
       UpperCamelCase_      				=     depths

       return config, expected_shape




def        lowerCAmelCase_       (							UpperCamelCase_     )   ->						str:
       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 "stages" in name:
              UpperCamelCase_      				=     name.replace("stages"    ,						"encoder.stages"     )
       if "norm" in name:
              UpperCamelCase_      				=     name.replace("norm"    ,						"layernorm"     )
       if "gamma" in name:
              UpperCamelCase_      				=     name.replace("gamma"    ,						"layer_scale_parameter"     )
       if "head" in name:
              UpperCamelCase_      				=     name.replace("head"    ,						"classifier"     )

       return name




def        lowerCAmelCase_       (							)   ->						Any:
       UpperCamelCase_      				=     "http://images.cocodataset.org/val2017/000000039769.jpg"
       UpperCamelCase_      				=     Image.open(requests.get(UpperCamelCase_    ,						stream=UpperCamelCase_     ).raw     )
       return im




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

       # load HuggingFace model
       UpperCamelCase_      				=     ConvNextForImageClassification(UpperCamelCase_     )
       model.load_state_dict(UpperCamelCase_     )
       model.eval()

       # Check outputs on an image, prepared by ConvNextImageProcessor
       UpperCamelCase_      				=     224 if "224" in checkpoint_url else 384
       UpperCamelCase_      				=     ConvNextImageProcessor(size=UpperCamelCase_     )
       UpperCamelCase_      				=     image_processor(images=prepare_img()    ,						return_tensors="pt"     ).pixel_values

       UpperCamelCase_      				=     model(UpperCamelCase_     ).logits

       # note: the logits below were obtained without center cropping
       if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
              UpperCamelCase_      				=     torch.tensor([-0.12_10, -0.66_05, 0.19_18]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
              UpperCamelCase_      				=     torch.tensor([-0.44_73, -0.18_47, -0.63_65]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
              UpperCamelCase_      				=     torch.tensor([0.45_25, 0.75_39, 0.03_08]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
              UpperCamelCase_      				=     torch.tensor([0.35_61, 0.63_50, -0.03_84]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
              UpperCamelCase_      				=     torch.tensor([0.41_74, -0.09_89, 0.14_89]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
              UpperCamelCase_      				=     torch.tensor([0.25_13, -0.13_49, -0.16_13]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
              UpperCamelCase_      				=     torch.tensor([1.29_80, 0.36_31, -0.11_98]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
              UpperCamelCase_      				=     torch.tensor([1.29_63, 0.12_27, 0.17_23]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
              UpperCamelCase_      				=     torch.tensor([1.79_56, 0.83_90, 0.28_20]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
              UpperCamelCase_      				=     torch.tensor([-0.28_22, -0.05_02, -0.08_78]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
              UpperCamelCase_      				=     torch.tensor([-0.56_72, -0.07_30, -0.43_48]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
              UpperCamelCase_      				=     torch.tensor([0.26_81, 0.23_65, 0.62_46]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
              UpperCamelCase_      				=     torch.tensor([-0.26_42, 0.39_31, 0.51_16]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
              UpperCamelCase_      				=     torch.tensor([-0.66_77, -0.18_73, -0.83_79]     )
       elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
              UpperCamelCase_      				=     torch.tensor([-0.77_49, -0.29_67, -0.64_44]     )
       else:
              raise ValueError(F'''Unknown URL: {checkpoint_url}'''     )

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

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

       print("Pushing model to the hub..."     )
       UpperCamelCase_      				=     "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(UpperCamelCase_    ,						UpperCamelCase_     )    ,						organization="nielsr"    ,						commit_message="Add model"    ,						)


if __name__ == "__main__":
     _UpperCAmelCase				       = 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.',
     )

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








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_UpperCAmelCase				       = logging.get_logger(__name__)

_UpperCAmelCase				       = {
    'MCG-NJU/videomae-base': 'https://huggingface.co/MCG-NJU/videomae-base/resolve/main/config.json',
}



class 							_UpperCamelCase    (		lowerCAmelCase_   ):
       _UpperCamelCase :    List[Any]   			=      '''videomae'''

       def __init__(						self:       Optional[Any]    ,      _SCREAMING_SNAKE_CASE:       List[str]=224    ,      _SCREAMING_SNAKE_CASE:       Dict=16    ,      _SCREAMING_SNAKE_CASE:       List[Any]=3    ,      _SCREAMING_SNAKE_CASE:       Tuple=16    ,      _SCREAMING_SNAKE_CASE:       Tuple=2    ,      _SCREAMING_SNAKE_CASE:       Optional[Any]=768    ,      _SCREAMING_SNAKE_CASE:       str=12    ,      _SCREAMING_SNAKE_CASE:       List[Any]=12    ,      _SCREAMING_SNAKE_CASE:       Dict=3072    ,      _SCREAMING_SNAKE_CASE:       Tuple="gelu"    ,      _SCREAMING_SNAKE_CASE:       List[Any]=0.0    ,      _SCREAMING_SNAKE_CASE:       List[Any]=0.0    ,      _SCREAMING_SNAKE_CASE:       Tuple=0.02    ,      _SCREAMING_SNAKE_CASE:       Optional[int]=1e-12    ,      _SCREAMING_SNAKE_CASE:       Dict=True    ,      _SCREAMING_SNAKE_CASE:       Dict=True    ,      _SCREAMING_SNAKE_CASE:       List[str]=6    ,      _SCREAMING_SNAKE_CASE:       Dict=384    ,      _SCREAMING_SNAKE_CASE:       Any=4    ,      _SCREAMING_SNAKE_CASE:       Dict=1536    ,      _SCREAMING_SNAKE_CASE:       Optional[int]=True    ,      **_SCREAMING_SNAKE_CASE:       str    ,      )				->       Dict:




              """simple docstring"""


              super().__init__(**_SCREAMING_SNAKE_CASE						)

              UpperCamelCase_      				=     image_size
              UpperCamelCase_      				=     patch_size
              UpperCamelCase_      				=     num_channels
              UpperCamelCase_      				=     num_frames
              UpperCamelCase_      				=     tubelet_size

              UpperCamelCase_      				=     hidden_size
              UpperCamelCase_      				=     num_hidden_layers
              UpperCamelCase_      				=     num_attention_heads
              UpperCamelCase_      				=     intermediate_size
              UpperCamelCase_      				=     hidden_act
              UpperCamelCase_      				=     hidden_dropout_prob
              UpperCamelCase_      				=     attention_probs_dropout_prob
              UpperCamelCase_      				=     initializer_range
              UpperCamelCase_      				=     layer_norm_eps
              UpperCamelCase_      				=     qkv_bias
              UpperCamelCase_      				=     use_mean_pooling

              UpperCamelCase_      				=     decoder_num_attention_heads
              UpperCamelCase_      				=     decoder_hidden_size
              UpperCamelCase_      				=     decoder_num_hidden_layers
              UpperCamelCase_      				=     decoder_intermediate_size
              UpperCamelCase_      				=     norm_pix_loss

