
"""simple docstring"""

import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
    from transformers.models.esm.modeling_esm import (
        ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
        EsmEmbeddings,
        create_position_ids_from_input_ids,
    )
class      _UpperCAmelCase							:




      def __init__(							self   ,       lowercase_   ,       lowercase_=1_3   ,       lowercase_=7   ,       lowercase_=False   ,       lowercase_=True   ,       lowercase_=False   ,       lowercase_=True   ,       lowercase_=3_3   ,       lowercase_=3_2   ,       lowercase_=5   ,       lowercase_=4   ,       lowercase_=3_7   ,       lowercase_="gelu"   ,       lowercase_=0.1   ,       lowercase_=0.1   ,       lowercase_=5_1_2   ,       lowercase_=1_6   ,       lowercase_=2   ,       lowercase_=0.0_2   ,       lowercase_=3   ,       lowercase_=4   ,       lowercase_=None   ,       )						->					Tuple:
          UpperCAmelCase									=					parent
          UpperCAmelCase									=					batch_size
          UpperCAmelCase									=					seq_length
          UpperCAmelCase									=					is_training
          UpperCAmelCase									=					use_input_mask
          UpperCAmelCase									=					use_token_type_ids
          UpperCAmelCase									=					use_labels
          UpperCAmelCase									=					vocab_size
          UpperCAmelCase									=					hidden_size
          UpperCAmelCase									=					num_hidden_layers
          UpperCAmelCase									=					num_attention_heads
          UpperCAmelCase									=					intermediate_size
          UpperCAmelCase									=					hidden_act
          UpperCAmelCase									=					hidden_dropout_prob
          UpperCAmelCase									=					attention_probs_dropout_prob
          UpperCAmelCase									=					max_position_embeddings
          UpperCAmelCase									=					type_vocab_size
          UpperCAmelCase									=					type_sequence_label_size
          UpperCAmelCase									=					initializer_range
          UpperCAmelCase									=					num_labels
          UpperCAmelCase									=					num_choices
          UpperCAmelCase									=					scope




      def     a_      (							self							)						->					List[Any]:
          UpperCAmelCase									=					ids_tensor([self.batch_size, self.seq_length]   ,       self.vocab_size							)

          UpperCAmelCase									=					None
          if self.use_input_mask:
              UpperCAmelCase									=					random_attention_mask([self.batch_size, self.seq_length]							)

          UpperCAmelCase									=					None
          UpperCAmelCase									=					None
          UpperCAmelCase									=					None
          if self.use_labels:
              UpperCAmelCase									=					ids_tensor([self.batch_size]   ,       self.type_sequence_label_size							)
              UpperCAmelCase									=					ids_tensor([self.batch_size, self.seq_length]   ,       self.num_labels							)
              UpperCAmelCase									=					ids_tensor([self.batch_size]   ,       self.num_choices							)

          UpperCAmelCase									=					self.get_config()

          return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels




      def     a_      (							self							)						->					Union[str, Any]:
          return EsmConfig(
              vocab_size=self.vocab_size   ,       hidden_size=self.hidden_size   ,       pad_token_id=1   ,       num_hidden_layers=self.num_hidden_layers   ,       num_attention_heads=self.num_attention_heads   ,       intermediate_size=self.intermediate_size   ,       hidden_act=self.hidden_act   ,       hidden_dropout_prob=self.hidden_dropout_prob   ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,       max_position_embeddings=self.max_position_embeddings   ,       type_vocab_size=self.type_vocab_size   ,       initializer_range=self.initializer_range   ,       )




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					int:
          UpperCAmelCase									=					EsmModel(config=__a							)
          model.to(__a							)
          model.eval()
          UpperCAmelCase									=					model(__a   ,       attention_mask=__a							)
          UpperCAmelCase									=					model(__a							)
          UpperCAmelCase									=					model(__a							)

          self.parent.assertEqual(result.last_hidden_state.shape   ,       (self.batch_size, self.seq_length, self.hidden_size)							)
          self.parent.assertEqual(result.pooler_output.shape   ,       (self.batch_size, self.hidden_size)							)




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					int:
          UpperCAmelCase									=					EsmForMaskedLM(config=__a							)
          model.to(__a							)
          model.eval()
          UpperCAmelCase									=					model(__a   ,       attention_mask=__a   ,       labels=__a							)
          self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.seq_length, self.vocab_size)							)




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					List[str]:
          UpperCAmelCase									=					self.num_labels
          UpperCAmelCase									=					EsmForTokenClassification(config=__a							)
          model.to(__a							)
          model.eval()
          UpperCAmelCase									=					model(__a   ,       attention_mask=__a   ,       labels=__a							)
          self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.seq_length, self.num_labels)							)




      def     a_      (							self							)						->					Optional[int]:
          UpperCAmelCase									=					self.prepare_config_and_inputs()
          (
              UpperCAmelCase			
          )									=					config_and_inputs
          UpperCAmelCase									=					{"""input_ids""": input_ids, """attention_mask""": input_mask}
          return config, inputs_dict

