






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__				:		Union[str, Any]      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   [2, 2, 6, 2] if 'tiny' in model_name else [2, 2, 18, 2]
				SCREAMING_SNAKE_CASE_   :	Dict													=   True if 'large' in model_name or 'huge' in model_name else False
				SCREAMING_SNAKE_CASE_   :	Any													=   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

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

				if "tiny" in model_name or "small" in model_name or "base" in model_name:
								SCREAMING_SNAKE_CASE_   :	List[str]													=   [3, 3, 3, 3]
								if "lrf" in model_name:
												SCREAMING_SNAKE_CASE_   :	Optional[int]													=   [3, 3, 3, 3]
								else:
												SCREAMING_SNAKE_CASE_   :	Any													=   [2, 2, 2, 2]

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

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

				SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   json.load(open(hf_hub_download(lowerCAmelCase__    ,    lowerCAmelCase__    ,    repo_type='dataset'      )    ,    'r'      )      )
				SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   {int(lowerCAmelCase__      ): v for k, v in idalabel.items()}
				SCREAMING_SNAKE_CASE_   :	int													=   {v: k for k, v in idalabel.items()}

				SCREAMING_SNAKE_CASE_   :	Tuple													=   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__				:		Optional[Any]      ):





				'''simple docstring'''







				if "patch_embed.proj" in name:
								SCREAMING_SNAKE_CASE_   :	Dict													=   name.replace('patch_embed.proj'    ,    'embeddings.patch_embeddings.projection'      )
				if "patch_embed.norm" in name:
								SCREAMING_SNAKE_CASE_   :	List[str]													=   name.replace('patch_embed.norm'    ,    'embeddings.norm'      )
				if "layers" in name:
								SCREAMING_SNAKE_CASE_   :	Optional[int]													=   '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_   :	Union[str, Any]													=   name.replace('downsample.proj'    ,    'downsample.projection'      )
				if "blocks" in name:
								SCREAMING_SNAKE_CASE_   :	Optional[int]													=   name.replace('blocks'    ,    'layers'      )
				if "modulation.f.weight" in name or "modulation.f.bias" in name:
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   name.replace('modulation.f'    ,    'modulation.projection_in'      )
				if "modulation.h.weight" in name or "modulation.h.bias" in name:
								SCREAMING_SNAKE_CASE_   :	int													=   name.replace('modulation.h'    ,    'modulation.projection_context'      )
				if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   name.replace('modulation.proj'    ,    'modulation.projection_out'      )

				if name == "norm.weight":
								SCREAMING_SNAKE_CASE_   :	List[str]													=   'layernorm.weight'
				if name == "norm.bias":
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   'layernorm.bias'

				if "head" in name:
								SCREAMING_SNAKE_CASE_   :	str													=   name.replace('head'    ,    'classifier'      )
				else:
								SCREAMING_SNAKE_CASE_   :	Tuple													=   'focalnet.' + name

				return name



