


import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


_A		:				Union[str, Any]									=       logging.get_logger(__name__)

enable_full_determinism()



class        a					(			_UpperCAmelCase     ,_UpperCAmelCase     ,unittest.TestCase   ):
     UpperCAmelCase__ :  List[Any]						      =	UNetaDModel
     UpperCAmelCase__ :  Tuple						      =	"sample"





     @property
     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]							):
        __lowerCamelCase:				List[Any]     =	4
        __lowerCamelCase:				Optional[int]     =	3
        __lowerCamelCase:				Dict     =	(32, 32)

        __lowerCamelCase:				Union[str, Any]     =	floats_tensor((batch_size, num_channels) + sizes							).to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Optional[Any]     =	torch.tensor([10]							).to(SCREAMING_SNAKE_CASE_							)

        return {"sample": noise, "timestep": time_step}





     @property
     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[Any]							):
        return (3, 32, 32)





     @property
     def  SCREAMING_SNAKE_CASE__       (    self					:				int							):
        return (3, 32, 32)







     def  SCREAMING_SNAKE_CASE__       (    self					:				Any							):
        __lowerCamelCase:				Dict     =	{
            """block_out_channels""": (32, 64),
            """down_block_types""": ("""DownBlock2D""", """AttnDownBlock2D"""),
            """up_block_types""": ("""AttnUpBlock2D""", """UpBlock2D"""),
            """attention_head_dim""": 3,
            """out_channels""": 3,
            """in_channels""": 3,
            """layers_per_block""": 2,
            """sample_size""": 32,
        }
        __lowerCamelCase:				str     =	self.dummy_input
        return init_dict, inputs_dict



class        a					(			_UpperCAmelCase     ,_UpperCAmelCase     ,unittest.TestCase   ):
     UpperCAmelCase__ :  List[Any]						      =	UNetaDModel
     UpperCAmelCase__ :  List[Any]						      =	"sample"





     @property
     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[Any]							):
        __lowerCamelCase:				str     =	4
        __lowerCamelCase:				str     =	4
        __lowerCamelCase:				str     =	(32, 32)

        __lowerCamelCase:				Dict     =	floats_tensor((batch_size, num_channels) + sizes							).to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				List[str]     =	torch.tensor([10]							).to(SCREAMING_SNAKE_CASE_							)

        return {"sample": noise, "timestep": time_step}





     @property
     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]							):
        return (4, 32, 32)





     @property
     def  SCREAMING_SNAKE_CASE__       (    self					:				Dict							):
        return (4, 32, 32)





     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]							):
        __lowerCamelCase:				Optional[Any]     =	{
            """sample_size""": 32,
            """in_channels""": 4,
            """out_channels""": 4,
            """layers_per_block""": 2,
            """block_out_channels""": (32, 64),
            """attention_head_dim""": 32,
            """down_block_types""": ("""DownBlock2D""", """DownBlock2D"""),
            """up_block_types""": ("""UpBlock2D""", """UpBlock2D"""),
        }
        __lowerCamelCase:				Any     =	self.dummy_input
        return init_dict, inputs_dict





     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[Any]							):
        __lowerCamelCase   ,     __lowerCamelCase:				List[Any]     =	UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""    ,	output_loading_info=SCREAMING_SNAKE_CASE_							)

        self.assertIsNotNone(SCREAMING_SNAKE_CASE_							)
        self.assertEqual(len(loading_info["""missing_keys"""]							)    ,	0							)

        model.to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				List[Any]     =	model(**self.dummy_input							).sample

        assert image is not None, "Make sure output is not None"





     @unittest.skipIf(torch_device != """cuda"""    ,	"""This test is supposed to run on GPU"""							)
     def  SCREAMING_SNAKE_CASE__       (    self					:				str							):
        __lowerCamelCase   ,     __lowerCamelCase:				str     =	UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""    ,	output_loading_info=SCREAMING_SNAKE_CASE_							)
        model.to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Any     =	model(**self.dummy_input							).sample

        assert image is not None, "Make sure output is not None"





     @unittest.skipIf(torch_device != """cuda"""    ,	"""This test is supposed to run on GPU"""							)
     def  SCREAMING_SNAKE_CASE__       (    self					:				Union[str, Any]							):
        # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
        __lowerCamelCase   ,     __lowerCamelCase:				Optional[Any]     =	UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""    ,	output_loading_info=SCREAMING_SNAKE_CASE_							)
        model_accelerate.to(SCREAMING_SNAKE_CASE_							)
        model_accelerate.eval()

        __lowerCamelCase:				Optional[Any]     =	torch.randn(
            1    ,	model_accelerate.config.in_channels    ,	model_accelerate.config.sample_size    ,	model_accelerate.config.sample_size    ,	generator=torch.manual_seed(0							)    ,	)
        __lowerCamelCase:				List[str]     =	noise.to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Union[str, Any]     =	torch.tensor([10] * noise.shape[0]							).to(SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				Optional[int]     =	model_accelerate(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							)["""sample"""]

        # two models don't need to stay in the device at the same time
        del model_accelerate
        torch.cuda.empty_cache()
        gc.collect()

        __lowerCamelCase   ,     __lowerCamelCase:				List[str]     =	UNetaDModel.from_pretrained(
            """fusing/unet-ldm-dummy-update"""    ,	output_loading_info=SCREAMING_SNAKE_CASE_    ,	low_cpu_mem_usage=SCREAMING_SNAKE_CASE_							)
        model_normal_load.to(SCREAMING_SNAKE_CASE_							)
        model_normal_load.eval()
        __lowerCamelCase:				List[Any]     =	model_normal_load(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							)["""sample"""]

        assert torch_all_close(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	rtol=1E-3							)







     def  SCREAMING_SNAKE_CASE__       (    self					:				int							):
        __lowerCamelCase:				List[str]     =	UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""							)
        model.eval()
        model.to(SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				int     =	torch.randn(
            1    ,	model.config.in_channels    ,	model.config.sample_size    ,	model.config.sample_size    ,	generator=torch.manual_seed(0							)    ,	)
        __lowerCamelCase:				List[Any]     =	noise.to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Union[str, Any]     =	torch.tensor([10] * noise.shape[0]							).to(SCREAMING_SNAKE_CASE_							)

        with torch.no_grad():
           __lowerCamelCase:				Optional[Any]     =	model(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							).sample

        __lowerCamelCase:				Optional[Any]     =	output[0, -1, -3:, -3:].flatten().cpu()
        # fmt: off
        __lowerCamelCase:				Any     =	torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800]							)
        # fmt: on

        self.assertTrue(torch_all_close(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	rtol=1E-3							)							)



