





'''simple docstring'''






import argparse
import json

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

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)
def 							_lowerCamelCase(  UpperCamelCase__    : str    ,       UpperCamelCase__    : List[Any]=False		)		->  List[Any]:
    A       :       Tuple							=						[]
    for i in range(config.num_hidden_layers		):
        # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
        rename_keys.append((F'''module.blocks.{i}.norm1.weight''', F'''vit.encoder.layer.{i}.layernorm_before.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.norm1.bias''', F'''vit.encoder.layer.{i}.layernorm_before.bias''')		)
        rename_keys.append(
            (F'''module.blocks.{i}.attn.proj.weight''', F'''vit.encoder.layer.{i}.attention.output.dense.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.attn.proj.bias''', F'''vit.encoder.layer.{i}.attention.output.dense.bias''')		)
        rename_keys.append((F'''module.blocks.{i}.norm2.weight''', F'''vit.encoder.layer.{i}.layernorm_after.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.norm2.bias''', F'''vit.encoder.layer.{i}.layernorm_after.bias''')		)
        rename_keys.append((F'''module.blocks.{i}.mlp.fc1.weight''', F'''vit.encoder.layer.{i}.intermediate.dense.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.mlp.fc1.bias''', F'''vit.encoder.layer.{i}.intermediate.dense.bias''')		)
        rename_keys.append((F'''module.blocks.{i}.mlp.fc2.weight''', F'''vit.encoder.layer.{i}.output.dense.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.mlp.fc2.bias''', F'''vit.encoder.layer.{i}.output.dense.bias''')		)

    # projection layer + position embeddings
    rename_keys.extend(
        [
            ('''module.cls_token''', '''vit.embeddings.cls_token'''),
            ('''module.patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
            ('''module.patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
            ('''module.pos_embed''', '''vit.embeddings.position_embeddings'''),
        ]		)

    if base_model:
        # layernorm + pooler
        rename_keys.extend(
            [
                ('''module.norm.weight''', '''layernorm.weight'''),
                ('''module.norm.bias''', '''layernorm.bias'''),
            ]		)

        # if just the base model, we should remove "vit" from all keys that start with "vit"
        A       :       Any							=						[(pair[0], pair[1][4:]) if pair[1].startswith('''vit'''		) else pair for pair in rename_keys]
    else:
        # layernorm + classification head
        rename_keys.extend(
            [
                ('''norm.weight''', '''vit.layernorm.weight'''),
                ('''norm.bias''', '''vit.layernorm.bias'''),
                ('''head.weight''', '''classifier.weight'''),
                ('''head.bias''', '''classifier.bias'''),
            ]		)

    return rename_keys
def 							_lowerCamelCase(  UpperCamelCase__    : Tuple    ,       UpperCamelCase__    : Optional[Any]    ,       UpperCamelCase__    : str=False		)		->  str:
    for i in range(config.num_hidden_layers		):
        if base_model:
            A       :       Dict							=						''''''
        else:
            A       :       Optional[int]							=						'''vit.'''
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        A       :       Optional[Any]							=						state_dict.pop(F'''module.blocks.{i}.attn.qkv.weight'''		)
        A       :       str							=						state_dict.pop(F'''module.blocks.{i}.attn.qkv.bias'''		)
        # next, add query, keys and values (in that order) to the state dict
        A       :       int							=						in_proj_weight[
            : config.hidden_size, :
        ]
        A       :       List[str]							=						in_proj_bias[: config.hidden_size]
        A       :       str							=						in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        A       :       Tuple							=						in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        A       :       Optional[Any]							=						in_proj_weight[
            -config.hidden_size :, :
        ]
        A       :       Tuple							=						in_proj_bias[-config.hidden_size :]
def 							_lowerCamelCase(  UpperCamelCase__    : str		)		->  Optional[int]:
    A       :       Optional[Any]							=						['''head.weight''', '''head.bias''']
    for k in ignore_keys:
        state_dict.pop(UpperCamelCase__    ,       UpperCamelCase__		)
def 							_lowerCamelCase(  UpperCamelCase__    : Optional[int]		)		->  str:
    # projection head is used in the self-supervised pre-training in MSN,
    # for downstream task it's not needed.
    A       :       List[Any]							=						[
        '''module.fc.fc1.weight''',
        '''module.fc.fc1.bias''',
        '''module.fc.bn1.weight''',
        '''module.fc.bn1.bias''',
        '''module.fc.bn1.running_mean''',
        '''module.fc.bn1.running_var''',
        '''module.fc.bn1.num_batches_tracked''',
        '''module.fc.fc2.weight''',
        '''module.fc.fc2.bias''',
        '''module.fc.bn2.weight''',
        '''module.fc.bn2.bias''',
        '''module.fc.bn2.running_mean''',
        '''module.fc.bn2.running_var''',
        '''module.fc.bn2.num_batches_tracked''',
        '''module.fc.fc3.weight''',
        '''module.fc.fc3.bias''',
    ]
    for k in ignore_keys:
        state_dict.pop(UpperCamelCase__    ,       UpperCamelCase__		)
def 							_lowerCamelCase(  UpperCamelCase__    : Union[str, Any]    ,       UpperCamelCase__    : Tuple    ,       UpperCamelCase__    : str		)		->  Optional[Any]:
    A       :       List[Any]							=						dct.pop(UpperCamelCase__		)
    A       :       Optional[int]							=						val
def 							_lowerCamelCase(  UpperCamelCase__    : Optional[Any]    ,       UpperCamelCase__    : Optional[Any]		)		->  List[str]:
    A       :       Optional[Any]							=						ViTMSNConfig()
    A       :       Optional[Any]							=						1_000

    A       :       Tuple							=						'''datasets/huggingface/label-files'''
    A       :       List[str]							=						'''imagenet-1k-id2label.json'''
    A       :       Any							=						json.load(open(hf_hub_download(UpperCamelCase__    ,       UpperCamelCase__		)    ,       '''r'''		)		)
    A       :       Any							=						{int(UpperCamelCase__		): v for k, v in idalabel.items()}
    A       :       str							=						idalabel
    A       :       Any							=						{v: k for k, v in idalabel.items()}

    if "s16" in checkpoint_url:
        A       :       Union[str, Any]							=						384
        A       :       Any							=						1_536
        A       :       Any							=						6
    elif "l16" in checkpoint_url:
        A       :       Dict							=						1_024
        A       :       Tuple							=						4_096
        A       :       int							=						24
        A       :       Dict							=						16
        A       :       int							=						0.1
    elif "b4" in checkpoint_url:
        A       :       str							=						4
    elif "l7" in checkpoint_url:
        A       :       Optional[int]							=						7
        A       :       Optional[int]							=						1_024
        A       :       int							=						4_096
        A       :       str							=						24
        A       :       Any							=						16
        A       :       Union[str, Any]							=						0.1

    A       :       Dict							=						ViTMSNModel(UpperCamelCase__		)

    A       :       Union[str, Any]							=						torch.hub.load_state_dict_from_url(UpperCamelCase__    ,       map_location='''cpu'''		)['''target_encoder''']

    A       :       int							=						ViTImageProcessor(size=config.image_size		)

    remove_projection_head(UpperCamelCase__		)
    A       :       Tuple							=						create_rename_keys(UpperCamelCase__    ,       base_model=UpperCamelCase__		)

    for src, dest in rename_keys:
        rename_key(UpperCamelCase__    ,       UpperCamelCase__    ,       UpperCamelCase__		)
    read_in_q_k_v(UpperCamelCase__    ,       UpperCamelCase__    ,       base_model=UpperCamelCase__		)

    model.load_state_dict(UpperCamelCase__		)
    model.eval()

    A       :       Union[str, Any]							=						'''http://images.cocodataset.org/val2017/000000039769.jpg'''

    A       :       Optional[Any]							=						Image.open(requests.get(UpperCamelCase__    ,       stream=UpperCamelCase__		).raw		)
    A       :       Dict							=						ViTImageProcessor(
        size=config.image_size    ,       image_mean=UpperCamelCase__    ,       image_std=UpperCamelCase__		)
    A       :       List[str]							=						image_processor(images=UpperCamelCase__    ,       return_tensors='''pt'''		)

    # forward pass
    torch.manual_seed(2		)
    A       :       Optional[Any]							=						model(**UpperCamelCase__		)
    A       :       int							=						outputs.last_hidden_state

    # The following Colab Notebook was used to generate these outputs:
    # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
    if "s16" in checkpoint_url:
        A       :       str							=						torch.tensor([[-1.0_9_1_5, -1.4_8_7_6, -1.1_8_0_9]]		)
    elif "b16" in checkpoint_url:
        A       :       int							=						torch.tensor([[14.2_889, -18.9_045, 11.7_281]]		)
    elif "l16" in checkpoint_url:
        A       :       Optional[Any]							=						torch.tensor([[41.5_028, -22.8_681, 45.6_475]]		)
    elif "b4" in checkpoint_url:
        A       :       Optional[Any]							=						torch.tensor([[-4.3_8_6_8, 5.2_9_3_2, -0.4_1_3_7]]		)
    else:
        A       :       str							=						torch.tensor([[-0.1_7_9_2, -0.6_4_6_5, 2.4_2_6_3]]		)

    # verify logits
    assert torch.allclose(last_hidden_state[:, 0, :3]    ,       UpperCamelCase__    ,       atol=1e-4		)

    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__		)


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

     snake_case_     =	parser.parse_args()
     convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)










