


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 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		Any	):





			"""simple docstring"""




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

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

			if "tiny" in model_name:
						UpperCamelCase         =					96
			elif "small" in model_name:
						UpperCamelCase         =					96
			elif "base" in model_name:
						UpperCamelCase         =					128
			elif "large" in model_name:
						UpperCamelCase         =					192
			elif "xlarge" in model_name:
						UpperCamelCase         =					256
			elif "huge" in model_name:
						UpperCamelCase         =					352

			# set label information
			UpperCamelCase         =					"""huggingface/label-files"""
			if "large" in model_name or "huge" in model_name:
						UpperCamelCase         =					"""imagenet-22k-id2label.json"""
			else:
						UpperCamelCase         =					"""imagenet-1k-id2label.json"""

			UpperCamelCase         =					json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  repo_type="""dataset"""	)       ,  """r"""	)	)
			UpperCamelCase         =					{int(SCREAMING_SNAKE_CASE_	): v for k, v in idalabel.items()}
			UpperCamelCase         =					{v: k for k, v in idalabel.items()}

			UpperCamelCase         =					FocalNetConfig(
			    embed_dim=SCREAMING_SNAKE_CASE_       ,  depths=SCREAMING_SNAKE_CASE_       ,  focal_levels=SCREAMING_SNAKE_CASE_       ,  focal_windows=SCREAMING_SNAKE_CASE_       ,  use_conv_embed=SCREAMING_SNAKE_CASE_       ,  idalabel=SCREAMING_SNAKE_CASE_       ,  labelaid=SCREAMING_SNAKE_CASE_       ,  use_post_layernorm=SCREAMING_SNAKE_CASE_       ,  use_layerscale=SCREAMING_SNAKE_CASE_       ,  )

			return config






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		Union[str, Any]	):





			"""simple docstring"""




			if "patch_embed.proj" in name:
						UpperCamelCase         =					name.replace("""patch_embed.proj"""       ,  """embeddings.patch_embeddings.projection"""	)
			if "patch_embed.norm" in name:
						UpperCamelCase         =					name.replace("""patch_embed.norm"""       ,  """embeddings.norm"""	)
			if "layers" in name:
						UpperCamelCase         =					"""encoder.""" + name
			if "encoder.layers" in name:
						UpperCamelCase         =					name.replace("""encoder.layers"""       ,  """encoder.stages"""	)
			if "downsample.proj" in name:
						UpperCamelCase         =					name.replace("""downsample.proj"""       ,  """downsample.projection"""	)
			if "blocks" in name:
						UpperCamelCase         =					name.replace("""blocks"""       ,  """layers"""	)
			if "modulation.f.weight" in name or "modulation.f.bias" in name:
						UpperCamelCase         =					name.replace("""modulation.f"""       ,  """modulation.projection_in"""	)
			if "modulation.h.weight" in name or "modulation.h.bias" in name:
						UpperCamelCase         =					name.replace("""modulation.h"""       ,  """modulation.projection_context"""	)
			if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
						UpperCamelCase         =					name.replace("""modulation.proj"""       ,  """modulation.projection_out"""	)

			if name == "norm.weight":
						UpperCamelCase         =					"""layernorm.weight"""
			if name == "norm.bias":
						UpperCamelCase         =					"""layernorm.bias"""

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

			return name






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int       ,  SCREAMING_SNAKE_CASE_   :		Any       ,  SCREAMING_SNAKE_CASE_   :		Union[str, Any]=False	):





			"""simple docstring"""




			UpperCamelCase         =					{
			    """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         =					model_name_to_url[model_name]
			print("""Checkpoint URL: """       ,  SCREAMING_SNAKE_CASE_	)
			UpperCamelCase         =					torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_       ,  map_location="""cpu"""	)["""model"""]

			# rename keys
			for key in state_dict.copy().keys():
						UpperCamelCase         =					state_dict.pop(SCREAMING_SNAKE_CASE_	)
						UpperCamelCase         =					val

			UpperCamelCase         =					get_focalnet_config(SCREAMING_SNAKE_CASE_	)
			UpperCamelCase         =					FocalNetForImageClassification(SCREAMING_SNAKE_CASE_	)
			model.eval()

			# load state dict
			model.load_state_dict(SCREAMING_SNAKE_CASE_	)

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

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

			UpperCamelCase         =					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]	),
			    ]	)

			UpperCamelCase         =					image_transforms(SCREAMING_SNAKE_CASE_	).unsqueeze(0	)

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

			UpperCamelCase         =					model(**SCREAMING_SNAKE_CASE_	)

			UpperCamelCase         =					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         =					torch.tensor([0.2166, -0.4368, 0.2191]	)
			elif model_name == "focalnet-tiny-lrf":
						UpperCamelCase         =					torch.tensor([1.1669, 0.0125, -0.1695]	)
			elif model_name == "focalnet-small":
						UpperCamelCase         =					torch.tensor([0.4917, -0.0430, 0.1341]	)
			elif model_name == "focalnet-small-lrf":
						UpperCamelCase         =					torch.tensor([-0.2588, -0.5342, -0.2331]	)
			elif model_name == "focalnet-base":
						UpperCamelCase         =					torch.tensor([-0.1655, -0.4090, -0.1730]	)
			elif model_name == "focalnet-base-lrf":
						UpperCamelCase         =					torch.tensor([0.5306, -0.0483, -0.3928]	)
			assert torch.allclose(outputs.logits[0, :3]       ,  SCREAMING_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(SCREAMING_SNAKE_CASE_	)
						processor.save_pretrained(SCREAMING_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__":
			__snake_case            =   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.",
			)

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








from ..utils import DummyObject, requires_backends




class    UpperCAmelCase					(       metaclass=__snake_case    ):
			lowercase						= ["""keras_nlp"""]




			def __init__(						self							:		List[str]       ,       *__magic_name__							:		str       ,       **__magic_name__							:		int						):







						"""simple docstring"""







						requires_backends(self       ,       ["""keras_nlp"""]						)




