


import argparse
import json
from pathlib import Path

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

from transformers import (
    BertTokenizer,
    ViltConfig,
    ViltForImageAndTextRetrieval,
    ViltForImagesAndTextClassification,
    ViltForMaskedLM,
    ViltForQuestionAnswering,
    ViltImageProcessor,
    ViltProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE :   Union[str, Any]							      =				logging.get_logger(__name__)



def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_=False							,				lowerCamelCase_=False							,				lowerCamelCase_=False				)					->		Tuple:
    _lowercase :							str   	=	[]
    for i in range(config.num_hidden_layers				):
        # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
        rename_keys.append((F'''transformer.blocks.{i}.norm1.weight''', F'''vilt.encoder.layer.{i}.layernorm_before.weight''')				)
        rename_keys.append((F'''transformer.blocks.{i}.norm1.bias''', F'''vilt.encoder.layer.{i}.layernorm_before.bias''')				)
        rename_keys.append(
            (F'''transformer.blocks.{i}.attn.proj.weight''', F'''vilt.encoder.layer.{i}.attention.output.dense.weight''')				)
        rename_keys.append(
            (F'''transformer.blocks.{i}.attn.proj.bias''', F'''vilt.encoder.layer.{i}.attention.output.dense.bias''')				)
        rename_keys.append((F'''transformer.blocks.{i}.norm2.weight''', F'''vilt.encoder.layer.{i}.layernorm_after.weight''')				)
        rename_keys.append((F'''transformer.blocks.{i}.norm2.bias''', F'''vilt.encoder.layer.{i}.layernorm_after.bias''')				)
        rename_keys.append(
            (F'''transformer.blocks.{i}.mlp.fc1.weight''', F'''vilt.encoder.layer.{i}.intermediate.dense.weight''')				)
        rename_keys.append((F'''transformer.blocks.{i}.mlp.fc1.bias''', F'''vilt.encoder.layer.{i}.intermediate.dense.bias''')				)
        rename_keys.append((F'''transformer.blocks.{i}.mlp.fc2.weight''', F'''vilt.encoder.layer.{i}.output.dense.weight''')				)
        rename_keys.append((F'''transformer.blocks.{i}.mlp.fc2.bias''', F'''vilt.encoder.layer.{i}.output.dense.bias''')				)

    # embeddings
    rename_keys.extend(
        [
            # text embeddings
            ('text_embeddings.word_embeddings.weight', 'vilt.embeddings.text_embeddings.word_embeddings.weight'),
            (
                'text_embeddings.position_embeddings.weight',
                'vilt.embeddings.text_embeddings.position_embeddings.weight',
            ),
            ('text_embeddings.position_ids', 'vilt.embeddings.text_embeddings.position_ids'),
            (
                'text_embeddings.token_type_embeddings.weight',
                'vilt.embeddings.text_embeddings.token_type_embeddings.weight',
            ),
            ('text_embeddings.LayerNorm.weight', 'vilt.embeddings.text_embeddings.LayerNorm.weight'),
            ('text_embeddings.LayerNorm.bias', 'vilt.embeddings.text_embeddings.LayerNorm.bias'),
            # patch embeddings
            ('transformer.cls_token', 'vilt.embeddings.cls_token'),
            ('transformer.patch_embed.proj.weight', 'vilt.embeddings.patch_embeddings.projection.weight'),
            ('transformer.patch_embed.proj.bias', 'vilt.embeddings.patch_embeddings.projection.bias'),
            ('transformer.pos_embed', 'vilt.embeddings.position_embeddings'),
            # token type embeddings
            ('token_type_embeddings.weight', 'vilt.embeddings.token_type_embeddings.weight'),
        ]				)

    # final layernorm + pooler
    rename_keys.extend(
        [
            ('transformer.norm.weight', 'vilt.layernorm.weight'),
            ('transformer.norm.bias', 'vilt.layernorm.bias'),
            ('pooler.dense.weight', 'vilt.pooler.dense.weight'),
            ('pooler.dense.bias', 'vilt.pooler.dense.bias'),
        ]				)

    # classifier head(s)
    if vqa_model:
        # classification head
        rename_keys.extend(
            [
                ('vqa_classifier.0.weight', 'classifier.0.weight'),
                ('vqa_classifier.0.bias', 'classifier.0.bias'),
                ('vqa_classifier.1.weight', 'classifier.1.weight'),
                ('vqa_classifier.1.bias', 'classifier.1.bias'),
                ('vqa_classifier.3.weight', 'classifier.3.weight'),
                ('vqa_classifier.3.bias', 'classifier.3.bias'),
            ]				)
    elif nlvr_model:
        # classification head
        rename_keys.extend(
            [
                ('nlvr2_classifier.0.weight', 'classifier.0.weight'),
                ('nlvr2_classifier.0.bias', 'classifier.0.bias'),
                ('nlvr2_classifier.1.weight', 'classifier.1.weight'),
                ('nlvr2_classifier.1.bias', 'classifier.1.bias'),
                ('nlvr2_classifier.3.weight', 'classifier.3.weight'),
                ('nlvr2_classifier.3.bias', 'classifier.3.bias'),
            ]				)
    else:
        pass

    return rename_keys



