



'''simple docstring'''





import unittest

from transformers import JukeboxTokenizer
from transformers.testing_utils import require_torch






class __A     (    unittest.TestCase		):



 '''simple docstring'''







 __lowerCamelCase       :			Tuple		=			JukeboxTokenizer
 __lowerCamelCase       :			Dict		=			{
     'artist': 'Zac Brown Band',
     'genres': 'Country',
     'lyrics': 'I met a traveller from an antique land,\n        Who said "Two vast and trunkless legs of stone\n        Stand in the desert. . . . Near them, on the sand,\n        Half sunk a shattered visage lies, whose frown,\n        And wrinkled lip, and sneer of cold command,\n        Tell that its sculptor well those passions read\n        Which yet survive, stamped on these lifeless things,\n        The hand that mocked them, and the heart that fed;\n        And on the pedestal, these words appear:\n        My name is Ozymandias, King of Kings;\n        Look on my Works, ye Mighty, and despair!\n        Nothing beside remains. Round the decay\n        Of that colossal Wreck, boundless and bare\n        The lone and level sands stretch far away\n        ',
 }



 @require_torch
 def    a__    (self			)       ->	Dict:


      """simple docstring"""



      import torch

      _a           =    JukeboxTokenizer.from_pretrained('''openai/jukebox-1b-lyrics'''			)
      _a           =    tokenizer(**self.metas			)['''input_ids''']
      # fmt: off
      _a           =    [
          torch.tensor([[
              0, 0, 0, 7_169, 507, 9, 76, 39, 31, 46, 76, 27,
              76, 46, 44, 27, 48, 31, 38, 38, 31, 44, 76, 32,
              44, 41, 39, 76, 27, 40, 76, 27, 40, 46, 35, 43,
              47, 31, 76, 38, 27, 40, 30, 64, 78, 76, 76, 76,
              76, 76, 76, 76, 76, 23, 34, 41, 76, 45, 27, 35,
              30, 76, 71, 20, 49, 41, 76, 48, 27, 45, 46, 76,
              27, 40, 30, 76, 46, 44, 47, 40, 37, 38, 31, 45,
              45, 76, 38, 31, 33, 45, 76, 41, 32, 76, 45, 46,
              41, 40, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
              19, 46, 27, 40, 30, 76, 35, 40, 76, 46, 34, 31,
              76, 30, 31, 45, 31, 44, 46, 63, 76, 63, 76, 63,
              76, 63, 76, 14, 31, 27, 44, 76, 46, 34, 31, 39,
              64, 76, 41, 40, 76, 46, 34, 31, 76, 45, 27, 40,
              30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76, 8,
              27, 38, 32, 76, 45, 47, 40, 37, 76, 27, 76, 45,
              34, 27, 46, 46, 31, 44, 31, 30, 76, 48, 35, 45,
              27, 33, 31, 76, 38, 35, 31, 45, 64, 76, 49, 34,
              41, 45, 31, 76, 32, 44, 41, 49, 40, 64, 78, 76,
              76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76, 49,
              44, 35, 40, 37, 38, 31, 30, 76, 38, 35, 42, 64,
              76, 27, 40, 30, 76, 45, 40, 31, 31, 44, 76, 41,
              32, 76, 29, 41, 38, 30, 76, 29, 41, 39, 39, 27,
              40, 30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76,
              20, 31, 38, 38, 76, 46, 34, 27, 46, 76, 35, 46,
              45, 76, 45, 29, 47, 38, 42, 46, 41, 44, 76, 49,
              31, 38, 38, 76, 46, 34, 41, 45, 31, 76, 42, 27,
              45, 45, 35, 41, 40, 45, 76, 44, 31, 27, 30, 78,
              76, 76, 76, 76, 76, 76, 76, 76, 23, 34, 35, 29,
              34, 76, 51, 31, 46, 76, 45, 47, 44, 48, 35, 48,
              31, 64, 76, 45, 46, 27, 39, 42, 31, 30, 76, 41,
              40, 76, 46, 34, 31, 45, 31, 76, 38, 35, 32, 31,
              38, 31, 45, 45, 76, 46, 34, 35, 40, 33, 45, 64,
              78, 76, 76, 76, 76, 76, 76, 76, 76, 20, 34, 31,
              76, 34, 27, 40, 30, 76, 46, 34, 27, 46, 76, 39,
              41, 29, 37, 31, 30, 76, 46, 34, 31, 39, 64, 76,
              27, 40, 30, 76, 46, 34, 31, 76, 34, 31, 27, 44,
              46, 76, 46, 34, 27, 46, 76, 32, 31, 30, 66, 78,
              76, 76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76,
              41, 40, 76, 46, 34, 31, 76, 42, 31, 30, 31, 45,
              46, 27, 38, 64, 76, 46, 34, 31, 45, 31, 76, 49,
              41, 44, 30, 45, 76, 27, 42, 42, 31, 27, 44, 65,
              78, 76, 76, 76, 76, 76, 76, 76, 76, 13, 51, 76,
              40, 27, 39, 31, 76, 35, 45, 76, 15, 52, 51, 39,
              27, 40, 30, 35, 27, 45, 64, 76, 11, 35, 40, 33,
              76, 41, 32, 76, 11, 35, 40, 33, 45, 66, 78, 76,
              76, 76, 76, 76, 76, 76, 76, 12, 41, 41, 37, 76,
              41, 40, 76, 39, 51, 76, 23, 41, 44, 37, 45, 64,
              76, 51, 31, 76, 13, 35, 33, 34, 46, 51, 64, 76,
              27, 40, 30, 76, 30, 31, 45, 42, 27, 35, 44, 67,
              78, 76, 76, 76, 76, 76, 76, 76, 76, 14, 41, 46,
              34, 35, 40, 33, 76, 28, 31, 45, 35, 30, 31, 76,
              44, 31, 39, 27, 35, 40, 45, 63, 76, 18, 41, 47,
              40, 30, 76, 46, 34, 31, 76, 30, 31, 29, 27, 51,
              78, 76, 76, 76, 76, 76, 76, 76, 76, 15, 32, 76,
              46, 34, 27, 46, 76, 29, 41, 38, 41, 45, 45, 27,
              38, 76, 23, 44, 31, 29, 37, 64, 76, 28, 41, 47,
              40, 30, 38, 31, 45, 45, 76, 27, 40, 30, 76, 28,
              27, 44, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
              20, 34, 31, 76, 38, 41, 40, 31, 76, 27, 40, 30,
              76, 38, 31, 48, 31, 38, 76, 45, 27, 40, 30, 45,
              76, 45, 46, 44, 31, 46, 29, 34, 76, 32, 27, 44,
              76, 27, 49, 27, 51, 78, 76, 76, 76, 76, 76, 76,
              76, 76]]			),
          torch.tensor([[0, 0, 0, 1_069, 11]]			),
          torch.tensor([[0, 0, 0, 1_069, 11]]			),
      ]
      # fmt: on
      self.assertTrue(torch.allclose(tokens[0]					,  EXPECTED_OUTPUT[0]			)			)
      self.assertTrue(torch.allclose(tokens[1]					,  EXPECTED_OUTPUT[1]			)			)
      self.assertTrue(torch.allclose(tokens[2]					,  EXPECTED_OUTPUT[2]			)			)



 @require_torch
 def    a__    (self			)       ->	Optional[int]:


      """simple docstring"""



      import torch

      _a           =    JukeboxTokenizer.from_pretrained('''openai/jukebox-5b-lyrics'''			)
      _a           =    tokenizer(**self.metas			)['''input_ids''']
      # fmt: off
      _a           =    [
          torch.tensor([[
              0, 0, 0, 1_069, 11, -1, -1, -1, -1, 9, 77, 39,
              31, 46, 77, 27, 77, 46, 44, 27, 48, 31, 38, 38,
              31, 44, 77, 32, 44, 41, 39, 77, 27, 40, 77, 27,
              40, 46, 35, 43, 47, 31, 77, 38, 27, 40, 30, 64,
              79, 77, 77, 77, 77, 77, 77, 77, 77, 23, 34, 41,
              77, 45, 27, 35, 30, 77, 72, 20, 49, 41, 77, 48,
              27, 45, 46, 77, 27, 40, 30, 77, 46, 44, 47, 40,
              37, 38, 31, 45, 45, 77, 38, 31, 33, 45, 77, 41,
              32, 77, 45, 46, 41, 40, 31, 79, 77, 77, 77, 77,
              77, 77, 77, 77, 19, 46, 27, 40, 30, 77, 35, 40,
              77, 46, 34, 31, 77, 30, 31, 45, 31, 44, 46, 63,
              77, 63, 77, 63, 77, 63, 77, 14, 31, 27, 44, 77,
              46, 34, 31, 39, 64, 77, 41, 40, 77, 46, 34, 31,
              77, 45, 27, 40, 30, 64, 79, 77, 77, 77, 77, 77,
              77, 77, 77, 8, 27, 38, 32, 77, 45, 47, 40, 37,
              77, 27, 77, 45, 34, 27, 46, 46, 31, 44, 31, 30,
              77, 48, 35, 45, 27, 33, 31, 77, 38, 35, 31, 45,
              64, 77, 49, 34, 41, 45, 31, 77, 32, 44, 41, 49,
              40, 64, 79, 77, 77, 77, 77, 77, 77, 77, 77, 1,
              40, 30, 77, 49, 44, 35, 40, 37, 38, 31, 30, 77,
              38, 35, 42, 64, 77, 27, 40, 30, 77, 45, 40, 31,
              31, 44, 77, 41, 32, 77, 29, 41, 38, 30, 77, 29,
              41, 39, 39, 27, 40, 30, 64, 79, 77, 77, 77, 77,
              77, 77, 77, 77, 20, 31, 38, 38, 77, 46, 34, 27,
              46, 77, 35, 46, 45, 77, 45, 29, 47, 38, 42, 46,
              41, 44, 77, 49, 31, 38, 38, 77, 46, 34, 41, 45,
              31, 77, 42, 27, 45, 45, 35, 41, 40, 45, 77, 44,
              31, 27, 30, 79, 77, 77, 77, 77, 77, 77, 77, 77,
              23, 34, 35, 29, 34, 77, 51, 31, 46, 77, 45, 47,
              44, 48, 35, 48, 31, 64, 77, 45, 46, 27, 39, 42,
              31, 30, 77, 41, 40, 77, 46, 34, 31, 45, 31, 77,
              38, 35, 32, 31, 38, 31, 45, 45, 77, 46, 34, 35,
              40, 33, 45, 64, 79, 77, 77, 77, 77, 77, 77, 77,
              77, 20, 34, 31, 77, 34, 27, 40, 30, 77, 46, 34,
              27, 46, 77, 39, 41, 29, 37, 31, 30, 77, 46, 34,
              31, 39, 64, 77, 27, 40, 30, 77, 46, 34, 31, 77,
              34, 31, 27, 44, 46, 77, 46, 34, 27, 46, 77, 32,
              31, 30, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77,
              1, 40, 30, 77, 41, 40, 77, 46, 34, 31, 77, 42,
              31, 30, 31, 45, 46, 27, 38, 64, 77, 46, 34, 31,
              45, 31, 77, 49, 41, 44, 30, 45, 77, 27, 42, 42,
              31, 27, 44, 65, 79, 77, 77, 77, 77, 77, 77, 77,
              77, 13, 51, 77, 40, 27, 39, 31, 77, 35, 45, 77,
              15, 52, 51, 39, 27, 40, 30, 35, 27, 45, 64, 77,
              11, 35, 40, 33, 77, 41, 32, 77, 11, 35, 40, 33,
              45, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77, 12,
              41, 41, 37, 77, 41, 40, 77, 39, 51, 77, 23, 41,
              44, 37, 45, 64, 77, 51, 31, 77, 13, 35, 33, 34,
              46, 51, 64, 77, 27, 40, 30, 77, 30, 31, 45, 42,
              27, 35, 44, 67, 79, 77, 77, 77, 77, 77, 77, 77,
              77, 14, 41, 46, 34, 35, 40, 33, 77, 28, 31, 45,
              35, 30, 31, 77, 44, 31, 39, 27, 35, 40, 45, 63,
              77, 18, 41, 47, 40, 30, 77, 46, 34, 31, 77, 30,
              31, 29, 27, 51, 79, 77, 77, 77, 77, 77, 77, 77,
              77, 15, 32, 77, 46, 34, 27, 46, 77, 29, 41, 38,
              41, 45, 45, 27, 38, 77, 23, 44, 31, 29, 37, 64,
              77, 28, 41, 47, 40, 30, 38, 31, 45, 45, 77, 27,
              40, 30, 77, 28, 27, 44, 31, 79, 77, 77, 77, 77,
              77, 77, 77, 77, 20, 34, 31, 77, 38, 41, 40, 31,
              77, 27, 40, 30, 77, 38, 31, 48, 31, 38, 77, 45,
              27, 40, 30, 45, 77, 45, 46, 44, 31, 46, 29, 34,
              77, 32, 27, 44, 77, 27, 49, 27, 51, 79, 77, 77,
              77, 77, 77, 77, 77, 77]]			),
          torch.tensor([[0, 0, 0, 1_069, 11, -1, -1, -1, -1]]			),
          torch.tensor([[0, 0, 0, 1_069, 11, -1, -1, -1, -1]]			),
      ]
      # fmt: on
      self.assertTrue(torch.allclose(tokens[0]					,  EXPECTED_OUTPUT[0]			)			)
      self.assertTrue(torch.allclose(tokens[1]					,  EXPECTED_OUTPUT[1]			)			)
      self.assertTrue(torch.allclose(tokens[2]					,  EXPECTED_OUTPUT[2]			)			)











