





'''simple docstring'''






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   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=model.config

  _UpperCamelCase       			=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  ,      )
  _UpperCamelCase       			=MBartConfig(
      is_decoder=__SCREAMING_SNAKE_CASE  ,      is_encoder_decoder=__SCREAMING_SNAKE_CASE  ,      add_cross_attention=__SCREAMING_SNAKE_CASE  ,      decoder_layers=original_config.decoder_layer  ,      max_position_embeddings=original_config.max_position_embeddings  ,      vocab_size=len(
          model.decoder.tokenizer   )  ,      scale_embedding=__SCREAMING_SNAKE_CASE  ,      add_final_layer_norm=__SCREAMING_SNAKE_CASE  ,      )

  return encoder_config, decoder_config


def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  if "encoder.model" in name:
    _UpperCamelCase       			=name.replace('''encoder.model'''  ,      '''encoder'''   )
  if "decoder.model" in name:
    _UpperCamelCase       			=name.replace('''decoder.model'''  ,      '''decoder'''   )
  if "patch_embed.proj" in name:
    _UpperCamelCase       			=name.replace('''patch_embed.proj'''  ,      '''embeddings.patch_embeddings.projection'''   )
  if "patch_embed.norm" in name:
    _UpperCamelCase       			=name.replace('''patch_embed.norm'''  ,      '''embeddings.norm'''   )
  if name.startswith('''encoder'''   ):
    if "layers" in name:
      _UpperCamelCase       			='''encoder.''' + name
    if "attn.proj" in name:
      _UpperCamelCase       			=name.replace('''attn.proj'''  ,      '''attention.output.dense'''   )
    if "attn" in name and "mask" not in name:
      _UpperCamelCase       			=name.replace('''attn'''  ,      '''attention.self'''   )
    if "norm1" in name:
      _UpperCamelCase       			=name.replace('''norm1'''  ,      '''layernorm_before'''   )
    if "norm2" in name:
      _UpperCamelCase       			=name.replace('''norm2'''  ,      '''layernorm_after'''   )
    if "mlp.fc1" in name:
      _UpperCamelCase       			=name.replace('''mlp.fc1'''  ,      '''intermediate.dense'''   )
    if "mlp.fc2" in name:
      _UpperCamelCase       			=name.replace('''mlp.fc2'''  ,      '''output.dense'''   )

    if name == "encoder.norm.weight":
      _UpperCamelCase       			='''encoder.layernorm.weight'''
    if name == "encoder.norm.bias":
      _UpperCamelCase       			='''encoder.layernorm.bias'''

  return name


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  for key in orig_state_dict.copy().keys():
    _UpperCamelCase       			=orig_state_dict.pop(__SCREAMING_SNAKE_CASE   )

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

      if "weight" in key:
        _UpperCamelCase       			=val[:dim, :]
        _UpperCamelCase       			=val[dim : dim * 2, :]
        _UpperCamelCase       			=val[-dim:, :]
      else:
        _UpperCamelCase       			=val[:dim]
        _UpperCamelCase       			=val[dim : dim * 2]
        _UpperCamelCase       			=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:
      _UpperCamelCase       			=val

  return orig_state_dict


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE=None  ,      __SCREAMING_SNAKE_CASE=False   ):


  """simple docstring"""

  _UpperCamelCase       			=DonutModel.from_pretrained(__SCREAMING_SNAKE_CASE   ).eval()

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

  _UpperCamelCase       			=original_model.state_dict()
  _UpperCamelCase       			=convert_state_dict(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
  model.load_state_dict(__SCREAMING_SNAKE_CASE   )

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

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

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

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

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

  # verify decoder hidden states
  _UpperCamelCase       			=original_model(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ).logits
  _UpperCamelCase       			=model(__SCREAMING_SNAKE_CASE  ,      decoder_input_ids=__SCREAMING_SNAKE_CASE   ).logits
  assert torch.allclose(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      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(__SCREAMING_SNAKE_CASE   )
    processor.save_pretrained(__SCREAMING_SNAKE_CASE   )

  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__":
      __lowerCamelCase						: List[str]		=   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.',
      )

      __lowerCamelCase						: str		=   parser.parse_args()
      convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)













'''simple docstring'''






import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np







