import importlib
import inspect
import json
import os
import re
import shutil
import sys
from pathlib import Path
from typing import Dict, Optional, Union
from urllib import request

from huggingface_hub import HfFolder, cached_download, hf_hub_download, model_info
from packaging import version

from .. import __version__
from . import DIFFUSERS_DYNAMIC_MODULE_NAME, HF_MODULES_CACHE, logging


snake_case__				= (
    'https://raw.githubusercontent.com/huggingface/diffusers/{revision}/examples/community/{pipeline}.py'
)


snake_case__				= logging.get_logger(__name__)  # pylint: disable=invalid-name
def       __magic_name__(  )						->    str:



      '''simple docstring'''

      _lowerCamelCase			    =				'''https://pypi.org/pypi/diffusers/json'''
      _lowerCamelCase			    =				json.loads(request.urlopen(__UpperCAmelCase		).read()		)['''releases'''].keys()
      return sorted(__UpperCAmelCase ,						key=lambda __UpperCAmelCase		: version.Version(__UpperCAmelCase		)		)
def       __magic_name__(  )						->    int:



      '''simple docstring'''

      if HF_MODULES_CACHE in sys.path:
            return

      sys.path.append(__UpperCAmelCase		)
      os.makedirs(__UpperCAmelCase ,						exist_ok=__UpperCAmelCase		)
      _lowerCamelCase			    =				Path(__UpperCAmelCase		) / '''__init__.py'''
      if not init_path.exists():
            init_path.touch()
def       __magic_name__(  __UpperCAmelCase		)						->    Optional[Any]:



      '''simple docstring'''

      init_hf_modules()
      _lowerCamelCase			    =				Path(__UpperCAmelCase		) / name
      # If the parent module does not exist yet, recursively create it.
      if not dynamic_module_path.parent.exists():
            create_dynamic_module(dynamic_module_path.parent		)
      os.makedirs(__UpperCAmelCase ,						exist_ok=__UpperCAmelCase		)
      _lowerCamelCase			    =				dynamic_module_path / '''__init__.py'''
      if not init_path.exists():
            init_path.touch()
def       __magic_name__(  __UpperCAmelCase		)						->    Optional[Any]:



      '''simple docstring'''

      with open(__UpperCAmelCase ,						'''r''' ,						encoding='''utf-8'''		) as f:
            _lowerCamelCase			    =				f.read()

      # Imports of the form `import .xxx`
      _lowerCamelCase			    =				re.findall('''^\s*import\s+\.(\S+)\s*$''' ,						__UpperCAmelCase ,						flags=re.MULTILINE		)
      # Imports of the form `from .xxx import yyy`
      relative_imports += re.findall('''^\s*from\s+\.(\S+)\s+import''' ,						__UpperCAmelCase ,						flags=re.MULTILINE		)
      # Unique-ify
      return list(set(__UpperCAmelCase		)		)
def       __magic_name__(  __UpperCAmelCase		)						->    List[Any]:



      '''simple docstring'''

      _lowerCamelCase			    =				False
      _lowerCamelCase			    =				[module_file]
      _lowerCamelCase			    =				[]

      # Let's recurse through all relative imports
      while not no_change:
            _lowerCamelCase			    =				[]
            for f in files_to_check:
                  new_imports.extend(get_relative_imports(__UpperCAmelCase		)		)

            _lowerCamelCase			    =				Path(__UpperCAmelCase		).parent
            _lowerCamelCase			    =				[str(module_path / m		) for m in new_imports]
            _lowerCamelCase			    =				[f for f in new_import_files if f not in all_relative_imports]
            _lowerCamelCase			    =				[F'{f}.py' for f in new_import_files]

            _lowerCamelCase			    =				len(__UpperCAmelCase		) == 0
            all_relative_imports.extend(__UpperCAmelCase		)

      return all_relative_imports
