





'''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 YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


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



def 						__A						(   lowerCamelCase_     ):


		"""simple docstring"""



		SCREAMING_SNAKE_CASE      :							int						=       YolosConfig()

		# size of the architecture
		if "yolos_ti" in yolos_name:
				SCREAMING_SNAKE_CASE      :							Any						=       1_92
				SCREAMING_SNAKE_CASE      :							Dict						=       7_68
				SCREAMING_SNAKE_CASE      :							Dict						=       12
				SCREAMING_SNAKE_CASE      :							List[Any]						=       3
				SCREAMING_SNAKE_CASE      :							Dict						=       [8_00, 13_33]
				SCREAMING_SNAKE_CASE      :							List[str]						=       False
		elif yolos_name == "yolos_s_dWr":
				SCREAMING_SNAKE_CASE      :							Optional[int]						=       3_30
				SCREAMING_SNAKE_CASE      :							List[str]						=       14
				SCREAMING_SNAKE_CASE      :							str						=       6
				SCREAMING_SNAKE_CASE      :							List[str]						=       13_20
		elif "yolos_s" in yolos_name:
				SCREAMING_SNAKE_CASE      :							Union[str, Any]						=       3_84
				SCREAMING_SNAKE_CASE      :							Dict						=       15_36
				SCREAMING_SNAKE_CASE      :							int						=       12
				SCREAMING_SNAKE_CASE      :							List[str]						=       6
		elif "yolos_b" in yolos_name:
				SCREAMING_SNAKE_CASE      :							List[Any]						=       [8_00, 13_44]

		SCREAMING_SNAKE_CASE      :							Optional[Any]						=       91
		SCREAMING_SNAKE_CASE      :							str						=       '''huggingface/label-files'''
		SCREAMING_SNAKE_CASE      :							Optional[Any]						=       '''coco-detection-id2label.json'''
		SCREAMING_SNAKE_CASE      :							Dict						=       json.load(open(hf_hub_download(_UpperCamelCase       ,       _UpperCamelCase       ,       repo_type="""dataset"""     )       ,       """r"""     )     )
		SCREAMING_SNAKE_CASE      :							Dict						=       {int(_UpperCamelCase     ): v for k, v in idalabel.items()}
		SCREAMING_SNAKE_CASE      :							Optional[Any]						=       idalabel
		SCREAMING_SNAKE_CASE      :							List[Any]						=       {v: k for k, v in idalabel.items()}

		return config



def 						__A						(   lowerCamelCase_       ,       lowerCamelCase_       ,       lowerCamelCase_ = False     ):


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



def 						__A						(   lowerCamelCase_     ):


		"""simple docstring"""



		if "backbone" in name:
				SCREAMING_SNAKE_CASE      :							List[str]						=       name.replace("""backbone"""       ,       """vit"""     )
		if "cls_token" in name:
				SCREAMING_SNAKE_CASE      :							List[str]						=       name.replace("""cls_token"""       ,       """embeddings.cls_token"""     )
		if "det_token" in name:
				SCREAMING_SNAKE_CASE      :							int						=       name.replace("""det_token"""       ,       """embeddings.detection_tokens"""     )
		if "mid_pos_embed" in name:
				SCREAMING_SNAKE_CASE      :							Optional[int]						=       name.replace("""mid_pos_embed"""       ,       """encoder.mid_position_embeddings"""     )
		if "pos_embed" in name:
				SCREAMING_SNAKE_CASE      :							Tuple						=       name.replace("""pos_embed"""       ,       """embeddings.position_embeddings"""     )
		if "patch_embed.proj" in name:
				SCREAMING_SNAKE_CASE      :							Dict						=       name.replace("""patch_embed.proj"""       ,       """embeddings.patch_embeddings.projection"""     )
		if "blocks" in name:
				SCREAMING_SNAKE_CASE      :							Any						=       name.replace("""blocks"""       ,       """encoder.layer"""     )
		if "attn.proj" in name:
				SCREAMING_SNAKE_CASE      :							Optional[int]						=       name.replace("""attn.proj"""       ,       """attention.output.dense"""     )
		if "attn" in name:
				SCREAMING_SNAKE_CASE      :							int						=       name.replace("""attn"""       ,       """attention.self"""     )
		if "norm1" in name:
				SCREAMING_SNAKE_CASE      :							str						=       name.replace("""norm1"""       ,       """layernorm_before"""     )
		if "norm2" in name:
				SCREAMING_SNAKE_CASE      :							Dict						=       name.replace("""norm2"""       ,       """layernorm_after"""     )
		if "mlp.fc1" in name:
				SCREAMING_SNAKE_CASE      :							int						=       name.replace("""mlp.fc1"""       ,       """intermediate.dense"""     )
		if "mlp.fc2" in name:
				SCREAMING_SNAKE_CASE      :							Any						=       name.replace("""mlp.fc2"""       ,       """output.dense"""     )
		if "class_embed" in name:
				SCREAMING_SNAKE_CASE      :							Any						=       name.replace("""class_embed"""       ,       """class_labels_classifier"""     )
		if "bbox_embed" in name:
				SCREAMING_SNAKE_CASE      :							List[Any]						=       name.replace("""bbox_embed"""       ,       """bbox_predictor"""     )
		if "vit.norm" in name:
				SCREAMING_SNAKE_CASE      :							Union[str, Any]						=       name.replace("""vit.norm"""       ,       """vit.layernorm"""     )

		return name



def 						__A						(   lowerCamelCase_       ,       lowerCamelCase_     ):


		"""simple docstring"""



		for key in orig_state_dict.copy().keys():
				SCREAMING_SNAKE_CASE      :							List[str]						=       orig_state_dict.pop(_UpperCamelCase     )

				if "qkv" in key:
						SCREAMING_SNAKE_CASE      :							Optional[int]						=       key.split("""."""     )
						SCREAMING_SNAKE_CASE      :							Optional[int]						=       int(key_split[2]     )
						SCREAMING_SNAKE_CASE      :							int						=       model.vit.encoder.layer[layer_num].attention.attention.all_head_size
						if "weight" in key:
								SCREAMING_SNAKE_CASE      :							Dict						=       val[:dim, :]
								SCREAMING_SNAKE_CASE      :							Dict						=       val[
								    dim : dim * 2, :
								]
								SCREAMING_SNAKE_CASE      :							List[str]						=       val[-dim:, :]
						else:
								SCREAMING_SNAKE_CASE      :							int						=       val[:dim]
								SCREAMING_SNAKE_CASE      :							Optional[int]						=       val[dim : dim * 2]
								SCREAMING_SNAKE_CASE      :							Optional[int]						=       val[-dim:]
				else:
						SCREAMING_SNAKE_CASE      :							List[Any]						=       val

		return orig_state_dict



def 						__A						(   ):


		"""simple docstring"""



		SCREAMING_SNAKE_CASE      :							Optional[Any]						=       '''http://images.cocodataset.org/val2017/000000039769.jpg'''
		SCREAMING_SNAKE_CASE      :							int						=       Image.open(requests.get(_UpperCamelCase       ,       stream=_UpperCamelCase     ).raw     )
		return im



