





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 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							List[str]    ):
		SCREAMING_SNAKE_CASE__			=      DPTConfig()

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

		if "ade" in checkpoint_url:
				SCREAMING_SNAKE_CASE__			=      True

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

		return config, expected_shape






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							Optional[int]    ):
		SCREAMING_SNAKE_CASE__			=      ["""pretrained.model.head.weight""", """pretrained.model.head.bias"""]
		for k in ignore_keys:
				state_dict.pop(UpperCamelCase__					,   UpperCamelCase__    )






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

		return name






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






def 				SCREAMING_SNAKE_CASE__							(       ):
		SCREAMING_SNAKE_CASE__			=      """http://images.cocodataset.org/val2017/000000039769.jpg"""
		SCREAMING_SNAKE_CASE__			=      Image.open(requests.get(UpperCamelCase__					,   stream=UpperCamelCase__    ).raw    )
		return im






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

		# load HuggingFace model
		SCREAMING_SNAKE_CASE__			=      DPTForSemanticSegmentation(UpperCamelCase__    ) if """ade""" in checkpoint_url else DPTForDepthEstimation(UpperCamelCase__    )
		model.load_state_dict(UpperCamelCase__    )
		model.eval()

		# Check outputs on an image
		SCREAMING_SNAKE_CASE__			=      480 if """ade""" in checkpoint_url else 384
		SCREAMING_SNAKE_CASE__			=      DPTImageProcessor(size=UpperCamelCase__    )

		SCREAMING_SNAKE_CASE__			=      prepare_img()
		SCREAMING_SNAKE_CASE__			=      image_processor(UpperCamelCase__					,   return_tensors="""pt"""    )

		# forward pass
		SCREAMING_SNAKE_CASE__			=      model(**UpperCamelCase__    ).logits if """ade""" in checkpoint_url else model(**UpperCamelCase__    ).predicted_depth

		# Assert logits
		SCREAMING_SNAKE_CASE__			=      torch.tensor([[6.3_1_9_9, 6.3_6_2_9, 6.4_1_4_8], [6.3_8_5_0, 6.3_6_1_5, 6.4_1_6_6], [6.3_5_1_9, 6.3_1_7_6, 6.3_5_7_5]]    )
		if "ade" in checkpoint_url:
				SCREAMING_SNAKE_CASE__			=      torch.tensor([[4.0_4_8_0, 4.2_4_2_0, 4.4_3_6_0], [4.3_1_2_4, 4.5_6_9_3, 4.8_2_6_1], [4.5_7_6_8, 4.8_9_6_5, 5.2_1_6_3]]    )
		assert outputs.shape == torch.Size(UpperCamelCase__    )
		assert (
		    torch.allclose(outputs[0, 0, :3, :3]					,   UpperCamelCase__					,   atol=1e-4    )
		    if "ade" in checkpoint_url
		    else torch.allclose(outputs[0, :3, :3]					,   UpperCamelCase__    )
		)

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

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


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)






