






"""simple docstring"""
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       __UpperCamelCase						(		snake_case__       ):
	A_		:    List[Any]						     = [2, 2, 6, 2] if """tiny""" in model_name else [2, 2, 18, 2]
	A_		:    Any						     = True if """large""" in model_name or """huge""" in model_name else False
	A_		:    Union[str, Any]						     = True if """large""" in model_name or """huge""" in model_name else False
	A_		:    Optional[int]						     = 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:
			A_		:    List[str]						     = [3, 3, 3, 3]
			A_		:    Tuple						     = [5, 5, 5, 5]
		elif "fl4" in model_name:
			A_		:    Union[str, Any]						     = [4, 4, 4, 4]
			A_		:    Union[str, Any]						     = [3, 3, 3, 3]

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

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

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

	A_		:    Any						     = json.load(open(hf_hub_download(snake_case__ , snake_case__ , repo_type="""dataset"""       ) , """r"""       )       )
	A_		:    Any						     = {int(snake_case__       ): v for k, v in idalabel.items()}
	A_		:    Tuple						     = {v: k for k, v in idalabel.items()}

	A_		:    str						     = FocalNetConfig(
	    embed_dim=snake_case__ , depths=snake_case__ , focal_levels=snake_case__ , focal_windows=snake_case__ , use_conv_embed=snake_case__ , idalabel=snake_case__ , labelaid=snake_case__ , use_post_layernorm=snake_case__ , use_layerscale=snake_case__ , )

	return config






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

	if name == "norm.weight":
		A_		:    Any						     = """layernorm.weight"""
	if name == "norm.bias":
		A_		:    List[str]						     = """layernorm.bias"""

	if "head" in name:
		A_		:    str						     = name.replace("""head""" , """classifier"""       )
	else:
		A_		:    int						     = """focalnet.""" + name

	return name






def       __UpperCamelCase						(		snake_case__ , snake_case__ , snake_case__=False       ):
	# fmt: off
	A_		:    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

	A_		:    Dict						     = model_name_to_url[model_name]
	print("""Checkpoint URL: """ , snake_case__       )
	A_		:    Optional[int]						     = torch.hub.load_state_dict_from_url(snake_case__ , map_location="""cpu"""       )["""model"""]

	# rename keys
	for key in state_dict.copy().keys():
		A_		:    Union[str, Any]						     = state_dict.pop(snake_case__       )
		A_		:    List[str]						     = val

	A_		:    int						     = get_focalnet_config(snake_case__       )
	A_		:    List[Any]						     = FocalNetForImageClassification(snake_case__       )
	model.eval()

	# load state dict
	model.load_state_dict(snake_case__       )

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

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

	A_		:    Optional[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]       ),
	    ]       )

	A_		:    str						     = image_transforms(snake_case__       ).unsqueeze(0       )

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

	A_		:    List[Any]						     = model(**snake_case__       )

	A_		:    Optional[Any]						     = 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":
		A_		:    int						     = torch.tensor([0.2_166, -0.4_368, 0.2_191]       )
	elif model_name == "focalnet-tiny-lrf":
		A_		:    Tuple						     = torch.tensor([1.1_669, 0.0_125, -0.1_695]       )
	elif model_name == "focalnet-small":
		A_		:    Optional[int]						     = torch.tensor([0.4_917, -0.0_430, 0.1_341]       )
	elif model_name == "focalnet-small-lrf":
		A_		:    Tuple						     = torch.tensor([-0.2_588, -0.5_342, -0.2_331]       )
	elif model_name == "focalnet-base":
		A_		:    Any						     = torch.tensor([-0.1_655, -0.4_090, -0.1_730]       )
	elif model_name == "focalnet-base-lrf":
		A_		:    str						     = torch.tensor([0.5_306, -0.0_483, -0.3_928]       )
	assert torch.allclose(outputs.logits[0, :3] , snake_case__ , 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(snake_case__       )
		processor.save_pretrained(snake_case__       )

	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)









"""simple docstring"""
from typing import Dict

from .base import GenericTensor, Pipeline



class   SCREAMING_SNAKE_CASE	(					_SCREAMING_SNAKE_CASE    ):





	"""simple docstring"""




	def  lowerCamelCase(self  ,  lowerCAmelCase_=None  ,  lowerCAmelCase_=None  ,  lowerCAmelCase_=None  ,  **lowerCAmelCase_				):
		if tokenize_kwargs is None:
			A_		:    Dict						     = {}

		if truncation is not None:
			if "truncation" in tokenize_kwargs:
				raise ValueError(
				    """truncation parameter defined twice (given as keyword argument as well as in tokenize_kwargs)"""				)
			A_		:    Optional[Any]						     = truncation

		A_		:    Optional[int]						     = tokenize_kwargs

		A_		:    Optional[Any]						     = {}
		if return_tensors is not None:
			A_		:    Optional[Any]						     = return_tensors

		return preprocess_params, {}, postprocess_params





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  **lowerCAmelCase_				):
		A_		:    List[Any]						     = self.framework
		A_		:    Optional[int]						     = self.tokenizer(lowerCAmelCase_  ,  return_tensors=lowerCAmelCase_  ,  **lowerCAmelCase_				)
		return model_inputs





	def  lowerCamelCase(self  ,  lowerCAmelCase_				):
		A_		:    List[Any]						     = self.model(**lowerCAmelCase_				)
		return model_outputs





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_=False				):
		# [0] is the first available tensor, logits or last_hidden_state.
		if return_tensors:
			return model_outputs[0]
		if self.framework == "pt":
			return model_outputs[0].tolist()
		elif self.framework == "tf":
			return model_outputs[0].numpy().tolist()





	def __call__(self  ,  *lowerCAmelCase_  ,  **lowerCAmelCase_				):
		return super().__call__(*lowerCAmelCase_  ,  **lowerCAmelCase_				)

