




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

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

			if "tiny" in model_name:
						__snake_case			=		96
			elif "small" in model_name:
						__snake_case			=		96
			elif "base" in model_name:
						__snake_case			=		1_28
			elif "large" in model_name:
						__snake_case			=		1_92
			elif "xlarge" in model_name:
						__snake_case			=		2_56
			elif "huge" in model_name:
						__snake_case			=		3_52

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

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

			__snake_case			=		FocalNetConfig(
			    embed_dim=_UpperCAmelCase		,				depths=_UpperCAmelCase		,				focal_levels=_UpperCAmelCase		,				focal_windows=_UpperCAmelCase		,				use_conv_embed=_UpperCAmelCase		,				idalabel=_UpperCAmelCase		,				labelaid=_UpperCAmelCase		,				use_post_layernorm=_UpperCAmelCase		,				use_layerscale=_UpperCAmelCase		,				)

			return config


def __UpperCAmelCase	(   _UpperCAmelCase		:   Tuple						)      ->      Optional[Any]:
			if "patch_embed.proj" in name:
						__snake_case			=		name.replace("patch_embed.proj"		,				"embeddings.patch_embeddings.projection"						)
			if "patch_embed.norm" in name:
						__snake_case			=		name.replace("patch_embed.norm"		,				"embeddings.norm"						)
			if "layers" in name:
						__snake_case			=		"encoder." + name
			if "encoder.layers" in name:
						__snake_case			=		name.replace("encoder.layers"		,				"encoder.stages"						)
			if "downsample.proj" in name:
						__snake_case			=		name.replace("downsample.proj"		,				"downsample.projection"						)
			if "blocks" in name:
						__snake_case			=		name.replace("blocks"		,				"layers"						)
			if "modulation.f.weight" in name or "modulation.f.bias" in name:
						__snake_case			=		name.replace("modulation.f"		,				"modulation.projection_in"						)
			if "modulation.h.weight" in name or "modulation.h.bias" in name:
						__snake_case			=		name.replace("modulation.h"		,				"modulation.projection_context"						)
			if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
						__snake_case			=		name.replace("modulation.proj"		,				"modulation.projection_out"						)

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

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

			return name