import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import PoolFormerConfig, PoolFormerForImageClassification, PoolFormerImageProcessor
from transformers.utils import logging


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






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							Union[str, Any]					,   UpperCamelCase__:							List[Any]					,   UpperCamelCase__:							Optional[Any]					,   UpperCamelCase__:							Optional[Any]    ):
		SCREAMING_SNAKE_CASE__			=      original_name.split("""."""    )[0]
		SCREAMING_SNAKE_CASE__			=      key.split("""."""    )
		SCREAMING_SNAKE_CASE__			=      int(key_list[key_list.index(UpperCamelCase__    ) - 2]    )
		SCREAMING_SNAKE_CASE__			=      int(key_list[key_list.index(UpperCamelCase__    ) - 1]    )
		SCREAMING_SNAKE_CASE__			=      orig_block_num - offset

		SCREAMING_SNAKE_CASE__			=      key.replace(f'''{orig_block_num}.{layer_num}.{original_name}'''					,   f'''block.{new_block_num}.{layer_num}.{new_name}'''    )
		return key






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							int    ):
		SCREAMING_SNAKE_CASE__			=      OrderedDict()
		SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      0, 0
		for key, value in state_dict.items():
				if key.startswith("""network"""    ):
						SCREAMING_SNAKE_CASE__			=      key.replace("""network"""					,   """poolformer.encoder"""    )
				if "proj" in key:
						# Works for the first embedding as well as the internal embedding layers
						if key.endswith("""bias"""    ) and "patch_embed" not in key:
								patch_emb_offset += 1
						SCREAMING_SNAKE_CASE__			=      key[: key.find("""proj"""    )]
						SCREAMING_SNAKE_CASE__			=      key.replace(UpperCamelCase__					,   f'''patch_embeddings.{total_embed_found}.'''    )
						SCREAMING_SNAKE_CASE__			=      key.replace("""proj"""					,   """projection"""    )
						if key.endswith("""bias"""    ):
								total_embed_found += 1
				if "patch_embeddings" in key:
						SCREAMING_SNAKE_CASE__			=      """poolformer.encoder.""" + key
				if "mlp.fc1" in key:
						SCREAMING_SNAKE_CASE__			=      replace_key_with_offset(UpperCamelCase__					,   UpperCamelCase__					,   """mlp.fc1"""					,   """output.conv1"""    )
				if "mlp.fc2" in key:
						SCREAMING_SNAKE_CASE__			=      replace_key_with_offset(UpperCamelCase__					,   UpperCamelCase__					,   """mlp.fc2"""					,   """output.conv2"""    )
				if "norm1" in key:
						SCREAMING_SNAKE_CASE__			=      replace_key_with_offset(UpperCamelCase__					,   UpperCamelCase__					,   """norm1"""					,   """before_norm"""    )
				if "norm2" in key:
						SCREAMING_SNAKE_CASE__			=      replace_key_with_offset(UpperCamelCase__					,   UpperCamelCase__					,   """norm2"""					,   """after_norm"""    )
				if "layer_scale_1" in key:
						SCREAMING_SNAKE_CASE__			=      replace_key_with_offset(UpperCamelCase__					,   UpperCamelCase__					,   """layer_scale_1"""					,   """layer_scale_1"""    )
				if "layer_scale_2" in key:
						SCREAMING_SNAKE_CASE__			=      replace_key_with_offset(UpperCamelCase__					,   UpperCamelCase__					,   """layer_scale_2"""					,   """layer_scale_2"""    )
				if "head" in key:
						SCREAMING_SNAKE_CASE__			=      key.replace("""head"""					,   """classifier"""    )
				SCREAMING_SNAKE_CASE__			=      value
		return new_state_dict






def 				SCREAMING_SNAKE_CASE__							(       ):
		SCREAMING_SNAKE_CASE__			=      """http://images.cocodataset.org/val2017/000000039769.jpg"""
		SCREAMING_SNAKE_CASE__			=      Image.open(requests.get(UpperCamelCase__					,   stream=UpperCamelCase__    ).raw    )

		return image