class 		UpperCAmelCase   :


  """simple docstring"""







  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Union[str, Any]:
    torch.manual_seed(0  )
    _UpperCamelCase       			=TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=UNetaDConditionModel(
        sample_size=32							,	layers_per_block=1							,	block_out_channels=[32, 64]							,	down_block_types=[
            '''ResnetDownsampleBlock2D''',
            '''SimpleCrossAttnDownBlock2D''',
        ]							,	mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,	up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,	in_channels=3							,	out_channels=6							,	cross_attention_dim=32							,	encoder_hid_dim=32							,	attention_head_dim=8							,	addition_embed_type='''text'''							,	addition_embed_type_num_heads=2							,	cross_attention_norm='''group_norm'''							,	resnet_time_scale_shift='''scale_shift'''							,	act_fn='''gelu'''							,	)
    unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    torch.manual_seed(0  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	thresholding=UpperCamelCase__							,	dynamic_thresholding_ratio=0.95							,	sample_max_value=1.0							,	prediction_type='''epsilon'''							,	variance_type='''learned_range'''							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=IFWatermarker()

    return {
        "text_encoder": text_encoder,
        "tokenizer": tokenizer,
        "unet": unet,
        "scheduler": scheduler,
        "watermarker": watermarker,
        "safety_checker": None,
        "feature_extractor": None,
    }


  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Optional[int]:
    torch.manual_seed(0  )
    _UpperCamelCase       			=TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=UNetaDConditionModel(
        sample_size=32							,	layers_per_block=[1, 2]							,	block_out_channels=[32, 64]							,	down_block_types=[
            '''ResnetDownsampleBlock2D''',
            '''SimpleCrossAttnDownBlock2D''',
        ]							,	mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,	up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,	in_channels=6							,	out_channels=6							,	cross_attention_dim=32							,	encoder_hid_dim=32							,	attention_head_dim=8							,	addition_embed_type='''text'''							,	addition_embed_type_num_heads=2							,	cross_attention_norm='''group_norm'''							,	resnet_time_scale_shift='''scale_shift'''							,	act_fn='''gelu'''							,	class_embed_type='''timestep'''							,	mid_block_scale_factor=1.414							,	time_embedding_act_fn='''gelu'''							,	time_embedding_dim=32							,	)
    unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    torch.manual_seed(0  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	thresholding=UpperCamelCase__							,	dynamic_thresholding_ratio=0.95							,	sample_max_value=1.0							,	prediction_type='''epsilon'''							,	variance_type='''learned_range'''							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=IFWatermarker()

    return {
        "text_encoder": text_encoder,
        "tokenizer": tokenizer,
        "unet": unet,
        "scheduler": scheduler,
        "image_noising_scheduler": image_noising_scheduler,
        "watermarker": watermarker,
        "safety_checker": None,
        "feature_extractor": None,
    }


  def   UpperCamelCase__      (				self    :					Any  )     ->						Any:
    _UpperCamelCase       			=self.get_dummy_components()
    _UpperCamelCase       			=self.pipeline_class(**UpperCamelCase__  )
    pipe.to(UpperCamelCase__  )
    pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )

    _UpperCamelCase       			=inputs['''prompt''']
    _UpperCamelCase       			=inputs['''generator''']
    _UpperCamelCase       			=inputs['''num_inference_steps''']
    _UpperCamelCase       			=inputs['''output_type''']

    if "image" in inputs:
      _UpperCamelCase       			=inputs['''image''']
    else:
      _UpperCamelCase       			=None

    if "mask_image" in inputs:
      _UpperCamelCase       			=inputs['''mask_image''']
    else:
      _UpperCamelCase       			=None

    if "original_image" in inputs:
      _UpperCamelCase       			=inputs['''original_image''']
    else:
      _UpperCamelCase       			=None

    _UpperCamelCase			,     _UpperCamelCase       			=pipe.encode_prompt(UpperCamelCase__  )

    # inputs with prompt converted to embeddings
    _UpperCamelCase       			={
        '''prompt_embeds''': prompt_embeds,
        '''negative_prompt_embeds''': negative_prompt_embeds,
        '''generator''': generator,
        '''num_inference_steps''': num_inference_steps,
        '''output_type''': output_type,
    }

    if image is not None:
      _UpperCamelCase       			=image

    if mask_image is not None:
      _UpperCamelCase       			=mask_image

    if original_image is not None:
      _UpperCamelCase       			=original_image

    # set all optional components to None
    for optional_component in pipe._optional_components:
      setattr(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )

    _UpperCamelCase       			=pipe(**UpperCamelCase__  )[0]

    with tempfile.TemporaryDirectory() as tmpdir:
      pipe.save_pretrained(UpperCamelCase__  )
      _UpperCamelCase       			=self.pipeline_class.from_pretrained(UpperCamelCase__  )
      pipe_loaded.to(UpperCamelCase__  )
      pipe_loaded.set_progress_bar_config(disable=UpperCamelCase__  )

    pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    for optional_component in pipe._optional_components:
      self.assertTrue(
          getattr(UpperCamelCase__							,	UpperCamelCase__  ) is None							,	F'''`{optional_component}` did not stay set to None after loading.'''							,	)

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )

    _UpperCamelCase       			=inputs['''generator''']
    _UpperCamelCase       			=inputs['''num_inference_steps''']
    _UpperCamelCase       			=inputs['''output_type''']

    # inputs with prompt converted to embeddings
    _UpperCamelCase       			={
        '''prompt_embeds''': prompt_embeds,
        '''negative_prompt_embeds''': negative_prompt_embeds,
        '''generator''': generator,
        '''num_inference_steps''': num_inference_steps,
        '''output_type''': output_type,
    }

    if image is not None:
      _UpperCamelCase       			=image

    if mask_image is not None:
      _UpperCamelCase       			=mask_image

    if original_image is not None:
      _UpperCamelCase       			=original_image

    _UpperCamelCase       			=pipe_loaded(**UpperCamelCase__  )[0]

    _UpperCamelCase       			=np.abs(to_np(UpperCamelCase__  ) - to_np(UpperCamelCase__  )  ).max()
    self.assertLess(UpperCamelCase__							,	1E-4  )




  def   UpperCamelCase__      (				self    :					int  )     ->						str:
    _UpperCamelCase       			=self.get_dummy_components()
    _UpperCamelCase       			=self.pipeline_class(**UpperCamelCase__  )
    pipe.to(UpperCamelCase__  )
    pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )
    _UpperCamelCase       			=pipe(**UpperCamelCase__  )[0]

    with tempfile.TemporaryDirectory() as tmpdir:
      pipe.save_pretrained(UpperCamelCase__  )
      _UpperCamelCase       			=self.pipeline_class.from_pretrained(UpperCamelCase__  )
      pipe_loaded.to(UpperCamelCase__  )
      pipe_loaded.set_progress_bar_config(disable=UpperCamelCase__  )

    pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )
    _UpperCamelCase       			=pipe_loaded(**UpperCamelCase__  )[0]

    _UpperCamelCase       			=np.abs(to_np(UpperCamelCase__  ) - to_np(UpperCamelCase__  )  ).max()
    self.assertLess(UpperCamelCase__							,	1E-4  )






