






from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block


@dataclass
class       __lowercase      (	lowercase_  ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				42


class       __lowercase      (	nn.Module  ):







     '''simple docstring'''




     def __init__(				self       :			Union[str, Any] ,	UpperCamelCase_       :			int=3 ,	UpperCamelCase_       :			List[Any]=3 ,	UpperCamelCase_       :			List[Any]=("DownEncoderBlock2D",) ,	UpperCamelCase_       :			int=(64,) ,	UpperCamelCase_       :			List[str]=2 ,	UpperCamelCase_       :			Optional[Any]=32 ,	UpperCamelCase_       :			Dict="silu" ,	UpperCamelCase_       :			Optional[int]=True ,	):


         """simple docstring"""



         super().__init__()
         __A         =	layers_per_block

         __A         =	torch.nn.Convad(
             UpperCamelCase_ ,	block_out_channels[0] ,	kernel_size=3 ,	stride=1 ,	padding=1 ,	)

         __A         =	None
         __A         =	nn.ModuleList([]				)

         # down
         __A         =	block_out_channels[0]
         for i, down_block_type in enumerate(UpperCamelCase_				):
             __A         =	output_channel
             __A         =	block_out_channels[i]
             __A         =	i == len(UpperCamelCase_				) - 1

             __A         =	get_down_block(
                 UpperCamelCase_ ,	num_layers=self.layers_per_block ,	in_channels=UpperCamelCase_ ,	out_channels=UpperCamelCase_ ,	add_downsample=not is_final_block ,	resnet_eps=1e-6 ,	downsample_padding=0 ,	resnet_act_fn=UpperCamelCase_ ,	resnet_groups=UpperCamelCase_ ,	attention_head_dim=UpperCamelCase_ ,	temb_channels=UpperCamelCase_ ,	)
             self.down_blocks.append(UpperCamelCase_				)

         # mid
         __A         =	UNetMidBlockaD(
             in_channels=block_out_channels[-1] ,	resnet_eps=1e-6 ,	resnet_act_fn=UpperCamelCase_ ,	output_scale_factor=1 ,	resnet_time_scale_shift="""default""" ,	attention_head_dim=block_out_channels[-1] ,	resnet_groups=UpperCamelCase_ ,	temb_channels=UpperCamelCase_ ,	)

         # out
         __A         =	nn.GroupNorm(num_channels=block_out_channels[-1] ,	num_groups=UpperCamelCase_ ,	eps=1e-6				)
         __A         =	nn.SiLU()

         __A         =	2 * out_channels if double_z else out_channels
         __A         =	nn.Convad(block_out_channels[-1] ,	UpperCamelCase_ ,	3 ,	padding=1				)

         __A         =	False





     def 		lowerCAmelCase_     (				self       :			int ,	UpperCamelCase_       :			List[str]				):


         """simple docstring"""



         __A         =	x
         __A         =	self.conv_in(UpperCamelCase_				)

         if self.training and self.gradient_checkpointing:

             def create_custom_forward(UpperCamelCase_       :			Optional[Any]				):
                 def custom_forward(*UpperCamelCase_       :			str				):
                     return module(*UpperCamelCase_				)

                 return custom_forward

             # down
             if is_torch_version(""">=""" ,	"""1.11.0"""				):
                 for down_block in self.down_blocks:
                     __A         =	torch.utils.checkpoint.checkpoint(
                         create_custom_forward(UpperCamelCase_				) ,	UpperCamelCase_ ,	use_reentrant=UpperCamelCase_				)
                 # middle
                 __A         =	torch.utils.checkpoint.checkpoint(
                     create_custom_forward(self.mid_block				) ,	UpperCamelCase_ ,	use_reentrant=UpperCamelCase_				)
             else:
                 for down_block in self.down_blocks:
                     __A         =	torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCamelCase_				) ,	UpperCamelCase_				)
                 # middle
                 __A         =	torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block				) ,	UpperCamelCase_				)

         else:
             # down
             for down_block in self.down_blocks:
                 __A         =	down_block(UpperCamelCase_				)

             # middle
             __A         =	self.mid_block(UpperCamelCase_				)

         # post-process
         __A         =	self.conv_norm_out(UpperCamelCase_				)
         __A         =	self.conv_act(UpperCamelCase_				)
         __A         =	self.conv_out(UpperCamelCase_				)

         return sample