@torch.no_grad()
def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							List[str]					,   UpperCamelCase__:							Optional[int]					,   UpperCamelCase__:							Any    ):
		SCREAMING_SNAKE_CASE__			=      PoolFormerConfig()

		# set attributes based on model_name
		SCREAMING_SNAKE_CASE__			=      """huggingface/label-files"""
		SCREAMING_SNAKE_CASE__			=      model_name[-3:]
		SCREAMING_SNAKE_CASE__			=      1_000
		SCREAMING_SNAKE_CASE__			=      """imagenet-1k-id2label.json"""
		SCREAMING_SNAKE_CASE__			=      (1, 1_000)

		# set config attributes
		SCREAMING_SNAKE_CASE__			=      json.load(open(hf_hub_download(UpperCamelCase__					,   UpperCamelCase__					,   repo_type="""dataset"""    )					,   """r"""    )    )
		SCREAMING_SNAKE_CASE__			=      {int(UpperCamelCase__    ): v for k, v in idalabel.items()}
		SCREAMING_SNAKE_CASE__			=      idalabel
		SCREAMING_SNAKE_CASE__			=      {v: k for k, v in idalabel.items()}
		if size == "s12":
				SCREAMING_SNAKE_CASE__			=      [2, 2, 6, 2]
				SCREAMING_SNAKE_CASE__			=      [64, 128, 320, 512]
				SCREAMING_SNAKE_CASE__			=      4.0
				SCREAMING_SNAKE_CASE__			=      0.9
		elif size == "s24":
				SCREAMING_SNAKE_CASE__			=      [4, 4, 12, 4]
				SCREAMING_SNAKE_CASE__			=      [64, 128, 320, 512]
				SCREAMING_SNAKE_CASE__			=      4.0
				SCREAMING_SNAKE_CASE__			=      0.9
		elif size == "s36":
				SCREAMING_SNAKE_CASE__			=      [6, 6, 18, 6]
				SCREAMING_SNAKE_CASE__			=      [64, 128, 320, 512]
				SCREAMING_SNAKE_CASE__			=      4.0
				SCREAMING_SNAKE_CASE__			=      1e-6
				SCREAMING_SNAKE_CASE__			=      0.9
		elif size == "m36":
				SCREAMING_SNAKE_CASE__			=      [6, 6, 18, 6]
				SCREAMING_SNAKE_CASE__			=      [96, 192, 384, 768]
				SCREAMING_SNAKE_CASE__			=      4.0
				SCREAMING_SNAKE_CASE__			=      1e-6
				SCREAMING_SNAKE_CASE__			=      0.9_5
		elif size == "m48":
				SCREAMING_SNAKE_CASE__			=      [8, 8, 24, 8]
				SCREAMING_SNAKE_CASE__			=      [96, 192, 384, 768]
				SCREAMING_SNAKE_CASE__			=      4.0
				SCREAMING_SNAKE_CASE__			=      1e-6
				SCREAMING_SNAKE_CASE__			=      0.9_5
		else:
				raise ValueError(f'''Size {size} not supported'''    )

		# load image processor
		SCREAMING_SNAKE_CASE__			=      PoolFormerImageProcessor(crop_pct=UpperCamelCase__    )

		# Prepare image
		SCREAMING_SNAKE_CASE__			=      prepare_img()
		SCREAMING_SNAKE_CASE__			=      image_processor(images=UpperCamelCase__					,   return_tensors="""pt"""    ).pixel_values

		logger.info(f'''Converting model {model_name}...'''    )

		# load original state dict
		SCREAMING_SNAKE_CASE__			=      torch.load(UpperCamelCase__					,   map_location=torch.device("""cpu"""    )    )

		# rename keys
		SCREAMING_SNAKE_CASE__			=      rename_keys(UpperCamelCase__    )

		# create HuggingFace model and load state dict
		SCREAMING_SNAKE_CASE__			=      PoolFormerForImageClassification(UpperCamelCase__    )
		model.load_state_dict(UpperCamelCase__    )
		model.eval()

		# Define image processor
		SCREAMING_SNAKE_CASE__			=      PoolFormerImageProcessor(crop_pct=UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      image_processor(images=prepare_img()					,   return_tensors="""pt"""    ).pixel_values

		# forward pass
		SCREAMING_SNAKE_CASE__			=      model(UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      outputs.logits

		# define expected logit slices for different models
		if size == "s12":
				SCREAMING_SNAKE_CASE__			=      torch.tensor([-0.3_0_4_5, -0.6_7_5_8, -0.4_8_6_9]    )
		elif size == "s24":
				SCREAMING_SNAKE_CASE__			=      torch.tensor([0.4_4_0_2, -0.1_3_7_4, -0.8_0_4_5]    )
		elif size == "s36":
				SCREAMING_SNAKE_CASE__			=      torch.tensor([-0.6_0_8_0, -0.5_1_3_3, -0.5_8_9_8]    )
		elif size == "m36":
				SCREAMING_SNAKE_CASE__			=      torch.tensor([0.3_9_5_2, 0.2_2_6_3, -1.2_6_6_8]    )
		elif size == "m48":
				SCREAMING_SNAKE_CASE__			=      torch.tensor([0.1_1_6_7, -0.0_6_5_6, -0.3_4_2_3]    )
		else:
				raise ValueError(f'''Size {size} not supported'''    )

		# verify logits
		assert logits.shape == expected_shape
		assert torch.allclose(logits[0, :3]					,   UpperCamelCase__					,   atol=1e-2    )

		# finally, save model and image processor
		logger.info(f'''Saving PyTorch model and image processor to {pytorch_dump_folder_path}...'''    )
		Path(UpperCamelCase__    ).mkdir(exist_ok=UpperCamelCase__    )
		model.save_pretrained(UpperCamelCase__    )
		print(f'''Saving image processor to {pytorch_dump_folder_path}'''    )
		image_processor.save_pretrained(UpperCamelCase__    )


if __name__ == "__main__":
	_lowerCamelCase   				=			argparse.ArgumentParser()

	parser.add_argument(
	    '--model_name',
	    default='poolformer_s12',
	    type=str,
	    help='Name of the model you\'d like to convert.',
	)
	parser.add_argument(
	    '--checkpoint_path', default=None, type=str, help='Path 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.'
	)
	_lowerCamelCase   				=			parser.parse_args()
	convert_poolformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)