



import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


logging.set_verbosity_info()
_lowercase		:				str      		=logging.get_logger(__name__)

def      lowerCAmelCase_					(  _lowercase    :       List[str]						,       _lowercase    :       List[str])							->	Optional[Any]:







			"""simple docstring"""
			a__       :						Tuple							=  []
			for i in range(encoder_config.num_hidden_layers):
						# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
						rename_keys.append(
						    (F'''encoder.deit.blocks.{i}.norm1.weight''', F'''encoder.encoder.layer.{i}.layernorm_before.weight'''))
						rename_keys.append((F'''encoder.deit.blocks.{i}.norm1.bias''', F'''encoder.encoder.layer.{i}.layernorm_before.bias'''))
						rename_keys.append(
						    (F'''encoder.deit.blocks.{i}.attn.proj.weight''', F'''encoder.encoder.layer.{i}.attention.output.dense.weight'''))
						rename_keys.append(
						    (F'''encoder.deit.blocks.{i}.attn.proj.bias''', F'''encoder.encoder.layer.{i}.attention.output.dense.bias'''))
						rename_keys.append(
						    (F'''encoder.deit.blocks.{i}.norm2.weight''', F'''encoder.encoder.layer.{i}.layernorm_after.weight'''))
						rename_keys.append((F'''encoder.deit.blocks.{i}.norm2.bias''', F'''encoder.encoder.layer.{i}.layernorm_after.bias'''))
						rename_keys.append(
						    (F'''encoder.deit.blocks.{i}.mlp.fc1.weight''', F'''encoder.encoder.layer.{i}.intermediate.dense.weight'''))
						rename_keys.append(
						    (F'''encoder.deit.blocks.{i}.mlp.fc1.bias''', F'''encoder.encoder.layer.{i}.intermediate.dense.bias'''))
						rename_keys.append(
						    (F'''encoder.deit.blocks.{i}.mlp.fc2.weight''', F'''encoder.encoder.layer.{i}.output.dense.weight'''))
						rename_keys.append((F'''encoder.deit.blocks.{i}.mlp.fc2.bias''', F'''encoder.encoder.layer.{i}.output.dense.bias'''))

			# cls token, position embeddings and patch embeddings of encoder
			rename_keys.extend(
			    [
			        ("""encoder.deit.cls_token""", """encoder.embeddings.cls_token"""),
			        ("""encoder.deit.pos_embed""", """encoder.embeddings.position_embeddings"""),
			        ("""encoder.deit.patch_embed.proj.weight""", """encoder.embeddings.patch_embeddings.projection.weight"""),
			        ("""encoder.deit.patch_embed.proj.bias""", """encoder.embeddings.patch_embeddings.projection.bias"""),
			        ("""encoder.deit.norm.weight""", """encoder.layernorm.weight"""),
			        ("""encoder.deit.norm.bias""", """encoder.layernorm.bias"""),
			    ])

			return rename_keys

def      lowerCAmelCase_					(  _lowercase    :       Union[str, Any]						,       _lowercase    :       Any)							->	Union[str, Any]:







			"""simple docstring"""
			for i in range(encoder_config.num_hidden_layers):
						# queries, keys and values (only weights, no biases)
						a__       :						List[Any]							=  state_dict.pop(F'''encoder.deit.blocks.{i}.attn.qkv.weight''')

						a__       :						int							=  in_proj_weight[
						    : encoder_config.hidden_size, :
						]
						a__       :						Union[str, Any]							=  in_proj_weight[
						    encoder_config.hidden_size : encoder_config.hidden_size * 2, :
						]
						a__       :						Union[str, Any]							=  in_proj_weight[
						    -encoder_config.hidden_size :, :
						]

def      lowerCAmelCase_					(  _lowercase    :       List[str]						,       _lowercase    :       Tuple						,       _lowercase    :       Any)							->	int:







			"""simple docstring"""
			a__       :						Any							=  dct.pop(_lowercase)
			a__       :						str							=  val

def      lowerCAmelCase_					(  _lowercase    :       int)							->	Dict:







			"""simple docstring"""
			if "handwritten" in checkpoint_url:
						a__       :						List[str]							=  """https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg"""  # industry
						# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-12.jpg" # have
						# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-10.jpg" # let
						# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"  #
						# url = "https://fki.tic.heia-fr.ch/static/img/a01-122.jpg"
			elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
						a__       :						Dict							=  """https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg"""
			a__       :						Any							=  Image.open(requests.get(_lowercase						,       stream=_lowercase).raw).convert("""RGB""")
			return im