def 					_lowerCamelCase(			lowerCAmelCase__				:		str    ,    lowerCAmelCase__				:		Optional[int]    ,    lowerCAmelCase__				:		Dict=False      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	List[str]													=   {
				    '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_   :	Tuple													=   model_name_to_url[model_name]
				print('Checkpoint URL: '    ,    lowerCAmelCase__      )
				SCREAMING_SNAKE_CASE_   :	List[Any]													=   torch.hub.load_state_dict_from_url(lowerCAmelCase__    ,    map_location='cpu'      )['model']

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

				SCREAMING_SNAKE_CASE_   :	List[Any]													=   get_focalnet_config(lowerCAmelCase__      )
				SCREAMING_SNAKE_CASE_   :	Tuple													=   FocalNetForImageClassification(lowerCAmelCase__      )
				model.eval()

				# load state dict
				model.load_state_dict(lowerCAmelCase__      )

				# verify conversion
				SCREAMING_SNAKE_CASE_   :	Dict													=   'http://images.cocodataset.org/val2017/000000039769.jpg'

				SCREAMING_SNAKE_CASE_   :	Union[str, 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_   :	Optional[Any]													=   Image.open(requests.get(lowerCAmelCase__    ,    stream=lowerCAmelCase__      ).raw      )
				SCREAMING_SNAKE_CASE_   :	List[str]													=   processor(images=lowerCAmelCase__    ,    return_tensors='pt'      )

				SCREAMING_SNAKE_CASE_   :	Dict													=   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_   :	str													=   image_transforms(lowerCAmelCase__      ).unsqueeze(0      )

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

				SCREAMING_SNAKE_CASE_   :	Any													=   model(**lowerCAmelCase__      )

				SCREAMING_SNAKE_CASE_   :	Union[str, 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":
								SCREAMING_SNAKE_CASE_   :	Dict													=   torch.tensor([0.2_166, -0.4_368, 0.2_191]      )
				elif model_name == "focalnet-tiny-lrf":
								SCREAMING_SNAKE_CASE_   :	Optional[int]													=   torch.tensor([1.1_669, 0.0_125, -0.1_695]      )
				elif model_name == "focalnet-small":
								SCREAMING_SNAKE_CASE_   :	Any													=   torch.tensor([0.4_917, -0.0_430, 0.1_341]      )
				elif model_name == "focalnet-small-lrf":
								SCREAMING_SNAKE_CASE_   :	Tuple													=   torch.tensor([-0.2_588, -0.5_342, -0.2_331]      )
				elif model_name == "focalnet-base":
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   torch.tensor([-0.1_655, -0.4_090, -0.1_730]      )
				elif model_name == "focalnet-base-lrf":
								SCREAMING_SNAKE_CASE_   :	Any													=   torch.tensor([0.5_306, -0.0_483, -0.3_928]      )
				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        =			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        =			parser.parse_args()
	convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union

import pyarrow as pa


if TYPE_CHECKING:
				from .features import FeatureType


@dataclass
class 							__a							:





		'''simple docstring'''



		UpperCAmelCase__      :				List[str]
		UpperCAmelCase__      :				Optional[str]				    =						None
		# Automatically constructed
		UpperCAmelCase__      :				ClassVar[str]				    =						"dict"
		UpperCAmelCase__      :				ClassVar[Any]				    =						None
		UpperCAmelCase__      :				str				    =						field(default="""Translation""" ,	init=__A ,	repr=__A						)


		def __call__(						self       ):
						return pa.struct({lang: pa.string() for lang in sorted(self.languages       )}       )






		def   __snake_case			(						self       ):
						from .features import Value

						return {k: Value('string'       ) for k in sorted(self.languages       )}






@dataclass
class 							__a							:





		'''simple docstring'''



		UpperCAmelCase__      :				Optional[List]				    =						None
		UpperCAmelCase__      :				Optional[int]				    =						None
		UpperCAmelCase__      :				Optional[str]				    =						None
		# Automatically constructed
		UpperCAmelCase__      :				ClassVar[str]				    =						"dict"
		UpperCAmelCase__      :				ClassVar[Any]				    =						None
		UpperCAmelCase__      :				str				    =						field(default="""TranslationVariableLanguages""" ,	init=__A ,	repr=__A						)


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Any													=   sorted(set(self.languages       )       ) if self.languages else None
						SCREAMING_SNAKE_CASE_   :	Dict													=   len(self.languages       ) if self.languages else None


		def __call__(						self       ):
						return pa.struct({'language': pa.list_(pa.string()       ), 'translation': pa.list_(pa.string()       )}       )


		def   __snake_case			(						self ,    UpperCamelCase__       ):
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   set(self.languages       )
						if self.languages and set(UpperCamelCase__       ) - lang_set:
										raise ValueError(
										    F'''Some languages in example ({', '.join(sorted(set(UpperCamelCase__       ) - lang_set       )       )}) are not in valid set ({', '.join(UpperCamelCase__       )}).'''       )

						# Convert dictionary into tuples, splitting out cases where there are
						# multiple translations for a single language.
						SCREAMING_SNAKE_CASE_   :	List[str]													=   []
						for lang, text in translation_dict.items():
										if isinstance(UpperCamelCase__ ,    UpperCamelCase__       ):
														translation_tuples.append((lang, text)       )
										else:
														translation_tuples.extend([(lang, el) for el in text]       )

        # Ensure translations are in ascending order by language code.
						SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_   :	Dict													=   zip(*sorted(UpperCamelCase__       )       )

						return {"language": languages, "translation": translations}






		def   __snake_case			(						self       ):
						from .features import Sequence, Value

						return {
						    "language": Sequence(Value('string'       )       ),
						    "translation": Sequence(Value('string'       )       ),
						}