


'''simple docstring'''


import argparse
import json
from pathlib import Path

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

from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


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






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase=False      ):
		_UpperCAmelCase						:				Optional[Any]	    =					[]
		for i in range(config.num_hidden_layers      ):
				# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
				rename_keys.append((F"""blocks.{i}.norm1.weight""", F"""vit.encoder.layer.{i}.layernorm_before.weight""")      )
				rename_keys.append((F"""blocks.{i}.norm1.bias""", F"""vit.encoder.layer.{i}.layernorm_before.bias""")      )
				rename_keys.append((F"""blocks.{i}.attn.proj.weight""", F"""vit.encoder.layer.{i}.attention.output.dense.weight""")      )
				rename_keys.append((F"""blocks.{i}.attn.proj.bias""", F"""vit.encoder.layer.{i}.attention.output.dense.bias""")      )
				rename_keys.append((F"""blocks.{i}.norm2.weight""", F"""vit.encoder.layer.{i}.layernorm_after.weight""")      )
				rename_keys.append((F"""blocks.{i}.norm2.bias""", F"""vit.encoder.layer.{i}.layernorm_after.bias""")      )
				rename_keys.append((F"""blocks.{i}.mlp.fc1.weight""", F"""vit.encoder.layer.{i}.intermediate.dense.weight""")      )
				rename_keys.append((F"""blocks.{i}.mlp.fc1.bias""", F"""vit.encoder.layer.{i}.intermediate.dense.bias""")      )
				rename_keys.append((F"""blocks.{i}.mlp.fc2.weight""", F"""vit.encoder.layer.{i}.output.dense.weight""")      )
				rename_keys.append((F"""blocks.{i}.mlp.fc2.bias""", F"""vit.encoder.layer.{i}.output.dense.bias""")      )

		# projection layer + position embeddings
		rename_keys.extend(
		    [
		        ("cls_token", "vit.embeddings.cls_token"),
		        ("patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight"),
		        ("patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias"),
		        ("pos_embed", "vit.embeddings.position_embeddings"),
		    ]      )

		if base_model:
				# layernorm + pooler
				rename_keys.extend(
				    [
				        ("norm.weight", "layernorm.weight"),
				        ("norm.bias", "layernorm.bias"),
				    ]      )

				# if just the base model, we should remove "vit" from all keys that start with "vit"
				_UpperCAmelCase						:				Any	    =					[(pair[0], pair[1][4:]) if pair[1].startswith("vit"      ) else pair for pair in rename_keys]
		else:
				# layernorm + classification head
				rename_keys.extend(
				    [
				        ("norm.weight", "vit.layernorm.weight"),
				        ("norm.bias", "vit.layernorm.bias"),
				        ("head.weight", "classifier.weight"),
				        ("head.bias", "classifier.bias"),
				    ]      )

		return rename_keys






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






