






"""simple docstring"""
import argparse
import json
from pathlib import Path

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

from transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
snake_case : List[str]     		=		logging.get_logger(__name__)


def 							A						(				__snake_case: Tuple ,     __snake_case: Any=False  )    ->		str:

       """simple docstring"""





       __magic_name__      				=    []
       for i in range(config.num_hidden_layers  ):
              # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
              rename_keys.append((F"""blocks.{i}.norm1.weight""", F"""vit.encoder.layer.{i}.layernorm_before.weight""")  )
              rename_keys.append((F"""blocks.{i}.norm1.bias""", F"""vit.encoder.layer.{i}.layernorm_before.bias""")  )
              rename_keys.append((F"""blocks.{i}.attn.proj.weight""", F"""vit.encoder.layer.{i}.attention.output.dense.weight""")  )
              rename_keys.append((F"""blocks.{i}.attn.proj.bias""", F"""vit.encoder.layer.{i}.attention.output.dense.bias""")  )
              rename_keys.append((F"""blocks.{i}.norm2.weight""", F"""vit.encoder.layer.{i}.layernorm_after.weight""")  )
              rename_keys.append((F"""blocks.{i}.norm2.bias""", F"""vit.encoder.layer.{i}.layernorm_after.bias""")  )
              rename_keys.append((F"""blocks.{i}.mlp.fc1.weight""", F"""vit.encoder.layer.{i}.intermediate.dense.weight""")  )
              rename_keys.append((F"""blocks.{i}.mlp.fc1.bias""", F"""vit.encoder.layer.{i}.intermediate.dense.bias""")  )
              rename_keys.append((F"""blocks.{i}.mlp.fc2.weight""", F"""vit.encoder.layer.{i}.output.dense.weight""")  )
              rename_keys.append((F"""blocks.{i}.mlp.fc2.bias""", F"""vit.encoder.layer.{i}.output.dense.bias""")  )

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

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

              # if just the base model, we should remove "vit" from all keys that start with "vit"
              __magic_name__      				=    [(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						(				__snake_case: List[Any] ,     __snake_case: str ,     __snake_case: str=False  )    ->		Tuple:

       """simple docstring"""





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


def 							A						(				__snake_case: List[str]  )    ->		Optional[Any]:

       """simple docstring"""





       __magic_name__      				=    ['head.weight', 'head.bias']
       for k in ignore_keys:
              state_dict.pop(lowerCAmelCase_ ,     lowerCAmelCase_  )


def 							A						(				__snake_case: List[Any] ,     __snake_case: str ,     __snake_case: Union[str, Any]  )    ->		Tuple:

       """simple docstring"""





       __magic_name__      				=    dct.pop(lowerCAmelCase_  )
       __magic_name__      				=    val


def 							A						(				)    ->		Union[str, Any]:

       """simple docstring"""





       __magic_name__      				=    'http://images.cocodataset.org/val2017/000000039769.jpg'
       __magic_name__      				=    Image.open(requests.get(lowerCAmelCase_ ,     stream=lowerCAmelCase_  ).raw  )
       return im


@torch.no_grad()
def 							A						(				__snake_case: Tuple ,     __snake_case: List[str]  )    ->		Union[str, Any]:

       """simple docstring"""





       __magic_name__      				=    ViTConfig()
       __magic_name__      				=    False
       # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
       if vit_name[-5:] == "in21k":
              __magic_name__      				=    True
              __magic_name__      				=    int(vit_name[-1_2:-1_0]  )
              __magic_name__      				=    int(vit_name[-9:-6]  )
       else:
              __magic_name__      				=    1_0_0_0
              __magic_name__      				=    'huggingface/label-files'
              __magic_name__      				=    'imagenet-1k-id2label.json'
              __magic_name__      				=    json.load(open(hf_hub_download(lowerCAmelCase_ ,     lowerCAmelCase_ ,     repo_type='dataset'  ) ,     'r'  )  )
              __magic_name__      				=    {int(lowerCAmelCase_  ): v for k, v in idalabel.items()}
              __magic_name__      				=    idalabel
              __magic_name__      				=    {v: k for k, v in idalabel.items()}
              __magic_name__      				=    int(vit_name[-6:-4]  )
              __magic_name__      				=    int(vit_name[-3:]  )
       # size of the architecture
       if "deit" in vit_name:
              if vit_name[9:].startswith('tiny'  ):
                     __magic_name__      				=    1_9_2
                     __magic_name__      				=    7_6_8
                     __magic_name__      				=    1_2
                     __magic_name__      				=    3
              elif vit_name[9:].startswith('small'  ):
                     __magic_name__      				=    3_8_4
                     __magic_name__      				=    1_5_3_6
                     __magic_name__      				=    1_2
                     __magic_name__      				=    6
              else:
                     pass
       else:
              if vit_name[4:].startswith('small'  ):
                     __magic_name__      				=    7_6_8
                     __magic_name__      				=    2_3_0_4
                     __magic_name__      				=    8
                     __magic_name__      				=    8
              elif vit_name[4:].startswith('base'  ):
                     pass
              elif vit_name[4:].startswith('large'  ):
                     __magic_name__      				=    1_0_2_4
                     __magic_name__      				=    4_0_9_6
                     __magic_name__      				=    2_4
                     __magic_name__      				=    1_6
              elif vit_name[4:].startswith('huge'  ):
                     __magic_name__      				=    1_2_8_0
                     __magic_name__      				=    5_1_2_0
                     __magic_name__      				=    3_2
                     __magic_name__      				=    1_6

    # load original model from timm
       __magic_name__      				=    timm.create_model(lowerCAmelCase_ ,     pretrained=lowerCAmelCase_  )
       timm_model.eval()

       # load state_dict of original model, remove and rename some keys
       __magic_name__      				=    timm_model.state_dict()
       if base_model:
              remove_classification_head_(lowerCAmelCase_  )
       __magic_name__      				=    create_rename_keys(lowerCAmelCase_ ,     lowerCAmelCase_  )
       for src, dest in rename_keys:
              rename_key(lowerCAmelCase_ ,     lowerCAmelCase_ ,     lowerCAmelCase_  )
       read_in_q_k_v(lowerCAmelCase_ ,     lowerCAmelCase_ ,     lowerCAmelCase_  )

       # load HuggingFace model
       if vit_name[-5:] == "in21k":
              __magic_name__      				=    ViTModel(lowerCAmelCase_  ).eval()
       else:
              __magic_name__      				=    ViTForImageClassification(lowerCAmelCase_  ).eval()
       model.load_state_dict(lowerCAmelCase_  )

       # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
       if "deit" in vit_name:
              __magic_name__      				=    DeiTImageProcessor(size=config.image_size  )
       else:
              __magic_name__      				=    ViTImageProcessor(size=config.image_size  )
       __magic_name__      				=    image_processor(images=prepare_img() ,     return_tensors='pt'  )
       __magic_name__      				=    encoding['pixel_values']
       __magic_name__      				=    model(lowerCAmelCase_  )

       if base_model:
              __magic_name__      				=    timm_model.forward_features(lowerCAmelCase_  )
              assert timm_pooled_output.shape == outputs.pooler_output.shape
              assert torch.allclose(lowerCAmelCase_ ,     outputs.pooler_output ,     atol=1E-3  )
       else:
              __magic_name__      				=    timm_model(lowerCAmelCase_  )
              assert timm_logits.shape == outputs.logits.shape
              assert torch.allclose(lowerCAmelCase_ ,     outputs.logits ,     atol=1E-3  )

       Path(lowerCAmelCase_  ).mkdir(exist_ok=lowerCAmelCase_  )
       print(F"""Saving model {vit_name} to {pytorch_dump_folder_path}"""  )
       model.save_pretrained(lowerCAmelCase_  )
       print(F"""Saving image processor to {pytorch_dump_folder_path}"""  )
       image_processor.save_pretrained(lowerCAmelCase_  )


if __name__ == "__main__":
 snake_case : Any     		=		argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     """--vit_name""",
     default="""vit_base_patch16_224""",
     type=str,
     help="""Name of the ViT timm model 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 : Dict     		=		parser.parse_args()
 convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)



import unittest

from diffusers.models.unet_ad_blocks import *  # noqa F403
from diffusers.utils import torch_device

from .test_unet_blocks_common import UNetBlockTesterMixin


class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							DownBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'down'



  def        lowerCAmelCase__	(  self							:							int	) ->			Tuple:




     """simple docstring"""
     snake_case_            =       [-0.0_232, -0.9_869, 0.8_054, -0.0_637, -0.1_688, -1.4_264, 0.4_470, -1.3_394, 0.0_904]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							ResnetDownsampleBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'down'



  def        lowerCAmelCase__	(  self							:							Any	) ->			Union[str, Any]:




     """simple docstring"""
     snake_case_            =       [0.0_710, 0.2_410, -0.7_320, -1.0_757, -1.1_343, 0.3_540, -0.0_133, -0.2_576, 0.0_948]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							AttnDownBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'down'



  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			Tuple:




     """simple docstring"""
     snake_case_            =       [0.0_636, 0.8_964, -0.6_234, -1.0_131, 0.0_844, 0.4_935, 0.3_437, 0.0_911, -0.2_957]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							CrossAttnDownBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'down'



  def        lowerCAmelCase__	(  self							:							str	) ->			Tuple:




     """simple docstring"""
     snake_case_     ,     snake_case_            =       super().prepare_init_args_and_inputs_for_common()
     snake_case_            =       3_2
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							List[Any]	) ->			str:




     """simple docstring"""
     snake_case_            =       [0.2_238, -0.7_396, -0.2_255, -0.3_829, 0.1_925, 1.1_665, 0.0_603, -0.7_295, 0.1_983]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							SimpleCrossAttnDownBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'down'



  @property
  def        lowerCAmelCase__	(  self							:							Dict	) ->			Any:




     """simple docstring"""
     return super().get_dummy_input(include_encoder_hidden_states=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							List[Any]	) ->			Optional[int]:




     """simple docstring"""
     snake_case_     ,     snake_case_            =       super().prepare_init_args_and_inputs_for_common()
     snake_case_            =       3_2
     return init_dict, inputs_dict



  @unittest.skipIf(torch_device == "mps"     ,					"MPS result is not consistent"	)
  def        lowerCAmelCase__	(  self							:							int	) ->			Optional[Any]:




     """simple docstring"""
     snake_case_            =       [0.7_921, -0.0_992, -0.1_962, -0.7_695, -0.4_242, 0.7_804, 0.4_737, 0.2_765, 0.3_338]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							SkipDownBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'down'



  @property
  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			Tuple:




     """simple docstring"""
     return super().get_dummy_input(include_skip_sample=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							Optional[int]	) ->			int:




     """simple docstring"""
     snake_case_            =       [-0.0_845, -0.2_087, -0.2_465, 0.0_971, 0.1_900, -0.0_484, 0.2_664, 0.4_179, 0.5_069]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							AttnSkipDownBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'down'



  @property
  def        lowerCAmelCase__	(  self							:							Any	) ->			Dict:




     """simple docstring"""
     return super().get_dummy_input(include_skip_sample=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			Any:




     """simple docstring"""
     snake_case_            =       [0.5_539, 0.1_609, 0.4_924, 0.0_537, -0.1_995, 0.4_050, 0.0_979, -0.2_721, -0.0_642]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							DownEncoderBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'down'



  @property
  def        lowerCAmelCase__	(  self							:							int	) ->			Optional[Any]:




     """simple docstring"""
     return super().get_dummy_input(include_temb=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			Dict:




     """simple docstring"""
     snake_case_            =       {
         "in_channels": 3_2,
         "out_channels": 3_2,
     }
     snake_case_            =       self.dummy_input
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							Optional[int]	) ->			List[Any]:




     """simple docstring"""
     snake_case_            =       [1.1_102, 0.5_302, 0.4_872, -0.0_023, -0.8_042, 0.0_483, -0.3_489, -0.5_632, 0.7_626]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							AttnDownEncoderBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'down'



  @property
  def        lowerCAmelCase__	(  self							:							str	) ->			Optional[Any]:




     """simple docstring"""
     return super().get_dummy_input(include_temb=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							int	) ->			Dict:




     """simple docstring"""
     snake_case_            =       {
         "in_channels": 3_2,
         "out_channels": 3_2,
     }
     snake_case_            =       self.dummy_input
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							List[str]	) ->			Any:




     """simple docstring"""
     snake_case_            =       [0.8_966, -0.1_486, 0.8_568, 0.8_141, -0.9_046, -0.1_342, -0.0_972, -0.7_417, 0.1_538]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							UNetMidBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'mid'



  def        lowerCAmelCase__	(  self							:							Any	) ->			List[Any]:




     """simple docstring"""
     snake_case_            =       {
         "in_channels": 3_2,
         "temb_channels": 1_2_8,
     }
     snake_case_            =       self.dummy_input
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			int:




     """simple docstring"""
     snake_case_            =       [-0.1_062, 1.7_248, 0.3_494, 1.4_569, -0.0_910, -1.2_421, -0.9_984, 0.6_736, 1.0_028]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							UNetMidBlockaDCrossAttn  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'mid'



  def        lowerCAmelCase__	(  self							:							int	) ->			List[Any]:




     """simple docstring"""
     snake_case_     ,     snake_case_            =       super().prepare_init_args_and_inputs_for_common()
     snake_case_            =       3_2
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							Tuple	) ->			Any:




     """simple docstring"""
     snake_case_            =       [0.0_187, 2.4_220, 0.4_484, 1.1_203, -0.6_121, -1.5_122, -0.8_270, 0.7_851, 1.8_335]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							UNetMidBlockaDSimpleCrossAttn  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'mid'



  @property
  def        lowerCAmelCase__	(  self							:							Tuple	) ->			Any:




     """simple docstring"""
     return super().get_dummy_input(include_encoder_hidden_states=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							Optional[Any]	) ->			Union[str, Any]:




     """simple docstring"""
     snake_case_     ,     snake_case_            =       super().prepare_init_args_and_inputs_for_common()
     snake_case_            =       3_2
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			Optional[Any]:




     """simple docstring"""
     snake_case_            =       [0.7_143, 1.9_974, 0.5_448, 1.3_977, 0.1_282, -1.1_237, -1.4_238, 0.5_530, 0.8_880]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							UpBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'up'



  @property
  def        lowerCAmelCase__	(  self							:							Dict	) ->			List[Any]:




     """simple docstring"""
     return super().get_dummy_input(include_res_hidden_states_tuple=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							str	) ->			Any:




     """simple docstring"""
     snake_case_            =       [-0.2_041, -0.4_165, -0.3_022, 0.0_041, -0.6_628, -0.7_053, 0.1_928, -0.0_325, 0.0_523]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							ResnetUpsampleBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'up'



  @property
  def        lowerCAmelCase__	(  self							:							str	) ->			int:




     """simple docstring"""
     return super().get_dummy_input(include_res_hidden_states_tuple=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							Any	) ->			Optional[Any]:




     """simple docstring"""
     snake_case_            =       [0.2_287, 0.3_549, -0.1_346, 0.4_797, -0.1_715, -0.9_649, 0.7_305, -0.5_864, -0.6_244]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							CrossAttnUpBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'up'



  @property
  def        lowerCAmelCase__	(  self							:							List[str]	) ->			List[str]:




     """simple docstring"""
     return super().get_dummy_input(include_res_hidden_states_tuple=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							List[str]	) ->			Optional[Any]:




     """simple docstring"""
     snake_case_     ,     snake_case_            =       super().prepare_init_args_and_inputs_for_common()
     snake_case_            =       3_2
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			List[Any]:




     """simple docstring"""
     snake_case_            =       [-0.1_403, -0.3_515, -0.0_420, -0.1_425, 0.3_167, 0.5_094, -0.2_181, 0.5_931, 0.5_582]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							SimpleCrossAttnUpBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'up'



  @property
  def        lowerCAmelCase__	(  self							:							int	) ->			List[Any]:




     """simple docstring"""
     return super().get_dummy_input(include_res_hidden_states_tuple=_lowerCAmelCase     ,					include_encoder_hidden_states=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							Tuple	) ->			Any:




     """simple docstring"""
     snake_case_     ,     snake_case_            =       super().prepare_init_args_and_inputs_for_common()
     snake_case_            =       3_2
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			Optional[int]:




     """simple docstring"""
     snake_case_            =       [0.2_645, 0.1_480, 0.0_909, 0.8_044, -0.9_758, -0.9_083, 0.0_994, -1.1_453, -0.7_402]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							AttnUpBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'up'



  @property
  def        lowerCAmelCase__	(  self							:							List[Any]	) ->			Tuple:




     """simple docstring"""
     return super().get_dummy_input(include_res_hidden_states_tuple=_lowerCAmelCase	)



  @unittest.skipIf(torch_device == "mps"     ,					"MPS result is not consistent"	)
  def        lowerCAmelCase__	(  self							:							List[Any]	) ->			Optional[int]:




     """simple docstring"""
     snake_case_            =       [0.0_979, 0.1_326, 0.0_021, 0.0_659, 0.2_249, 0.0_059, 0.1_132, 0.5_952, 0.1_033]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							SkipUpBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'up'



  @property
  def        lowerCAmelCase__	(  self							:							List[str]	) ->			str:




     """simple docstring"""
     return super().get_dummy_input(include_res_hidden_states_tuple=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							Tuple	) ->			str:




     """simple docstring"""
     snake_case_            =       [-0.0_893, -0.1_234, -0.1_506, -0.0_332, 0.0_123, -0.0_211, 0.0_566, 0.0_143, 0.0_362]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							AttnSkipUpBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'up'



  @property
  def        lowerCAmelCase__	(  self							:							int	) ->			List[str]:




     """simple docstring"""
     return super().get_dummy_input(include_res_hidden_states_tuple=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							Optional[int]	) ->			Dict:




     """simple docstring"""
     snake_case_            =       [0.0_361, 0.0_617, 0.2_787, -0.0_350, 0.0_342, 0.3_421, -0.0_843, 0.0_913, 0.3_015]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							UpDecoderBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'up'



  @property
  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			Optional[int]:




     """simple docstring"""
     return super().get_dummy_input(include_temb=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							Any	) ->			Any:




     """simple docstring"""
     snake_case_            =       {"in_channels": 3_2, "out_channels": 3_2}

     snake_case_            =       self.dummy_input
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							str	) ->			List[str]:




     """simple docstring"""
     snake_case_            =       [0.4_404, 0.1_998, -0.9_886, -0.3_320, -0.3_128, -0.7_034, -0.6_955, -0.2_338, -0.3_137]
     super().test_output(_lowerCAmelCase	)






class      __lowerCAmelCase  (			a			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							AttnUpDecoderBlockaD  # noqa F405
  _SCREAMING_SNAKE_CASE		     =							'up'



  @property
  def        lowerCAmelCase__	(  self							:							Tuple	) ->			int:




     """simple docstring"""
     return super().get_dummy_input(include_temb=_lowerCAmelCase	)



  def        lowerCAmelCase__	(  self							:							int	) ->			Dict:




     """simple docstring"""
     snake_case_            =       {"in_channels": 3_2, "out_channels": 3_2}

     snake_case_            =       self.dummy_input
     return init_dict, inputs_dict



  def        lowerCAmelCase__	(  self							:							Optional[int]	) ->			Optional[int]:




     """simple docstring"""
     snake_case_            =       [0.6_738, 0.4_491, 0.1_055, 1.0_710, 0.7_316, 0.3_339, 0.3_352, 0.1_023, 0.3_568]
     super().test_output(_lowerCAmelCase	)