class       __lowercase      (	nn.Module  ):







     '''simple docstring'''




     def __init__(				self       :			Union[str, Any] ,	UpperCamelCase_       :			List[str]=3 ,	UpperCamelCase_       :			Union[str, Any]=3 ,	UpperCamelCase_       :			Union[str, Any]=("UpDecoderBlock2D",) ,	UpperCamelCase_       :			Tuple=(64,) ,	UpperCamelCase_       :			Union[str, Any]=2 ,	UpperCamelCase_       :			Tuple=32 ,	UpperCamelCase_       :			int="silu" ,	UpperCamelCase_       :			Tuple="group" ,	):


         """simple docstring"""



         super().__init__()
         __A         =	layers_per_block

         __A         =	nn.Convad(
             UpperCamelCase_ ,	block_out_channels[-1] ,	kernel_size=3 ,	stride=1 ,	padding=1 ,	)

         __A         =	None
         __A         =	nn.ModuleList([]				)

         __A         =	in_channels if norm_type == """spatial""" else None

         # mid
         __A         =	UNetMidBlockaD(
             in_channels=block_out_channels[-1] ,	resnet_eps=1e-6 ,	resnet_act_fn=UpperCamelCase_ ,	output_scale_factor=1 ,	resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type ,	attention_head_dim=block_out_channels[-1] ,	resnet_groups=UpperCamelCase_ ,	temb_channels=UpperCamelCase_ ,	)

         # up
         __A         =	list(reversed(UpperCamelCase_				)				)
         __A         =	reversed_block_out_channels[0]
         for i, up_block_type in enumerate(UpperCamelCase_				):
             __A         =	output_channel
             __A         =	reversed_block_out_channels[i]

             __A         =	i == len(UpperCamelCase_				) - 1

             __A         =	get_up_block(
                 UpperCamelCase_ ,	num_layers=self.layers_per_block + 1 ,	in_channels=UpperCamelCase_ ,	out_channels=UpperCamelCase_ ,	prev_output_channel=UpperCamelCase_ ,	add_upsample=not is_final_block ,	resnet_eps=1e-6 ,	resnet_act_fn=UpperCamelCase_ ,	resnet_groups=UpperCamelCase_ ,	attention_head_dim=UpperCamelCase_ ,	temb_channels=UpperCamelCase_ ,	resnet_time_scale_shift=UpperCamelCase_ ,	)
             self.up_blocks.append(UpperCamelCase_				)
             __A         =	output_channel

         # out
         if norm_type == "spatial":
             __A         =	SpatialNorm(block_out_channels[0] ,	UpperCamelCase_				)
         else:
             __A         =	nn.GroupNorm(num_channels=block_out_channels[0] ,	num_groups=UpperCamelCase_ ,	eps=1e-6				)
         __A         =	nn.SiLU()
         __A         =	nn.Convad(block_out_channels[0] ,	UpperCamelCase_ ,	3 ,	padding=1				)

         __A         =	False





     def 		lowerCAmelCase_     (				self       :			str ,	UpperCamelCase_       :			int ,	UpperCamelCase_       :			int=None				):


         """simple docstring"""



         __A         =	z
         __A         =	self.conv_in(UpperCamelCase_				)

         __A         =	next(iter(self.up_blocks.parameters()				)				).dtype
         if self.training and self.gradient_checkpointing:

             def create_custom_forward(UpperCamelCase_       :			Optional[Any]				):
                 def custom_forward(*UpperCamelCase_       :			int				):
                     return module(*UpperCamelCase_				)

                 return custom_forward

             if is_torch_version(""">=""" ,	"""1.11.0"""				):
                 # middle
                 __A         =	torch.utils.checkpoint.checkpoint(
                     create_custom_forward(self.mid_block				) ,	UpperCamelCase_ ,	UpperCamelCase_ ,	use_reentrant=UpperCamelCase_				)
                 __A         =	sample.to(UpperCamelCase_				)

                 # up
                 for up_block in self.up_blocks:
                     __A         =	torch.utils.checkpoint.checkpoint(
                         create_custom_forward(UpperCamelCase_				) ,	UpperCamelCase_ ,	UpperCamelCase_ ,	use_reentrant=UpperCamelCase_				)
             else:
                 # middle
                 __A         =	torch.utils.checkpoint.checkpoint(
                     create_custom_forward(self.mid_block				) ,	UpperCamelCase_ ,	UpperCamelCase_				)
                 __A         =	sample.to(UpperCamelCase_				)

                 # up
                 for up_block in self.up_blocks:
                     __A         =	torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCamelCase_				) ,	UpperCamelCase_ ,	UpperCamelCase_				)
         else:
             # middle
             __A         =	self.mid_block(UpperCamelCase_ ,	UpperCamelCase_				)
             __A         =	sample.to(UpperCamelCase_				)

             # up
             for up_block in self.up_blocks:
                 __A         =	up_block(UpperCamelCase_ ,	UpperCamelCase_				)

        # post-process
         if latent_embeds is None:
             __A         =	self.conv_norm_out(UpperCamelCase_				)
         else:
             __A         =	self.conv_norm_out(UpperCamelCase_ ,	UpperCamelCase_				)
         __A         =	self.conv_act(UpperCamelCase_				)
         __A         =	self.conv_out(UpperCamelCase_				)

         return sample