@torch.no_grad()
def      lowerCAmelCase_					(  _lowercase    :       Union[str, Any]						,       _lowercase    :       Any)							->	List[str]:







			"""simple docstring"""
			a__       :						int							=  ViTConfig(image_size=384						,       qkv_bias=_lowercase)
			a__       :						str							=  TrOCRConfig()

			# size of the architecture
			if "base" in checkpoint_url:
						a__       :						Optional[int]							=  768
			elif "large" in checkpoint_url:
						# use ViT-large encoder
						a__       :						Optional[Any]							=  1024
						a__       :						Optional[Any]							=  4096
						a__       :						Union[str, Any]							=  24
						a__       :						int							=  16
						a__       :						Dict							=  1024
			else:
						raise ValueError("""Should either find 'base' or 'large' in checkpoint URL""")

			# the large-printed + stage1 checkpoints uses sinusoidal position embeddings, no layernorm afterwards
			if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
						a__       :						List[Any]							=  False
						a__       :						Any							=  """relu"""
						a__       :						str							=  1024
						a__       :						Optional[int]							=  True
						a__       :						int							=  False
						a__       :						int							=  False

			# load HuggingFace model
			a__       :						Any							=  ViTModel(_lowercase						,       add_pooling_layer=_lowercase)
			a__       :						Optional[int]							=  TrOCRForCausalLM(_lowercase)
			a__       :						Dict							=  VisionEncoderDecoderModel(encoder=_lowercase						,       decoder=_lowercase)
			model.eval()

			# load state_dict of original model, rename some keys
			a__       :						Any							=  torch.hub.load_state_dict_from_url(_lowercase						,       map_location="""cpu"""						,       check_hash=_lowercase)["""model"""]

			a__       :						Tuple							=  create_rename_keys(_lowercase						,       _lowercase)
			for src, dest in rename_keys:
						rename_key(_lowercase						,       _lowercase						,       _lowercase)
			read_in_q_k_v(_lowercase						,       _lowercase)

			# remove parameters we don't need
			del state_dict["encoder.deit.head.weight"]
			del state_dict["encoder.deit.head.bias"]
			del state_dict["decoder.version"]

			# add prefix to decoder keys
			for key, val in state_dict.copy().items():
						a__       :						Tuple							=  state_dict.pop(_lowercase)
						if key.startswith("""decoder""") and "output_projection" not in key:
									a__       :						str							=  val
						else:
									a__       :						Optional[int]							=  val

    # load state dict
			model.load_state_dict(_lowercase)

			# Check outputs on an image
			a__       :						List[Any]							=  ViTImageProcessor(size=encoder_config.image_size)
			a__       :						List[str]							=  RobertaTokenizer.from_pretrained("""roberta-large""")
			a__       :						List[str]							=  TrOCRProcessor(_lowercase						,       _lowercase)

			a__       :						str							=  processor(images=prepare_img(_lowercase)						,       return_tensors="""pt""").pixel_values

			# verify logits
			a__       :						Tuple							=  torch.tensor([[model.config.decoder.decoder_start_token_id]])
			a__       :						str							=  model(pixel_values=_lowercase						,       decoder_input_ids=_lowercase)
			a__       :						Optional[Any]							=  outputs.logits

			a__       :						Optional[int]							=  torch.Size([1, 1, 5_0265])
			if "trocr-base-handwritten" in checkpoint_url:
						a__       :						Any							=  torch.tensor(
						    [-1.4502, -4.6683, -0.5347, -2.9291, 9.1435, -3.0571, 8.9764, 1.7560, 8.7358, -1.5311])
			elif "trocr-large-handwritten" in checkpoint_url:
						a__       :						Optional[int]							=  torch.tensor(
						    [-2.6437, -1.3129, -2.2596, -5.3455, 6.3539, 1.7604, 5.4991, 1.4702, 5.6113, 2.0170])
			elif "trocr-base-printed" in checkpoint_url:
						a__       :						List[Any]							=  torch.tensor(
						    [-5.6816, -5.8388, 1.1398, -6.9034, 6.8505, -2.4393, 1.2284, -1.0232, -1.9661, -3.9210])
			elif "trocr-large-printed" in checkpoint_url:
						a__       :						str							=  torch.tensor(
						    [-6.0162, -7.0959, 4.4155, -5.1063, 7.0468, -3.1631, 2.6466, -0.3081, -0.8106, -1.7535])

			if "stage1" not in checkpoint_url:
						assert logits.shape == expected_shape, "Shape of logits not as expected"
						assert torch.allclose(logits[0, 0, :10]						,       _lowercase						,       atol=1e-3), "First elements of logits not as expected"

			Path(_lowercase).mkdir(exist_ok=_lowercase)
			print(F'''Saving model to {pytorch_dump_folder_path}''')
			model.save_pretrained(_lowercase)
			print(F'''Saving processor to {pytorch_dump_folder_path}''')
			processor.save_pretrained(_lowercase)


