




from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, logging


if is_torch_available():
  import torch


_lowerCamelCase     :		Tuple								=       logging.get_logger(__name__)


class     UpperCamelCase_   (			UpperCAmelCase__					):

    '''simple docstring'''






    UpperCAmelCase__							  =				['''pixel_values''']


    def __init__(      self				:     Dict ,			UpperCAmelCase__				:     bool = True ,			UpperCAmelCase__				:     Optional[Dict[str, int]] = None ,			UpperCAmelCase__				:     PILImageResampling = PILImageResampling.BILINEAR ,			UpperCAmelCase__				:     bool = True ,			UpperCAmelCase__				:     Dict[str, int] = None ,			UpperCAmelCase__				:     bool = True ,			UpperCAmelCase__				:     Union[int, float] = 1 / 255 ,			UpperCAmelCase__				:     bool = True ,			UpperCAmelCase__				:     Optional[Union[float, List[float]]] = None ,			UpperCAmelCase__				:     Optional[Union[float, List[float]]] = None ,			**UpperCAmelCase__				:     List[str] ,			)						->None:





           '''simple docstring'''

           super().__init__(**UpperCAmelCase__)
           A__							=   size if size is not None else {'''shortest_edge''': 256}
           A__							=   get_size_dict(UpperCAmelCase__ ,			default_to_square=UpperCAmelCase__)
           A__							=   crop_size if crop_size is not None else {'''height''': 224, '''width''': 224}
           A__							=   get_size_dict(UpperCAmelCase__ ,			param_name='''crop_size''')
           A__							=   do_resize
           A__							=   size
           A__							=   resample
           A__							=   do_center_crop
           A__							=   crop_size
           A__							=   do_rescale
           A__							=   rescale_factor
           A__							=   do_normalize
           A__							=   image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
           A__							=   image_std if image_std is not None else IMAGENET_STANDARD_STD


    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any] ,			UpperCAmelCase__				:     np.ndarray ,			UpperCAmelCase__				:     Dict[str, int] ,			UpperCAmelCase__				:     PILImageResampling = PILImageResampling.BICUBIC ,			UpperCAmelCase__				:     Optional[Union[str, ChannelDimension]] = None ,			**UpperCAmelCase__				:     int ,			)						->np.ndarray:





           '''simple docstring'''

           A__							=   get_size_dict(UpperCAmelCase__ ,			default_to_square=UpperCAmelCase__)
           if "shortest_edge" not in size:
                  raise ValueError(f"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}""")
           A__							=   get_resize_output_image_size(UpperCAmelCase__ ,			size=size['''shortest_edge'''] ,			default_to_square=UpperCAmelCase__)
           return resize(UpperCAmelCase__ ,			size=UpperCAmelCase__ ,			resample=UpperCAmelCase__ ,			data_format=UpperCAmelCase__ ,			**UpperCAmelCase__)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any] ,			UpperCAmelCase__				:     np.ndarray ,			UpperCAmelCase__				:     Dict[str, int] ,			UpperCAmelCase__				:     Optional[Union[str, ChannelDimension]] = None ,			**UpperCAmelCase__				:     Optional[int] ,			)						->np.ndarray:





           '''simple docstring'''

           A__							=   get_size_dict(UpperCAmelCase__)
           if "height" not in size or "width" not in size:
                  raise ValueError(f"""The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}""")
           return center_crop(UpperCAmelCase__ ,			size=(size['''height'''], size['''width''']) ,			data_format=UpperCAmelCase__ ,			**UpperCAmelCase__)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Tuple ,			UpperCAmelCase__				:     np.ndarray ,			UpperCAmelCase__				:     float ,			UpperCAmelCase__				:     Optional[Union[str, ChannelDimension]] = None ,			**UpperCAmelCase__				:     Optional[int])						->np.ndarray:





           '''simple docstring'''

           return rescale(UpperCAmelCase__ ,			scale=UpperCAmelCase__ ,			data_format=UpperCAmelCase__ ,			**UpperCAmelCase__)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Any ,			UpperCAmelCase__				:     np.ndarray ,			UpperCAmelCase__				:     Union[float, List[float]] ,			UpperCAmelCase__				:     Union[float, List[float]] ,			UpperCAmelCase__				:     Optional[Union[str, ChannelDimension]] = None ,			**UpperCAmelCase__				:     Optional[Any] ,			)						->np.ndarray:





           '''simple docstring'''

           return normalize(UpperCAmelCase__ ,			mean=UpperCAmelCase__ ,			std=UpperCAmelCase__ ,			data_format=UpperCAmelCase__ ,			**UpperCAmelCase__)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Union[str, Any] ,			UpperCAmelCase__				:     ImageInput ,			UpperCAmelCase__				:     Optional[bool] = None ,			UpperCAmelCase__				:     Dict[str, int] = None ,			UpperCAmelCase__				:     PILImageResampling = None ,			UpperCAmelCase__				:     bool = None ,			UpperCAmelCase__				:     Dict[str, int] = None ,			UpperCAmelCase__				:     Optional[bool] = None ,			UpperCAmelCase__				:     Optional[float] = None ,			UpperCAmelCase__				:     Optional[bool] = None ,			UpperCAmelCase__				:     Optional[Union[float, List[float]]] = None ,			UpperCAmelCase__				:     Optional[Union[float, List[float]]] = None ,			UpperCAmelCase__				:     Optional[Union[str, TensorType]] = None ,			UpperCAmelCase__				:     Union[str, ChannelDimension] = ChannelDimension.FIRST ,			**UpperCAmelCase__				:     Dict ,			)						->Any:





           '''simple docstring'''

           A__							=   do_resize if do_resize is not None else self.do_resize
           A__							=   size if size is not None else self.size
           A__							=   get_size_dict(UpperCAmelCase__ ,			default_to_square=UpperCAmelCase__)
           A__							=   resample if resample is not None else self.resample
           A__							=   do_center_crop if do_center_crop is not None else self.do_center_crop
           A__							=   crop_size if crop_size is not None else self.crop_size
           A__							=   get_size_dict(UpperCAmelCase__ ,			param_name='''crop_size''')
           A__							=   do_rescale if do_rescale is not None else self.do_rescale
           A__							=   rescale_factor if rescale_factor is not None else self.rescale_factor
           A__							=   do_normalize if do_normalize is not None else self.do_normalize
           A__							=   image_mean if image_mean is not None else self.image_mean
           A__							=   image_std if image_std is not None else self.image_std

           A__							=   make_list_of_images(UpperCAmelCase__)

           if not valid_images(UpperCAmelCase__):
                  raise ValueError(
                      '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
                      '''torch.Tensor, tf.Tensor or jax.ndarray.''')

           if do_resize and size is None:
                  raise ValueError('''Size must be specified if do_resize is True.''')

           if do_center_crop and crop_size is None:
                  raise ValueError('''Crop size must be specified if do_center_crop is True.''')

           if do_rescale and rescale_factor is None:
                  raise ValueError('''Rescale factor must be specified if do_rescale is True.''')

           if do_normalize and (image_mean is None or image_std is None):
                  raise ValueError('''Image mean and std must be specified if do_normalize is True.''')

           # All transformations expect numpy arrays.
           A__							=   [to_numpy_array(UpperCAmelCase__) for image in images]

           if do_resize:
                  A__							=   [self.resize(image=UpperCAmelCase__ ,			size=UpperCAmelCase__ ,			resample=UpperCAmelCase__) for image in images]

           if do_center_crop:
                  A__							=   [self.center_crop(image=UpperCAmelCase__ ,			size=UpperCAmelCase__) for image in images]

           if do_rescale:
                  A__							=   [self.rescale(image=UpperCAmelCase__ ,			scale=UpperCAmelCase__) for image in images]

           if do_normalize:
                  A__							=   [self.normalize(image=UpperCAmelCase__ ,			mean=UpperCAmelCase__ ,			std=UpperCAmelCase__) for image in images]

           A__							=   [to_channel_dimension_format(UpperCAmelCase__ ,			UpperCAmelCase__) for image in images]

           A__							=   {'''pixel_values''': images}
           return BatchFeature(data=UpperCAmelCase__ ,			tensor_type=UpperCAmelCase__)





    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any] ,			UpperCAmelCase__				:     int ,			UpperCAmelCase__				:     List[Tuple] = None)						->Dict:





           '''simple docstring'''

           A__							=   outputs.logits

           # Resize logits and compute semantic segmentation maps
           if target_sizes is not None:
                  if len(UpperCAmelCase__) != len(UpperCAmelCase__):
                         raise ValueError(
                             '''Make sure that you pass in as many target sizes as the batch dimension of the logits''')

                  if is_torch_tensor(UpperCAmelCase__):
                         A__							=   target_sizes.numpy()

                  A__							=   []

                  for idx in range(len(UpperCAmelCase__)):
                         A__							=   torch.nn.functional.interpolate(
                             logits[idx].unsqueeze(dim=0) ,			size=target_sizes[idx] ,			mode='''bilinear''' ,			align_corners=UpperCAmelCase__)
                         A__							=   resized_logits[0].argmax(dim=0)
                         semantic_segmentation.append(UpperCAmelCase__)
           else:
                  A__							=   logits.argmax(dim=1)
                  A__							=   [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

           return semantic_segmentation








import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	str:



       """simple docstring"""

       if "cls_token" in name:
              A__							=   name.replace('''cls_token'''			,    '''vit.embeddings.cls_token'''						)
       if "mask_token" in name:
              A__							=   name.replace('''mask_token'''			,    '''decoder.mask_token'''						)
       if "decoder_pos_embed" in name:
              A__							=   name.replace('''decoder_pos_embed'''			,    '''decoder.decoder_pos_embed'''						)
       if "pos_embed" in name and "decoder" not in name:
              A__							=   name.replace('''pos_embed'''			,    '''vit.embeddings.position_embeddings'''						)
       if "patch_embed.proj" in name:
              A__							=   name.replace('''patch_embed.proj'''			,    '''vit.embeddings.patch_embeddings.projection'''						)
       if "patch_embed.norm" in name:
              A__							=   name.replace('''patch_embed.norm'''			,    '''vit.embeddings.norm'''						)
       if "decoder_blocks" in name:
              A__							=   name.replace('''decoder_blocks'''			,    '''decoder.decoder_layers'''						)
       if "blocks" in name:
              A__							=   name.replace('''blocks'''			,    '''vit.encoder.layer'''						)
       if "attn.proj" in name:
              A__							=   name.replace('''attn.proj'''			,    '''attention.output.dense'''						)
       if "attn" in name:
              A__							=   name.replace('''attn'''			,    '''attention.self'''						)
       if "norm1" in name:
              A__							=   name.replace('''norm1'''			,    '''layernorm_before'''						)
       if "norm2" in name:
              A__							=   name.replace('''norm2'''			,    '''layernorm_after'''						)
       if "mlp.fc1" in name:
              A__							=   name.replace('''mlp.fc1'''			,    '''intermediate.dense'''						)
       if "mlp.fc2" in name:
              A__							=   name.replace('''mlp.fc2'''			,    '''output.dense'''						)
       if "decoder_embed" in name:
              A__							=   name.replace('''decoder_embed'''			,    '''decoder.decoder_embed'''						)
       if "decoder_norm" in name:
              A__							=   name.replace('''decoder_norm'''			,    '''decoder.decoder_norm'''						)
       if "decoder_pred" in name:
              A__							=   name.replace('''decoder_pred'''			,    '''decoder.decoder_pred'''						)
       if "norm.weight" in name and "decoder" not in name:
              A__							=   name.replace('''norm.weight'''			,    '''vit.layernorm.weight'''						)
       if "norm.bias" in name and "decoder" not in name:
              A__							=   name.replace('''norm.bias'''			,    '''vit.layernorm.bias'''						)

       return name







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	str:



       """simple docstring"""

       for key in orig_state_dict.copy().keys():
              A__							=   orig_state_dict.pop(lowercase_						)

              if "qkv" in key:
                     A__							=   key.split('''.'''						)
                     A__							=   int(key_split[1]						)
                     if "decoder_blocks" in key:
                            A__							=   config.decoder_hidden_size
                            A__							=   '''decoder.decoder_layers.'''
                            if "weight" in key:
                                   A__							=   val[:dim, :]
                                   A__							=   val[dim : dim * 2, :]
                                   A__							=   val[-dim:, :]
                            elif "bias" in key:
                                   A__							=   val[:dim]
                                   A__							=   val[dim : dim * 2]
                                   A__							=   val[-dim:]
                     else:
                            A__							=   config.hidden_size
                            A__							=   '''vit.encoder.layer.'''
                            if "weight" in key:
                                   A__							=   val[:dim, :]
                                   A__							=   val[dim : dim * 2, :]
                                   A__							=   val[-dim:, :]
                            elif "bias" in key:
                                   A__							=   val[:dim]
                                   A__							=   val[dim : dim * 2]
                                   A__							=   val[-dim:]

              else:
                     A__							=   val

       return orig_state_dict







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	Tuple:



       """simple docstring"""

       A__							=   ViTMAEConfig()
       if "large" in checkpoint_url:
              A__							=   1_024
              A__							=   4_096
              A__							=   24
              A__							=   16
       elif "huge" in checkpoint_url:
              A__							=   14
              A__							=   1_280
              A__							=   5_120
              A__							=   32
              A__							=   16

       A__							=   ViTMAEForPreTraining(lowercase_						)

       A__							=   torch.hub.load_state_dict_from_url(lowercase_			,    map_location='''cpu'''						)['''model''']

       A__							=   ViTMAEImageProcessor(size=config.image_size						)

       A__							=   convert_state_dict(lowercase_			,    lowercase_						)

       model.load_state_dict(lowercase_						)
       model.eval()

       A__							=   '''https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg'''

       A__							=   Image.open(requests.get(lowercase_			,    stream=lowercase_						).raw						)
       A__							=   ViTMAEImageProcessor(size=config.image_size						)
       A__							=   image_processor(images=lowercase_			,    return_tensors='''pt'''						)

       # forward pass
       torch.manual_seed(2						)
       A__							=   model(**lowercase_						)
       A__							=   outputs.logits

       if "large" in checkpoint_url:
              A__							=   torch.tensor(
                  [[-0.73_09, -0.71_28, -1.01_69], [-1.01_61, -0.90_58, -1.18_78], [-1.04_78, -0.94_11, -1.19_11]]						)
       elif "huge" in checkpoint_url:
              A__							=   torch.tensor(
                  [[-1.15_99, -0.91_99, -1.22_21], [-1.19_52, -0.92_69, -1.23_07], [-1.21_43, -0.93_37, -1.22_62]]						)
       else:
              A__							=   torch.tensor(
                  [[-0.91_92, -0.84_81, -1.12_59], [-1.13_49, -1.00_34, -1.25_99], [-1.17_57, -1.04_29, -1.27_26]]						)

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

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

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


if __name__ == "__main__":
  _lowerCamelCase     :		Any								=       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--checkpoint_url""",
      default="""https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth""",
      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."""
  )

  _lowerCamelCase     :		Optional[int]								=       parser.parse_args()
  convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