class       __lowercase      (	nn.Module  ):







     '''simple docstring'''




     def __init__(				self       :			Optional[Any] ,	UpperCamelCase_       :			int ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			str ,	UpperCamelCase_       :			Tuple=None ,	UpperCamelCase_       :			Union[str, Any]="random" ,	UpperCamelCase_       :			Optional[Any]=False ,	UpperCamelCase_       :			List[Any]=True				):


         """simple docstring"""



         super().__init__()
         __A         =	n_e
         __A         =	vq_embed_dim
         __A         =	beta
         __A         =	legacy

         __A         =	nn.Embedding(self.n_e ,	self.vq_embed_dim				)
         self.embedding.weight.data.uniform_(-1.0 / self.n_e ,	1.0 / self.n_e				)

         __A         =	remap
         if self.remap is not None:
             self.register_buffer("""used""" ,	torch.tensor(np.load(self.remap				)				)				)
             __A         =	self.used.shape[0]
             __A         =	unknown_index  # "random" or "extra" or integer
             if self.unknown_index == "extra":
                 __A         =	self.re_embed
                 __A         =	self.re_embed + 1
             print(
                 F"Remapping {self.n_e} indices to {self.re_embed} indices. "
                 F"Using {self.unknown_index} for unknown indices."				)
         else:
             __A         =	n_e

         __A         =	sane_index_shape




     def 		lowerCAmelCase_     (				self       :			List[Any] ,	UpperCamelCase_       :			Dict				):


         """simple docstring"""



         __A         =	inds.shape
         assert len(UpperCamelCase_				) > 1
         __A         =	inds.reshape(ishape[0] ,	-1				)
         __A         =	self.used.to(UpperCamelCase_				)
         __A         =	(inds[:, :, None] == used[None, None, ...]).long()
         __A         =	match.argmax(-1				)
         __A         =	match.sum(2				) < 1
         if self.unknown_index == "random":
             __A         =	torch.randint(0 ,	self.re_embed ,	size=new[unknown].shape				).to(device=new.device				)
         else:
             __A         =	self.unknown_index
         return new.reshape(UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			Tuple ,	UpperCamelCase_       :			Tuple				):


         """simple docstring"""



         __A         =	inds.shape
         assert len(UpperCamelCase_				) > 1
         __A         =	inds.reshape(ishape[0] ,	-1				)
         __A         =	self.used.to(UpperCamelCase_				)
         if self.re_embed > self.used.shape[0]:  # extra token
             __A         =	0  # simply set to zero
         __A         =	torch.gather(used[None, :][inds.shape[0] * [0], :] ,	1 ,	UpperCamelCase_				)
         return back.reshape(UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			Optional[Any] ,	UpperCamelCase_       :			Dict				):


         """simple docstring"""



         __A         =	z.permute(0 ,	2 ,	3 ,	1				).contiguous()
         __A         =	z.view(-1 ,	self.vq_embed_dim				)

         # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
         __A         =	torch.argmin(torch.cdist(UpperCamelCase_ ,	self.embedding.weight				) ,	dim=1				)

         __A         =	self.embedding(UpperCamelCase_				).view(z.shape				)
         __A         =	None
         __A         =	None

         # compute loss for embedding
         if not self.legacy:
             __A         =	self.beta * torch.mean((z_q.detach() - z) ** 2				) + torch.mean((z_q - z.detach()) ** 2				)
         else:
             __A         =	torch.mean((z_q.detach() - z) ** 2				) + self.beta * torch.mean((z_q - z.detach()) ** 2				)

         # preserve gradients
         __A         =	z + (z_q - z).detach()

         # reshape back to match original input shape
         __A         =	z_q.permute(0 ,	3 ,	1 ,	2				).contiguous()

         if self.remap is not None:
             __A         =	min_encoding_indices.reshape(z.shape[0] ,	-1				)  # add batch axis
             __A         =	self.remap_to_used(UpperCamelCase_				)
             __A         =	min_encoding_indices.reshape(-1 ,	1				)  # flatten

         if self.sane_index_shape:
             __A         =	min_encoding_indices.reshape(z_q.shape[0] ,	z_q.shape[2] ,	z_q.shape[3]				)

         return z_q, loss, (perplexity, min_encodings, min_encoding_indices)





     def 		lowerCAmelCase_     (				self       :			Optional[Any] ,	UpperCamelCase_       :			List[str] ,	UpperCamelCase_       :			Union[str, Any]				):


         """simple docstring"""



         if self.remap is not None:
             __A         =	indices.reshape(shape[0] ,	-1				)  # add batch axis
             __A         =	self.unmap_to_all(UpperCamelCase_				)
             __A         =	indices.reshape(-1				)  # flatten again

         # get quantized latent vectors
         __A         =	self.embedding(UpperCamelCase_				)

         if shape is not None:
             __A         =	z_q.view(UpperCamelCase_				)
             # reshape back to match original input shape
             __A         =	z_q.permute(0 ,	3 ,	1 ,	2				).contiguous()

         return z_q