@torch.no_grad()
def 						__A						(   lowerCamelCase_       ,       lowerCamelCase_       ,       lowerCamelCase_       ,       lowerCamelCase_ = False     ):


		"""simple docstring"""



		SCREAMING_SNAKE_CASE      :							Optional[int]						=       get_yolos_config(_UpperCamelCase     )

		# load original state_dict
		SCREAMING_SNAKE_CASE      :							Union[str, Any]						=       torch.load(_UpperCamelCase       ,       map_location="""cpu"""     )['''model''']

		# load 🤗 model
		SCREAMING_SNAKE_CASE      :							int						=       YolosForObjectDetection(_UpperCamelCase     )
		model.eval()
		SCREAMING_SNAKE_CASE      :							str						=       convert_state_dict(_UpperCamelCase       ,       _UpperCamelCase     )
		model.load_state_dict(_UpperCamelCase     )

		# Check outputs on an image, prepared by YolosImageProcessor
		SCREAMING_SNAKE_CASE      :							str						=       8_00 if yolos_name != '''yolos_ti''' else 5_12
		SCREAMING_SNAKE_CASE      :							Tuple						=       YolosImageProcessor(format="""coco_detection"""       ,       size=_UpperCamelCase     )
		SCREAMING_SNAKE_CASE      :							Any						=       image_processor(images=prepare_img()       ,       return_tensors="""pt"""     )
		SCREAMING_SNAKE_CASE      :							Tuple						=       model(**_UpperCamelCase     )
		SCREAMING_SNAKE_CASE      :							Optional[Any]						=       outputs.logits, outputs.pred_boxes

		SCREAMING_SNAKE_CASE      :							List[Any]						=       None, None
		if yolos_name == "yolos_ti":
				SCREAMING_SNAKE_CASE      :							Optional[int]						=       torch.tensor(
				    [[-39.5_022, -11.9_820, -17.6_888], [-29.9_574, -9.9_769, -17.7_691], [-42.3_281, -20.7_200, -30.6_294]]     )
				SCREAMING_SNAKE_CASE      :							Union[str, Any]						=       torch.tensor(
				    [[0.4_021, 0.0_836, 0.7_979], [0.0_184, 0.2_609, 0.0_364], [0.1_781, 0.2_004, 0.2_095]]     )
		elif yolos_name == "yolos_s_200_pre":
				SCREAMING_SNAKE_CASE      :							str						=       torch.tensor(
				    [[-24.0_248, -10.3_024, -14.8_290], [-42.0_392, -16.8_200, -27.4_334], [-27.2_743, -11.8_154, -18.7_148]]     )
				SCREAMING_SNAKE_CASE      :							int						=       torch.tensor(
				    [[0.2_559, 0.5_455, 0.4_706], [0.2_989, 0.7_279, 0.1_875], [0.7_732, 0.4_017, 0.4_462]]     )
		elif yolos_name == "yolos_s_300_pre":
				SCREAMING_SNAKE_CASE      :							str						=       torch.tensor(
				    [[-36.2_220, -14.4_385, -23.5_457], [-35.6_970, -14.7_583, -21.3_935], [-31.5_939, -13.6_042, -16.8_049]]     )
				SCREAMING_SNAKE_CASE      :							int						=       torch.tensor(
				    [[0.7_614, 0.2_316, 0.4_728], [0.7_168, 0.4_495, 0.3_855], [0.4_996, 0.1_466, 0.9_996]]     )
		elif yolos_name == "yolos_s_dWr":
				SCREAMING_SNAKE_CASE      :							List[str]						=       torch.tensor(
				    [[-42.8_668, -24.1_049, -41.1_690], [-34.7_456, -14.1_274, -24.9_194], [-33.7_898, -12.1_946, -25.6_495]]     )
				SCREAMING_SNAKE_CASE      :							Optional[int]						=       torch.tensor(
				    [[0.5_587, 0.2_773, 0.0_605], [0.5_004, 0.3_014, 0.9_994], [0.4_999, 0.1_548, 0.9_994]]     )
		elif yolos_name == "yolos_base":
				SCREAMING_SNAKE_CASE      :							Optional[int]						=       torch.tensor(
				    [[-40.6_064, -24.3_084, -32.6_447], [-55.1_990, -30.7_719, -35.5_877], [-51.4_311, -33.3_507, -35.6_462]]     )
				SCREAMING_SNAKE_CASE      :							Optional[Any]						=       torch.tensor(
				    [[0.5_555, 0.2_794, 0.0_655], [0.9_049, 0.2_664, 0.1_894], [0.9_183, 0.1_984, 0.1_635]]     )
		else:
				raise ValueError(f'''Unknown yolos_name: {yolos_name}'''     )

		assert torch.allclose(logits[0, :3, :3]       ,       _UpperCamelCase       ,       atol=1E-4     )
		assert torch.allclose(pred_boxes[0, :3, :3]       ,       _UpperCamelCase       ,       atol=1E-4     )

		Path(_UpperCamelCase     ).mkdir(exist_ok=_UpperCamelCase     )
		print(f'''Saving model {yolos_name} 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:
				SCREAMING_SNAKE_CASE      :							Optional[Any]						=       {
				    '''yolos_ti''': '''yolos-tiny''',
				    '''yolos_s_200_pre''': '''yolos-small''',
				    '''yolos_s_300_pre''': '''yolos-small-300''',
				    '''yolos_s_dWr''': '''yolos-small-dwr''',
				    '''yolos_base''': '''yolos-base''',
				}

				print("""Pushing to the hub..."""     )
				SCREAMING_SNAKE_CASE      :							List[str]						=       model_mapping[yolos_name]
				image_processor.push_to_hub(_UpperCamelCase       ,       organization="""hustvl"""     )
				model.push_to_hub(_UpperCamelCase       ,       organization="""hustvl"""     )


if __name__ == "__main__":
							__UpperCAmelCase			    =    argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--yolos_name""",
							    default="""yolos_s_200_pre""",
							    type=str,
							    help=(
							        """Name of the YOLOS model you\'d like to convert. Should be one of \'yolos_ti\', \'yolos_s_200_pre\',"""
							        """ \'yolos_s_300_pre\', \'yolos_s_dWr\', \'yolos_base\'."""
							    ),
							)
							parser.add_argument(
							    """--checkpoint_path""", default=None, type=str, help="""Path to the original state dict (.pth file)."""
							)
							parser.add_argument(
							    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
							)
							parser.add_argument(
							    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
							)

							__UpperCAmelCase			    =    parser.parse_args()
							convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)










from math import factorial

lowerCAmelCase_    =				{str(digit): factorial(digit) for digit in range(1_0)}






def      lowerCamelCase_						(				_UpperCamelCase )		->						int:






				"""simple docstring"""


				if not isinstance(_UpperCamelCase ,				_UpperCamelCase ):
								raise TypeError('''Parameter number must be int''' )

				if number < 0:
								raise ValueError('''Parameter number must be greater than or equal to 0''' )

				# Converts number in string to iterate on its digits and adds its factorial.
				return sum(DIGIT_FACTORIAL[digit] for digit in str(_UpperCamelCase ) )






def      lowerCamelCase_						(				_UpperCamelCase = 60 ,				_UpperCamelCase = 1_000_000 )		->						int:






				"""simple docstring"""


				if not isinstance(_UpperCamelCase ,				_UpperCamelCase ) or not isinstance(_UpperCamelCase ,				_UpperCamelCase ):
								raise TypeError('''Parameters chain_length and number_limit must be int''' )

				if chain_length <= 0 or number_limit <= 0:
								raise ValueError(
								    '''Parameters chain_length and number_limit must be greater than 0''' )

				# the counter for the chains with the exact desired length
				snake_case_						:    Optional[Any]				=  0
				# the cached sizes of the previous chains
				snake_case_						:    dict[int, int]				=  {}

				for start_chain_element in range(1 ,				_UpperCamelCase ):
								# The temporary set will contain the elements of the chain
								snake_case_						:    List[str]				=  set()
								snake_case_						:    List[Any]				=  0

								# Stop computing the chain when you find a cached size, a repeating item or the
								# length is greater then the desired one.
								snake_case_						:    Any				=  start_chain_element
								while (
								    chain_element not in chain_sets_lengths
								    and chain_element not in chain_set
								    and chain_set_length <= chain_length
								):
												chain_set.add(_UpperCamelCase )
												chain_set_length += 1
												snake_case_						:    List[Any]				=  digit_factorial_sum(_UpperCamelCase )

								if chain_element in chain_sets_lengths:
												chain_set_length += chain_sets_lengths[chain_element]

								snake_case_						:    List[str]				=  chain_set_length

								# If chain contains the exact amount of elements increase the counter
								if chain_set_length == chain_length:
												chains_counter += 1

				return chains_counter


if __name__ == "__main__":
	import doctest

	doctest.testmod()
	print(F'''{solution()}''')