if __name__ == "__main__":
		_lowercase		:				Optional[int]      		=argparse.ArgumentParser()

		parser.add_argument(
		    "--checkpoint_url",
		    default="https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt",
		    type=str,
		    help="URL to the original PyTorch checkpoint (.pth file).",
		)
		parser.add_argument(
		    "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
		)
		_lowercase		:				Optional[Any]      		=parser.parse_args()
		convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









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


_lowercase		:				int      		=(
    "https://raw.githubusercontent.com/huggingface/diffusers/{revision}/examples/community/{pipeline}.py"
)


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

def      lowerCAmelCase_					(  )							->	Union[str, Any]:







			"""simple docstring"""
			a__       :						List[Any]							=  """https://pypi.org/pypi/diffusers/json"""
			a__       :						Tuple							=  json.loads(request.urlopen(_lowercase).read())["""releases"""].keys()
			return sorted(_lowercase						,       key=lambda _lowercase: version.Version(_lowercase))

def      lowerCAmelCase_					(  )							->	Optional[Any]:







			"""simple docstring"""
			# This function has already been executed if HF_MODULES_CACHE already is in the Python path.
			if HF_MODULES_CACHE in sys.path:
						return

			sys.path.append(_lowercase)
			os.makedirs(_lowercase						,       exist_ok=_lowercase)
			a__       :						Tuple							=  Path(_lowercase) / """__init__.py"""
			if not init_path.exists():
						init_path.touch()

def      lowerCAmelCase_					(  _lowercase    :       Union[str, os.PathLike])							->	Optional[Any]:







			"""simple docstring"""
			init_hf_modules()
			a__       :						Optional[int]							=  Path(_lowercase) / 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(_lowercase						,       exist_ok=_lowercase)
			a__       :						Any							=  dynamic_module_path / """__init__.py"""
			if not init_path.exists():
						init_path.touch()

def      lowerCAmelCase_					(  _lowercase    :       Union[str, Any])							->	Optional[int]:







			"""simple docstring"""
			with open(_lowercase						,       """r"""						,       encoding="""utf-8""") as f:
						a__       :						Union[str, Any]							=  f.read()

			# Imports of the form `import .xxx`
			a__       :						Optional[Any]							=  re.findall("""^\s*import\s+\.(\S+)\s*$"""						,       _lowercase						,       flags=re.MULTILINE)
			# Imports of the form `from .xxx import yyy`
			relative_imports += re.findall("""^\s*from\s+\.(\S+)\s+import"""						,       _lowercase						,       flags=re.MULTILINE)
			# Unique-ify
			return list(set(_lowercase))

def      lowerCAmelCase_					(  _lowercase    :       List[Any])							->	Dict:







			"""simple docstring"""
			a__       :						Dict							=  False
			a__       :						str							=  [module_file]
			a__       :						List[Any]							=  []

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

						a__       :						Dict							=  Path(_lowercase).parent
						a__       :						Any							=  [str(module_path / m) for m in new_imports]
						a__       :						Any							=  [f for f in new_import_files if f not in all_relative_imports]
						a__       :						List[Any]							=  [F'''{f}.py''' for f in new_import_files]

						a__       :						List[Any]							=  len(_lowercase) == 0
						all_relative_imports.extend(_lowercase)

			return all_relative_imports

def      lowerCAmelCase_					(  _lowercase    :       Optional[int])							->	Any:







			"""simple docstring"""
			with open(_lowercase						,       """r"""						,       encoding="""utf-8""") as f:
						a__       :						Optional[Any]							=  f.read()

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

			# Unique-ify and test we got them all
			a__       :						Optional[Any]							=  list(set(_lowercase))
			a__       :						Union[str, Any]							=  []
			for imp in imports:
						try:
									importlib.import_module(_lowercase)
						except ImportError:
									missing_packages.append(_lowercase)

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

			return get_relative_imports(_lowercase)

def      lowerCAmelCase_					(  _lowercase    :       int						,       _lowercase    :       Tuple)							->	Union[str, Any]:







			"""simple docstring"""
			a__       :						Dict							=  module_path.replace(os.path.sep						,       """.""")
			a__       :						List[Any]							=  importlib.import_module(_lowercase)

			if class_name is None:
						return find_pipeline_class(_lowercase)
			return getattr(_lowercase						,       _lowercase)

