




import argparse

import torch
from datasets import load_dataset
from donut import DonutModel

from transformers import (
    DonutImageProcessor,
    DonutProcessor,
    DonutSwinConfig,
    DonutSwinModel,
    MBartConfig,
    MBartForCausalLM,
    VisionEncoderDecoderModel,
    XLMRobertaTokenizerFast,
)





def 		__A		(	__lowerCamelCase    )      ->       Dict:
	a										=							model.config

	a										=							DonutSwinConfig(
	    image_size=original_config.input_size							,			patch_size=4							,			depths=original_config.encoder_layer							,			num_heads=[4, 8, 16, 32]							,			window_size=original_config.window_size							,			embed_dim=128							,			)
	a										=							MBartConfig(
	    is_decoder=__lowerCamelCase							,			is_encoder_decoder=__lowerCamelCase							,			add_cross_attention=__lowerCamelCase							,			decoder_layers=original_config.decoder_layer							,			max_position_embeddings=original_config.max_position_embeddings							,			vocab_size=len(
	        model.decoder.tokenizer    )							,			scale_embedding=__lowerCamelCase							,			add_final_layer_norm=__lowerCamelCase							,			)

	return encoder_config, decoder_config





def 		__A		(	__lowerCamelCase    )      ->       List[str]:
	if "encoder.model" in name:
		a										=							name.replace("""encoder.model"""							,			"""encoder"""    )
	if "decoder.model" in name:
		a										=							name.replace("""decoder.model"""							,			"""decoder"""    )
	if "patch_embed.proj" in name:
		a										=							name.replace("""patch_embed.proj"""							,			"""embeddings.patch_embeddings.projection"""    )
	if "patch_embed.norm" in name:
		a										=							name.replace("""patch_embed.norm"""							,			"""embeddings.norm"""    )
	if name.startswith("""encoder"""    ):
		if "layers" in name:
			a										=							"""encoder.""" + name
		if "attn.proj" in name:
			a										=							name.replace("""attn.proj"""							,			"""attention.output.dense"""    )
		if "attn" in name and "mask" not 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 name == "encoder.norm.weight":
			a										=							"""encoder.layernorm.weight"""
		if name == "encoder.norm.bias":
			a										=							"""encoder.layernorm.bias"""

	return name





def 		__A		(	__lowerCamelCase							,			__lowerCamelCase    )      ->       Any:
	for key in orig_state_dict.copy().keys():
		a										=							orig_state_dict.pop(__lowerCamelCase    )

		if "qkv" in key:
			a										=							key.split("""."""    )
			a										=							int(key_split[3]    )
			a										=							int(key_split[5]    )
			a										=							model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

			if "weight" in key:
				a										=							val[:dim, :]
				a										=							val[dim : dim * 2, :]
				a										=							val[-dim:, :]
			else:
				a										=							val[:dim]
				a										=							val[dim : dim * 2]
				a										=							val[-dim:]
		elif "attn_mask" in key or key in ["encoder.model.norm.weight", "encoder.model.norm.bias"]:
			# HuggingFace implementation doesn't use attn_mask buffer
			# and model doesn't use final LayerNorms for the encoder
			pass
		else:
			a										=							val

	return orig_state_dict





