





from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


_snake_case           =			logging.get_logger(__name__)

# General docstring
_snake_case           =			"""ResNetConfig"""

# Base docstring
_snake_case           =			"""microsoft/resnet-50"""
_snake_case           =			[1, 2048, 7, 7]

# Image classification docstring
_snake_case           =			"""microsoft/resnet-50"""
_snake_case           =			"""tiger cat"""

_snake_case           =			[
    """microsoft/resnet-50""",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :Optional[int]   ,							_lowercase  :int   ,							_lowercase  :int   ,							_lowercase  :int = 3   ,							_lowercase  :int = 1   ,							_lowercase  :str = "relu" ):






     '''simple docstring'''





     super().__init__()
     lowercase__			       =  nn.Convad(
         _lowercase   ,							_lowercase   ,							kernel_size=_lowercase   ,							stride=_lowercase   ,							padding=kernel_size // 2   ,							bias=_lowercase )
     lowercase__			       =  nn.BatchNormad(_lowercase )
     lowercase__			       =  ACTaFN[activation] if activation is not None else nn.Identity()







  def   UpperCAmelCase					(      self  :Optional[Any]   ,							_lowercase  :Tensor ):






     '''simple docstring'''





     lowercase__			       =  self.convolution(_lowercase )
     lowercase__			       =  self.normalization(_lowercase )
     lowercase__			       =  self.activation(_lowercase )
     return hidden_state





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :Optional[Any]   ,							_lowercase  :ResNetConfig ):






     '''simple docstring'''





     super().__init__()
     lowercase__			       =  ResNetConvLayer(
         config.num_channels   ,							config.embedding_size   ,							kernel_size=7   ,							stride=2   ,							activation=config.hidden_act )
     lowercase__			       =  nn.MaxPoolad(kernel_size=3   ,							stride=2   ,							padding=1 )
     lowercase__			       =  config.num_channels







  def   UpperCAmelCase					(      self  :Optional[Any]   ,							_lowercase  :Tensor ):






     '''simple docstring'''





     lowercase__			       =  pixel_values.shape[1]
     if num_channels != self.num_channels:
        raise ValueError(
            "Make sure that the channel dimension of the pixel values match with the one set in the configuration." )
     lowercase__			       =  self.embedder(_lowercase )
     lowercase__			       =  self.pooler(_lowercase )
     return embedding





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :List[Any]   ,							_lowercase  :int   ,							_lowercase  :int   ,							_lowercase  :int = 2 ):






     '''simple docstring'''





     super().__init__()
     lowercase__			       =  nn.Convad(_lowercase   ,							_lowercase   ,							kernel_size=1   ,							stride=_lowercase   ,							bias=_lowercase )
     lowercase__			       =  nn.BatchNormad(_lowercase )







  def   UpperCAmelCase					(      self  :List[str]   ,							_lowercase  :Tensor ):






     '''simple docstring'''





     lowercase__			       =  self.convolution(_lowercase )
     lowercase__			       =  self.normalization(_lowercase )
     return hidden_state





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :List[str]   ,							_lowercase  :int   ,							_lowercase  :int   ,							_lowercase  :int = 1   ,							_lowercase  :str = "relu" ):






     '''simple docstring'''





     super().__init__()
     lowercase__			       =  in_channels != out_channels or stride != 1
     lowercase__			       =  (
         ResNetShortCut(_lowercase   ,							_lowercase   ,							stride=_lowercase ) if should_apply_shortcut else nn.Identity()
     )
     lowercase__			       =  nn.Sequential(
         ResNetConvLayer(_lowercase   ,							_lowercase   ,							stride=_lowercase )   ,							ResNetConvLayer(_lowercase   ,							_lowercase   ,							activation=_lowercase )   ,							)
     lowercase__			       =  ACTaFN[activation]







  def   UpperCAmelCase					(      self  :int   ,							_lowercase  :Tuple ):






     '''simple docstring'''





     lowercase__			       =  hidden_state
     lowercase__			       =  self.layer(_lowercase )
     lowercase__			       =  self.shortcut(_lowercase )
     hidden_state += residual
     lowercase__			       =  self.activation(_lowercase )
     return hidden_state





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :Optional[int]   ,							_lowercase  :int   ,							_lowercase  :int   ,							_lowercase  :int = 1   ,							_lowercase  :str = "relu"   ,							_lowercase  :int = 4 ):






     '''simple docstring'''





     super().__init__()
     lowercase__			       =  in_channels != out_channels or stride != 1
     lowercase__			       =  out_channels // reduction
     lowercase__			       =  (
         ResNetShortCut(_lowercase   ,							_lowercase   ,							stride=_lowercase ) if should_apply_shortcut else nn.Identity()
     )
     lowercase__			       =  nn.Sequential(
         ResNetConvLayer(_lowercase   ,							_lowercase   ,							kernel_size=1 )   ,							ResNetConvLayer(_lowercase   ,							_lowercase   ,							stride=_lowercase )   ,							ResNetConvLayer(_lowercase   ,							_lowercase   ,							kernel_size=1   ,							activation=_lowercase )   ,							)
     lowercase__			       =  ACTaFN[activation]







  def   UpperCAmelCase					(      self  :int   ,							_lowercase  :Tuple ):






     '''simple docstring'''





     lowercase__			       =  hidden_state
     lowercase__			       =  self.layer(_lowercase )
     lowercase__			       =  self.shortcut(_lowercase )
     hidden_state += residual
     lowercase__			       =  self.activation(_lowercase )
     return hidden_state





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :List[str]   ,							_lowercase  :ResNetConfig   ,							_lowercase  :int   ,							_lowercase  :int   ,							_lowercase  :int = 2   ,							_lowercase  :int = 2   ,							):






     '''simple docstring'''





     super().__init__()

     lowercase__			       =  ResNetBottleNeckLayer if config.layer_type == "bottleneck" else ResNetBasicLayer

     lowercase__			       =  nn.Sequential(
         # downsampling is done in the first layer with stride of 2
         layer(_lowercase   ,							_lowercase   ,							stride=_lowercase   ,							activation=config.hidden_act )   ,							*[layer(_lowercase   ,							_lowercase   ,							activation=config.hidden_act ) for _ in range(depth - 1 )]   ,							)







  def   UpperCAmelCase					(      self  :int   ,							_lowercase  :Tensor ):






     '''simple docstring'''





     lowercase__			       =  input
     for layer in self.layers:
        lowercase__			       =  layer(_lowercase )
     return hidden_state





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :List[str]   ,							_lowercase  :ResNetConfig ):






     '''simple docstring'''





     super().__init__()
     lowercase__			       =  nn.ModuleList([] )
     # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
     self.stages.append(
         ResNetStage(
             _lowercase   ,							config.embedding_size   ,							config.hidden_sizes[0]   ,							stride=2 if config.downsample_in_first_stage else 1   ,							depth=config.depths[0]   ,							) )
     lowercase__			       =  zip(config.hidden_sizes   ,							config.hidden_sizes[1:] )
     for (in_channels, out_channels), depth in zip(_lowercase   ,							config.depths[1:] ):
        self.stages.append(ResNetStage(_lowercase   ,							_lowercase   ,							_lowercase   ,							depth=_lowercase ) )







  def   UpperCAmelCase					(      self  :Optional[Any]   ,							_lowercase  :Tensor   ,							_lowercase  :bool = False   ,							_lowercase  :bool = True ):






     '''simple docstring'''





     lowercase__			       =  () if output_hidden_states else None

     for stage_module in self.stages:
        if output_hidden_states:
           lowercase__			       =  hidden_states + (hidden_state,)

        lowercase__			       =  stage_module(_lowercase )

     if output_hidden_states:
        lowercase__			       =  hidden_states + (hidden_state,)

     if not return_dict:
        return tuple(v for v in [hidden_state, hidden_states] if v is not None )

     return BaseModelOutputWithNoAttention(
         last_hidden_state=_lowercase   ,							hidden_states=_lowercase   ,							)