def       __magic_name__(  __UpperCAmelCase		)						->    Dict:



      '''simple docstring'''

      with open(__UpperCAmelCase ,						'''r''' ,						encoding='''utf-8'''		) as f:
            _lowerCamelCase			    =				f.read()

      # Imports of the form `import xxx`
      _lowerCamelCase			    =				re.findall('''^\s*import\s+(\S+)\s*$''' ,						__UpperCAmelCase ,						flags=re.MULTILINE		)
      # Imports of the form `from xxx import yyy`
      imports += re.findall('''^\s*from\s+(\S+)\s+import''' ,						__UpperCAmelCase ,						flags=re.MULTILINE		)
      # Only keep the top-level module
      _lowerCamelCase			    =				[imp.split('''.'''		)[0] for imp in imports if not imp.startswith('''.'''		)]

      # Unique-ify and test we got them all
      _lowerCamelCase			    =				list(set(__UpperCAmelCase		)		)
      _lowerCamelCase			    =				[]
      for imp in imports:
            try:
                  importlib.import_module(__UpperCAmelCase		)
            except ImportError:
                  missing_packages.append(__UpperCAmelCase		)

      if len(__UpperCAmelCase		) > 0:
            raise ImportError(
                '''This modeling file requires the following packages that were not found in your environment: '''
                F'{", ".join(__UpperCAmelCase		)}. Run `pip install {" ".join(__UpperCAmelCase		)}`'		)

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



      '''simple docstring'''

      _lowerCamelCase			    =				module_path.replace(os.path.sep ,						'''.'''		)
      _lowerCamelCase			    =				importlib.import_module(__UpperCAmelCase		)

      if class_name is None:
            return find_pipeline_class(__UpperCAmelCase		)
      return getattr(__UpperCAmelCase ,						__UpperCAmelCase		)