def      lowerCAmelCase_					(  _lowercase    :       Union[str, Any])							->	Optional[int]:







			"""simple docstring"""
			from ..pipelines import DiffusionPipeline

			a__       :						Any							=  dict(inspect.getmembers(_lowercase						,       inspect.isclass))

			a__       :						str							=  None
			for cls_name, cls in cls_members.items():
						if (
						    cls_name != DiffusionPipeline.__name__
						    and issubclass(cls						,       _lowercase)
						    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}.''')
									a__       :						str							=  cls

			return pipeline_class

def      lowerCAmelCase_					(  _lowercase    :       Union[str, os.PathLike]						,       _lowercase    :       str						,       _lowercase    :       Optional[Union[str, os.PathLike]] = None						,       _lowercase    :       bool = False						,       _lowercase    :       bool = False						,       _lowercase    :       Optional[Dict[str, str]] = None						,       _lowercase    :       Optional[Union[bool, str]] = None						,       _lowercase    :       Optional[str] = None						,       _lowercase    :       bool = False						,       )							->	Optional[int]:







			"""simple docstring"""
			a__       :						Union[str, Any]							=  str(_lowercase)

			a__       :						Union[str, Any]							=  os.path.join(_lowercase						,       _lowercase)

			if os.path.isfile(_lowercase):
						a__       :						Dict							=  module_file_or_url
						a__       :						List[str]							=  """local"""
			elif pretrained_model_name_or_path.count("""/""") == 0:
						a__       :						Optional[Any]							=  get_diffusers_versions()
						# cut ".dev0"
						a__       :						Union[str, Any]							=  """v""" + """.""".join(__version__.split(""".""")[:3])

						# retrieve github version that matches
						if revision is None:
									a__       :						List[str]							=  latest_version if latest_version[1:] in available_versions else """main"""
									logger.info(F'''Defaulting to latest_version: {revision}.''')
						elif revision in available_versions:
									a__       :						str							=  F'''v{revision}'''
						elif revision == "main":
									a__       :						List[str]							=  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
						a__       :						Any							=  COMMUNITY_PIPELINES_URL.format(revision=_lowercase						,       pipeline=_lowercase)
						try:
									a__       :						Optional[int]							=  cached_download(
									    _lowercase						,       cache_dir=_lowercase						,       force_download=_lowercase						,       proxies=_lowercase						,       resume_download=_lowercase						,       local_files_only=_lowercase						,       use_auth_token=_lowercase						,       )
									a__       :						Any							=  """git"""
									a__       :						int							=  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
									a__       :						Any							=  hf_hub_download(
									    _lowercase						,       _lowercase						,       cache_dir=_lowercase						,       force_download=_lowercase						,       proxies=_lowercase						,       resume_download=_lowercase						,       local_files_only=_lowercase						,       use_auth_token=_lowercase						,       )
									a__       :						Optional[Any]							=  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
			a__       :						List[str]							=  check_imports(_lowercase)

			# Now we move the module inside our cached dynamic modules.
			a__       :						int							=  DIFFUSERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule
			create_dynamic_module(_lowercase)
			a__       :						Tuple							=  Path(_lowercase) / 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(_lowercase						,       submodule_path / module_file)
						for module_needed in modules_needed:
									a__       :						Dict							=  F'''{module_needed}.py'''
									shutil.copy(os.path.join(_lowercase						,       _lowercase)						,       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(_lowercase						,       _lowercase):
									a__       :						Optional[Any]							=  use_auth_token
						elif use_auth_token is True:
									a__       :						Dict							=  HfFolder.get_token()
						else:
									a__       :						str							=  None

						a__       :						Any							=  model_info(_lowercase						,       revision=_lowercase						,       token=_lowercase).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.
						a__       :						Optional[int]							=  submodule_path / commit_hash
						a__       :						Tuple							=  full_submodule + os.path.sep + commit_hash
						create_dynamic_module(_lowercase)

						if not (submodule_path / module_file).exists():
									shutil.copy(_lowercase						,       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(
												    _lowercase						,       F'''{module_needed}.py'''						,       cache_dir=_lowercase						,       force_download=_lowercase						,       resume_download=_lowercase						,       proxies=_lowercase						,       use_auth_token=_lowercase						,       revision=_lowercase						,       local_files_only=_lowercase						,       )
			return os.path.join(_lowercase						,       _lowercase)





def      lowerCAmelCase_					(  _lowercase    :       Union[str, os.PathLike]						,       _lowercase    :       str						,       _lowercase    :       Optional[str] = None						,       _lowercase    :       Optional[Union[str, os.PathLike]] = None						,       _lowercase    :       bool = False						,       _lowercase    :       bool = False						,       _lowercase    :       Optional[Dict[str, str]] = None						,       _lowercase    :       Optional[Union[bool, str]] = None						,       _lowercase    :       Optional[str] = None						,       _lowercase    :       bool = False						,       **_lowercase    :       Union[str, Any]						,       )							->	List[str]:







			"""simple docstring"""
			a__       :						int							=  get_cached_module_file(
			    _lowercase						,       _lowercase						,       cache_dir=_lowercase						,       force_download=_lowercase						,       resume_download=_lowercase						,       proxies=_lowercase						,       use_auth_token=_lowercase						,       revision=_lowercase						,       local_files_only=_lowercase						,       )
			return get_class_in_module(_lowercase						,       final_module.replace(""".py"""						,       """"""))