class  lowerCAmelCase							(       lowercase_ ):
  __lowerCamelCase		=      ResNetConfig
  __lowerCamelCase		=      'resnet'
  __lowerCamelCase		=      'pixel_values'
  __lowerCamelCase		=      True







  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :Optional[Any] ):






     '''simple docstring'''





     if isinstance(_lowercase   ,							nn.Convad ):
        nn.init.kaiming_normal_(module.weight   ,							mode="fan_out"   ,							nonlinearity="relu" )
     elif isinstance(_lowercase   ,							(nn.BatchNormad, nn.GroupNorm) ):
        nn.init.constant_(module.weight   ,							1 )
        nn.init.constant_(module.bias   ,							0 )







  def   UpperCAmelCase					(      self  :Optional[Any]   ,							_lowercase  :List[Any]   ,							_lowercase  :str=False ):






     '''simple docstring'''





     if isinstance(_lowercase   ,							_lowercase ):
        lowercase__			       =  value
_snake_case           =			R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

_snake_case           =			R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""





@add_start_docstrings(
    'The bare ResNet model outputting raw features without any specific head on top.'							,							lowercase_							,							)
class  lowerCAmelCase							(       lowercase_ ):







  def __init__(      self  :List[str]   ,							_lowercase  :str ):






     '''simple docstring'''





     super().__init__(_lowercase )
     lowercase__			       =  config
     lowercase__			       =  ResNetEmbeddings(_lowercase )
     lowercase__			       =  ResNetEncoder(_lowercase )
     lowercase__			       =  nn.AdaptiveAvgPoolad((1, 1) )
     # Initialize weights and apply final processing
     self.post_init()







  @add_start_docstrings_to_model_forward(_lowercase )
  @add_code_sample_docstrings(
      checkpoint=_CHECKPOINT_FOR_DOC   ,							output_type=_lowercase   ,							config_class=_CONFIG_FOR_DOC   ,							modality="vision"   ,							expected_output=_EXPECTED_OUTPUT_SHAPE   ,							)
  def   UpperCAmelCase					(      self  :Tuple   ,							_lowercase  :Tensor   ,							_lowercase  :Optional[bool] = None   ,							_lowercase  :Optional[bool] = None ):






     '''simple docstring'''





     lowercase__			       =  (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     lowercase__			       =  return_dict if return_dict is not None else self.config.use_return_dict

     lowercase__			       =  self.embedder(_lowercase )

     lowercase__			       =  self.encoder(
         _lowercase   ,							output_hidden_states=_lowercase   ,							return_dict=_lowercase )

     lowercase__			       =  encoder_outputs[0]

     lowercase__			       =  self.pooler(_lowercase )

     if not return_dict:
        return (last_hidden_state, pooled_output) + encoder_outputs[1:]

     return BaseModelOutputWithPoolingAndNoAttention(
         last_hidden_state=_lowercase   ,							pooler_output=_lowercase   ,							hidden_states=encoder_outputs.hidden_states   ,							)





@add_start_docstrings(
    '\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '							,							lowercase_							,							)
class  lowerCAmelCase							(       lowercase_ ):







  def __init__(      self  :Optional[Any]   ,							_lowercase  :Tuple ):






     '''simple docstring'''





     super().__init__(_lowercase )
     lowercase__			       =  config.num_labels
     lowercase__			       =  ResNetModel(_lowercase )
     # classification head
     lowercase__			       =  nn.Sequential(
         nn.Flatten()   ,							nn.Linear(config.hidden_sizes[-1]   ,							config.num_labels ) if config.num_labels > 0 else nn.Identity()   ,							)
     # initialize weights and apply final processing
     self.post_init()







  @add_start_docstrings_to_model_forward(_lowercase )
  @add_code_sample_docstrings(
      checkpoint=_IMAGE_CLASS_CHECKPOINT   ,							output_type=_lowercase   ,							config_class=_CONFIG_FOR_DOC   ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,							)
  def   UpperCAmelCase					(      self  :List[str]   ,							_lowercase  :Optional[torch.FloatTensor] = None   ,							_lowercase  :Optional[torch.LongTensor] = None   ,							_lowercase  :Optional[bool] = None   ,							_lowercase  :Optional[bool] = None   ,							):






     '''simple docstring'''





     lowercase__			       =  return_dict if return_dict is not None else self.config.use_return_dict

     lowercase__			       =  self.resnet(_lowercase   ,							output_hidden_states=_lowercase   ,							return_dict=_lowercase )

     lowercase__			       =  outputs.pooler_output if return_dict else outputs[1]

     lowercase__			       =  self.classifier(_lowercase )

     lowercase__			       =  None

     if labels is not None:
        if self.config.problem_type is None:
           if self.num_labels == 1:
              lowercase__			       =  "regression"
           elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
              lowercase__			       =  "single_label_classification"
           else:
              lowercase__			       =  "multi_label_classification"
        if self.config.problem_type == "regression":
           lowercase__			       =  MSELoss()
           if self.num_labels == 1:
              lowercase__			       =  loss_fct(logits.squeeze()   ,							labels.squeeze() )
           else:
              lowercase__			       =  loss_fct(_lowercase   ,							_lowercase )
        elif self.config.problem_type == "single_label_classification":
           lowercase__			       =  CrossEntropyLoss()
           lowercase__			       =  loss_fct(logits.view(-1   ,							self.num_labels )   ,							labels.view(-1 ) )
        elif self.config.problem_type == "multi_label_classification":
           lowercase__			       =  BCEWithLogitsLoss()
           lowercase__			       =  loss_fct(_lowercase   ,							_lowercase )

     if not return_dict:
        lowercase__			       =  (logits,) + outputs[2:]
        return (loss,) + output if loss is not None else output

     return ImageClassifierOutputWithNoAttention(loss=_lowercase   ,							logits=_lowercase   ,							hidden_states=outputs.hidden_states )





@add_start_docstrings(
    '\n    ResNet backbone, to be used with frameworks like DETR and MaskFormer.\n    '							,							lowercase_							,							)
class  lowerCAmelCase							(       lowercase_							,							lowercase_ ):







  def __init__(      self  :List[Any]   ,							_lowercase  :List[str] ):






     '''simple docstring'''





     super().__init__(_lowercase )
     super()._init_backbone(_lowercase )

     lowercase__			       =  [config.embedding_size] + config.hidden_sizes
     lowercase__			       =  ResNetEmbeddings(_lowercase )
     lowercase__			       =  ResNetEncoder(_lowercase )

     # initialize weights and apply final processing
     self.post_init()







  @add_start_docstrings_to_model_forward(_lowercase )
  @replace_return_docstrings(output_type=_lowercase   ,							config_class=_CONFIG_FOR_DOC )
  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :Tensor   ,							_lowercase  :Optional[bool] = None   ,							_lowercase  :Optional[bool] = None ):






     '''simple docstring'''





     lowercase__			       =  return_dict if return_dict is not None else self.config.use_return_dict
     lowercase__			       =  (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )

     lowercase__			       =  self.embedder(_lowercase )

     lowercase__			       =  self.encoder(_lowercase   ,							output_hidden_states=_lowercase   ,							return_dict=_lowercase )

     lowercase__			       =  outputs.hidden_states

     lowercase__			       =  ()
     for idx, stage in enumerate(self.stage_names ):
        if stage in self.out_features:
           feature_maps += (hidden_states[idx],)

     if not return_dict:
        lowercase__			       =  (feature_maps,)
        if output_hidden_states:
           output += (outputs.hidden_states,)
        return output

     return BackboneOutput(
         feature_maps=_lowercase   ,							hidden_states=outputs.hidden_states if output_hidden_states else None   ,							attentions=_lowercase   ,							)













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 _A       (   __magic_name__  , __magic_name__=False	):
   lowercase__			       =  []
   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"
      lowercase__			       =  [(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 _A       (   __magic_name__  , __magic_name__  , __magic_name__=False	):
   for i in range(config.num_hidden_layers	):
      if base_model:
         lowercase__			       =  ""
      else:
         lowercase__			       =  "vit."
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      lowercase__			       =  state_dict.pop(f'''module.blocks.{i}.attn.qkv.weight'''	)
      lowercase__			       =  state_dict.pop(f'''module.blocks.{i}.attn.qkv.bias'''	)
      # next, add query, keys and values (in that order) to the state dict
      lowercase__			       =  in_proj_weight[
          : config.hidden_size, :
      ]
      lowercase__			       =  in_proj_bias[: config.hidden_size]
      lowercase__			       =  in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      lowercase__			       =  in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      lowercase__			       =  in_proj_weight[
          -config.hidden_size :, :
      ]
      lowercase__			       =  in_proj_bias[-config.hidden_size :]


def _A       (   __magic_name__	):
   lowercase__			       =  ["head.weight", "head.bias"]
   for k in ignore_keys:
      state_dict.pop(__magic_name__  , __magic_name__	)


def _A       (   __magic_name__	):
   # projection head is used in the self-supervised pre-training in MSN,
   # for downstream task it's not needed.
   lowercase__			       =  [
       "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(__magic_name__  , __magic_name__	)


def _A       (   __magic_name__  , __magic_name__  , __magic_name__	):
   lowercase__			       =  dct.pop(__magic_name__	)
   lowercase__			       =  val


def _A       (   __magic_name__  , __magic_name__	):
   lowercase__			       =  ViTMSNConfig()
   lowercase__			       =  1000

   lowercase__			       =  "datasets/huggingface/label-files"
   lowercase__			       =  "imagenet-1k-id2label.json"
   lowercase__			       =  json.load(open(hf_hub_download(__magic_name__  , __magic_name__	)  , "r"	)	)
   lowercase__			       =  {int(__magic_name__	): v for k, v in idalabel.items()}
   lowercase__			       =  idalabel
   lowercase__			       =  {v: k for k, v in idalabel.items()}

   if "s16" in checkpoint_url:
      lowercase__			       =  384
      lowercase__			       =  1536
      lowercase__			       =  6
   elif "l16" in checkpoint_url:
      lowercase__			       =  1024
      lowercase__			       =  4096
      lowercase__			       =  24
      lowercase__			       =  16
      lowercase__			       =  0.1
   elif "b4" in checkpoint_url:
      lowercase__			       =  4
   elif "l7" in checkpoint_url:
      lowercase__			       =  7
      lowercase__			       =  1024
      lowercase__			       =  4096
      lowercase__			       =  24
      lowercase__			       =  16
      lowercase__			       =  0.1

   lowercase__			       =  ViTMSNModel(__magic_name__	)

   lowercase__			       =  torch.hub.load_state_dict_from_url(__magic_name__  , map_location="cpu"	)["target_encoder"]

   lowercase__			       =  ViTImageProcessor(size=config.image_size	)

   remove_projection_head(__magic_name__	)
   lowercase__			       =  create_rename_keys(__magic_name__  , base_model=__magic_name__	)

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

   model.load_state_dict(__magic_name__	)
   model.eval()

   lowercase__			       =  "http://images.cocodataset.org/val2017/000000039769.jpg"

   lowercase__			       =  Image.open(requests.get(__magic_name__  , stream=__magic_name__	).raw	)
   lowercase__			       =  ViTImageProcessor(
       size=config.image_size  , image_mean=__magic_name__  , image_std=__magic_name__	)
   lowercase__			       =  image_processor(images=__magic_name__  , return_tensors="pt"	)

   # forward pass
   torch.manual_seed(2	)
   lowercase__			       =  model(**__magic_name__	)
   lowercase__			       =  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:
      lowercase__			       =  torch.tensor([[-1.0_915, -1.4_876, -1.1_809]]	)
   elif "b16" in checkpoint_url:
      lowercase__			       =  torch.tensor([[14.2_889, -18.9_045, 11.7_281]]	)
   elif "l16" in checkpoint_url:
      lowercase__			       =  torch.tensor([[41.5_028, -22.8_681, 45.6_475]]	)
   elif "b4" in checkpoint_url:
      lowercase__			       =  torch.tensor([[-4.3_868, 5.2_932, -0.4_137]]	)
   else:
      lowercase__			       =  torch.tensor([[-0.1_792, -0.6_465, 2.4_263]]	)

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

   print(f'''Saving model to {pytorch_dump_folder_path}'''	)
   model.save_pretrained(__magic_name__	)

   print(f'''Saving image processor to {pytorch_dump_folder_path}'''	)
   image_processor.save_pretrained(__magic_name__	)


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)