def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_				)					->		List[Any]:
    for i in range(config.num_hidden_layers				):
        _lowercase :							Tuple   	=	'vilt.'
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        _lowercase :							Tuple   	=	state_dict.pop(F'''transformer.blocks.{i}.attn.qkv.weight'''				)
        _lowercase :							int   	=	state_dict.pop(F'''transformer.blocks.{i}.attn.qkv.bias'''				)
        # next, add query, keys and values (in that order) to the state dict
        _lowercase :							Any   	=	in_proj_weight[
            : config.hidden_size, :
        ]
        _lowercase :							int   	=	in_proj_bias[: config.hidden_size]
        _lowercase :							int   	=	in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        _lowercase :							List[Any]   	=	in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        _lowercase :							str   	=	in_proj_weight[
            -config.hidden_size :, :
        ]
        _lowercase :							str   	=	in_proj_bias[-config.hidden_size :]



def        UpperCamelCase_(   lowerCamelCase_				)					->		Any:
    _lowercase :							Dict   	=	['head.weight', 'head.bias']
    for k in ignore_keys:
        state_dict.pop(lowerCamelCase_							,				lowerCamelCase_				)



def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_							,				lowerCamelCase_				)					->		Any:
    _lowercase :							List[str]   	=	dct.pop(lowerCamelCase_				)
    _lowercase :							str   	=	val



