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

       if "tiny" in model_name or "small" in model_name or "base" in model_name:
              lowercase       :							Union[str, Any]									=				[3, 3, 3, 3]
              if "lrf" in model_name:
                     lowercase       :							str									=				[3, 3, 3, 3]
              else:
                     lowercase       :							Any									=				[2, 2, 2, 2]

       if "tiny" in model_name:
              lowercase       :							str									=				96
       elif "small" in model_name:
              lowercase       :							Any									=				96
       elif "base" in model_name:
              lowercase       :							Any									=				128
       elif "large" in model_name:
              lowercase       :							Any									=				192
       elif "xlarge" in model_name:
              lowercase       :							List[Any]									=				256
       elif "huge" in model_name:
              lowercase       :							List[Any]									=				352

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

       lowercase       :							Union[str, Any]									=				json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE__	,					SCREAMING_SNAKE_CASE__	,					repo_type="dataset"    )	,					"r"    )    )
       lowercase       :							str									=				{int(SCREAMING_SNAKE_CASE__    ): v for k, v in idalabel.items()}
       lowercase       :							Tuple									=				{v: k for k, v in idalabel.items()}

       lowercase       :							Optional[Any]									=				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     __magic_name__      ( __snake_case							:			Union[str, Any]    )					->							Union[str, Any]:





       if "patch_embed.proj" in name:
              lowercase       :							Optional[int]									=				name.replace("patch_embed.proj"	,					"embeddings.patch_embeddings.projection"    )
       if "patch_embed.norm" in name:
              lowercase       :							List[Any]									=				name.replace("patch_embed.norm"	,					"embeddings.norm"    )
       if "layers" in name:
              lowercase       :							Optional[int]									=				'encoder.' + name
       if "encoder.layers" in name:
              lowercase       :							Union[str, Any]									=				name.replace("encoder.layers"	,					"encoder.stages"    )
       if "downsample.proj" in name:
              lowercase       :							Tuple									=				name.replace("downsample.proj"	,					"downsample.projection"    )
       if "blocks" in name:
              lowercase       :							Any									=				name.replace("blocks"	,					"layers"    )
       if "modulation.f.weight" in name or "modulation.f.bias" in name:
              lowercase       :							Dict									=				name.replace("modulation.f"	,					"modulation.projection_in"    )
       if "modulation.h.weight" in name or "modulation.h.bias" in name:
              lowercase       :							Optional[Any]									=				name.replace("modulation.h"	,					"modulation.projection_context"    )
       if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
              lowercase       :							int									=				name.replace("modulation.proj"	,					"modulation.projection_out"    )

       if name == "norm.weight":
              lowercase       :							List[str]									=				'layernorm.weight'
       if name == "norm.bias":
              lowercase       :							List[str]									=				'layernorm.bias'

       if "head" in name:
              lowercase       :							Dict									=				name.replace("head"	,					"classifier"    )
       else:
              lowercase       :							Any									=				'focalnet.' + name

       return name