def 		__A		(	__lowerCamelCase							,			__lowerCamelCase=None							,			__lowerCamelCase=False    )      ->       List[str]:
	# load original model
	a										=							DonutModel.from_pretrained(__lowerCamelCase    ).eval()

	# load HuggingFace model
	a				, a										=							get_configs(__lowerCamelCase    )
	a										=							DonutSwinModel(__lowerCamelCase    )
	a										=							MBartForCausalLM(__lowerCamelCase    )
	a										=							VisionEncoderDecoderModel(encoder=__lowerCamelCase							,			decoder=__lowerCamelCase    )
	model.eval()

	a										=							original_model.state_dict()
	a										=							convert_state_dict(__lowerCamelCase							,			__lowerCamelCase    )
	model.load_state_dict(__lowerCamelCase    )

	# verify results on scanned document
	a										=							load_dataset("""hf-internal-testing/example-documents"""    )
	a										=							dataset["""test"""][0]["""image"""].convert("""RGB"""    )

	a										=							XLMRobertaTokenizerFast.from_pretrained(__lowerCamelCase							,			from_slow=__lowerCamelCase    )
	a										=							DonutImageProcessor(
	    do_align_long_axis=original_model.config.align_long_axis							,			size=original_model.config.input_size[::-1]    )
	a										=							DonutProcessor(__lowerCamelCase							,			__lowerCamelCase    )
	a										=							processor(__lowerCamelCase							,			return_tensors="""pt"""    ).pixel_values

	if model_name == "naver-clova-ix/donut-base-finetuned-docvqa":
		a										=							"""<s_docvqa><s_question>{user_input}</s_question><s_answer>"""
		a										=							"""When is the coffee break?"""
		a										=							task_prompt.replace("""{user_input}"""							,			__lowerCamelCase    )
	elif model_name == "naver-clova-ix/donut-base-finetuned-rvlcdip":
		a										=							"""<s_rvlcdip>"""
	elif model_name in [
	    "naver-clova-ix/donut-base-finetuned-cord-v1",
	    "naver-clova-ix/donut-base-finetuned-cord-v1-2560",
	]:
		a										=							"""<s_cord>"""
	elif model_name == "naver-clova-ix/donut-base-finetuned-cord-v2":
		a										=							"""s_cord-v2>"""
	elif model_name == "naver-clova-ix/donut-base-finetuned-zhtrainticket":
		a										=							"""<s_zhtrainticket>"""
	elif model_name in ["naver-clova-ix/donut-proto", "naver-clova-ix/donut-base"]:
		# use a random prompt
		a										=							"""hello world"""
	else:
		raise ValueError("""Model name not supported"""    )
	a										=							original_model.decoder.tokenizer(__lowerCamelCase							,			add_special_tokens=__lowerCamelCase							,			return_tensors="""pt"""    )[
	    """input_ids"""
	]

	a										=							original_model.encoder.model.patch_embed(__lowerCamelCase    )
	a				, a										=							model.encoder.embeddings(__lowerCamelCase    )
	assert torch.allclose(__lowerCamelCase							,			__lowerCamelCase							,			atol=1E-3    )

	# verify encoder hidden states
	a										=							original_model.encoder(__lowerCamelCase    )
	a										=							model.encoder(__lowerCamelCase    ).last_hidden_state
	assert torch.allclose(__lowerCamelCase							,			__lowerCamelCase							,			atol=1E-2    )

	# verify decoder hidden states
	a										=							original_model(__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase    ).logits
	a										=							model(__lowerCamelCase							,			decoder_input_ids=__lowerCamelCase    ).logits
	assert torch.allclose(__lowerCamelCase							,			__lowerCamelCase							,			atol=1E-3    )
	print("""Looks ok!"""    )

	if pytorch_dump_folder_path is not None:
		print(f'Saving model and processor to {pytorch_dump_folder_path}'    )
		model.save_pretrained(__lowerCamelCase    )
		processor.save_pretrained(__lowerCamelCase    )

	if push_to_hub:
		model.push_to_hub("""nielsr/""" + model_name.split("""/"""    )[-1]							,			commit_message="""Update model"""    )
		processor.push_to_hub("""nielsr/""" + model_name.split("""/"""    )[-1]							,			commit_message="""Update model"""    )


if __name__ == "__main__":
		__UpperCamelCase					:       Optional[Any]    					=      argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--model_name",
		    default="naver-clova-ix/donut-base-finetuned-docvqa",
		    required=False,
		    type=str,
		    help="Name of the original model you'd like to convert.",
		)
		parser.add_argument(
		    "--pytorch_dump_folder_path",
		    default=None,
		    required=False,
		    type=str,
		    help="Path to the output PyTorch model directory.",
		)
		parser.add_argument(
		    "--push_to_hub",
		    action="store_true",
		    help="Whether or not to push the converted model and processor to the 🤗 hub.",
		)

		__UpperCamelCase					:       Optional[Any]    					=      parser.parse_args()
		convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


__UpperCamelCase					:       Optional[Any]    					=      False

