from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .text_encoder import MultilingualCLIP


snake_case__				= logging.get_logger(__name__)  # pylint: disable=invalid-name

snake_case__				= '\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/Kandinsky-2-1-prior")\n        >>> pipe_prior.to("cuda")\n\n        >>> prompt = "red cat, 4k photo"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> negative_image_emb = out.negative_image_embeds\n\n        >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1")\n        >>> pipe.to("cuda")\n\n        >>> image = pipe(\n        ...     prompt,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=negative_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ... ).images\n\n        >>> image[0].save("cat.png")\n        ```\n'
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase=8		)						->    str:



      '''simple docstring'''

      _lowerCamelCase			    =				h // scale_factor**2
      if h % scale_factor**2 != 0:
            new_h += 1
      _lowerCamelCase			    =				w // scale_factor**2
      if w % scale_factor**2 != 0:
            new_w += 1
      return new_h * scale_factor, new_w * scale_factor




class  UpperCamelCase				( _UpperCAmelCase ):

     '''simple docstring'''

     def __init__(					self				,				A_				,				A_				,				A_				,				A_				,				A_				,				) ->				str:






           """simple docstring"""
           super().__init__()

           self.register_modules(
               text_encoder=lowercase__				,				tokenizer=lowercase__				,				unet=lowercase__				,				scheduler=lowercase__				,				movq=lowercase__				,				)
           _lowerCamelCase			    =				2 ** (len(self.movq.config.block_out_channels						) - 1)



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_				,				A_				,				A_				,				A_						) ->				List[Any]:






           """simple docstring"""






           if latents is None:
                 _lowerCamelCase			    =				randn_tensor(lowercase__				,				generator=lowercase__				,				device=lowercase__				,				dtype=lowercase__						)
           else:
                 if latents.shape != shape:
                       raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {shape}'						)
                 _lowerCamelCase			    =				latents.to(lowercase__						)

           _lowerCamelCase			    =				latents * scheduler.init_noise_sigma
           return latents



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_				,				A_				,				A_=None				,				) ->				Optional[int]:






           """simple docstring"""
           _lowerCamelCase			    =				len(lowercase__						) if isinstance(lowercase__				,				lowercase__						) else 1
           # get prompt text embeddings
           _lowerCamelCase			    =				self.tokenizer(
               lowercase__				,				padding='''max_length'''				,				truncation=lowercase__				,				max_length=77				,				return_attention_mask=lowercase__				,				add_special_tokens=lowercase__				,				return_tensors='''pt'''				,				)

           _lowerCamelCase			    =				text_inputs.input_ids
           _lowerCamelCase			    =				self.tokenizer(lowercase__				,				padding='''longest'''				,				return_tensors='''pt'''						).input_ids

           if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(lowercase__				,				lowercase__						):
                 _lowerCamelCase			    =				self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]						)
                 logger.warning(
                     '''The following part of your input was truncated because CLIP can only handle sequences up to'''
                     F' {self.tokenizer.model_max_length} tokens: {removed_text}'						)

           _lowerCamelCase			    =				text_input_ids.to(lowercase__						)
           _lowerCamelCase			    =				text_inputs.attention_mask.to(lowercase__						)

           _lowerCamelCase			    =				self.text_encoder(
               input_ids=lowercase__				,				attention_mask=lowercase__						)

           _lowerCamelCase			    =				prompt_embeds.repeat_interleave(lowercase__				,				dim=0						)
           _lowerCamelCase			    =				text_encoder_hidden_states.repeat_interleave(lowercase__				,				dim=0						)
           _lowerCamelCase			    =				text_mask.repeat_interleave(lowercase__				,				dim=0						)

           if do_classifier_free_guidance:
                 _lowerCamelCase			    =				42
                 if negative_prompt is None:
                       _lowerCamelCase			    =				[""] * batch_size
                 elif type(lowercase__						) is not type(lowercase__						):
                       raise TypeError(
                           F'`negative_prompt` should be the same type to `prompt`, but got {type(lowercase__						)} !='
                           F' {type(lowercase__						)}.'						)
                 elif isinstance(lowercase__				,				lowercase__						):
                       _lowerCamelCase			    =				[negative_prompt]
                 elif batch_size != len(lowercase__						):
                       raise ValueError(
                           F'`negative_prompt`: {negative_prompt} has batch size {len(lowercase__						)}, but `prompt`:'
                           F' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'
                           ''' the batch size of `prompt`.'''						)
                 else:
                       _lowerCamelCase			    =				negative_prompt

                 _lowerCamelCase			    =				self.tokenizer(
                     lowercase__				,				padding='''max_length'''				,				max_length=77				,				truncation=lowercase__				,				return_attention_mask=lowercase__				,				add_special_tokens=lowercase__				,				return_tensors='''pt'''				,				)
                 _lowerCamelCase			    =				uncond_input.input_ids.to(lowercase__						)
                 _lowerCamelCase			    =				uncond_input.attention_mask.to(lowercase__						)

                 _lowerCamelCase			    =				self.text_encoder(
                     input_ids=lowercase__				,				attention_mask=lowercase__						)

                 # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

                 _lowerCamelCase			    =				negative_prompt_embeds.shape[1]
                 _lowerCamelCase			    =				negative_prompt_embeds.repeat(1				,				lowercase__						)
                 _lowerCamelCase			    =				negative_prompt_embeds.view(batch_size * num_images_per_prompt				,				lowercase__						)

                 _lowerCamelCase			    =				uncond_text_encoder_hidden_states.shape[1]
                 _lowerCamelCase			    =				uncond_text_encoder_hidden_states.repeat(1				,				lowercase__				,				1						)
                 _lowerCamelCase			    =				uncond_text_encoder_hidden_states.view(
                     batch_size * num_images_per_prompt				,				lowercase__				,				-1						)
                 _lowerCamelCase			    =				uncond_text_mask.repeat_interleave(lowercase__				,				dim=0						)

                 # done duplicates

                 # For classifier free guidance, we need to do two forward passes.
                 # Here we concatenate the unconditional and text embeddings into a single batch
                 # to avoid doing two forward passes
                 _lowerCamelCase			    =				torch.cat([negative_prompt_embeds, prompt_embeds]						)
                 _lowerCamelCase			    =				torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]						)

                 _lowerCamelCase			    =				torch.cat([uncond_text_mask, text_mask]						)

           return prompt_embeds, text_encoder_hidden_states, text_mask



     def 		UpperCamelCase_						(					self				,				A_=0						) ->				List[str]:






           """simple docstring"""






           if is_accelerate_available():
                 from accelerate import cpu_offload
           else:
                 raise ImportError('''Please install accelerate via `pip install accelerate`'''						)

           _lowerCamelCase			    =				torch.device(F'cuda:{gpu_id}'						)

           _lowerCamelCase			    =				[
               self.unet,
               self.text_encoder,
               self.movq,
           ]
           for cpu_offloaded_model in models:
                 if cpu_offloaded_model is not None:
                       cpu_offload(lowercase__				,				lowercase__						)



     def 		UpperCamelCase_						(					self				,				A_=0						) ->				Optional[int]:






           """simple docstring"""






           if is_accelerate_available() and is_accelerate_version('''>='''				,				'''0.17.0.dev0'''						):
                 from accelerate import cpu_offload_with_hook
           else:
                 raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''						)

           _lowerCamelCase			    =				torch.device(F'cuda:{gpu_id}'						)

           if self.device.type != "cpu":
                 self.to('''cpu'''				,				silence_dtype_warnings=lowercase__						)
                 torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

           _lowerCamelCase			    =				None
           for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
                 _lowerCamelCase			    =				cpu_offload_with_hook(lowercase__				,				lowercase__				,				prev_module_hook=lowercase__						)

           if self.safety_checker is not None:
                 _lowerCamelCase			    =				cpu_offload_with_hook(self.safety_checker				,				lowercase__				,				prev_module_hook=lowercase__						)

           # We'll offload the last model manually.
           _lowerCamelCase			    =				hook



     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
     def 		UpperCamelCase_						(					self						) ->				Optional[int]:






           """simple docstring"""






           if not hasattr(self.unet				,				'''_hf_hook'''						):
                 return self.device
           for module in self.unet.modules():
                 if (
                     hasattr(lowercase__				,				'''_hf_hook'''						)
                     and hasattr(module._hf_hook				,				'''execution_device'''						)
                     and module._hf_hook.execution_device is not None
                 ):
                       return torch.device(module._hf_hook.execution_device						)
           return self.device





     @torch.no_grad()
     @replace_example_docstring(lowercase__						)
     def __call__(					self				,				A_				,				A_				,				A_				,				A_ = None				,				A_ = 5_12				,				A_ = 5_12				,				A_ = 1_00				,				A_ = 4.0				,				A_ = 1				,				A_ = None				,				A_ = None				,				A_ = "pil"				,				A_ = True				,				) ->				Dict:






           """simple docstring"""






           if isinstance(lowercase__				,				lowercase__						):
                 _lowerCamelCase			    =				1
           elif isinstance(lowercase__				,				lowercase__						):
                 _lowerCamelCase			    =				len(lowercase__						)
           else:
                 raise ValueError(F'`prompt` has to be of type `str` or `list` but is {type(lowercase__						)}'						)

           _lowerCamelCase			    =				self._execution_device

           _lowerCamelCase			    =				batch_size * num_images_per_prompt
           _lowerCamelCase			    =				guidance_scale > 1.0

           _lowerCamelCase			    =				self._encode_prompt(
               lowercase__				,				lowercase__				,				lowercase__				,				lowercase__				,				lowercase__						)

           if isinstance(lowercase__				,				lowercase__						):
                 _lowerCamelCase			    =				torch.cat(lowercase__				,				dim=0						)
           if isinstance(lowercase__				,				lowercase__						):
                 _lowerCamelCase			    =				torch.cat(lowercase__				,				dim=0						)

           if do_classifier_free_guidance:
                 _lowerCamelCase			    =				image_embeds.repeat_interleave(lowercase__				,				dim=0						)
                 _lowerCamelCase			    =				negative_image_embeds.repeat_interleave(lowercase__				,				dim=0						)

           _lowerCamelCase			    =				torch.cat([negative_image_embeds, image_embeds]				,				dim=0						).to(
               dtype=prompt_embeds.dtype				,				device=lowercase__						)

           self.scheduler.set_timesteps(lowercase__				,				device=lowercase__						)
           _lowerCamelCase			    =				self.scheduler.timesteps

           _lowerCamelCase			    =				self.unet.config.in_channels

           _lowerCamelCase			    =				get_new_h_w(lowercase__				,				lowercase__				,				self.movq_scale_factor						)

           # create initial latent
           _lowerCamelCase			    =				self.prepare_latents(
               (batch_size, num_channels_latents, height, width)				,				text_encoder_hidden_states.dtype				,				lowercase__				,				lowercase__				,				lowercase__				,				self.scheduler				,				)

           for i, t in enumerate(self.progress_bar(lowercase__						)						):
                 # expand the latents if we are doing classifier free guidance
                 _lowerCamelCase			    =				torch.cat([latents] * 2						) if do_classifier_free_guidance else latents

                 _lowerCamelCase			    =				{"text_embeds": prompt_embeds, "image_embeds": image_embeds}
                 _lowerCamelCase			    =				self.unet(
                     sample=lowercase__				,				timestep=lowercase__				,				encoder_hidden_states=lowercase__				,				added_cond_kwargs=lowercase__				,				return_dict=lowercase__				,				)[0]

                 if do_classifier_free_guidance:
                       _lowerCamelCase			    =				noise_pred.split(latents.shape[1]				,				dim=1						)
                       _lowerCamelCase			    =				noise_pred.chunk(2						)
                       _lowerCamelCase			    =				variance_pred.chunk(2						)
                       _lowerCamelCase			    =				noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                       _lowerCamelCase			    =				torch.cat([noise_pred, variance_pred_text]				,				dim=1						)

                 if not (
                     hasattr(self.scheduler.config				,				'''variance_type'''						)
                     and self.scheduler.config.variance_type in ["learned", "learned_range"]
                 ):
                       _lowerCamelCase			    =				noise_pred.split(latents.shape[1]				,				dim=1						)

                 # compute the previous noisy sample x_t -> x_t-1
                 _lowerCamelCase			    =				self.scheduler.step(
                     lowercase__				,				lowercase__				,				lowercase__				,				generator=lowercase__				,				).prev_sample
           # post-processing
           _lowerCamelCase			    =				self.movq.decode(lowercase__				,				force_not_quantize=lowercase__						)["sample"]

           if output_type not in ["pt", "np", "pil"]:
                 raise ValueError(F'Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}'						)

           if output_type in ["np", "pil"]:
                 _lowerCamelCase			    =				image * 0.5 + 0.5
                 _lowerCamelCase			    =				image.clamp(0				,				1						)
                 _lowerCamelCase			    =				image.cpu().permute(0				,				2				,				3				,				1						).float().numpy()

           if output_type == "pil":
                 _lowerCamelCase			    =				self.numpy_to_pil(lowercase__						)

           if not return_dict:
                 return (image,)

           return ImagePipelineOutput(images=lowercase__						)
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       __magic_name__(  __UpperCAmelCase		)						->    str:



      '''simple docstring'''

      _lowerCamelCase			    =				model.config

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

      return encoder_config, decoder_config
def       __magic_name__(  __UpperCAmelCase		)						->    Tuple:



      '''simple docstring'''

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

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

      return name
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    Dict:



      '''simple docstring'''

      for key in orig_state_dict.copy().keys():
            _lowerCamelCase			    =				orig_state_dict.pop(__UpperCAmelCase		)

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

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

      return orig_state_dict
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase=None ,						__UpperCAmelCase=False		)						->    int:



      '''simple docstring'''

      _lowerCamelCase			    =				DonutModel.from_pretrained(__UpperCAmelCase		).eval()

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

      _lowerCamelCase			    =				original_model.state_dict()
      _lowerCamelCase			    =				convert_state_dict(__UpperCAmelCase ,						__UpperCAmelCase		)
      model.load_state_dict(__UpperCAmelCase		)

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

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

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

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

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

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

      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__":
     snake_case__				= 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.',
     )

     snake_case__				= parser.parse_args()
     convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)