def     __magic_name__      ( __snake_case							:			List[str]	,					__snake_case							:			Optional[Any]	,					__snake_case							:			Optional[int]=False    )					->							Any:
       lowercase       :							Optional[int]									=				{
           '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

       lowercase       :							Dict									=				model_name_to_url[model_name]
       print("Checkpoint URL: "	,					SCREAMING_SNAKE_CASE__    )
       lowercase       :							int									=				torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE__	,					map_location="cpu"    )['model']

       # rename keys
       for key in state_dict.copy().keys():
              lowercase       :							List[str]									=				state_dict.pop(SCREAMING_SNAKE_CASE__    )
              lowercase       :							str									=				val

       lowercase       :							int									=				get_focalnet_config(SCREAMING_SNAKE_CASE__    )
       lowercase       :							Tuple									=				FocalNetForImageClassification(SCREAMING_SNAKE_CASE__    )
       model.eval()

       # load state dict
       model.load_state_dict(SCREAMING_SNAKE_CASE__    )

       # verify conversion
       lowercase       :							Tuple									=				'http://images.cocodataset.org/val2017/000000039769.jpg'

       lowercase       :							Optional[int]									=				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__	,					)
       lowercase       :							Optional[int]									=				Image.open(requests.get(SCREAMING_SNAKE_CASE__	,					stream=SCREAMING_SNAKE_CASE__    ).raw    )
       lowercase       :							Optional[int]									=				processor(images=SCREAMING_SNAKE_CASE__	,					return_tensors="pt"    )

       lowercase       :							Dict									=				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]    ),
           ]    )

       lowercase       :							Any									=				image_transforms(SCREAMING_SNAKE_CASE__    ).unsqueeze(0    )

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

       lowercase       :							int									=				model(**SCREAMING_SNAKE_CASE__    )

       lowercase       :							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":
              lowercase       :							Optional[Any]									=				torch.tensor([0.21_66, -0.43_68, 0.21_91]    )
       elif model_name == "focalnet-tiny-lrf":
              lowercase       :							str									=				torch.tensor([1.16_69, 0.01_25, -0.16_95]    )
       elif model_name == "focalnet-small":
              lowercase       :							List[str]									=				torch.tensor([0.49_17, -0.04_30, 0.13_41]    )
       elif model_name == "focalnet-small-lrf":
              lowercase       :							Any									=				torch.tensor([-0.25_88, -0.53_42, -0.23_31]    )
       elif model_name == "focalnet-base":
              lowercase       :							List[str]									=				torch.tensor([-0.16_55, -0.40_90, -0.17_30]    )
       elif model_name == "focalnet-base-lrf":
              lowercase       :							Tuple									=				torch.tensor([0.53_06, -0.04_83, -0.39_28]    )
       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__":
   _A     :			Union[str, 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.""",
   )

   _A     :			Tuple  						=			parser.parse_args()
   convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






import time
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
 import torch

 from transformers.generation import (
     MaxLengthCriteria,
     MaxNewTokensCriteria,
     MaxTimeCriteria,
     StoppingCriteriaList,
     validate_stopping_criteria,
 )




@require_torch
class A    (  unittest.TestCase	):






      def 			lowerCamelCase			(   self		:	Optional[int]     ,  lowercase_		:	int					)   ->		Optional[int]:
             """simple docstring"""




             _lowerCamelCase			:						List[str]            =3
             _lowerCamelCase			:						Dict            =250

             _lowerCamelCase			:						Tuple            =ids_tensor((batch_size, length)     ,  lowercase_					)
             _lowerCamelCase			:						str            =torch.ones((batch_size, length)     ,  device=lowercase_     ,  dtype=torch.float					) / length
             return input_ids, scores






      def 			lowerCamelCase			(   self		:	Optional[int]					)   ->		Optional[Any]:
             """simple docstring"""




             _lowerCamelCase      ,	_lowerCamelCase			:						Optional[Any]            =self._get_tensors(5					)

             _lowerCamelCase			:						Dict            =StoppingCriteriaList(
                 [
                     MaxLengthCriteria(max_length=10					),
                     MaxTimeCriteria(max_time=0.1					),
                 ]					)

             self.assertFalse(criteria(lowercase_     ,  lowercase_					)					)

             _lowerCamelCase      ,	_lowerCamelCase			:						List[str]            =self._get_tensors(9					)
             self.assertFalse(criteria(lowercase_     ,  lowercase_					)					)

             _lowerCamelCase      ,	_lowerCamelCase			:						Optional[int]            =self._get_tensors(10					)
             self.assertTrue(criteria(lowercase_     ,  lowercase_					)					)






      def 			lowerCamelCase			(   self		:	Optional[int]					)   ->		Tuple:
             """simple docstring"""




             _lowerCamelCase			:						Any            =MaxLengthCriteria(max_length=10					)

             _lowerCamelCase      ,	_lowerCamelCase			:						List[str]            =self._get_tensors(5					)
             self.assertFalse(criteria(lowercase_     ,  lowercase_					)					)

             _lowerCamelCase      ,	_lowerCamelCase			:						Optional[int]            =self._get_tensors(9					)
             self.assertFalse(criteria(lowercase_     ,  lowercase_					)					)

             _lowerCamelCase      ,	_lowerCamelCase			:						Optional[int]            =self._get_tensors(10					)
             self.assertTrue(criteria(lowercase_     ,  lowercase_					)					)






      def 			lowerCamelCase			(   self		:	List[str]					)   ->		Tuple:
             """simple docstring"""




             _lowerCamelCase			:						Union[str, Any]            =MaxNewTokensCriteria(start_length=5     ,  max_new_tokens=5					)

             _lowerCamelCase      ,	_lowerCamelCase			:						str            =self._get_tensors(5					)
             self.assertFalse(criteria(lowercase_     ,  lowercase_					)					)

             _lowerCamelCase      ,	_lowerCamelCase			:						Union[str, Any]            =self._get_tensors(9					)
             self.assertFalse(criteria(lowercase_     ,  lowercase_					)					)

             _lowerCamelCase      ,	_lowerCamelCase			:						Dict            =self._get_tensors(10					)
             self.assertTrue(criteria(lowercase_     ,  lowercase_					)					)

             _lowerCamelCase			:						Optional[Any]            =StoppingCriteriaList([criteria]					)
             self.assertEqual(criteria_list.max_length     ,  10					)






      def 			lowerCamelCase			(   self		:	Optional[Any]					)   ->		str:
             """simple docstring"""




             _lowerCamelCase      ,	_lowerCamelCase			:						Optional[Any]            =self._get_tensors(5					)

             _lowerCamelCase			:						Tuple            =MaxTimeCriteria(max_time=0.1					)
             self.assertFalse(criteria(lowercase_     ,  lowercase_					)					)

             _lowerCamelCase			:						int            =MaxTimeCriteria(max_time=0.1     ,  initial_timestamp=time.time() - 0.2					)
             self.assertTrue(criteria(lowercase_     ,  lowercase_					)					)






      def 			lowerCamelCase			(   self		:	Optional[int]					)   ->		Optional[int]:
             """simple docstring"""




             validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10					)]					)     ,  10					)

             with self.assertWarns(lowercase_					):
                    validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10					)]					)     ,  11					)

             _lowerCamelCase			:						Optional[Any]            =validate_stopping_criteria(StoppingCriteriaList()     ,  11					)

             self.assertEqual(len(lowercase_					)     ,  1					)