@require_torch
class      _UpperCAmelCase							(		__lowercase     ,			__lowercase     ,			unittest.TestCase							):
      __SCREAMING_SNAKE_CASE   :							Any							      =	False

      __SCREAMING_SNAKE_CASE   :							Dict							      =	(
          (
              EsmForMaskedLM,
              EsmModel,
              EsmForSequenceClassification,
              EsmForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      __SCREAMING_SNAKE_CASE   :							Dict							      =	()
      __SCREAMING_SNAKE_CASE   :							Optional[int]							      =	(
          {
              '''feature-extraction''': EsmModel,
              '''fill-mask''': EsmForMaskedLM,
              '''text-classification''': EsmForSequenceClassification,
              '''token-classification''': EsmForTokenClassification,
              '''zero-shot''': EsmForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __SCREAMING_SNAKE_CASE   :							Any							      =	True




      def     a_      (							self							)						->					Optional[int]:
          UpperCAmelCase									=					EsmModelTester(self							)
          UpperCAmelCase									=					ConfigTester(self   ,       config_class=__a   ,       hidden_size=3_7							)




      def     a_      (							self							)						->					Optional[Any]:
          self.config_tester.run_common_tests()




      def     a_      (							self							)						->					Optional[Any]:
          UpperCAmelCase									=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__a							)




      def     a_      (							self							)						->					Any:
          UpperCAmelCase									=					self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
              UpperCAmelCase									=					type
              self.model_tester.create_and_check_model(*__a							)




      def     a_      (							self							)						->					Tuple:
          UpperCAmelCase									=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*__a							)




      def     a_      (							self							)						->					Optional[Any]:
          UpperCAmelCase									=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*__a							)




      @slow
      def     a_      (							self							)						->					List[Any]:
          for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              UpperCAmelCase									=					EsmModel.from_pretrained(__a							)
              self.assertIsNotNone(__a							)




      def     a_      (							self							)						->					List[str]:
          UpperCAmelCase									=					self.model_tester.prepare_config_and_inputs()[0]
          UpperCAmelCase									=					EsmEmbeddings(config=__a							)

          UpperCAmelCase									=					torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]							)
          UpperCAmelCase									=					torch.as_tensor(
              [
                  [
                      0 + model.padding_idx + 1,
                      1 + model.padding_idx + 1,
                      2 + model.padding_idx + 1,
                      model.padding_idx,
                  ]
              ]							)
          UpperCAmelCase									=					create_position_ids_from_input_ids(__a   ,       model.padding_idx							)
          self.assertEqual(position_ids.shape   ,       expected_positions.shape							)
          self.assertTrue(torch.all(torch.eq(__a   ,       __a							)							)							)




      def     a_      (							self							)						->					int:
          UpperCAmelCase									=					self.model_tester.prepare_config_and_inputs()[0]
          UpperCAmelCase									=					EsmEmbeddings(config=__a							)

          UpperCAmelCase									=					torch.empty(2   ,       4   ,       3_0							)
          UpperCAmelCase									=					[
              0 + embeddings.padding_idx + 1,
              1 + embeddings.padding_idx + 1,
              2 + embeddings.padding_idx + 1,
              3 + embeddings.padding_idx + 1,
          ]
          UpperCAmelCase									=					torch.as_tensor([expected_single_positions, expected_single_positions]							)
          UpperCAmelCase									=					embeddings.create_position_ids_from_inputs_embeds(__a							)
          self.assertEqual(position_ids.shape   ,       expected_positions.shape							)
          self.assertTrue(torch.all(torch.eq(__a   ,       __a							)							)							)




      @unittest.skip('Esm does not support embedding resizing'							)
      def     a_      (							self							)						->					Optional[Any]:
          pass




      @unittest.skip('Esm does not support embedding resizing'							)
      def     a_      (							self							)						->					Dict:
          pass




      @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'							)
      def     a_      (							self							)						->					Dict:
          pass