"""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()
_lowerCAmelCase       =				logging.get_logger(__name__)






def       __UpperCamelCase						(		snake_case__       ):
	A_		:    List[str]						     = DPTConfig()

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

	if "ade" in checkpoint_url:
		A_		:    Tuple						     = True

		A_		:    Optional[Any]						     = 150
		A_		:    str						     = """huggingface/label-files"""
		A_		:    List[Any]						     = """ade20k-id2label.json"""
		A_		:    List[str]						     = json.load(open(cached_download(hf_hub_url(snake_case__ , snake_case__ , repo_type="""dataset"""       )       ) , """r"""       )       )
		A_		:    Dict						     = {int(snake_case__       ): v for k, v in idalabel.items()}
		A_		:    int						     = idalabel
		A_		:    Dict						     = {v: k for k, v in idalabel.items()}
		A_		:    Optional[Any]						     = [1, 150, 480, 480]

	return config, expected_shape






def       __UpperCamelCase						(		snake_case__       ):
	A_		:    List[str]						     = ["""pretrained.model.head.weight""", """pretrained.model.head.bias"""]
	for k in ignore_keys:
		state_dict.pop(snake_case__ , snake_case__       )






def       __UpperCamelCase						(		snake_case__       ):
	if (
	    "pretrained.model" in name
	    and "cls_token" not in name
	    and "pos_embed" not in name
	    and "patch_embed" not in name
	):
		A_		:    Any						     = name.replace("""pretrained.model""" , """dpt.encoder"""       )
	if "pretrained.model" in name:
		A_		:    int						     = name.replace("""pretrained.model""" , """dpt.embeddings"""       )
	if "patch_embed" in name:
		A_		:    List[Any]						     = name.replace("""patch_embed""" , """patch_embeddings"""       )
	if "pos_embed" in name:
		A_		:    Optional[int]						     = name.replace("""pos_embed""" , """position_embeddings"""       )
	if "attn.proj" in name:
		A_		:    str						     = name.replace("""attn.proj""" , """attention.output.dense"""       )
	if "proj" in name and "project" not in name:
		A_		:    Tuple						     = name.replace("""proj""" , """projection"""       )
	if "blocks" in name:
		A_		:    Dict						     = name.replace("""blocks""" , """layer"""       )
	if "mlp.fc1" in name:
		A_		:    List[str]						     = name.replace("""mlp.fc1""" , """intermediate.dense"""       )
	if "mlp.fc2" in name:
		A_		:    Any						     = name.replace("""mlp.fc2""" , """output.dense"""       )
	if "norm1" in name:
		A_		:    Tuple						     = name.replace("""norm1""" , """layernorm_before"""       )
	if "norm2" in name:
		A_		:    Optional[int]						     = name.replace("""norm2""" , """layernorm_after"""       )
	if "scratch.output_conv" in name:
		A_		:    Dict						     = name.replace("""scratch.output_conv""" , """head"""       )
	if "scratch" in name:
		A_		:    int						     = name.replace("""scratch""" , """neck"""       )
	if "layer1_rn" in name:
		A_		:    Any						     = name.replace("""layer1_rn""" , """convs.0"""       )
	if "layer2_rn" in name:
		A_		:    List[str]						     = name.replace("""layer2_rn""" , """convs.1"""       )
	if "layer3_rn" in name:
		A_		:    int						     = name.replace("""layer3_rn""" , """convs.2"""       )
	if "layer4_rn" in name:
		A_		:    Optional[Any]						     = name.replace("""layer4_rn""" , """convs.3"""       )
	if "refinenet" in name:
		A_		:    Dict						     = 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
		A_		:    Tuple						     = name.replace(F"""refinenet{layer_idx}""" , F"""fusion_stage.layers.{abs(layer_idx-4       )}"""       )
	if "out_conv" in name:
		A_		:    int						     = name.replace("""out_conv""" , """projection"""       )
	if "resConfUnit1" in name:
		A_		:    str						     = name.replace("""resConfUnit1""" , """residual_layer1"""       )
	if "resConfUnit2" in name:
		A_		:    Optional[int]						     = name.replace("""resConfUnit2""" , """residual_layer2"""       )
	if "conv1" in name:
		A_		:    Any						     = name.replace("""conv1""" , """convolution1"""       )
	if "conv2" in name:
		A_		:    Dict						     = name.replace("""conv2""" , """convolution2"""       )
	# readout blocks
	if "pretrained.act_postprocess1.0.project.0" in name:
		A_		:    List[str]						     = name.replace("""pretrained.act_postprocess1.0.project.0""" , """neck.reassemble_stage.readout_projects.0.0"""       )
	if "pretrained.act_postprocess2.0.project.0" in name:
		A_		:    List[str]						     = name.replace("""pretrained.act_postprocess2.0.project.0""" , """neck.reassemble_stage.readout_projects.1.0"""       )
	if "pretrained.act_postprocess3.0.project.0" in name:
		A_		:    int						     = name.replace("""pretrained.act_postprocess3.0.project.0""" , """neck.reassemble_stage.readout_projects.2.0"""       )
	if "pretrained.act_postprocess4.0.project.0" in name:
		A_		:    Union[str, Any]						     = name.replace("""pretrained.act_postprocess4.0.project.0""" , """neck.reassemble_stage.readout_projects.3.0"""       )
	# resize blocks
	if "pretrained.act_postprocess1.3" in name:
		A_		:    Union[str, Any]						     = name.replace("""pretrained.act_postprocess1.3""" , """neck.reassemble_stage.layers.0.projection"""       )
	if "pretrained.act_postprocess1.4" in name:
		A_		:    Optional[Any]						     = name.replace("""pretrained.act_postprocess1.4""" , """neck.reassemble_stage.layers.0.resize"""       )
	if "pretrained.act_postprocess2.3" in name:
		A_		:    List[str]						     = name.replace("""pretrained.act_postprocess2.3""" , """neck.reassemble_stage.layers.1.projection"""       )
	if "pretrained.act_postprocess2.4" in name:
		A_		:    Optional[int]						     = name.replace("""pretrained.act_postprocess2.4""" , """neck.reassemble_stage.layers.1.resize"""       )
	if "pretrained.act_postprocess3.3" in name:
		A_		:    Any						     = name.replace("""pretrained.act_postprocess3.3""" , """neck.reassemble_stage.layers.2.projection"""       )
	if "pretrained.act_postprocess4.3" in name:
		A_		:    int						     = name.replace("""pretrained.act_postprocess4.3""" , """neck.reassemble_stage.layers.3.projection"""       )
	if "pretrained.act_postprocess4.4" in name:
		A_		:    Dict						     = name.replace("""pretrained.act_postprocess4.4""" , """neck.reassemble_stage.layers.3.resize"""       )
	if "pretrained" in name:
		A_		:    List[str]						     = name.replace("""pretrained""" , """dpt"""       )
	if "bn" in name:
		A_		:    str						     = name.replace("""bn""" , """batch_norm"""       )
	if "head" in name:
		A_		:    Any						     = name.replace("""head""" , """head.head"""       )
	if "encoder.norm" in name:
		A_		:    Optional[Any]						     = name.replace("""encoder.norm""" , """layernorm"""       )
	if "auxlayer" in name:
		A_		:    List[Any]						     = name.replace("""auxlayer""" , """auxiliary_head.head"""       )

	return name






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






def       __UpperCamelCase						(		):
	A_		:    Optional[Any]						     = """http://images.cocodataset.org/val2017/000000039769.jpg"""
	A_		:    Optional[Any]						     = Image.open(requests.get(snake_case__ , stream=snake_case__       ).raw       )
	return im