'''simple docstring'''






import importlib.util
import os
import platform
from argparse import ArgumentParser

import huggingface_hub

from .. import __version__ as version
from ..utils import (
    is_accelerate_available,
    is_flax_available,
    is_safetensors_available,
    is_tf_available,
    is_torch_available,
)
from . import BaseTransformersCLICommand
def 							_lowerCamelCase(  UpperCamelCase__    : Optional[int]		)		->  List[str]:
    return EnvironmentCommand()
def 							_lowerCamelCase(  UpperCamelCase__    : List[Any]		)		->  Tuple:
    return EnvironmentCommand(args.accelerate_config_file		)





class   _lowercase						(							a				):


   @staticmethod
   def 					snake_case				( _UpperCAmelCase					):
       A       :       Optional[Any]							=						parser.add_parser('''env'''					)
       download_parser.set_defaults(func=_UpperCAmelCase					)
       download_parser.add_argument(
           '''--accelerate-config_file'''    ,   default=_UpperCAmelCase    ,   help='''The accelerate config file to use for the default values in the launching script.'''    ,   )
       download_parser.set_defaults(func=_UpperCAmelCase					)


   def __init__( self    ,   _UpperCAmelCase    ,   *_UpperCAmelCase					):
       A       :       Any							=						accelerate_config_file


   def 					snake_case				( self					):
       A       :       Any							=						'''not installed'''
       if is_safetensors_available():
           import safetensors

           A       :       Dict							=						safetensors.__version__
       elif importlib.util.find_spec('''safetensors'''					) is not None:
           import safetensors

           A       :       int							=						f'''{safetensors.__version__} but is ignored because of PyTorch version too old.'''

       A       :       Tuple							=						'''not installed'''
       A       :       Tuple							=						'''not found'''
       if is_accelerate_available():
           import accelerate
           from accelerate.commands.config import default_config_file, load_config_from_file

           A       :       Any							=						accelerate.__version__
           # Get the default from the config file.
           if self._accelerate_config_file is not None or os.path.isfile(_UpperCAmelCase					):
               A       :       int							=						load_config_from_file(self._accelerate_config_file					).to_dict()

           A       :       Tuple							=						(
               '''\n'''.join([f'''\t- {prop}: {val}''' for prop, val in accelerate_config.items()]					)
               if isinstance(_UpperCAmelCase    ,   _UpperCAmelCase					)
               else f'''\t{accelerate_config}'''
           )

       A       :       str							=						'''not installed'''
       A       :       Optional[Any]							=						'''NA'''
       if is_torch_available():
           import torch

           A       :       Optional[int]							=						torch.__version__
           A       :       List[Any]							=						torch.cuda.is_available()

       A       :       Dict							=						'''not installed'''
       A       :       Any							=						'''NA'''
       if is_tf_available():
           import tensorflow as tf

           A       :       Optional[int]							=						tf.__version__
           try:
               # deprecated in v2.1
               A       :       int							=						tf.test.is_gpu_available()
           except AttributeError:
               # returns list of devices, convert to bool
               A       :       Optional[int]							=						bool(tf.config.list_physical_devices('''GPU'''					)					)

       A       :       str							=						'''not installed'''
       A       :       Dict							=						'''not installed'''
       A       :       int							=						'''not installed'''
       A       :       Tuple							=						'''NA'''
       if is_flax_available():
           import flax
           import jax
           import jaxlib

           A       :       Any							=						flax.__version__
           A       :       Tuple							=						jax.__version__
           A       :       int							=						jaxlib.__version__
           A       :       List[Any]							=						jax.lib.xla_bridge.get_backend().platform

       A       :       Union[str, Any]							=						{
           '''`transformers` version''': version,
           '''Platform''': platform.platform(),
           '''Python version''': platform.python_version(),
           '''Huggingface_hub version''': huggingface_hub.__version__,
           '''Safetensors version''': f'''{safetensors_version}''',
           '''Accelerate version''': f'''{accelerate_version}''',
           '''Accelerate config''': f'''{accelerate_config_str}''',
           '''PyTorch version (GPU?)''': f'''{pt_version} ({pt_cuda_available})''',
           '''Tensorflow version (GPU?)''': f'''{tf_version} ({tf_cuda_available})''',
           '''Flax version (CPU?/GPU?/TPU?)''': f'''{flax_version} ({jax_backend})''',
           '''Jax version''': f'''{jax_version}''',
           '''JaxLib version''': f'''{jaxlib_version}''',
           '''Using GPU in script?''': '''<fill in>''',
           '''Using distributed or parallel set-up in script?''': '''<fill in>''',
       }

       print('''\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n'''					)
       print(self.format_dict(_UpperCAmelCase					)					)

       return info


   @staticmethod
   def 					snake_case				( _UpperCAmelCase					):
       return "\n".join([f'''- {prop}: {val}''' for prop, val in d.items()]					) + "\n"
