



import argparse
import json

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

from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling







def       lowerCAmelCase_     (lowerCAmelCase__:						Tuple   ):





			"""simple docstring"""


			UpperCAmelCase_:						Tuple					=			[2, 2, 6, 2] if """tiny""" in model_name else [2, 2, 1_8, 2]
			UpperCAmelCase_:						List[str]					=			True if """large""" in model_name or """huge""" in model_name else False
			UpperCAmelCase_:						List[Any]					=			True if """large""" in model_name or """huge""" in model_name else False
			UpperCAmelCase_:						Tuple					=			True if """large""" in model_name or """huge""" in model_name else False

			if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
						if "fl3" in model_name:
									UpperCAmelCase_:						Union[str, Any]					=			[3, 3, 3, 3]
									UpperCAmelCase_:						List[str]					=			[5, 5, 5, 5]
						elif "fl4" in model_name:
									UpperCAmelCase_:						List[str]					=			[4, 4, 4, 4]
									UpperCAmelCase_:						Optional[int]					=			[3, 3, 3, 3]

			if "tiny" in model_name or "small" in model_name or "base" in model_name:
						UpperCAmelCase_:						Any					=			[3, 3, 3, 3]
						if "lrf" in model_name:
									UpperCAmelCase_:						Optional[Any]					=			[3, 3, 3, 3]
						else:
									UpperCAmelCase_:						Any					=			[2, 2, 2, 2]

			if "tiny" in model_name:
						UpperCAmelCase_:						List[Any]					=			9_6
			elif "small" in model_name:
						UpperCAmelCase_:						Optional[int]					=			9_6
			elif "base" in model_name:
						UpperCAmelCase_:						Dict					=			1_2_8
			elif "large" in model_name:
						UpperCAmelCase_:						Dict					=			1_9_2
			elif "xlarge" in model_name:
						UpperCAmelCase_:						int					=			2_5_6
			elif "huge" in model_name:
						UpperCAmelCase_:						Union[str, Any]					=			3_5_2

			# set label information
			UpperCAmelCase_:						int					=			"""huggingface/label-files"""
			if "large" in model_name or "huge" in model_name:
						UpperCAmelCase_:						Optional[Any]					=			"""imagenet-22k-id2label.json"""
			else:
						UpperCAmelCase_:						int					=			"""imagenet-1k-id2label.json"""

			UpperCAmelCase_:						Optional[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_:						Union[str, Any]					=			{v: k for k, v in idalabel.items()}

			UpperCAmelCase_:						List[str]					=			FocalNetConfig(
			    embed_dim=lowerCAmelCase__						,		depths=lowerCAmelCase__						,		focal_levels=lowerCAmelCase__						,		focal_windows=lowerCAmelCase__						,		use_conv_embed=lowerCAmelCase__						,		idalabel=lowerCAmelCase__						,		labelaid=lowerCAmelCase__						,		use_post_layernorm=lowerCAmelCase__						,		use_layerscale=lowerCAmelCase__						,		)

			return config







def       lowerCAmelCase_     (lowerCAmelCase__:						Dict   ):





			"""simple docstring"""


			if "patch_embed.proj" in name:
						UpperCAmelCase_:						Union[str, Any]					=			name.replace("""patch_embed.proj"""						,		"""embeddings.patch_embeddings.projection"""   )
			if "patch_embed.norm" in name:
						UpperCAmelCase_:						List[str]					=			name.replace("""patch_embed.norm"""						,		"""embeddings.norm"""   )
			if "layers" in name:
						UpperCAmelCase_:						Union[str, Any]					=			"""encoder.""" + name
			if "encoder.layers" in name:
						UpperCAmelCase_:						str					=			name.replace("""encoder.layers"""						,		"""encoder.stages"""   )
			if "downsample.proj" in name:
						UpperCAmelCase_:						Dict					=			name.replace("""downsample.proj"""						,		"""downsample.projection"""   )
			if "blocks" in name:
						UpperCAmelCase_:						Optional[int]					=			name.replace("""blocks"""						,		"""layers"""   )
			if "modulation.f.weight" in name or "modulation.f.bias" in name:
						UpperCAmelCase_:						List[str]					=			name.replace("""modulation.f"""						,		"""modulation.projection_in"""   )
			if "modulation.h.weight" in name or "modulation.h.bias" in name:
						UpperCAmelCase_:						List[str]					=			name.replace("""modulation.h"""						,		"""modulation.projection_context"""   )
			if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
						UpperCAmelCase_:						Tuple					=			name.replace("""modulation.proj"""						,		"""modulation.projection_out"""   )

			if name == "norm.weight":
						UpperCAmelCase_:						List[Any]					=			"""layernorm.weight"""
			if name == "norm.bias":
						UpperCAmelCase_:						int					=			"""layernorm.bias"""

			if "head" in name:
						UpperCAmelCase_:						Optional[Any]					=			name.replace("""head"""						,		"""classifier"""   )
			else:
						UpperCAmelCase_:						Optional[int]					=			"""focalnet.""" + name

			return name







def       lowerCAmelCase_     (lowerCAmelCase__:						Any						,		lowerCAmelCase__:						Optional[Any]						,		lowerCAmelCase__:						Tuple=False   ):





			"""simple docstring"""


			UpperCAmelCase_:						Optional[Any]					=			{
			    """focalnet-tiny""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth""",
			    """focalnet-tiny-lrf""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth""",
			    """focalnet-small""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth""",
			    """focalnet-small-lrf""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth""",
			    """focalnet-base""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth""",
			    """focalnet-base-lrf""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth""",
			    """focalnet-large-lrf-fl3""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth""",
			    """focalnet-large-lrf-fl4""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth""",
			    """focalnet-xlarge-lrf-fl3""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth""",
			    """focalnet-xlarge-lrf-fl4""": """https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth""",
			}
			# fmt: on

			UpperCAmelCase_:						Optional[int]					=			model_name_to_url[model_name]
			print("""Checkpoint URL: """						,		lowerCAmelCase__   )
			UpperCAmelCase_:						Tuple					=			torch.hub.load_state_dict_from_url(lowerCAmelCase__						,		map_location="""cpu"""   )["""model"""]

			# rename keys
			for key in state_dict.copy().keys():
						UpperCAmelCase_:						List[Any]					=			state_dict.pop(lowerCAmelCase__   )
						UpperCAmelCase_:						Tuple					=			val

			UpperCAmelCase_:						Optional[Any]					=			get_focalnet_config(lowerCAmelCase__   )
			UpperCAmelCase_:						Dict					=			FocalNetForImageClassification(lowerCAmelCase__   )
			model.eval()

			# load state dict
			model.load_state_dict(lowerCAmelCase__   )

			# verify conversion
			UpperCAmelCase_:						str					=			"""http://images.cocodataset.org/val2017/000000039769.jpg"""

			UpperCAmelCase_:						Optional[int]					=			BitImageProcessor(
			    do_resize=lowerCAmelCase__						,		size={"""shortest_edge""": 2_5_6}						,		resample=PILImageResampling.BILINEAR						,		do_center_crop=lowerCAmelCase__						,		crop_size=2_2_4						,		do_normalize=lowerCAmelCase__						,		image_mean=lowerCAmelCase__						,		image_std=lowerCAmelCase__						,		)
			UpperCAmelCase_:						Any					=			Image.open(requests.get(lowerCAmelCase__						,		stream=lowerCAmelCase__   ).raw   )
			UpperCAmelCase_:						Dict					=			processor(images=lowerCAmelCase__						,		return_tensors="""pt"""   )

			UpperCAmelCase_:						Optional[int]					=			transforms.Compose(
			    [
			        transforms.Resize(2_5_6   ),
			        transforms.CenterCrop(2_2_4   ),
			        transforms.ToTensor(),
			        transforms.Normalize(mean=[0.485, 0.456, 0.406]						,		std=[0.229, 0.224, 0.225]   ),
			    ]   )

			UpperCAmelCase_:						List[Any]					=			image_transforms(lowerCAmelCase__   ).unsqueeze(0   )

			# verify pixel_values
			assert torch.allclose(inputs.pixel_values						,		lowerCAmelCase__						,		atol=1e-4   )

			UpperCAmelCase_:						int					=			model(**lowerCAmelCase__   )

			UpperCAmelCase_:						Optional[int]					=			outputs.logits.argmax(-1   ).item()
			print("""Predicted class:"""						,		model.config.idalabel[predicted_class_idx]   )

			print("""First values of logits:"""						,		outputs.logits[0, :3]   )

			if model_name == "focalnet-tiny":
						UpperCAmelCase_:						List[str]					=			torch.tensor([0.2166, -0.4368, 0.2191]   )
			elif model_name == "focalnet-tiny-lrf":
						UpperCAmelCase_:						Union[str, Any]					=			torch.tensor([1.1669, 0.0125, -0.1695]   )
			elif model_name == "focalnet-small":
						UpperCAmelCase_:						Union[str, Any]					=			torch.tensor([0.4917, -0.0430, 0.1341]   )
			elif model_name == "focalnet-small-lrf":
						UpperCAmelCase_:						Optional[int]					=			torch.tensor([-0.2588, -0.5342, -0.2331]   )
			elif model_name == "focalnet-base":
						UpperCAmelCase_:						Optional[int]					=			torch.tensor([-0.1655, -0.4090, -0.1730]   )
			elif model_name == "focalnet-base-lrf":
						UpperCAmelCase_:						Optional[int]					=			torch.tensor([0.5306, -0.0483, -0.3928]   )
			assert torch.allclose(outputs.logits[0, :3]						,		lowerCAmelCase__						,		atol=1e-4   )
			print("""Looks ok!"""   )

			if pytorch_dump_folder_path is not None:
						print(F'Saving model and processor of {model_name} to {pytorch_dump_folder_path}'   )
						model.save_pretrained(lowerCAmelCase__   )
						processor.save_pretrained(lowerCAmelCase__   )

			if push_to_hub:
						print(F'Pushing model and processor of {model_name} to the hub...'   )
						model.push_to_hub(F'{model_name}'   )
						processor.push_to_hub(F'{model_name}'   )


if __name__ == "__main__":
	a	:   Optional[int]       			=       argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '--model_name',
	    default='focalnet-tiny',
	    type=str,
	    help='Name of the FocalNet model 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(
	    '--push_to_hub',
	    action='store_true',
	    help='Whether to push the model and processor to the hub.',
	)

	a	:   str       			=       parser.parse_args()
	convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)