@torch.no_grad()
def       __UpperCamelCase						(		snake_case__ , snake_case__ , snake_case__ , snake_case__       ):
	A_     ,  A_		:    int						     = get_dpt_config(snake_case__       )
	# load original state_dict from URL
	A_		:    Dict						     = torch.hub.load_state_dict_from_url(snake_case__ , map_location="""cpu"""       )
	# remove certain keys
	remove_ignore_keys_(snake_case__       )
	# rename keys
	for key in state_dict.copy().keys():
		A_		:    str						     = state_dict.pop(snake_case__       )
		A_		:    str						     = val
	# read in qkv matrices
	read_in_q_k_v(snake_case__ , snake_case__       )

	# load HuggingFace model
	A_		:    Union[str, Any]						     = DPTForSemanticSegmentation(snake_case__       ) if """ade""" in checkpoint_url else DPTForDepthEstimation(snake_case__       )
	model.load_state_dict(snake_case__       )
	model.eval()

	# Check outputs on an image
	A_		:    str						     = 480 if """ade""" in checkpoint_url else 384
	A_		:    List[str]						     = DPTImageProcessor(size=snake_case__       )

	A_		:    Dict						     = prepare_img()
	A_		:    List[Any]						     = image_processor(snake_case__ , return_tensors="""pt"""       )

	# forward pass
	A_		:    Tuple						     = model(**snake_case__       ).logits if """ade""" in checkpoint_url else model(**snake_case__       ).predicted_depth

	# Assert logits
	A_		:    List[str]						     = torch.tensor([[6.3_199, 6.3_629, 6.4_148], [6.3_850, 6.3_615, 6.4_166], [6.3_519, 6.3_176, 6.3_575]]       )
	if "ade" in checkpoint_url:
		A_		:    str						     = torch.tensor([[4.0_480, 4.2_420, 4.4_360], [4.3_124, 4.5_693, 4.8_261], [4.5_768, 4.8_965, 5.2_163]]       )
	assert outputs.shape == torch.Size(snake_case__       )
	assert (
	    torch.allclose(outputs[0, 0, :3, :3] , snake_case__ , atol=1E-4       )
	    if "ade" in checkpoint_url
	    else torch.allclose(outputs[0, :3, :3] , snake_case__       )
	)

	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 image processor to {pytorch_dump_folder_path}"""       )
	image_processor.save_pretrained(snake_case__       )

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


if __name__ == "__main__":
							_lowerCAmelCase       =				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.",
							)

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

