


"""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 _lowerCamelCase      (						_UpperCamelCase ):

			'''simple docstring'''




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

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

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

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

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

			__lowerCAmelCase							       =     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 _lowerCamelCase      (						_UpperCamelCase ):

			'''simple docstring'''




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

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

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

			return name


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase=False ):

			'''simple docstring'''




			__lowerCAmelCase							       =     {
			    "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

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

			# rename keys
			for key in state_dict.copy().keys():
						__lowerCAmelCase							       =     state_dict.pop(_UpperCamelCase )
						__lowerCAmelCase							       =     val

			__lowerCAmelCase							       =     get_focalnet_config(_UpperCamelCase )
			__lowerCAmelCase							       =     FocalNetForImageClassification(_UpperCamelCase )
			model.eval()

			# load state dict
			model.load_state_dict(_UpperCamelCase )

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

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

			__lowerCAmelCase							       =     transforms.Compose(
			    [
			        transforms.Resize(256 ),
			        transforms.CenterCrop(224 ),
			        transforms.ToTensor(),
			        transforms.Normalize(mean=[0.4_85, 0.4_56, 0.4_06]       ,					std=[0.2_29, 0.2_24, 0.2_25] ),
			    ] )

			__lowerCAmelCase							       =     image_transforms(_UpperCamelCase ).unsqueeze(0 )

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

			__lowerCAmelCase							       =     model(**_UpperCamelCase )

			__lowerCAmelCase							       =     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":
						__lowerCAmelCase							       =     torch.tensor([0.21_66, -0.43_68, 0.21_91] )
			elif model_name == "focalnet-tiny-lrf":
						__lowerCAmelCase							       =     torch.tensor([1.16_69, 0.01_25, -0.16_95] )
			elif model_name == "focalnet-small":
						__lowerCAmelCase							       =     torch.tensor([0.49_17, -0.04_30, 0.13_41] )
			elif model_name == "focalnet-small-lrf":
						__lowerCAmelCase							       =     torch.tensor([-0.25_88, -0.53_42, -0.23_31] )
			elif model_name == "focalnet-base":
						__lowerCAmelCase							       =     torch.tensor([-0.16_55, -0.40_90, -0.17_30] )
			elif model_name == "focalnet-base-lrf":
						__lowerCAmelCase							       =     torch.tensor([0.53_06, -0.04_83, -0.39_28] )
			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	:  Optional[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	:  Optional[int]            =     parser.parse_args()
	convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








"""simple docstring"""


from pathlib import PurePosixPath
from typing import Optional

import fsspec
from fsspec import AbstractFileSystem
from huggingface_hub.hf_api import DatasetInfo

from ..utils.file_utils import get_authentication_headers_for_url
from ..utils.hub import hf_hub_url
class     _UpperCamelCase    (  lowerCAmelCase__						):
							'''simple docstring'''







							__UpperCAmelCase		:       Union[str, Any]							 =""""""
							__UpperCAmelCase		:       Optional[int]							 ="""hf-legacy"""  # "hf://"" is reserved for hffs

							def __init__(	self       ,  __a = None       ,  __a = None       ,  **__a       ,  ):
										super().__init__(self       ,  **__a      )
										__lowerCAmelCase							       =     repo_info
										__lowerCAmelCase							       =     token
										__lowerCAmelCase							       =     None

							def      snake_case							(	self      ):
										if self.dir_cache is None:
													__lowerCAmelCase							       =     {}
													for hf_file in self.repo_info.siblings:
																# TODO(QL): add sizes
																__lowerCAmelCase							       =     {
																    "name": hf_file.rfilename,
																    "size": None,
																    "type": "file",
																}
																self.dir_cache.update(
																    {
																        str(__a      ): {"name": str(__a      ), "size": None, "type": "directory"}
																        for d in list(PurePosixPath(hf_file.rfilename      ).parents      )[:-1]
																    }      )

							def      snake_case							(	self       ,  __a       ,  __a = "rb"       ,  **__a       ,  ):
										if not isinstance(self.repo_info       ,  __a      ):
													raise NotImplementedError(f"Open is only implemented for dataset repositories, but got {self.repo_info}"      )
										__lowerCAmelCase							       =     hf_hub_url(self.repo_info.id       ,  __a       ,  revision=self.repo_info.sha      )
										return fsspec.open(
										    __a       ,  mode=__a       ,  headers=get_authentication_headers_for_url(__a       ,  use_auth_token=self.token      )       ,  client_kwargs={"trust_env": True}       ,  ).open()

							def      snake_case							(	self       ,  __a       ,  **__a      ):
										self._get_dirs()
										__lowerCAmelCase							       =     self._strip_protocol(__a      )
										if path in self.dir_cache:
													return self.dir_cache[path]
										else:
													raise FileNotFoundError(__a      )



							def      snake_case							(	self       ,  __a       ,  __a=False       ,  **__a      ):
										self._get_dirs()
										__lowerCAmelCase							       =     PurePosixPath(path.strip("/"      )      )
										__lowerCAmelCase							       =     {}
										for p, f in self.dir_cache.items():
													__lowerCAmelCase							       =     PurePosixPath(p.strip("/"      )      )
													__lowerCAmelCase							       =     p.parent
													if root == path:
																__lowerCAmelCase							       =     f
										__lowerCAmelCase							       =     list(paths.values()      )
										if detail:
													return out
										else:
													return sorted(f["name"] for f in out      )