class       __lowercase      (	lowercase_  ):







     '''simple docstring'''




     def __init__(				self       :			Tuple ,	UpperCamelCase_       :			Union[str, Any] ,	UpperCamelCase_       :			List[str]=False				):


         """simple docstring"""



         __A         =	parameters
         __A    ,	__A         =	torch.chunk(UpperCamelCase_ ,	2 ,	dim=1				)
         __A         =	torch.clamp(self.logvar ,	-30.0 ,	20.0				)
         __A         =	deterministic
         __A         =	torch.exp(0.5 * self.logvar				)
         __A         =	torch.exp(self.logvar				)
         if self.deterministic:
             __A         =	__A         =	torch.zeros_like(
                 self.mean ,	device=self.parameters.device ,	dtype=self.parameters.dtype				)




     def 		lowerCAmelCase_     (				self       :			List[Any] ,	UpperCamelCase_       :			Optional[torch.Generator] = None				):


         """simple docstring"""



         __A         =	randn_tensor(
             self.mean.shape ,	generator=UpperCamelCase_ ,	device=self.parameters.device ,	dtype=self.parameters.dtype				)
         __A         =	self.mean + self.std * sample
         return x




     def 		lowerCAmelCase_     (				self       :			Optional[int] ,	UpperCamelCase_       :			Tuple=None				):


         """simple docstring"""



         if self.deterministic:
             return torch.Tensor([0.0]				)
         else:
             if other is None:
                 return 0.5 * torch.sum(torch.pow(self.mean ,	2				) + self.var - 1.0 - self.logvar ,	dim=[1, 2, 3]				)
             else:
                 return 0.5 * torch.sum(
                     torch.pow(self.mean - other.mean ,	2				) / other.var
                     + self.var / other.var
                     - 1.0
                     - self.logvar
                     + other.logvar ,	dim=[1, 2, 3] ,	)




     def 		lowerCAmelCase_     (				self       :			Any ,	UpperCamelCase_       :			Dict ,	UpperCamelCase_       :			Optional[Any]=[1, 2, 3]				):


         """simple docstring"""



         if self.deterministic:
             return torch.Tensor([0.0]				)
         __A         =	np.log(2.0 * np.pi				)
         return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean ,	2				) / self.var ,	dim=UpperCamelCase_				)





     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         return self.mean












import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			str    )						->      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   :			Any			,				__lowercase   :			Dict    )						->      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   :			Tuple			,				__lowercase   :			str    )						->      Optional[Any]:





    """simple docstring"""

    __A         =	ViTMAEConfig()
    if "large" in checkpoint_url:
        __A         =	1_0_2_4
        __A         =	4_0_9_6
        __A         =	2_4
        __A         =	1_6
    elif "huge" in checkpoint_url:
        __A         =	1_4
        __A         =	1_2_8_0
        __A         =	5_1_2_0
        __A         =	3_2
        __A         =	1_6

    __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.7_309, -0.7_128, -1.0_169], [-1.0_161, -0.9_058, -1.1_878], [-1.0_478, -0.9_411, -1.1_911]]    )
    elif "huge" in checkpoint_url:
        __A         =	torch.tensor(
            [[-1.1_599, -0.9_199, -1.2_221], [-1.1_952, -0.9_269, -1.2_307], [-1.2_143, -0.9_337, -1.2_262]]    )
    else:
        __A         =	torch.tensor(
            [[-0.9_192, -0.8_481, -1.1_259], [-1.1_349, -1.0_034, -1.2_599], [-1.1_757, -1.0_429, -1.2_726]]    )

    # 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__":
  __a      :       List[str]				  =   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."
  )

  __a      :       List[str]				  =   parser.parse_args()
  convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