def       __magic_name__(  __UpperCAmelCase		)						->    int:



      '''simple docstring'''

      from ..pipelines import DiffusionPipeline

      _lowerCamelCase			    =				dict(inspect.getmembers(__UpperCAmelCase ,						inspect.isclass		)		)

      _lowerCamelCase			    =				None
      for cls_name, cls in cls_members.items():
            if (
                cls_name != DiffusionPipeline.__name__
                and issubclass(cls ,						__UpperCAmelCase		)
                and cls.__module__.split('''.'''		)[0] != "diffusers"
            ):
                  if pipeline_class is not None:
                        raise ValueError(
                            F'Multiple classes that inherit from {DiffusionPipeline.__name__} have been found:'
                            F' {pipeline_class.__name__}, and {cls_name}. Please make sure to define only one in'
                            F' {loaded_module}.'		)
                  _lowerCamelCase			    =				cls

      return pipeline_class
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase = None ,						__UpperCAmelCase = False ,						__UpperCAmelCase = False ,						__UpperCAmelCase = None ,						__UpperCAmelCase = None ,						__UpperCAmelCase = None ,						__UpperCAmelCase = False ,						)						->    Any:



      '''simple docstring'''

      _lowerCamelCase			    =				str(__UpperCAmelCase		)

      _lowerCamelCase			    =				os.path.join(__UpperCAmelCase ,						__UpperCAmelCase		)

      if os.path.isfile(__UpperCAmelCase		):
            _lowerCamelCase			    =				module_file_or_url
            _lowerCamelCase			    =				'''local'''
      elif pretrained_model_name_or_path.count('''/'''		) == 0:
            _lowerCamelCase			    =				get_diffusers_versions()
            # cut ".dev0"
            _lowerCamelCase			    =				'''v''' + '''.'''.join(__version__.split('''.'''		)[:3]		)

            # retrieve github version that matches
            if revision is None:
                  _lowerCamelCase			    =				latest_version if latest_version[1:] in available_versions else '''main'''
                  logger.info(F'Defaulting to latest_version: {revision}.'		)
            elif revision in available_versions:
                  _lowerCamelCase			    =				F'v{revision}'
            elif revision == "main":
                  _lowerCamelCase			    =				revision
            else:
                  raise ValueError(
                      F'`custom_revision`: {revision} does not exist. Please make sure to choose one of'
                      F' {", ".join(available_versions + ["main"]		)}.'		)

            # community pipeline on GitHub
            _lowerCamelCase			    =				COMMUNITY_PIPELINES_URL.format(revision=__UpperCAmelCase ,						pipeline=__UpperCAmelCase		)
            try:
                  _lowerCamelCase			    =				cached_download(
                      __UpperCAmelCase ,						cache_dir=__UpperCAmelCase ,						force_download=__UpperCAmelCase ,						proxies=__UpperCAmelCase ,						resume_download=__UpperCAmelCase ,						local_files_only=__UpperCAmelCase ,						use_auth_token=__UpperCAmelCase ,						)
                  _lowerCamelCase			    =				'''git'''
                  _lowerCamelCase			    =				pretrained_model_name_or_path + '''.py'''
            except EnvironmentError:
                  logger.error(F'Could not locate the {module_file} inside {pretrained_model_name_or_path}.'		)
                  raise
      else:
            try:
                  # Load from URL or cache if already cached
                  _lowerCamelCase			    =				hf_hub_download(
                      __UpperCAmelCase ,						__UpperCAmelCase ,						cache_dir=__UpperCAmelCase ,						force_download=__UpperCAmelCase ,						proxies=__UpperCAmelCase ,						resume_download=__UpperCAmelCase ,						local_files_only=__UpperCAmelCase ,						use_auth_token=__UpperCAmelCase ,						)
                  _lowerCamelCase			    =				os.path.join('''local''' ,						'''--'''.join(pretrained_model_name_or_path.split('''/'''		)		)		)
            except EnvironmentError:
                  logger.error(F'Could not locate the {module_file} inside {pretrained_model_name_or_path}.'		)
                  raise

    # Check we have all the requirements in our environment
      _lowerCamelCase			    =				check_imports(__UpperCAmelCase		)

      # Now we move the module inside our cached dynamic modules.
      _lowerCamelCase			    =				DIFFUSERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule
      create_dynamic_module(__UpperCAmelCase		)
      _lowerCamelCase			    =				Path(__UpperCAmelCase		) / full_submodule
      if submodule == "local" or submodule == "git":
            # We always copy local files (we could hash the file to see if there was a change, and give them the name of
            # that hash, to only copy when there is a modification but it seems overkill for now).
            # The only reason we do the copy is to avoid putting too many folders in sys.path.
            shutil.copy(__UpperCAmelCase ,						submodule_path / module_file		)
            for module_needed in modules_needed:
                  _lowerCamelCase			    =				F'{module_needed}.py'
                  shutil.copy(os.path.join(__UpperCAmelCase ,						__UpperCAmelCase		) ,						submodule_path / module_needed		)
      else:
            # Get the commit hash
            # TODO: we will get this info in the etag soon, so retrieve it from there and not here.
            if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		):
                  _lowerCamelCase			    =				use_auth_token
            elif use_auth_token is True:
                  _lowerCamelCase			    =				HfFolder.get_token()
            else:
                  _lowerCamelCase			    =				None

            _lowerCamelCase			    =				model_info(__UpperCAmelCase ,						revision=__UpperCAmelCase ,						token=__UpperCAmelCase		).sha

            # The module file will end up being placed in a subfolder with the git hash of the repo. This way we get the
            # benefit of versioning.
            _lowerCamelCase			    =				submodule_path / commit_hash
            _lowerCamelCase			    =				full_submodule + os.path.sep + commit_hash
            create_dynamic_module(__UpperCAmelCase		)

            if not (submodule_path / module_file).exists():
                  shutil.copy(__UpperCAmelCase ,						submodule_path / module_file		)
            # Make sure we also have every file with relative
            for module_needed in modules_needed:
                  if not (submodule_path / module_needed).exists():
                        get_cached_module_file(
                            __UpperCAmelCase ,						F'{module_needed}.py' ,						cache_dir=__UpperCAmelCase ,						force_download=__UpperCAmelCase ,						resume_download=__UpperCAmelCase ,						proxies=__UpperCAmelCase ,						use_auth_token=__UpperCAmelCase ,						revision=__UpperCAmelCase ,						local_files_only=__UpperCAmelCase ,						)
      return os.path.join(__UpperCAmelCase ,						__UpperCAmelCase		)




def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase = None ,						__UpperCAmelCase = None ,						__UpperCAmelCase = False ,						__UpperCAmelCase = False ,						__UpperCAmelCase = None ,						__UpperCAmelCase = None ,						__UpperCAmelCase = None ,						__UpperCAmelCase = False ,						**__UpperCAmelCase ,						)						->    Any:



      '''simple docstring'''

      _lowerCamelCase			    =				get_cached_module_file(
          __UpperCAmelCase ,						__UpperCAmelCase ,						cache_dir=__UpperCAmelCase ,						force_download=__UpperCAmelCase ,						resume_download=__UpperCAmelCase ,						proxies=__UpperCAmelCase ,						use_auth_token=__UpperCAmelCase ,						revision=__UpperCAmelCase ,						local_files_only=__UpperCAmelCase ,						)
      return get_class_in_module(__UpperCAmelCase ,						final_module.replace('''.py''' ,						''''''		)		)
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)