@torch.no_grad()
def        UpperCamelCase_(   lowerCamelCase_							,				lowerCamelCase_				)					->		int:
    _lowercase :							int   	=	ViltConfig(image_size=384							,				patch_size=32							,				tie_word_embeddings=lowerCamelCase_				)
    _lowercase :							List[str]   	=	False
    _lowercase :							Optional[int]   	=	False
    _lowercase :							str   	=	False
    _lowercase :							Any   	=	False
    if "vqa" in checkpoint_url:
        _lowercase :							Optional[Any]   	=	True
        _lowercase :							Dict   	=	3129
        _lowercase :							str   	=	'huggingface/label-files'
        _lowercase :							List[str]   	=	'vqa2-id2label.json'
        _lowercase :							List[Any]   	=	json.load(open(hf_hub_download(lowerCamelCase_							,				lowerCamelCase_							,				repo_type='dataset'				)							,				'r'				)				)
        _lowercase :							Optional[int]   	=	{int(lowerCamelCase_				): v for k, v in idalabel.items()}
        _lowercase :							Any   	=	idalabel
        _lowercase :							Dict   	=	{v: k for k, v in idalabel.items()}
        _lowercase :							Union[str, Any]   	=	ViltForQuestionAnswering(lowerCamelCase_				)
    elif "nlvr" in checkpoint_url:
        _lowercase :							Any   	=	True
        _lowercase :							List[str]   	=	2
        _lowercase :							List[str]   	=	{0: 'False', 1: 'True'}
        _lowercase :							Tuple   	=	{v: k for k, v in config.idalabel.items()}
        _lowercase :							List[str]   	=	3
        _lowercase :							Dict   	=	ViltForImagesAndTextClassification(lowerCamelCase_				)
    elif "irtr" in checkpoint_url:
        _lowercase :							int   	=	True
        _lowercase :							Union[str, Any]   	=	ViltForImageAndTextRetrieval(lowerCamelCase_				)
    elif "mlm_itm" in checkpoint_url:
        _lowercase :							Optional[int]   	=	True
        _lowercase :							str   	=	ViltForMaskedLM(lowerCamelCase_				)
    else:
        raise ValueError('Unknown model type'				)

    # load state_dict of original model, remove and rename some keys
    _lowercase :							List[Any]   	=	torch.hub.load_state_dict_from_url(lowerCamelCase_							,				map_location='cpu'				)['state_dict']
    _lowercase :							List[Any]   	=	create_rename_keys(lowerCamelCase_							,				lowerCamelCase_							,				lowerCamelCase_							,				lowerCamelCase_				)
    for src, dest in rename_keys:
        rename_key(lowerCamelCase_							,				lowerCamelCase_							,				lowerCamelCase_				)
    read_in_q_k_v(lowerCamelCase_							,				lowerCamelCase_				)
    if mlm_model or irtr_model:
        _lowercase :							Optional[int]   	=	['itm_score.fc.weight', 'itm_score.fc.bias']
        for k in ignore_keys:
            state_dict.pop(lowerCamelCase_							,				lowerCamelCase_				)

    # load state dict into HuggingFace model
    model.eval()
    if mlm_model:
        _lowercase	,       _lowercase :							Dict   	=	model.load_state_dict(lowerCamelCase_							,				strict=lowerCamelCase_				)
        assert missing_keys == ["mlm_score.decoder.bias"]
    else:
        model.load_state_dict(lowerCamelCase_				)

    # Define processor
    _lowercase :							Any   	=	ViltImageProcessor(size=384				)
    _lowercase :							Tuple   	=	BertTokenizer.from_pretrained('bert-base-uncased'				)
    _lowercase :							Optional[int]   	=	ViltProcessor(lowerCamelCase_							,				lowerCamelCase_				)

    # Forward pass on example inputs (image + text)
    if nlvr_model:
        _lowercase :							List[str]   	=	Image.open(requests.get('https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'							,				stream=lowerCamelCase_				).raw				)
        _lowercase :							Any   	=	Image.open(requests.get('https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'							,				stream=lowerCamelCase_				).raw				)
        _lowercase :							Any   	=	(
            'The left image contains twice the number of dogs as the right image, and at least two dogs in total are'
            ' standing.'
        )
        _lowercase :							Union[str, Any]   	=	processor(lowerCamelCase_							,				lowerCamelCase_							,				return_tensors='pt'				)
        _lowercase :							Union[str, Any]   	=	processor(lowerCamelCase_							,				lowerCamelCase_							,				return_tensors='pt'				)
        _lowercase :							Optional[Any]   	=	model(
            input_ids=encoding_a.input_ids							,				pixel_values=encoding_a.pixel_values							,				pixel_values_a=encoding_a.pixel_values							,				)
    else:
        _lowercase :							Dict   	=	Image.open(requests.get('http://images.cocodataset.org/val2017/000000039769.jpg'							,				stream=lowerCamelCase_				).raw				)
        if mlm_model:
            _lowercase :							List[Any]   	=	'a bunch of [MASK] laying on a [MASK].'
        else:
            _lowercase :							Dict   	=	'How many cats are there?'
        _lowercase :							Union[str, Any]   	=	processor(lowerCamelCase_							,				lowerCamelCase_							,				return_tensors='pt'				)
        _lowercase :							List[Any]   	=	model(**lowerCamelCase_				)

    # Verify outputs
    if mlm_model:
        _lowercase :							Optional[Any]   	=	torch.Size([1, 11, 3_0522]				)
        _lowercase :							List[str]   	=	torch.tensor([-12.50_61, -12.51_23, -12.51_74]				)
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3]							,				lowerCamelCase_							,				atol=1e-4				)

        # verify masked token prediction equals "cats"
        _lowercase :							Union[str, Any]   	=	outputs.logits[0, 4, :].argmax(-1				).item()
        assert tokenizer.decode([predicted_id]				) == "cats"
    elif vqa_model:
        _lowercase :							Dict   	=	torch.Size([1, 3129]				)
        _lowercase :							Optional[Any]   	=	torch.tensor([-15.94_95, -18.14_72, -10.30_41]				)
        assert torch.allclose(outputs.logits[0, :3]							,				lowerCamelCase_							,				atol=1e-4				)
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3]							,				lowerCamelCase_							,				atol=1e-4				)

        # verify vqa prediction equals "2"
        _lowercase :							Optional[int]   	=	outputs.logits.argmax(-1				).item()
        assert model.config.idalabel[predicted_idx] == "2"
    elif nlvr_model:
        _lowercase :							Optional[int]   	=	torch.Size([1, 2]				)
        _lowercase :							str   	=	torch.tensor([-2.87_21, 2.12_91]				)
        assert torch.allclose(outputs.logits[0, :3]							,				lowerCamelCase_							,				atol=1e-4				)
        assert outputs.logits.shape == expected_shape

    Path(lowerCamelCase_				).mkdir(exist_ok=lowerCamelCase_				)
    print(F'''Saving model and processor to {pytorch_dump_folder_path}'''				)
    model.save_pretrained(lowerCamelCase_				)
    processor.save_pretrained(lowerCamelCase_				)


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE :   str							      =				argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--checkpoint_url",
     default="https://github.com/dandelin/ViLT/releases/download/200k/vilt_200k_mlm_itm.ckpt",
     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."
 )

 SCREAMING_SNAKE_CASE :   Optional[int]							      =				parser.parse_args()
 convert_vilt_checkpoint(args.checkpoint_url, 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     ):
       lowercase_				:			str									=    DownBlockaD  # noqa F405
       lowercase_				:			List[str]									=    """down"""
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           _lowercase :							str   	=	[-0.0_2_3_2, -0.9_8_6_9, 0.8_0_5_4, -0.0_6_3_7, -0.1_6_8_8, -1.4_2_6_4, 0.4_4_7_0, -1.3_3_9_4, 0.0_9_0_4]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			Optional[Any]									=    ResnetDownsampleBlockaD  # noqa F405
       lowercase_				:			Optional[Any]									=    """down"""
       def     UpperCamelCase  (      self) ->     Any:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	[0.0_7_1_0, 0.2_4_1_0, -0.7_3_2_0, -1.0_7_5_7, -1.1_3_4_3, 0.3_5_4_0, -0.0_1_3_3, -0.2_5_7_6, 0.0_9_4_8]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			Any									=    AttnDownBlockaD  # noqa F405
       lowercase_				:			Dict									=    """down"""
       def     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           _lowercase :							Optional[int]   	=	[0.0_6_3_6, 0.8_9_6_4, -0.6_2_3_4, -1.0_1_3_1, 0.0_8_4_4, 0.4_9_3_5, 0.3_4_3_7, 0.0_9_1_1, -0.2_9_5_7]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			str									=    CrossAttnDownBlockaD  # noqa F405
       lowercase_				:			Optional[int]									=    """down"""
       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           _lowercase	,       _lowercase :							int   	=	super().prepare_init_args_and_inputs_for_common()
           _lowercase :							int   	=	32
           return init_dict, inputs_dict




       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           _lowercase :							List[str]   	=	[0.2_2_3_8, -0.7_3_9_6, -0.2_2_5_5, -0.3_8_2_9, 0.1_9_2_5, 1.1_6_6_5, 0.0_6_0_3, -0.7_2_9_5, 0.1_9_8_3]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			Dict									=    SimpleCrossAttnDownBlockaD  # noqa F405
       lowercase_				:			Optional[Any]									=    """down"""
       @property
       def     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           return super().get_dummy_input(include_encoder_hidden_states=lowerCamelCase)
       def     UpperCamelCase  (      self) ->     List[Any]:




           """simple docstring"""

           _lowercase	,       _lowercase :							Dict   	=	super().prepare_init_args_and_inputs_for_common()
           _lowercase :							Optional[Any]   	=	32
           return init_dict, inputs_dict




       @unittest.skipIf(torch_device == 'mps',   'MPS result is not consistent')
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase :							Tuple   	=	[0.7_9_2_1, -0.0_9_9_2, -0.1_9_6_2, -0.7_6_9_5, -0.4_2_4_2, 0.7_8_0_4, 0.4_7_3_7, 0.2_7_6_5, 0.3_3_3_8]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			str									=    SkipDownBlockaD  # noqa F405
       lowercase_				:			Union[str, Any]									=    """down"""
       @property
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           return super().get_dummy_input(include_skip_sample=lowerCamelCase)




       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           _lowercase :							int   	=	[-0.0_8_4_5, -0.2_0_8_7, -0.2_4_6_5, 0.0_9_7_1, 0.1_9_0_0, -0.0_4_8_4, 0.2_6_6_4, 0.4_1_7_9, 0.5_0_6_9]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			str									=    AttnSkipDownBlockaD  # noqa F405
       lowercase_				:			Optional[Any]									=    """down"""
       @property
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           return super().get_dummy_input(include_skip_sample=lowerCamelCase)




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




           """simple docstring"""

           _lowercase :							Optional[int]   	=	[0.5_5_3_9, 0.1_6_0_9, 0.4_9_2_4, 0.0_5_3_7, -0.1_9_9_5, 0.4_0_5_0, 0.0_9_7_9, -0.2_7_2_1, -0.0_6_4_2]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			str									=    DownEncoderBlockaD  # noqa F405
       lowercase_				:			Optional[Any]									=    """down"""
       @property
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           return super().get_dummy_input(include_temb=lowerCamelCase)
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase :							Dict   	=	{
               'in_channels': 32,
               'out_channels': 32,
           }
           _lowercase :							int   	=	self.dummy_input
           return init_dict, inputs_dict




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




           """simple docstring"""

           _lowercase :							List[str]   	=	[1.1_1_0_2, 0.5_3_0_2, 0.4_8_7_2, -0.0_0_2_3, -0.8_0_4_2, 0.0_4_8_3, -0.3_4_8_9, -0.5_6_3_2, 0.7_6_2_6]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			int									=    AttnDownEncoderBlockaD  # noqa F405
       lowercase_				:			Tuple									=    """down"""
       @property
       def     UpperCamelCase  (      self) ->     Tuple:




           """simple docstring"""

           return super().get_dummy_input(include_temb=lowerCamelCase)
       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           _lowercase :							Any   	=	{
               'in_channels': 32,
               'out_channels': 32,
           }
           _lowercase :							Optional[int]   	=	self.dummy_input
           return init_dict, inputs_dict




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




           """simple docstring"""

           _lowercase :							Union[str, Any]   	=	[0.8_9_6_6, -0.1_4_8_6, 0.8_5_6_8, 0.8_1_4_1, -0.9_0_4_6, -0.1_3_4_2, -0.0_9_7_2, -0.7_4_1_7, 0.1_5_3_8]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			int									=    UNetMidBlockaD  # noqa F405
       lowercase_				:			Optional[int]									=    """mid"""
       def     UpperCamelCase  (      self) ->     Optional[Any]:




           """simple docstring"""

           _lowercase :							int   	=	{
               'in_channels': 32,
               'temb_channels': 1_28,
           }
           _lowercase :							Optional[Any]   	=	self.dummy_input
           return init_dict, inputs_dict




       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           _lowercase :							Dict   	=	[-0.1_0_6_2, 1.7_2_4_8, 0.3_4_9_4, 1.4_5_6_9, -0.0_9_1_0, -1.2_4_2_1, -0.9_9_8_4, 0.6_7_3_6, 1.0_0_2_8]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			Dict									=    UNetMidBlockaDCrossAttn  # noqa F405
       lowercase_				:			Dict									=    """mid"""
       def     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           _lowercase	,       _lowercase :							Union[str, Any]   	=	super().prepare_init_args_and_inputs_for_common()
           _lowercase :							Any   	=	32
           return init_dict, inputs_dict




       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           _lowercase :							Dict   	=	[0.0_1_8_7, 2.4_2_2_0, 0.4_4_8_4, 1.1_2_0_3, -0.6_1_2_1, -1.5_1_2_2, -0.8_2_7_0, 0.7_8_5_1, 1.8_3_3_5]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			List[Any]									=    UNetMidBlockaDSimpleCrossAttn  # noqa F405
       lowercase_				:			Any									=    """mid"""
       @property
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           return super().get_dummy_input(include_encoder_hidden_states=lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           _lowercase	,       _lowercase :							str   	=	super().prepare_init_args_and_inputs_for_common()
           _lowercase :							List[str]   	=	32
           return init_dict, inputs_dict




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




           """simple docstring"""

           _lowercase :							int   	=	[0.7_1_4_3, 1.9_9_7_4, 0.5_4_4_8, 1.3_9_7_7, 0.1_2_8_2, -1.1_2_3_7, -1.4_2_3_8, 0.5_5_3_0, 0.8_8_8_0]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			List[Any]									=    UpBlockaD  # noqa F405
       lowercase_				:			List[str]									=    """up"""
       @property
       def     UpperCamelCase  (      self) ->     List[Any]:




           """simple docstring"""

           return super().get_dummy_input(include_res_hidden_states_tuple=lowerCamelCase)




       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           _lowercase :							Any   	=	[-0.2_0_4_1, -0.4_1_6_5, -0.3_0_2_2, 0.0_0_4_1, -0.6_6_2_8, -0.7_0_5_3, 0.1_9_2_8, -0.0_3_2_5, 0.0_5_2_3]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			List[str]									=    ResnetUpsampleBlockaD  # noqa F405
       lowercase_				:			Optional[int]									=    """up"""
       @property
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           return super().get_dummy_input(include_res_hidden_states_tuple=lowerCamelCase)




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




           """simple docstring"""

           _lowercase :							List[Any]   	=	[0.2_2_8_7, 0.3_5_4_9, -0.1_3_4_6, 0.4_7_9_7, -0.1_7_1_5, -0.9_6_4_9, 0.7_3_0_5, -0.5_8_6_4, -0.6_2_4_4]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			Dict									=    CrossAttnUpBlockaD  # noqa F405
       lowercase_				:			List[str]									=    """up"""
       @property
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           return super().get_dummy_input(include_res_hidden_states_tuple=lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Union[str, Any]:




           """simple docstring"""

           _lowercase	,       _lowercase :							Union[str, Any]   	=	super().prepare_init_args_and_inputs_for_common()
           _lowercase :							Tuple   	=	32
           return init_dict, inputs_dict




       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           _lowercase :							Dict   	=	[-0.1_4_0_3, -0.3_5_1_5, -0.0_4_2_0, -0.1_4_2_5, 0.3_1_6_7, 0.5_0_9_4, -0.2_1_8_1, 0.5_9_3_1, 0.5_5_8_2]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			Optional[int]									=    SimpleCrossAttnUpBlockaD  # noqa F405
       lowercase_				:			List[Any]									=    """up"""
       @property
       def     UpperCamelCase  (      self) ->     Any:




           """simple docstring"""

           return super().get_dummy_input(include_res_hidden_states_tuple=lowerCamelCase,   include_encoder_hidden_states=lowerCamelCase)
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           _lowercase	,       _lowercase :							int   	=	super().prepare_init_args_and_inputs_for_common()
           _lowercase :							Dict   	=	32
           return init_dict, inputs_dict




       def     UpperCamelCase  (      self) ->     Dict:




           """simple docstring"""

           _lowercase :							List[str]   	=	[0.2_6_4_5, 0.1_4_8_0, 0.0_9_0_9, 0.8_0_4_4, -0.9_7_5_8, -0.9_0_8_3, 0.0_9_9_4, -1.1_4_5_3, -0.7_4_0_2]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			List[str]									=    AttnUpBlockaD  # noqa F405
       lowercase_				:			Optional[int]									=    """up"""
       @property
       def     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           return super().get_dummy_input(include_res_hidden_states_tuple=lowerCamelCase)




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




           """simple docstring"""

           _lowercase :							Union[str, Any]   	=	[0.0_9_7_9, 0.1_3_2_6, 0.0_0_2_1, 0.0_6_5_9, 0.2_2_4_9, 0.0_0_5_9, 0.1_1_3_2, 0.5_9_5_2, 0.1_0_3_3]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			List[str]									=    SkipUpBlockaD  # noqa F405
       lowercase_				:			Optional[Any]									=    """up"""
       @property
       def     UpperCamelCase  (      self) ->     List[str]:




           """simple docstring"""

           return super().get_dummy_input(include_res_hidden_states_tuple=lowerCamelCase)




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




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	[-0.0_8_9_3, -0.1_2_3_4, -0.1_5_0_6, -0.0_3_3_2, 0.0_1_2_3, -0.0_2_1_1, 0.0_5_6_6, 0.0_1_4_3, 0.0_3_6_2]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			int									=    AttnSkipUpBlockaD  # noqa F405
       lowercase_				:			Any									=    """up"""
       @property
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           return super().get_dummy_input(include_res_hidden_states_tuple=lowerCamelCase)




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




           """simple docstring"""

           _lowercase :							Tuple   	=	[0.0_3_6_1, 0.0_6_1_7, 0.2_7_8_7, -0.0_3_5_0, 0.0_3_4_2, 0.3_4_2_1, -0.0_8_4_3, 0.0_9_1_3, 0.3_0_1_5]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			List[str]									=    UpDecoderBlockaD  # noqa F405
       lowercase_				:			int									=    """up"""
       @property
       def     UpperCamelCase  (      self) ->     List[Any]:




           """simple docstring"""

           return super().get_dummy_input(include_temb=lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Optional[int]:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	{'in_channels': 32, 'out_channels': 32}

           _lowercase :							Optional[Any]   	=	self.dummy_input
           return init_dict, inputs_dict




       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           _lowercase :							List[Any]   	=	[0.4_4_0_4, 0.1_9_9_8, -0.9_8_8_6, -0.3_3_2_0, -0.3_1_2_8, -0.7_0_3_4, -0.6_9_5_5, -0.2_3_3_8, -0.3_1_3_7]
           super().test_output(lowerCamelCase)






class _lowerCamelCase(		_a,				unittest.TestCase     ):
       lowercase_				:			Any									=    AttnUpDecoderBlockaD  # noqa F405
       lowercase_				:			str									=    """up"""
       @property
       def     UpperCamelCase  (      self) ->     int:




           """simple docstring"""

           return super().get_dummy_input(include_temb=lowerCamelCase)
       def     UpperCamelCase  (      self) ->     Any:




           """simple docstring"""

           _lowercase :							Optional[Any]   	=	{'in_channels': 32, 'out_channels': 32}

           _lowercase :							List[Any]   	=	self.dummy_input
           return init_dict, inputs_dict




       def     UpperCamelCase  (      self) ->     str:




           """simple docstring"""

           _lowercase :							Optional[int]   	=	[0.6_7_3_8, 0.4_4_9_1, 0.1_0_5_5, 1.0_7_1_0, 0.7_3_1_6, 0.3_3_3_9, 0.3_3_5_2, 0.1_0_2_3, 0.3_5_6_8]
           super().test_output(lowerCamelCase)


