



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   :     str     ):



	"""simple docstring"""





	__SCREAMING_SNAKE_CASE       :    List[str]       =	[2, 2, 6, 2] if "tiny" in model_name else [2, 2, 18, 2]
	__SCREAMING_SNAKE_CASE       :    int       =	True if "large" in model_name or "huge" in model_name else False
	__SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	True if "large" in model_name or "huge" in model_name else False
	__SCREAMING_SNAKE_CASE       :    str       =	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:
			__SCREAMING_SNAKE_CASE       :    int       =	[3, 3, 3, 3]
			__SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	[5, 5, 5, 5]
		elif "fl4" in model_name:
			__SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	[4, 4, 4, 4]
			__SCREAMING_SNAKE_CASE       :    List[Any]       =	[3, 3, 3, 3]

	if "tiny" in model_name or "small" in model_name or "base" in model_name:
		__SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	[3, 3, 3, 3]
		if "lrf" in model_name:
			__SCREAMING_SNAKE_CASE       :    Dict       =	[3, 3, 3, 3]
		else:
			__SCREAMING_SNAKE_CASE       :    List[Any]       =	[2, 2, 2, 2]

	if "tiny" in model_name:
		__SCREAMING_SNAKE_CASE       :    Dict       =	96
	elif "small" in model_name:
		__SCREAMING_SNAKE_CASE       :    Tuple       =	96
	elif "base" in model_name:
		__SCREAMING_SNAKE_CASE       :    Dict       =	128
	elif "large" in model_name:
		__SCREAMING_SNAKE_CASE       :    List[Any]       =	192
	elif "xlarge" in model_name:
		__SCREAMING_SNAKE_CASE       :    List[str]       =	256
	elif "huge" in model_name:
		__SCREAMING_SNAKE_CASE       :    Optional[int]       =	352

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

	__SCREAMING_SNAKE_CASE       :    Tuple       =	json.load(open(hf_hub_download(__lowerCamelCase		,	__lowerCamelCase		,	repo_type="dataset"     )		,	"r"     )     )
	__SCREAMING_SNAKE_CASE       :    Optional[int]       =	{int(__lowerCamelCase     ): v for k, v in idalabel.items()}
	__SCREAMING_SNAKE_CASE       :    List[str]       =	{v: k for k, v in idalabel.items()}

	__SCREAMING_SNAKE_CASE       :    Dict       =	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   :     Tuple     ):



	"""simple docstring"""





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

	if name == "norm.weight":
		__SCREAMING_SNAKE_CASE       :    Dict       =	"layernorm.weight"
	if name == "norm.bias":
		__SCREAMING_SNAKE_CASE       :    Optional[int]       =	"layernorm.bias"

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

	return name






def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[Any]		,	__lowerCamelCase   :     Optional[int]		,	__lowerCamelCase   :     Dict=False     ):



	"""simple docstring"""





	__SCREAMING_SNAKE_CASE       :    Dict       =	{
	    "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

	__SCREAMING_SNAKE_CASE       :    str       =	model_name_to_url[model_name]
	print("Checkpoint URL: "		,	__lowerCamelCase     )
	__SCREAMING_SNAKE_CASE       :    List[str]       =	torch.hub.load_state_dict_from_url(__lowerCamelCase		,	map_location="cpu"     )["model"]

	# rename keys
	for key in state_dict.copy().keys():
		__SCREAMING_SNAKE_CASE       :    List[str]       =	state_dict.pop(__lowerCamelCase     )
		__SCREAMING_SNAKE_CASE       :    Optional[int]       =	val

	__SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	get_focalnet_config(__lowerCamelCase     )
	__SCREAMING_SNAKE_CASE       :    str       =	FocalNetForImageClassification(__lowerCamelCase     )
	model.eval()

	# load state dict
	model.load_state_dict(__lowerCamelCase     )

	# verify conversion
	__SCREAMING_SNAKE_CASE       :    Any       =	"http://images.cocodataset.org/val2017/000000039769.jpg"

	__SCREAMING_SNAKE_CASE       :    Any       =	BitImageProcessor(
	    do_resize=__lowerCamelCase		,	size={"shortest_edge": 256}		,	resample=PILImageResampling.BILINEAR		,	do_center_crop=__lowerCamelCase		,	crop_size=224		,	do_normalize=__lowerCamelCase		,	image_mean=__lowerCamelCase		,	image_std=__lowerCamelCase		,	)
	__SCREAMING_SNAKE_CASE       :    Tuple       =	Image.open(requests.get(__lowerCamelCase		,	stream=__lowerCamelCase     ).raw     )
	__SCREAMING_SNAKE_CASE       :    str       =	processor(images=__lowerCamelCase		,	return_tensors="pt"     )

	__SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	transforms.Compose(
	    [
	        transforms.Resize(256     ),
	        transforms.CenterCrop(224     ),
	        transforms.ToTensor(),
	        transforms.Normalize(mean=[0.485, 0.456, 0.406]		,	std=[0.229, 0.224, 0.225]     ),
	    ]     )

	__SCREAMING_SNAKE_CASE       :    List[str]       =	image_transforms(__lowerCamelCase     ).unsqueeze(0     )

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

	__SCREAMING_SNAKE_CASE       :    str       =	model(**__lowerCamelCase     )

	__SCREAMING_SNAKE_CASE       :    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":
		__SCREAMING_SNAKE_CASE       :    List[str]       =	torch.tensor([0.2166, -0.4368, 0.2191]     )
	elif model_name == "focalnet-tiny-lrf":
		__SCREAMING_SNAKE_CASE       :    Any       =	torch.tensor([1.1669, 0.0125, -0.1695]     )
	elif model_name == "focalnet-small":
		__SCREAMING_SNAKE_CASE       :    Tuple       =	torch.tensor([0.4917, -0.0430, 0.1341]     )
	elif model_name == "focalnet-small-lrf":
		__SCREAMING_SNAKE_CASE       :    Tuple       =	torch.tensor([-0.2588, -0.5342, -0.2331]     )
	elif model_name == "focalnet-base":
		__SCREAMING_SNAKE_CASE       :    int       =	torch.tensor([-0.1655, -0.4090, -0.1730]     )
	elif model_name == "focalnet-base-lrf":
		__SCREAMING_SNAKE_CASE       :    Any       =	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__":
						_lowerCamelCase						    =  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.""",
						)

						_lowerCamelCase						    =  parser.parse_args()
						convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)











from math import sqrt






def 	_lowerCAmelCase						(				__lowerCamelCase   :     int = 1000000     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    int       =	0
 __SCREAMING_SNAKE_CASE       :    int       =	0
 __SCREAMING_SNAKE_CASE       :    int

 while num_cuboids <= limit:
  max_cuboid_size += 1
  for sum_shortest_sides in range(2		,	2 * max_cuboid_size + 1     ):
   if sqrt(sum_shortest_sides**2 + max_cuboid_size**2     ).is_integer():
    num_cuboids += (
        min(__lowerCamelCase		,	sum_shortest_sides // 2     )
        - max(1		,	sum_shortest_sides - max_cuboid_size     )
        + 1
    )

 return max_cuboid_size


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






