



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_    ) ->       Optional[Any]:
						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(lowerCamelCase_    ,   lowerCamelCase_    ,   repo_type="dataset"    )    ,   "r"    )    )
						UpperCAmelCase         =     {int(lowerCamelCase_    ): v for k, v in idalabel.items()}
						UpperCAmelCase         =     {v: k for k, v in idalabel.items()}

						UpperCAmelCase         =     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:
						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    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_=False    ) ->       Dict:
						# fmt: off
						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: "    ,   lowerCamelCase_    )
						UpperCAmelCase         =     torch.hub.load_state_dict_from_url(lowerCamelCase_    ,   map_location="cpu"    )["model"]

						# rename keys
						for key in state_dict.copy().keys():
												UpperCAmelCase         =     state_dict.pop(lowerCamelCase_    )
												UpperCAmelCase         =     val

						UpperCAmelCase         =     get_focalnet_config(lowerCamelCase_    )
						UpperCAmelCase         =     FocalNetForImageClassification(lowerCamelCase_    )
						model.eval()

						# load state dict
						model.load_state_dict(lowerCamelCase_    )

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

						UpperCAmelCase         =     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_    ,   )
						UpperCAmelCase         =     Image.open(requests.get(lowerCamelCase_    ,   stream=lowerCamelCase_    ).raw    )
						UpperCAmelCase         =     processor(images=lowerCamelCase_    ,   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(lowerCamelCase_    ).unsqueeze(0    )

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

						UpperCAmelCase         =     model(**lowerCamelCase_    )

						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]    ,   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     :							Optional[Any]  							=				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     :							Any  							=				parser.parse_args()
		convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








import inspect
import unittest

from transformers import RegNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.models.regnet.modeling_flax_regnet import FlaxRegNetForImageClassification, FlaxRegNetModel

if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor





class  __magic_name__   (							unittest.TestCase						):






		def __init__(						self      :     List[Any]				,		UpperCamelCase__      :     Dict				,		UpperCamelCase__      :     Union[str, Any]=3				,		UpperCamelCase__      :     Optional[Any]=32				,		UpperCamelCase__      :     Optional[int]=3				,		UpperCamelCase__      :     str=10				,		UpperCamelCase__      :     Optional[int]=[10, 20, 30, 40]				,		UpperCamelCase__      :     int=[1, 1, 2, 1]				,		UpperCamelCase__      :     Any=True				,		UpperCamelCase__      :     Any=True				,		UpperCamelCase__      :     Any="relu"				,		UpperCamelCase__      :     List[str]=3				,		UpperCamelCase__      :     str=None				,		)				->					Optional[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     parent
								UpperCAmelCase         =     batch_size
								UpperCAmelCase         =     image_size
								UpperCAmelCase         =     num_channels
								UpperCAmelCase         =     embeddings_size
								UpperCAmelCase         =     hidden_sizes
								UpperCAmelCase         =     depths
								UpperCAmelCase         =     is_training
								UpperCAmelCase         =     use_labels
								UpperCAmelCase         =     hidden_act
								UpperCAmelCase         =     num_labels
								UpperCAmelCase         =     scope
								UpperCAmelCase         =     len(UpperCamelCase__      )






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]      )				->					Optional[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

								UpperCAmelCase         =     self.get_config()

								return config, pixel_values






		def SCREAMING_SNAKE_CASE_			(						self      :     Any      )				->					Any:
								'''simple docstring'''







								return RegNetConfig(
								    num_channels=self.num_channels				,		embeddings_size=self.embeddings_size				,		hidden_sizes=self.hidden_sizes				,		depths=self.depths				,		hidden_act=self.hidden_act				,		num_labels=self.num_labels				,		image_size=self.image_size				,		)






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]				,		UpperCamelCase__      :     Union[str, Any]				,		UpperCamelCase__      :     List[str]      )				->					Any:
								'''simple docstring'''







								UpperCAmelCase         =     FlaxRegNetModel(config=UpperCamelCase__      )
								UpperCAmelCase         =     model(UpperCamelCase__      )

								# Output shape (b, c, h, w)
								self.parent.assertEqual(
								    result.last_hidden_state.shape				,		(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)				,		)






		def SCREAMING_SNAKE_CASE_			(						self      :     Union[str, Any]				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     str      )				->					Union[str, Any]:
								'''simple docstring'''







								UpperCAmelCase         =     self.num_labels
								UpperCAmelCase         =     FlaxRegNetForImageClassification(config=UpperCamelCase__      )
								UpperCAmelCase         =     model(UpperCamelCase__      )
								self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.num_labels)      )







		def SCREAMING_SNAKE_CASE_			(						self      :     str      )				->					List[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     self.prepare_config_and_inputs()
								UpperCAmelCase       ,		UpperCAmelCase         =     config_and_inputs
								UpperCAmelCase         =     {"pixel_values": pixel_values}
								return config, inputs_dict





@require_flax
class  __magic_name__   (							A__,  unittest.TestCase						):
		lowercase      :       List[str]          =(FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

		lowercase      :       Tuple          =False
		lowercase      :       Union[str, Any]          =False
		lowercase      :       Dict          =False






		def SCREAMING_SNAKE_CASE_			(						self      :     List[Any]      )				->					None:
								'''simple docstring'''







								UpperCAmelCase         =     FlaxRegNetModelTester(self      )
								UpperCAmelCase         =     ConfigTester(self				,		config_class=UpperCamelCase__				,		has_text_modality=UpperCamelCase__      )






		def SCREAMING_SNAKE_CASE_			(						self      :     List[str]      )				->					Optional[Any]:
								'''simple docstring'''







								self.create_and_test_config_common_properties()
								self.config_tester.create_and_test_config_to_json_string()
								self.config_tester.create_and_test_config_to_json_file()
								self.config_tester.create_and_test_config_from_and_save_pretrained()
								self.config_tester.create_and_test_config_with_num_labels()
								self.config_tester.check_config_can_be_init_without_params()
								self.config_tester.check_config_arguments_init()






		def SCREAMING_SNAKE_CASE_			(						self      :     int      )				->					Dict:
								'''simple docstring'''







								return






		def SCREAMING_SNAKE_CASE_			(						self      :     List[Any]      )				->					List[str]:
								'''simple docstring'''







								UpperCAmelCase         =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*UpperCamelCase__      )






		def SCREAMING_SNAKE_CASE_			(						self      :     str      )				->					Tuple:
								'''simple docstring'''







								UpperCAmelCase         =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__      )






		@unittest.skip(reason="RegNet does not use inputs_embeds"      )
		def SCREAMING_SNAKE_CASE_			(						self      :     str      )				->					Optional[int]:
								'''simple docstring'''







								pass






		@unittest.skip(reason="RegNet does not support input and output embeddings"      )
		def SCREAMING_SNAKE_CASE_			(						self      :     int      )				->					Optional[Any]:
								'''simple docstring'''







								pass






		def SCREAMING_SNAKE_CASE_			(						self      :     List[str]      )				->					Optional[Any]:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase         =     self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )
														UpperCAmelCase         =     inspect.signature(model.__call__      )
														# signature.parameters is an OrderedDict => so arg_names order is deterministic
														UpperCAmelCase         =     [*signature.parameters.keys()]

														UpperCAmelCase         =     ["pixel_values"]
														self.assertListEqual(arg_names[:1]				,		UpperCamelCase__      )






		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple      )				->					Dict:
								'''simple docstring'''







								def check_hidden_states_output(UpperCamelCase__      :     List[str]				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     Any      ):
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

														UpperCAmelCase         =     model(**self._prepare_for_class(UpperCamelCase__				,		UpperCamelCase__      )      )

														UpperCAmelCase         =     outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

														UpperCAmelCase         =     self.model_tester.num_stages
														self.assertEqual(len(UpperCamelCase__      )				,		expected_num_stages + 1      )

								UpperCAmelCase       ,		UpperCAmelCase         =     self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
														UpperCAmelCase         =     True
														check_hidden_states_output(UpperCamelCase__				,		UpperCamelCase__				,		UpperCamelCase__      )

														# check that output_hidden_states also work using config
														del inputs_dict["output_hidden_states"]
														UpperCAmelCase         =     True

														check_hidden_states_output(UpperCamelCase__				,		UpperCamelCase__				,		UpperCamelCase__      )







		def SCREAMING_SNAKE_CASE_			(						self      :     List[str]      )				->					Optional[int]:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase         =     self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__      ):
																				UpperCAmelCase         =     self._prepare_for_class(UpperCamelCase__				,		UpperCamelCase__      )
																				UpperCAmelCase         =     model_class(UpperCamelCase__      )

																				@jax.jit
																				def model_jitted(UpperCamelCase__      :     Optional[int]				,		**UpperCamelCase__      :     Union[str, Any]      ):
																										return model(pixel_values=UpperCamelCase__				,		**UpperCamelCase__      )

																				with self.subTest("JIT Enabled"      ):
																										UpperCAmelCase         =     model_jitted(**UpperCamelCase__      ).to_tuple()

																				with self.subTest("JIT Disabled"      ):
																										with jax.disable_jit():
																																UpperCAmelCase         =     model_jitted(**UpperCamelCase__      ).to_tuple()

																				self.assertEqual(len(UpperCamelCase__      )				,		len(UpperCamelCase__      )      )
																				for jitted_output, output in zip(UpperCamelCase__				,		UpperCamelCase__      ):
																										self.assertEqual(jitted_output.shape				,		output.shape      )

def    lowerCamelCase_() ->       Union[str, Any]:
						UpperCAmelCase         =     Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
						return image





@require_flax
class  __magic_name__   (							unittest.TestCase						):






		@cached_property
		def SCREAMING_SNAKE_CASE_			(						self      :     Union[str, Any]      )				->					Optional[Any]:
								'''simple docstring'''







								return AutoImageProcessor.from_pretrained("facebook/regnet-y-040"      ) if is_vision_available() else None







		@slow
		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[int]      )				->					Tuple:
								'''simple docstring'''







								UpperCAmelCase         =     FlaxRegNetForImageClassification.from_pretrained("facebook/regnet-y-040"      )

								UpperCAmelCase         =     self.default_image_processor
								UpperCAmelCase         =     prepare_img()
								UpperCAmelCase         =     image_processor(images=UpperCamelCase__				,		return_tensors="np"      )

								UpperCAmelCase         =     model(**UpperCamelCase__      )

								# verify the logits
								UpperCAmelCase         =     (1, 10_00)
								self.assertEqual(outputs.logits.shape				,		UpperCamelCase__      )

								UpperCAmelCase         =     jnp.array([-0.41_80, -1.50_51, -3.48_36]      )

								self.assertTrue(jnp.allclose(outputs.logits[0, :3]				,		UpperCamelCase__				,		atol=1e-4      )      )



