

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      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Any     ,					SCREAMING_SNAKE_CASE_		:   List[Any]=False      ):



     '''simple docstring'''





     lowercase__   :       List[str]				=	[]
     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__   :       Tuple				=	[(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      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Dict     ,					SCREAMING_SNAKE_CASE_		:   int     ,					SCREAMING_SNAKE_CASE_		:   Optional[int]=False      ):



     '''simple docstring'''





     for i in range(config.num_hidden_layers      ):
          if base_model:
               lowercase__   :       List[str]				=	''
          else:
               lowercase__   :       Dict				=	'vit.'
          # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
          lowercase__   :       Optional[int]				=	state_dict.pop(f"""module.blocks.{i}.attn.qkv.weight"""      )
          lowercase__   :       List[str]				=	state_dict.pop(f"""module.blocks.{i}.attn.qkv.bias"""      )
          # next, add query, keys and values (in that order) to the state dict
          lowercase__   :       Union[str, Any]				=	in_proj_weight[
              : config.hidden_size, :
          ]
          lowercase__   :       Optional[Any]				=	in_proj_bias[: config.hidden_size]
          lowercase__   :       str				=	in_proj_weight[
              config.hidden_size : config.hidden_size * 2, :
          ]
          lowercase__   :       Union[str, Any]				=	in_proj_bias[
              config.hidden_size : config.hidden_size * 2
          ]
          lowercase__   :       Dict				=	in_proj_weight[
              -config.hidden_size :, :
          ]
          lowercase__   :       Union[str, Any]				=	in_proj_bias[-config.hidden_size :]

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Any      ):



     '''simple docstring'''





     lowercase__   :       Optional[int]				=	['head.weight', 'head.bias']
     for k in ignore_keys:
          state_dict.pop(SCREAMING_SNAKE_CASE_     ,					SCREAMING_SNAKE_CASE_      )

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Tuple      ):



     '''simple docstring'''





     lowercase__   :       str				=	[
         '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(SCREAMING_SNAKE_CASE_     ,					SCREAMING_SNAKE_CASE_      )

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Union[str, Any]     ,					SCREAMING_SNAKE_CASE_		:   str     ,					SCREAMING_SNAKE_CASE_		:   List[Any]      ):



     '''simple docstring'''





     lowercase__   :       str				=	dct.pop(SCREAMING_SNAKE_CASE_      )
     lowercase__   :       str				=	val

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Optional[Any]     ,					SCREAMING_SNAKE_CASE_		:   str      ):



     '''simple docstring'''





     lowercase__   :       Union[str, Any]				=	ViTMSNConfig()
     lowercase__   :       Optional[int]				=	1_000

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

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

     lowercase__   :       int				=	ViTMSNModel(SCREAMING_SNAKE_CASE_      )

     lowercase__   :       Tuple				=	torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_     ,					map_location='cpu'      )['target_encoder']

     lowercase__   :       Optional[Any]				=	ViTImageProcessor(size=config.image_size      )

     remove_projection_head(SCREAMING_SNAKE_CASE_      )
     lowercase__   :       Union[str, Any]				=	create_rename_keys(SCREAMING_SNAKE_CASE_     ,					base_model=SCREAMING_SNAKE_CASE_      )

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

     model.load_state_dict(SCREAMING_SNAKE_CASE_      )
     model.eval()

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

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

     # forward pass
     torch.manual_seed(2      )
     lowercase__   :       str				=	model(**SCREAMING_SNAKE_CASE_      )
     lowercase__   :       Any				=	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__   :       int				=	torch.tensor([[-1.0915, -1.4876, -1.1809]]      )
     elif "b16" in checkpoint_url:
          lowercase__   :       int				=	torch.tensor([[14.2889, -18.9045, 11.7281]]      )
     elif "l16" in checkpoint_url:
          lowercase__   :       Union[str, Any]				=	torch.tensor([[41.5028, -22.8681, 45.6475]]      )
     elif "b4" in checkpoint_url:
          lowercase__   :       Dict				=	torch.tensor([[-4.3868, 5.2932, -0.4137]]      )
     else:
          lowercase__   :       int				=	torch.tensor([[-0.1792, -0.6465, 2.4263]]      )

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

     print(f"""Saving model to {pytorch_dump_folder_path}"""      )
     model.save_pretrained(SCREAMING_SNAKE_CASE_      )

     print(f"""Saving image processor to {pytorch_dump_folder_path}"""      )
     image_processor.save_pretrained(SCREAMING_SNAKE_CASE_      )


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)




from __future__ import annotations

from collections.abc import Callable

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Callable[[int | float], int | float]     ,					SCREAMING_SNAKE_CASE_		:   int | float     ,					SCREAMING_SNAKE_CASE_		:   int | float     ,					SCREAMING_SNAKE_CASE_		:   int = 100     ,					):



     '''simple docstring'''





     lowercase__   :       Tuple				=	x_start
     lowercase__   :       Tuple				=	fnc(SCREAMING_SNAKE_CASE_      )
     lowercase__   :       List[Any]				=	0.0
     for _ in range(SCREAMING_SNAKE_CASE_      ):
          # Approximates small segments of curve as linear and solve
          # for trapezoidal area
          lowercase__   :       Any				=	(x_end - x_start) / steps + xa
          lowercase__   :       Optional[Any]				=	fnc(SCREAMING_SNAKE_CASE_      )
          area += abs(fxa + fxa      ) * (xa - xa) / 2
          # Increment step
          lowercase__   :       Any				=	xa
          lowercase__   :       str				=	fxa
     return area


if __name__ == "__main__":

       def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Tuple      ):



            '''simple docstring'''





            return x**3 + x**2

       print('''f(x) = x^3 + x^2''')
       print('''The area between the curve, x = -5, x = 5 and the x axis is:''')
       snake_case_        =			10
       while i <= 100_000:
              print(F'''with {i} steps: {trapezoidal_area(f, -5, 5, i)}''')
              i *= 10