@require_torch
class      _UpperCAmelCase							(		__lowercase							):




      @slow
      def     a_      (							self							)						->					Union[str, Any]:
          with torch.no_grad():
              UpperCAmelCase									=					EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'							)
              model.eval()
              UpperCAmelCase									=					torch.tensor([[0, 1, 2, 3, 4, 5]]							)
              UpperCAmelCase									=					model(__a							)[0]

              UpperCAmelCase									=					3_3

              UpperCAmelCase									=					torch.Size((1, 6, vocab_size)							)
              self.assertEqual(output.shape   ,       __a							)

              UpperCAmelCase									=					torch.tensor(
                  [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]]							)
              self.assertTrue(torch.allclose(output[:, :3, :3]   ,       __a   ,       atol=1E-4							)							)




      @slow
      def     a_      (							self							)						->					str:
          with torch.no_grad():
              UpperCAmelCase									=					EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'							)
              model.eval()

              UpperCAmelCase									=					torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]							)
              UpperCAmelCase									=					model(__a							)[0]
              # compare the actual values for a slice.
              UpperCAmelCase									=					torch.tensor(
                  [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]]							)
              self.assertTrue(torch.allclose(output[:, :3, :3]   ,       __a   ,       atol=1E-4							)							)





"""simple docstring"""

import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_url
from PIL import Image

from transformers import DPTConfig, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE_    					=						logging.get_logger(__name__)




def     lowercase__						(					lowerCAmelCase				:					int ) ->      Any:






    """simple docstring"""




    UpperCAmelCase									=					DPTConfig()

    if "large" in checkpoint_url:
        UpperCAmelCase									=					1_024
        UpperCAmelCase									=					4_096
        UpperCAmelCase									=					24
        UpperCAmelCase									=					16
        UpperCAmelCase									=					[5, 11, 17, 23]
        UpperCAmelCase									=					[256, 512, 1_024, 1_024]
        UpperCAmelCase									=					(1, 384, 384)

    if "ade" in checkpoint_url:
        UpperCAmelCase									=					True

        UpperCAmelCase									=					150
        UpperCAmelCase									=					'huggingface/label-files'
        UpperCAmelCase									=					'ade20k-id2label.json'
        UpperCAmelCase									=					json.load(open(cached_download(hf_hub_url(lowerCAmelCase ,				lowerCAmelCase ,				repo_type='dataset' ) ) ,				'r' ) )
        UpperCAmelCase									=					{int(lowerCAmelCase ): v for k, v in idalabel.items()}
        UpperCAmelCase									=					idalabel
        UpperCAmelCase									=					{v: k for k, v in idalabel.items()}
        UpperCAmelCase									=					[1, 150, 480, 480]

    return config, expected_shape




def     lowercase__						(					lowerCAmelCase				:					str ) ->      Any:






    """simple docstring"""




    UpperCAmelCase									=					['pretrained.model.head.weight', 'pretrained.model.head.bias']
    for k in ignore_keys:
        state_dict.pop(lowerCAmelCase ,				lowerCAmelCase )