"""simple docstring"""
from math import isqrt, loga

def       _snake_case						(						UpperCamelCase						:		int						):
 UpperCAmelCase	:					Any	=	[True] * max_number
 for i in range(2			,		isqrt(max_number - 1						) + 1						):
  if is_prime[i]:
   for j in range(i**2			,		UpperCamelCase			,		UpperCamelCase						):
    UpperCAmelCase	:					str	=	False

 return [i for i in range(2			,		UpperCamelCase						) if is_prime[i]]

def       _snake_case						(						UpperCamelCase						:		int = 800800			,		UpperCamelCase						:		int = 800800						):
 UpperCAmelCase	:					Union[str, Any]	=	degree * loga(UpperCamelCase						)
 UpperCAmelCase	:					int	=	int(UpperCamelCase						)
 UpperCAmelCase	:					Union[str, Any]	=	calculate_prime_numbers(UpperCamelCase						)

 UpperCAmelCase	:					Dict	=	0
 UpperCAmelCase	:					Optional[int]	=	0
 UpperCAmelCase	:					Dict	=	len(UpperCamelCase						) - 1
 while left < right:
  while (
      prime_numbers[right] * loga(prime_numbers[left]						)
      + prime_numbers[left] * loga(prime_numbers[right]						)
      > upper_bound
  ):
   right -= 1
  hybrid_integers_count += right - left
  left += 1

 return hybrid_integers_count


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