def __UpperCAmelCase	(   _UpperCAmelCase		:   List[str]		,				_UpperCAmelCase		:   Optional[Any]		,				_UpperCAmelCase		:   List[str]=False						)      ->      Dict:
			# fmt: off
			__snake_case			=		{
			    "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

			__snake_case			=		model_name_to_url[model_name]
			print("Checkpoint URL: "		,				_UpperCAmelCase						)
			__snake_case			=		torch.hub.load_state_dict_from_url(_UpperCAmelCase		,				map_location="cpu"						)["model"]

			# rename keys
			for key in state_dict.copy().keys():
						__snake_case			=		state_dict.pop(_UpperCAmelCase						)
						__snake_case			=		val

			__snake_case			=		get_focalnet_config(_UpperCAmelCase						)
			__snake_case			=		FocalNetForImageClassification(_UpperCAmelCase						)
			model.eval()

			# load state dict
			model.load_state_dict(_UpperCAmelCase						)

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

			__snake_case			=		BitImageProcessor(
			    do_resize=_UpperCAmelCase		,				size={"shortest_edge": 2_56}		,				resample=PILImageResampling.BILINEAR		,				do_center_crop=_UpperCAmelCase		,				crop_size=2_24		,				do_normalize=_UpperCAmelCase		,				image_mean=_UpperCAmelCase		,				image_std=_UpperCAmelCase		,				)
			__snake_case			=		Image.open(requests.get(_UpperCAmelCase		,				stream=_UpperCAmelCase						).raw						)
			__snake_case			=		processor(images=_UpperCAmelCase		,				return_tensors="pt"						)

			__snake_case			=		transforms.Compose(
			    [
			        transforms.Resize(2_56						),
			        transforms.CenterCrop(2_24						),
			        transforms.ToTensor(),
			        transforms.Normalize(mean=[0.485, 0.456, 0.406]		,				std=[0.229, 0.224, 0.225]						),
			    ]						)

			__snake_case			=		image_transforms(_UpperCAmelCase						).unsqueeze(0						)

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

			__snake_case			=		model(**_UpperCAmelCase						)

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

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












"""simple docstring"""





import importlib.metadata
import operator
import re
import sys
from typing import Optional

from packaging import version


__A											=						{
    """<""": operator.lt,
    """<=""": operator.le,
    """==""": operator.eq,
    """!=""": operator.ne,
    """>=""": operator.ge,
    """>""": operator.gt,
}
def 					__A  (_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE     )   ->Union[str, Any]:





  """simple docstring"""


  if got_ver is None or want_ver is None:
    raise ValueError(
        F"Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider"
        F" reinstalling {pkg}."     )
  if not ops[op](version.parse(_SCREAMING_SNAKE_CASE     )	,  version.parse(_SCREAMING_SNAKE_CASE     )     ):
    raise ImportError(
        F"{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}"     )
def 					__A  (_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE = None     )   ->None:





  """simple docstring"""


  lowerCAmelCase__     :List[str]											=				F"\n{hint}" if hint is not None else ''

  # non-versioned check
  if re.match(r'^[\w_\-\d]+$'	,  _SCREAMING_SNAKE_CASE     ):
    lowerCAmelCase__				,							lowerCAmelCase__				,							lowerCAmelCase__     :Any											=				requirement, None, None
  else:
    lowerCAmelCase__     :List[str]											=				re.findall(r'^([^!=<>\s]+)([\s!=<>]{1,2}.+)'	,  _SCREAMING_SNAKE_CASE     )
    if not match:
      raise ValueError(
          'requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but'
          F" got {requirement}"     )
    lowerCAmelCase__				,							lowerCAmelCase__     :Union[str, Any]											=				match[0]
    lowerCAmelCase__     :List[Any]											=				want_full.split(','     )  # there could be multiple requirements
    lowerCAmelCase__     :Any											=				{}
    for w in want_range:
      lowerCAmelCase__     :Tuple											=				re.findall(r'^([\s!=<>]{1,2})(.+)'	,  _SCREAMING_SNAKE_CASE     )
      if not match:
        raise ValueError(
            'requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23,'
            F" but got {requirement}"     )
      lowerCAmelCase__				,							lowerCAmelCase__     :int											=				match[0]
      lowerCAmelCase__     :str											=				want_ver
      if op not in ops:
        raise ValueError(F"{requirement}: need one of {list(ops.keys()     )}, but got {op}"     )

    # special case
  if pkg == "python":
    lowerCAmelCase__     :Any											=				'.'.join([str(_SCREAMING_SNAKE_CASE     ) for x in sys.version_info[:3]]     )
    for op, want_ver in wanted.items():
      _compare_versions(_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE     )
    return

  # check if any version is installed
  try:
    lowerCAmelCase__     :List[Any]											=				importlib.metadata.version(_SCREAMING_SNAKE_CASE     )
  except importlib.metadata.PackageNotFoundError:
    raise importlib.metadata.PackageNotFoundError(
        F"The '{requirement}' distribution was not found and is required by this application. {hint}"     )

  # check that the right version is installed if version number or a range was provided
  if want_ver is not None:
    for op, want_ver in wanted.items():
      _compare_versions(_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE     )




def 					__A  (_SCREAMING_SNAKE_CASE     )   ->List[Any]:





  """simple docstring"""


  lowerCAmelCase__     :Optional[Any]											=				'Try: pip install transformers -U or pip install -e \'.[dev]\' if you\'re working with git main'
  return require_version(_SCREAMING_SNAKE_CASE	,  _SCREAMING_SNAKE_CASE     )
