




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  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   DownBlockaD  # noqa F405
		__snake_case									=   'down'



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [-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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   ResnetDownsampleBlockaD  # noqa F405
		__snake_case									=   'down'



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   AttnDownBlockaD  # noqa F405
		__snake_case									=   'down'



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   CrossAttnDownBlockaD  # noqa F405
		__snake_case									=   'down'



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_			,					lowerCAmelCase_      			=  super().prepare_init_args_and_inputs_for_common()
								lowerCAmelCase_      			=  32
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   SimpleCrossAttnDownBlockaD  # noqa F405
		__snake_case									=   'down'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_encoder_hidden_states=UpperCamelCase__   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_			,					lowerCAmelCase_      			=  super().prepare_init_args_and_inputs_for_common()
								lowerCAmelCase_      			=  32
								return init_dict, inputs_dict







		@unittest.skipIf(torch_device == '''mps''',	'''MPS result is not consistent'''   )
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   SkipDownBlockaD  # noqa F405
		__snake_case									=   'down'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_skip_sample=UpperCamelCase__   )







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [-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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   AttnSkipDownBlockaD  # noqa F405
		__snake_case									=   'down'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_skip_sample=UpperCamelCase__   )







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   DownEncoderBlockaD  # noqa F405
		__snake_case									=   'down'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_temb=UpperCamelCase__   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  {
								    '''in_channels''': 32,
								    '''out_channels''': 32,
								}
								lowerCAmelCase_      			=  self.dummy_input
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   AttnDownEncoderBlockaD  # noqa F405
		__snake_case									=   'down'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_temb=UpperCamelCase__   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  {
								    '''in_channels''': 32,
								    '''out_channels''': 32,
								}
								lowerCAmelCase_      			=  self.dummy_input
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   UNetMidBlockaD  # noqa F405
		__snake_case									=   'mid'



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  {
								    '''in_channels''': 32,
								    '''temb_channels''': 128,
								}
								lowerCAmelCase_      			=  self.dummy_input
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [-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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   UNetMidBlockaDCrossAttn  # noqa F405
		__snake_case									=   'mid'



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_			,					lowerCAmelCase_      			=  super().prepare_init_args_and_inputs_for_common()
								lowerCAmelCase_      			=  32
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   UNetMidBlockaDSimpleCrossAttn  # noqa F405
		__snake_case									=   'mid'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_encoder_hidden_states=UpperCamelCase__   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_			,					lowerCAmelCase_      			=  super().prepare_init_args_and_inputs_for_common()
								lowerCAmelCase_      			=  32
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   UpBlockaD  # noqa F405
		__snake_case									=   'up'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__   )







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [-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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   ResnetUpsampleBlockaD  # noqa F405
		__snake_case									=   'up'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__   )







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   CrossAttnUpBlockaD  # noqa F405
		__snake_case									=   'up'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_			,					lowerCAmelCase_      			=  super().prepare_init_args_and_inputs_for_common()
								lowerCAmelCase_      			=  32
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [-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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   SimpleCrossAttnUpBlockaD  # noqa F405
		__snake_case									=   'up'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__,	include_encoder_hidden_states=UpperCamelCase__   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_			,					lowerCAmelCase_      			=  super().prepare_init_args_and_inputs_for_common()
								lowerCAmelCase_      			=  32
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   AttnUpBlockaD  # noqa F405
		__snake_case									=   'up'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__   )







		@unittest.skipIf(torch_device == '''mps''',	'''MPS result is not consistent'''   )
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   SkipUpBlockaD  # noqa F405
		__snake_case									=   'up'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__   )







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [-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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   AttnSkipUpBlockaD  # noqa F405
		__snake_case									=   'up'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__   )







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   UpDecoderBlockaD  # noqa F405
		__snake_case									=   'up'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_temb=UpperCamelCase__   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  {'''in_channels''': 32, '''out_channels''': 32}

								lowerCAmelCase_      			=  self.dummy_input
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )







class  A		(  __UpperCAmelCase      ,				unittest.TestCase     ):
		__snake_case									=   AttnUpDecoderBlockaD  # noqa F405
		__snake_case									=   'up'



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return super().get_dummy_input(include_temb=UpperCamelCase__   )



		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  {'''in_channels''': 32, '''out_channels''': 32}

								lowerCAmelCase_      			=  self.dummy_input
								return init_dict, inputs_dict







		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								lowerCAmelCase_      			=  [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(UpperCamelCase__   )










import argparse
import os
import re

import torch
from flax.traverse_util import flatten_dict
from tax import checkpoints

from transformers import (
    AutoTokenizer,
    PixaStructConfig,
    PixaStructForConditionalGeneration,
    PixaStructImageProcessor,
    PixaStructProcessor,
    PixaStructTextConfig,
    PixaStructVisionConfig,
)





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  checkpoints.load_tax_checkpoint(_A  )
						lowerCAmelCase_      			=  flatten_dict(_A  )
						return flax_params





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  {}

						lowerCAmelCase_      			=  {
						    '''token_embedder''': '''embeddings''',
						    '''encoder_norm''': '''layernorm''',
						    '''kernel''': '''weight''',
						    '''.out''': '''.output''',
						    '''scale''': '''weight''',
						    '''embedders_0.pos_embedding''': '''row_embedder.weight''',
						    '''embedders_1.pos_embedding''': '''column_embedder.weight''',
						}

						lowerCAmelCase_      			=  {
						    '''query''': '''attention.query''',
						    '''key''': '''attention.key''',
						    '''value''': '''attention.value''',
						    '''output.dense''': '''output''',
						    '''encoder_decoder_attention.o''': '''encoder_decoder_attention.attention.o''',
						    '''pre_self_attention_layer_norm''': '''self_attention.layer_norm''',
						    '''pre_cross_attention_layer_norm''': '''encoder_decoder_attention.layer_norm''',
						    '''mlp.''': '''mlp.DenseReluDense.''',
						    '''pre_mlp_layer_norm''': '''mlp.layer_norm''',
						    '''self_attention.o''': '''self_attention.attention.o''',
						    '''decoder.embeddings.embedding''': '''decoder.embed_tokens.weight''',
						    '''decoder.relpos_bias.rel_embedding''': '''decoder.layer.0.self_attention.attention.relative_attention_bias.weight''',
						    '''decoder.decoder_norm.weight''': '''decoder.final_layer_norm.weight''',
						    '''decoder.logits_dense.weight''': '''decoder.lm_head.weight''',
						}

						for key in flax_dict.keys():
												if "target" in key:
																		# remove the first prefix from the key
																		lowerCAmelCase_      			=  '''.'''.join(key[1:]  )

																		# rename the key
																		for old, new in CONVERSION_MAPPING.items():
																								lowerCAmelCase_      			=  new_key.replace(_A ,  _A  )

																		if "decoder" in new_key:
																								for old, new in DECODER_CONVERSION_MAPPING.items():
																														lowerCAmelCase_      			=  new_key.replace(_A ,  _A  )

																		if "layers" in new_key and "decoder" not in new_key:
																								# use regex to replace the layer number
																								lowerCAmelCase_      			=  re.sub(r'''layers_(\d+)''' ,  r'''layer.\1''' ,  _A  )
																								lowerCAmelCase_      			=  new_key.replace('''encoder''' ,  '''encoder.encoder'''  )

																		elif "layers" in new_key and "decoder" in new_key:
																								# use regex to replace the layer number
																								lowerCAmelCase_      			=  re.sub(r'''layers_(\d+)''' ,  r'''layer.\1''' ,  _A  )

																		lowerCAmelCase_      			=  flax_dict[key]

						lowerCAmelCase_      			=  {}
						# convert converted_dict into torch format
						for key in converted_dict.keys():
												if ("embed_tokens" not in key) and ("embedder" not in key):
																		lowerCAmelCase_      			=  torch.from_numpy(converted_dict[key].T  )
												else:
																		lowerCAmelCase_      			=  torch.from_numpy(converted_dict[key]  )

						return converted_torch_dict





def __UpperCamelCase			(       _A ,  _A ,  _A=False ,  _A=False  ):
						lowerCAmelCase_      			=  get_flax_param(_A  )

						if not use_large:
												lowerCAmelCase_      			=  PixaStructVisionConfig()
												lowerCAmelCase_      			=  PixaStructTextConfig()
						else:
												lowerCAmelCase_      			=  PixaStructVisionConfig(
												    hidden_size=1536 ,  d_ff=3968 ,  num_attention_heads=24 ,  num_hidden_layers=18  )
												lowerCAmelCase_      			=  PixaStructTextConfig(hidden_size=1536 ,  d_ff=3968 ,  num_heads=24 ,  num_layers=18  )
						lowerCAmelCase_      			=  PixaStructConfig(
						    vision_config=encoder_config.to_dict() ,  text_config=decoder_config.to_dict() ,  is_vqa=_A  )

						lowerCAmelCase_      			=  PixaStructForConditionalGeneration(_A  )

						lowerCAmelCase_      			=  rename_and_convert_flax_params(_A  )
						model.load_state_dict(_A  )

						lowerCAmelCase_      			=  AutoTokenizer.from_pretrained('''ybelkada/test-pix2struct-tokenizer'''  )
						lowerCAmelCase_      			=  PixaStructImageProcessor()
						lowerCAmelCase_      			=  PixaStructProcessor(image_processor=_A ,  tokenizer=_A  )

						if use_large:
												lowerCAmelCase_      			=  4096

						lowerCAmelCase_      			=  True

						# mkdir if needed
						os.makedirs(_A ,  exist_ok=_A  )

						model.save_pretrained(_A  )
						processor.save_pretrained(_A  )

						print('''Model saved in {}'''.format(_A  )  )


if __name__ == "__main__":
						_A     =			argparse.ArgumentParser()
						parser.add_argument('''--t5x_checkpoint_path''', default=None, type=str, help='''Path to the original T5x checkpoint.''')
						parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
						parser.add_argument('''--use_large''', action='''store_true''', help='''Use large model.''')
						parser.add_argument('''--is_vqa''', action='''store_true''', help='''Use large model.''')
						_A     =			parser.parse_args()

						convert_pixastruct_original_pytorch_checkpoint_to_hf(
						    args.tax_checkpoint_path, args.pytorch_dump_folder_path, args.use_large
						)