def     lowercase__						(					lowerCAmelCase				:					Union[str, Any] ) ->      List[Any]:






    """simple docstring"""




    if (
        "pretrained.model" in name
        and "cls_token" not in name
        and "pos_embed" not in name
        and "patch_embed" not in name
    ):
        UpperCAmelCase									=					name.replace('pretrained.model' ,				'dpt.encoder' )
    if "pretrained.model" in name:
        UpperCAmelCase									=					name.replace('pretrained.model' ,				'dpt.embeddings' )
    if "patch_embed" in name:
        UpperCAmelCase									=					name.replace('patch_embed' ,				'patch_embeddings' )
    if "pos_embed" in name:
        UpperCAmelCase									=					name.replace('pos_embed' ,				'position_embeddings' )
    if "attn.proj" in name:
        UpperCAmelCase									=					name.replace('attn.proj' ,				'attention.output.dense' )
    if "proj" in name and "project" not in name:
        UpperCAmelCase									=					name.replace('proj' ,				'projection' )
    if "blocks" in name:
        UpperCAmelCase									=					name.replace('blocks' ,				'layer' )
    if "mlp.fc1" in name:
        UpperCAmelCase									=					name.replace('mlp.fc1' ,				'intermediate.dense' )
    if "mlp.fc2" in name:
        UpperCAmelCase									=					name.replace('mlp.fc2' ,				'output.dense' )
    if "norm1" in name:
        UpperCAmelCase									=					name.replace('norm1' ,				'layernorm_before' )
    if "norm2" in name:
        UpperCAmelCase									=					name.replace('norm2' ,				'layernorm_after' )
    if "scratch.output_conv" in name:
        UpperCAmelCase									=					name.replace('scratch.output_conv' ,				'head' )
    if "scratch" in name:
        UpperCAmelCase									=					name.replace('scratch' ,				'neck' )
    if "layer1_rn" in name:
        UpperCAmelCase									=					name.replace('layer1_rn' ,				'convs.0' )
    if "layer2_rn" in name:
        UpperCAmelCase									=					name.replace('layer2_rn' ,				'convs.1' )
    if "layer3_rn" in name:
        UpperCAmelCase									=					name.replace('layer3_rn' ,				'convs.2' )
    if "layer4_rn" in name:
        UpperCAmelCase									=					name.replace('layer4_rn' ,				'convs.3' )
    if "refinenet" in name:
        UpperCAmelCase									=					int(name[len('neck.refinenet' ) : len('neck.refinenet' ) + 1] )
        # tricky here: we need to map 4 to 0, 3 to 1, 2 to 2 and 1 to 3
        UpperCAmelCase									=					name.replace(F"refinenet{layer_idx}" ,				F"fusion_stage.layers.{abs(layer_idx-4 )}" )
    if "out_conv" in name:
        UpperCAmelCase									=					name.replace('out_conv' ,				'projection' )
    if "resConfUnit1" in name:
        UpperCAmelCase									=					name.replace('resConfUnit1' ,				'residual_layer1' )
    if "resConfUnit2" in name:
        UpperCAmelCase									=					name.replace('resConfUnit2' ,				'residual_layer2' )
    if "conv1" in name:
        UpperCAmelCase									=					name.replace('conv1' ,				'convolution1' )
    if "conv2" in name:
        UpperCAmelCase									=					name.replace('conv2' ,				'convolution2' )
    # readout blocks
    if "pretrained.act_postprocess1.0.project.0" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess1.0.project.0' ,				'neck.reassemble_stage.readout_projects.0.0' )
    if "pretrained.act_postprocess2.0.project.0" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess2.0.project.0' ,				'neck.reassemble_stage.readout_projects.1.0' )
    if "pretrained.act_postprocess3.0.project.0" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess3.0.project.0' ,				'neck.reassemble_stage.readout_projects.2.0' )
    if "pretrained.act_postprocess4.0.project.0" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess4.0.project.0' ,				'neck.reassemble_stage.readout_projects.3.0' )
    # resize blocks
    if "pretrained.act_postprocess1.3" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess1.3' ,				'neck.reassemble_stage.layers.0.projection' )
    if "pretrained.act_postprocess1.4" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess1.4' ,				'neck.reassemble_stage.layers.0.resize' )
    if "pretrained.act_postprocess2.3" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess2.3' ,				'neck.reassemble_stage.layers.1.projection' )
    if "pretrained.act_postprocess2.4" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess2.4' ,				'neck.reassemble_stage.layers.1.resize' )
    if "pretrained.act_postprocess3.3" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess3.3' ,				'neck.reassemble_stage.layers.2.projection' )
    if "pretrained.act_postprocess4.3" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess4.3' ,				'neck.reassemble_stage.layers.3.projection' )
    if "pretrained.act_postprocess4.4" in name:
        UpperCAmelCase									=					name.replace('pretrained.act_postprocess4.4' ,				'neck.reassemble_stage.layers.3.resize' )
    if "pretrained" in name:
        UpperCAmelCase									=					name.replace('pretrained' ,				'dpt' )
    if "bn" in name:
        UpperCAmelCase									=					name.replace('bn' ,				'batch_norm' )
    if "head" in name:
        UpperCAmelCase									=					name.replace('head' ,				'head.head' )
    if "encoder.norm" in name:
        UpperCAmelCase									=					name.replace('encoder.norm' ,				'layernorm' )
    if "auxlayer" in name:
        UpperCAmelCase									=					name.replace('auxlayer' ,				'auxiliary_head.head' )

    return name




def     lowercase__						(					lowerCAmelCase				:					List[Any] ,				lowerCAmelCase				:					int ) ->      int:






    """simple docstring"""




    for i in range(config.num_hidden_layers ):
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        UpperCAmelCase									=					state_dict.pop(F"dpt.encoder.layer.{i}.attn.qkv.weight" )
        UpperCAmelCase									=					state_dict.pop(F"dpt.encoder.layer.{i}.attn.qkv.bias" )
        # next, add query, keys and values (in that order) to the state dict
        UpperCAmelCase									=					in_proj_weight[: config.hidden_size, :]
        UpperCAmelCase									=					in_proj_bias[: config.hidden_size]
        UpperCAmelCase									=					in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        UpperCAmelCase									=					in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        UpperCAmelCase									=					in_proj_weight[
            -config.hidden_size :, :
        ]
        UpperCAmelCase									=					in_proj_bias[-config.hidden_size :]