@skip_mps
class 		__lowerCAmelCase	(	__magic_name__      ,       __magic_name__      ,       __magic_name__      ,       unittest.TestCase ):
							UpperCamelCase__					     =  StableDiffusionAttendAndExcitePipeline
							UpperCamelCase__					     =  False
							UpperCamelCase__					     =  TEXT_TO_IMAGE_PARAMS
							UpperCamelCase__					     =  TEXT_TO_IMAGE_BATCH_PARAMS.union({'''token_indices'''} )
							UpperCamelCase__					     =  TEXT_TO_IMAGE_IMAGE_PARAMS
							UpperCamelCase__					     =  TEXT_TO_IMAGE_IMAGE_PARAMS
							@classmethod
							def       lowerCamelCase__							(						cls						:Union[str, Any]				):







								'''simple docstring'''







								super().setUpClass()
								torch.use_deterministic_algorithms(__magic_name__				)
							@classmethod
							def       lowerCamelCase__							(						cls						:Union[str, Any]				):







								'''simple docstring'''







								super().tearDownClass()
								torch.use_deterministic_algorithms(__magic_name__				)
							def       lowerCamelCase__							(						self						:Tuple				):







								'''simple docstring'''







								torch.manual_seed(0				)
								a										=							UNetaDConditionModel(
								    block_out_channels=(32, 64)		,      layers_per_block=1		,      sample_size=32		,      in_channels=4		,      out_channels=4		,      down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")		,      up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")		,      cross_attention_dim=32		,      attention_head_dim=(2, 4)		,      use_linear_projection=__magic_name__		,      )
								a										=							DDIMScheduler(
								    beta_start=0.00085		,      beta_end=0.012		,      beta_schedule="""scaled_linear"""		,      clip_sample=__magic_name__		,      set_alpha_to_one=__magic_name__		,      )
								torch.manual_seed(0				)
								a										=							AutoencoderKL(
								    block_out_channels=[32, 64]		,      in_channels=3		,      out_channels=3		,      down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]		,      up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]		,      latent_channels=4		,      sample_size=128		,      )
								torch.manual_seed(0				)
								a										=							CLIPTextConfig(
								    bos_token_id=0		,      eos_token_id=2		,      hidden_size=32		,      intermediate_size=37		,      layer_norm_eps=1E-0_5		,      num_attention_heads=4		,      num_hidden_layers=5		,      pad_token_id=1		,      vocab_size=1000		,      hidden_act="""gelu"""		,      projection_dim=512		,      )
								a										=							CLIPTextModel(__magic_name__				)
								a										=							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

								a										=							{
								    """unet""": unet,
								    """scheduler""": scheduler,
								    """vae""": vae,
								    """text_encoder""": text_encoder,
								    """tokenizer""": tokenizer,
								    """safety_checker""": None,
								    """feature_extractor""": None,
								}

								return components
							def       lowerCamelCase__							(						self						:Dict		,      __magic_name__						:Optional[int]		,      __magic_name__						:Optional[int]=0				):







								'''simple docstring'''







								if str(__magic_name__				).startswith("""mps"""				):
									a										=							torch.manual_seed(__magic_name__				)
								else:
									a										=							torch.Generator(device=__magic_name__				).manual_seed(__magic_name__				)
								a										=							a										=							{
								    """prompt""": """a cat and a frog""",
								    """token_indices""": [2, 5],
								    """generator""": generator,
								    """num_inference_steps""": 1,
								    """guidance_scale""": 6.0,
								    """output_type""": """numpy""",
								    """max_iter_to_alter""": 2,
								    """thresholds""": {0: 0.7},
								}
								return inputs
							def       lowerCamelCase__							(						self						:str				):







								'''simple docstring'''







								a										=							"""cpu"""

								a										=							self.get_dummy_components()
								a										=							self.pipeline_class(**__magic_name__				)
								pipe.to(__magic_name__				)
								pipe.set_progress_bar_config(disable=__magic_name__				)

								a										=							self.get_dummy_inputs(__magic_name__				)
								a										=							pipe(**__magic_name__				).images
								a										=							image[0, -3:, -3:, -1]

								self.assertEqual(image.shape		,      (1, 64, 64, 3)				)
								a										=							np.array(
								    [0.63905364, 0.62897307, 0.48599017, 0.5133624, 0.5550048, 0.45769516, 0.50326973, 0.5023139, 0.45384496]				)
								a										=							np.abs(image_slice.flatten() - expected_slice				).max()
								self.assertLessEqual(__magic_name__		,      1E-3				)
							def       lowerCamelCase__							(						self						:Tuple				):







								'''simple docstring'''







								super().test_cpu_offload_forward_pass(expected_max_diff=5E-4				)
							def       lowerCamelCase__							(						self						:List[Any]				):







								'''simple docstring'''







								self._test_inference_batch_consistent(batch_sizes=[1, 2]				)
							def       lowerCamelCase__							(						self						:List[str]				):







								'''simple docstring'''







								self._test_inference_batch_single_identical(batch_size=2		,      expected_max_diff=7E-4				)
							def       lowerCamelCase__							(						self						:List[str]				):







								'''simple docstring'''







								super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3				)
							def       lowerCamelCase__							(						self						:Dict				):







								'''simple docstring'''







								super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5E-4				)
							def       lowerCamelCase__							(						self						:List[str]				):







								'''simple docstring'''







								super().test_save_load_local(expected_max_difference=5E-4				)





							def       lowerCamelCase__							(						self						:str				):







								'''simple docstring'''







								super().test_save_load_optional_components(expected_max_difference=4E-4				)







@require_torch_gpu
@slow
class 		__lowerCAmelCase	(	unittest.TestCase ):
							@classmethod
							def       lowerCamelCase__							(						cls						:int				):







								'''simple docstring'''







								super().setUpClass()
								torch.use_deterministic_algorithms(__magic_name__				)
							@classmethod
							def       lowerCamelCase__							(						cls						:Any				):







								'''simple docstring'''







								super().tearDownClass()
								torch.use_deterministic_algorithms(__magic_name__				)
							def       lowerCamelCase__							(						self						:Union[str, Any]				):







								'''simple docstring'''







								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()





							def       lowerCamelCase__							(						self						:int				):







								'''simple docstring'''







								a										=							torch.manual_seed(51				)

								a										=							StableDiffusionAttendAndExcitePipeline.from_pretrained(
								    """CompVis/stable-diffusion-v1-4"""		,      safety_checker=__magic_name__		,      torch_dtype=torch.floataa				)
								pipe.to("""cuda"""				)

								a										=							"""a painting of an elephant with glasses"""
								a										=							[5, 7]

								a										=							pipe(
								    prompt=__magic_name__		,      token_indices=__magic_name__		,      guidance_scale=7.5		,      generator=__magic_name__		,      num_inference_steps=5		,      max_iter_to_alter=5		,      output_type="""numpy"""		,      ).images[0]

								a										=							load_numpy(
								    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy"""				)
								assert np.abs((expected_image - image).max()				) < 5E-1

