

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

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   List[Any]      ):



     '''simple docstring'''





     lowercase__   :       int				=	DPTConfig()

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

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

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

     return config, expected_shape

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   int      ):



     '''simple docstring'''





     lowercase__   :       Optional[Any]				=	['pretrained.model.head.weight', 'pretrained.model.head.bias']
     for k in ignore_keys:
          state_dict.pop(SCREAMING_SNAKE_CASE_     ,					SCREAMING_SNAKE_CASE_      )

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Optional[int]      ):



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

     return name

def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   Any     ,					SCREAMING_SNAKE_CASE_		:   Optional[Any]      ):



     '''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)
          lowercase__   :       Any				=	state_dict.pop(f"""dpt.encoder.layer.{i}.attn.qkv.weight"""      )
          lowercase__   :       Tuple				=	state_dict.pop(f"""dpt.encoder.layer.{i}.attn.qkv.bias"""      )
          # next, add query, keys and values (in that order) to the state dict
          lowercase__   :       List[Any]				=	in_proj_weight[: config.hidden_size, :]
          lowercase__   :       Any				=	in_proj_bias[: config.hidden_size]
          lowercase__   :       Any				=	in_proj_weight[
              config.hidden_size : config.hidden_size * 2, :
          ]
          lowercase__   :       List[Any]				=	in_proj_bias[
              config.hidden_size : config.hidden_size * 2
          ]
          lowercase__   :       List[Any]				=	in_proj_weight[
              -config.hidden_size :, :
          ]
          lowercase__   :       List[str]				=	in_proj_bias[-config.hidden_size :]

def      snake_case__      (   ):



     '''simple docstring'''





     lowercase__   :       Dict				=	'http://images.cocodataset.org/val2017/000000039769.jpg'
     lowercase__   :       List[Any]				=	Image.open(requests.get(SCREAMING_SNAKE_CASE_     ,					stream=SCREAMING_SNAKE_CASE_      ).raw      )
     return im

@torch.no_grad()
def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   str     ,					SCREAMING_SNAKE_CASE_		:   Optional[Any]     ,					SCREAMING_SNAKE_CASE_		:   List[Any]     ,					SCREAMING_SNAKE_CASE_		:   Optional[int]      ):



     '''simple docstring'''





     lowercase__				,			lowercase__   :       Any				=	get_dpt_config(SCREAMING_SNAKE_CASE_      )
     # load original state_dict from URL
     lowercase__   :       List[Any]				=	torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_     ,					map_location='cpu'      )
     # remove certain keys
     remove_ignore_keys_(SCREAMING_SNAKE_CASE_      )
     # rename keys
     for key in state_dict.copy().keys():
          lowercase__   :       str				=	state_dict.pop(SCREAMING_SNAKE_CASE_      )
          lowercase__   :       List[str]				=	val
     # read in qkv matrices
     read_in_q_k_v(SCREAMING_SNAKE_CASE_     ,					SCREAMING_SNAKE_CASE_      )

     # load HuggingFace model
     lowercase__   :       List[Any]				=	DPTForSemanticSegmentation(SCREAMING_SNAKE_CASE_      ) if 'ade' in checkpoint_url else DPTForDepthEstimation(SCREAMING_SNAKE_CASE_      )
     model.load_state_dict(SCREAMING_SNAKE_CASE_      )
     model.eval()

     # Check outputs on an image
     lowercase__   :       Tuple				=	480 if 'ade' in checkpoint_url else 384
     lowercase__   :       Tuple				=	DPTImageProcessor(size=SCREAMING_SNAKE_CASE_      )

     lowercase__   :       Optional[int]				=	prepare_img()
     lowercase__   :       Optional[Any]				=	image_processor(SCREAMING_SNAKE_CASE_     ,					return_tensors='pt'      )

     # forward pass
     lowercase__   :       Optional[int]				=	model(**SCREAMING_SNAKE_CASE_      ).logits if 'ade' in checkpoint_url else model(**SCREAMING_SNAKE_CASE_      ).predicted_depth

     # Assert logits
     lowercase__   :       Tuple				=	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:
          lowercase__   :       List[Any]				=	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(SCREAMING_SNAKE_CASE_      )
     assert (
         torch.allclose(outputs[0, 0, :3, :3]     ,					SCREAMING_SNAKE_CASE_     ,					atol=1E-4      )
         if "ade" in checkpoint_url
         else torch.allclose(outputs[0, :3, :3]     ,					SCREAMING_SNAKE_CASE_      )
     )

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

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


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

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




def      snake_case__      (   SCREAMING_SNAKE_CASE_		:   int = 50_000_000      ):



     '''simple docstring'''





     lowercase__   :       List[Any]				=	set()
     lowercase__   :       Any				=	int((limit - 24) ** (1 / 2)      )

     lowercase__   :       Any				=	set(range(3     ,					prime_square_limit + 1     ,					2      )      )
     primes.add(2      )
     for p in range(3     ,					prime_square_limit + 1     ,					2      ):
          if p not in primes:
               continue
          primes.difference_update(set(range(p * p     ,					prime_square_limit + 1     ,					SCREAMING_SNAKE_CASE_      )      )      )

     for primea in primes:
          lowercase__   :       Any				=	primea * primea
          for primea in primes:
               lowercase__   :       Optional[int]				=	primea * primea * primea
               if square + cube >= limit - 16:
                    break
               for primea in primes:
                    lowercase__   :       Dict				=	primea * primea * primea * primea
                    lowercase__   :       List[Any]				=	square + cube + tetr
                    if total >= limit:
                         break
                    ret.add(SCREAMING_SNAKE_CASE_      )

     return len(SCREAMING_SNAKE_CASE_      )


if __name__ == "__main__":
       print(F'''{solution() = }''')