class        a					(			_UpperCAmelCase     ,_UpperCAmelCase     ,unittest.TestCase   ):
     UpperCAmelCase__ :  int						      =	UNetaDModel
     UpperCAmelCase__ :  int						      =	"sample"





     @property
     def  SCREAMING_SNAKE_CASE__       (    self					:				Dict    ,	SCREAMING_SNAKE_CASE_					:				List[Any]=(32, 32)							):
        __lowerCamelCase:				Optional[int]     =	4
        __lowerCamelCase:				int     =	3

        __lowerCamelCase:				Tuple     =	floats_tensor((batch_size, num_channels) + sizes							).to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Tuple     =	torch.tensor(batch_size * [10]							).to(dtype=torch.intaa    ,	device=SCREAMING_SNAKE_CASE_							)

        return {"sample": noise, "timestep": time_step}





     @property
     def  SCREAMING_SNAKE_CASE__       (    self					:				Dict							):
        return (3, 32, 32)





     @property
     def  SCREAMING_SNAKE_CASE__       (    self					:				Union[str, Any]							):
        return (3, 32, 32)





     def  SCREAMING_SNAKE_CASE__       (    self					:				int							):
        __lowerCamelCase:				List[Any]     =	{
            """block_out_channels""": [32, 64, 64, 64],
            """in_channels""": 3,
            """layers_per_block""": 1,
            """out_channels""": 3,
            """time_embedding_type""": """fourier""",
            """norm_eps""": 1E-6,
            """mid_block_scale_factor""": math.sqrt(2.0							),
            """norm_num_groups""": None,
            """down_block_types""": [
                """SkipDownBlock2D""",
                """AttnSkipDownBlock2D""",
                """SkipDownBlock2D""",
                """SkipDownBlock2D""",
            ],
            """up_block_types""": [
                """SkipUpBlock2D""",
                """SkipUpBlock2D""",
                """AttnSkipUpBlock2D""",
                """SkipUpBlock2D""",
            ],
        }
        __lowerCamelCase:				Dict     =	self.dummy_input
        return init_dict, inputs_dict





     @slow
     def  SCREAMING_SNAKE_CASE__       (    self					:				Tuple							):
        __lowerCamelCase   ,     __lowerCamelCase:				Any     =	UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""    ,	output_loading_info=SCREAMING_SNAKE_CASE_							)
        self.assertIsNotNone(SCREAMING_SNAKE_CASE_							)
        self.assertEqual(len(loading_info["""missing_keys"""]							)    ,	0							)

        model.to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				List[str]     =	self.dummy_input
        __lowerCamelCase:				int     =	floats_tensor((4, 3) + (256, 256)							).to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Union[str, Any]     =	noise
        __lowerCamelCase:				str     =	model(**SCREAMING_SNAKE_CASE_							)

        assert image is not None, "Make sure output is not None"





     @slow
     def  SCREAMING_SNAKE_CASE__       (    self					:				str							):
        __lowerCamelCase:				Optional[int]     =	UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""							)
        model.to(SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				str     =	4
        __lowerCamelCase:				List[str]     =	3
        __lowerCamelCase:				str     =	(256, 256)

        __lowerCamelCase:				int     =	torch.ones((batch_size, num_channels) + sizes							).to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				int     =	torch.tensor(batch_size * [1E-4]							).to(SCREAMING_SNAKE_CASE_							)

        with torch.no_grad():
           __lowerCamelCase:				List[Any]     =	model(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							).sample

        __lowerCamelCase:				Optional[Any]     =	output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        __lowerCamelCase:				int     =	torch.tensor([-4842.8691, -6499.6631, -3800.1953, -7978.2686, -1_0980.7129, -2_0028.8535, 8148.2822, 2342.2905, 567.7608]							)
        # fmt: on

        self.assertTrue(torch_all_close(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	rtol=1E-2							)							)





     def  SCREAMING_SNAKE_CASE__       (    self					:				Tuple							):
        __lowerCamelCase:				int     =	UNetaDModel.from_pretrained("""fusing/ncsnpp-ffhq-ve-dummy-update"""							)
        model.to(SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				Dict     =	4
        __lowerCamelCase:				Optional[int]     =	3
        __lowerCamelCase:				Optional[int]     =	(32, 32)

        __lowerCamelCase:				Dict     =	torch.ones((batch_size, num_channels) + sizes							).to(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				str     =	torch.tensor(batch_size * [1E-4]							).to(SCREAMING_SNAKE_CASE_							)

        with torch.no_grad():
           __lowerCamelCase:				List[str]     =	model(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							).sample

        __lowerCamelCase:				List[str]     =	output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        __lowerCamelCase:				List[Any]     =	torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256]							)
        # fmt: on

        self.assertTrue(torch_all_close(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	rtol=1E-2							)							)







     def  SCREAMING_SNAKE_CASE__       (    self					:				int							):
        # not required for this model
        pass






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()
_A		:				List[Any]									=       logging.get_logger(__name__)



def     __lowerCAmelCase						(							snake_case	:    Any						,      snake_case	:    int )       ->							Dict:
   __lowerCamelCase:				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						(							snake_case	:    str						,      snake_case	:    Union[str, Any] )       ->							Tuple:
   for i in range(encoder_config.num_hidden_layers ):
      # queries, keys and values (only weights, no biases)
      __lowerCamelCase:				Union[str, Any]     =	state_dict.pop(f'encoder.deit.blocks.{i}.attn.qkv.weight' )

      __lowerCamelCase:				Any     =	in_proj_weight[
          : encoder_config.hidden_size, :
      ]
      __lowerCamelCase:				Dict     =	in_proj_weight[
          encoder_config.hidden_size : encoder_config.hidden_size * 2, :
      ]
      __lowerCamelCase:				Any     =	in_proj_weight[
          -encoder_config.hidden_size :, :
      ]



def     __lowerCAmelCase						(							snake_case	:    Dict						,      snake_case	:    Union[str, Any]						,      snake_case	:    Optional[int] )       ->							int:
   __lowerCamelCase:				List[str]     =	dct.pop(snake_case )
   __lowerCamelCase:				Tuple     =	val



def     __lowerCAmelCase						(							snake_case	:    Dict )       ->							Any:
   if "handwritten" in checkpoint_url:
      __lowerCamelCase:				Tuple     =	"""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:
      __lowerCamelCase:				Union[str, Any]     =	"""https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg"""
   __lowerCamelCase:				Union[str, Any]     =	Image.open(requests.get(snake_case						,      stream=snake_case ).raw ).convert("""RGB""" )
   return im



@torch.no_grad()
def     __lowerCAmelCase						(							snake_case	:    List[Any]						,      snake_case	:    Any )       ->							List[str]:
   __lowerCamelCase:				Union[str, Any]     =	ViTConfig(image_size=384						,      qkv_bias=snake_case )
   __lowerCamelCase:				Optional[int]     =	TrOCRConfig()

   # size of the architecture
   if "base" in checkpoint_url:
      __lowerCamelCase:				Optional[int]     =	768
   elif "large" in checkpoint_url:
      # use ViT-large encoder
      __lowerCamelCase:				Optional[int]     =	1024
      __lowerCamelCase:				int     =	4096
      __lowerCamelCase:				Tuple     =	24
      __lowerCamelCase:				Optional[Any]     =	16
      __lowerCamelCase:				str     =	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:
      __lowerCamelCase:				Any     =	False
      __lowerCamelCase:				Optional[Any]     =	"""relu"""
      __lowerCamelCase:				List[Any]     =	1024
      __lowerCamelCase:				List[str]     =	True
      __lowerCamelCase:				int     =	False
      __lowerCamelCase:				int     =	False

   # load HuggingFace model
   __lowerCamelCase:				Any     =	ViTModel(snake_case						,      add_pooling_layer=snake_case )
   __lowerCamelCase:				Dict     =	TrOCRForCausalLM(snake_case )
   __lowerCamelCase:				Any     =	VisionEncoderDecoderModel(encoder=snake_case						,      decoder=snake_case )
   model.eval()

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

   __lowerCamelCase:				Optional[Any]     =	create_rename_keys(snake_case						,      snake_case )
   for src, dest in rename_keys:
      rename_key(snake_case						,      snake_case						,      snake_case )
   read_in_q_k_v(snake_case						,      snake_case )

   # 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():
      __lowerCamelCase:				int     =	state_dict.pop(snake_case )
      if key.startswith("""decoder""" ) and "output_projection" not in key:
         __lowerCamelCase:				Any     =	val
      else:
         __lowerCamelCase:				List[str]     =	val

    # load state dict
   model.load_state_dict(snake_case )

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

   __lowerCamelCase:				Optional[int]     =	processor(images=prepare_img(snake_case )						,      return_tensors="""pt""" ).pixel_values

   # verify logits
   __lowerCamelCase:				Tuple     =	torch.tensor([[model.config.decoder.decoder_start_token_id]] )
   __lowerCamelCase:				Dict     =	model(pixel_values=snake_case						,      decoder_input_ids=snake_case )
   __lowerCamelCase:				Any     =	outputs.logits

   __lowerCamelCase:				Dict     =	torch.Size([1, 1, 50265] )
   if "trocr-base-handwritten" in checkpoint_url:
      __lowerCamelCase:				Tuple     =	torch.tensor(
          [-1.4_502, -4.6_683, -0.5_347, -2.9_291, 9.1_435, -3.0_571, 8.9_764, 1.7_560, 8.7_358, -1.5_311] )
   elif "trocr-large-handwritten" in checkpoint_url:
      __lowerCamelCase:				Optional[int]     =	torch.tensor(
          [-2.6_437, -1.3_129, -2.2_596, -5.3_455, 6.3_539, 1.7_604, 5.4_991, 1.4_702, 5.6_113, 2.0_170] )
   elif "trocr-base-printed" in checkpoint_url:
      __lowerCamelCase:				Union[str, Any]     =	torch.tensor(
          [-5.6_816, -5.8_388, 1.1_398, -6.9_034, 6.8_505, -2.4_393, 1.2_284, -1.0_232, -1.9_661, -3.9_210] )
   elif "trocr-large-printed" in checkpoint_url:
      __lowerCamelCase:				Union[str, Any]     =	torch.tensor(
          [-6.0_162, -7.0_959, 4.4_155, -5.1_063, 7.0_468, -3.1_631, 2.6_466, -0.3_081, -0.8_106, -1.7_535] )

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

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


if __name__ == "__main__":
   _A		:				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.'''
   )
   _A		:				Optional[int]									=       parser.parse_args()
   convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