def     lowercase__						(					) ->      Optional[Any]:






    """simple docstring"""




    UpperCAmelCase									=					'http://images.cocodataset.org/val2017/000000039769.jpg'
    UpperCAmelCase									=					Image.open(requests.get(lowerCAmelCase ,				stream=lowerCAmelCase ).raw )
    return im




@torch.no_grad()
def     lowercase__						(					lowerCAmelCase				:					Dict ,				lowerCAmelCase				:					Optional[int] ,				lowerCAmelCase				:					List[str] ,				lowerCAmelCase				:					Union[str, Any] ) ->      Union[str, Any]:






    """simple docstring"""




    UpperCAmelCase			,							UpperCAmelCase									=					get_dpt_config(lowerCAmelCase )
    # load original state_dict from URL
    UpperCAmelCase									=					torch.hub.load_state_dict_from_url(lowerCAmelCase ,				map_location='cpu' )
    # remove certain keys
    remove_ignore_keys_(lowerCAmelCase )
    # rename keys
    for key in state_dict.copy().keys():
        UpperCAmelCase									=					state_dict.pop(lowerCAmelCase )
        UpperCAmelCase									=					val
    # read in qkv matrices
    read_in_q_k_v(lowerCAmelCase ,				lowerCAmelCase )

    # load HuggingFace model
    UpperCAmelCase									=					DPTForSemanticSegmentation(lowerCAmelCase ) if 'ade' in checkpoint_url else DPTForDepthEstimation(lowerCAmelCase )
    model.load_state_dict(lowerCAmelCase )
    model.eval()

    # Check outputs on an image
    UpperCAmelCase									=					480 if 'ade' in checkpoint_url else 384
    UpperCAmelCase									=					DPTImageProcessor(size=lowerCAmelCase )

    UpperCAmelCase									=					prepare_img()
    UpperCAmelCase									=					image_processor(lowerCAmelCase ,				return_tensors='pt' )

    # forward pass
    UpperCAmelCase									=					model(**lowerCAmelCase ).logits if 'ade' in checkpoint_url else model(**lowerCAmelCase ).predicted_depth

    # Assert logits
    UpperCAmelCase									=					torch.tensor([[6.3199, 6.3629, 6.4148], [6.3850, 6.3615, 6.4166], [6.3519, 6.3176, 6.3575]] )
    if "ade" in checkpoint_url:
        UpperCAmelCase									=					torch.tensor([[4.0480, 4.2420, 4.4360], [4.3124, 4.5693, 4.8261], [4.5768, 4.8965, 5.2163]] )
    assert outputs.shape == torch.Size(lowerCAmelCase )
    assert (
        torch.allclose(outputs[0, 0, :3, :3] ,				lowerCAmelCase ,				atol=1E-4 )
        if "ade" in checkpoint_url
        else torch.allclose(outputs[0, :3, :3] ,				lowerCAmelCase )
    )

    Path(lowerCAmelCase ).mkdir(exist_ok=lowerCAmelCase )
    print(F"Saving model to {pytorch_dump_folder_path}" )
    model.save_pretrained(lowerCAmelCase )
    print(F"Saving image processor to {pytorch_dump_folder_path}" )
    image_processor.save_pretrained(lowerCAmelCase )

    if push_to_hub:
        print('Pushing model to hub...' )
        model.push_to_hub(
            repo_path_or_name=Path(lowerCAmelCase ,				lowerCAmelCase ) ,				organization='nielsr' ,				commit_message='Add model' ,				use_temp_dir=lowerCAmelCase ,				)
        image_processor.push_to_hub(
            repo_path_or_name=Path(lowerCAmelCase ,				lowerCAmelCase ) ,				organization='nielsr' ,				commit_message='Add image processor' ,				use_temp_dir=lowerCAmelCase ,				)


if __name__ == "__main__":
   SCREAMING_SNAKE_CASE_    					=						argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '''--checkpoint_url''',
       default='''https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt''',
       type=str,
       help='''URL of the original DPT checkpoint you\'d like to convert.''',
   )
   parser.add_argument(
       '''--pytorch_dump_folder_path''',
       default=None,
       type=str,
       required=True,
       help='''Path to the output PyTorch model directory.''',
   )
   parser.add_argument(
       '''--push_to_hub''',
       action='''store_true''',
   )
   parser.add_argument(
       '''--model_name''',
       default='''dpt-large''',
       type=str,
       help='''Name of the model, in case you\'re pushing to the hub.''',
   )

   SCREAMING_SNAKE_CASE_    					=						parser.parse_args()
   convert_dpt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub, args.model_name)