def        __lowerCAmelCase       (__lowerCAmelCase      ):
		_UpperCAmelCase						:				Tuple	    =					["head.weight", "head.bias"]
		for k in ignore_keys:
				state_dict.pop(__lowerCAmelCase					,       __lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				Union[str, Any]	    =					dct.pop(__lowerCAmelCase      )
		_UpperCAmelCase						:				Optional[int]	    =					val






def        __lowerCAmelCase       ():
		_UpperCAmelCase						:				int	    =					"http://images.cocodataset.org/val2017/000000039769.jpg"
		_UpperCAmelCase						:				Union[str, Any]	    =					Image.open(requests.get(__lowerCAmelCase					,       stream=__lowerCAmelCase      ).raw      )
		return im






@torch.no_grad()
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase=True      ):
		_UpperCAmelCase						:				Tuple	    =					ViTConfig()
		# patch_size
		if model_name[-1] == "8":
				_UpperCAmelCase						:				Optional[Any]	    =					8
		# set labels if required
		if not base_model:
				_UpperCAmelCase						:				str	    =					1_000
				_UpperCAmelCase						:				str	    =					"huggingface/label-files"
				_UpperCAmelCase						:				Any	    =					"imagenet-1k-id2label.json"
				_UpperCAmelCase						:				Union[str, Any]	    =					json.load(open(hf_hub_download(__lowerCAmelCase					,       __lowerCAmelCase					,       repo_type="dataset"      )					,       "r"      )      )
				_UpperCAmelCase						:				List[str]	    =					{int(__lowerCAmelCase      ): v for k, v in idalabel.items()}
				_UpperCAmelCase						:				Tuple	    =					idalabel
				_UpperCAmelCase						:				Any	    =					{v: k for k, v in idalabel.items()}
		# size of the architecture
		if model_name in ["dino_vits8", "dino_vits16"]:
				_UpperCAmelCase						:				int	    =					384
				_UpperCAmelCase						:				str	    =					1_536
				_UpperCAmelCase						:				Optional[Any]	    =					12
				_UpperCAmelCase						:				Optional[int]	    =					6

		# load original model from torch hub
		_UpperCAmelCase						:				Any	    =					torch.hub.load("facebookresearch/dino:main"					,       __lowerCAmelCase      )
		original_model.eval()

		# load state_dict of original model, remove and rename some keys
		_UpperCAmelCase						:				List[str]	    =					original_model.state_dict()
		if base_model:
				remove_classification_head_(__lowerCAmelCase      )
		_UpperCAmelCase						:				Any	    =					create_rename_keys(__lowerCAmelCase					,       base_model=__lowerCAmelCase      )
		for src, dest in rename_keys:
				rename_key(__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      )
		read_in_q_k_v(__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      )

		# load HuggingFace model
		if base_model:
				_UpperCAmelCase						:				List[Any]	    =					ViTModel(__lowerCAmelCase					,       add_pooling_layer=__lowerCAmelCase      ).eval()
		else:
				_UpperCAmelCase						:				int	    =					ViTForImageClassification(__lowerCAmelCase      ).eval()
		model.load_state_dict(__lowerCAmelCase      )

		# Check outputs on an image, prepared by ViTImageProcessor
		_UpperCAmelCase						:				List[str]	    =					ViTImageProcessor()
		_UpperCAmelCase						:				Tuple	    =					image_processor(images=prepare_img()					,       return_tensors="pt"      )
		_UpperCAmelCase						:				int	    =					encoding["pixel_values"]
		_UpperCAmelCase						:				Union[str, Any]	    =					model(__lowerCAmelCase      )

		if base_model:
				_UpperCAmelCase						:				Optional[int]	    =					original_model(__lowerCAmelCase      )
				assert torch.allclose(__lowerCAmelCase					,       outputs.last_hidden_state[:, 0, :]					,       atol=1e-1      )
		else:
				_UpperCAmelCase						:				List[Any]	    =					original_model(__lowerCAmelCase      )
				assert logits.shape == outputs.logits.shape
				assert torch.allclose(__lowerCAmelCase					,       outputs.logits					,       atol=1e-3      )

		Path(__lowerCAmelCase      ).mkdir(exist_ok=__lowerCAmelCase      )
		print(F"""Saving model {model_name} 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 __name__ == "__main__":
						lowerCamelCase__					=   argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--model_name',
						    default='dino_vitb16',
						    type=str,
						    help='Name of the model trained with DINO you\'d like to convert.',
						)
						parser.add_argument(
						    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
						)
						parser.add_argument(
						    '--base_model',
						    action='store_true',
						    help='Whether to only convert the base model (no projection head weights).',
						)

						parser.set_defaults(base_model=True)
						lowerCamelCase__					=   parser.parse_args()
						convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)







'''simple docstring'''


def        __lowerCAmelCase       (__lowerCAmelCase = 50      ):
		_UpperCAmelCase						:				Tuple	    =					[[0] * 3 for _ in range(length + 1      )]

		for row_length in range(length + 1      ):
				for tile_length in range(2					,       5      ):
						for tile_start in range(row_length - tile_length + 1      ):
								different_colour_ways_number[row_length][tile_length - 2] += (
								    different_colour_ways_number[row_length - tile_start - tile_length][
								        tile_length - 2
								    ]
								    + 1
								)

		return sum(different_colour_ways_number[length]      )


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



