






'''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  _a(  UpperCamelCase__				:						Tuple							):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__     :				Tuple									=[2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 1_8, 2]
      SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=True if '''large''' in model_name or '''huge''' in model_name else False
      SCREAMING_SNAKE_CASE__     :				int									=True if '''large''' in model_name or '''huge''' in model_name else False
      SCREAMING_SNAKE_CASE__     :				Optional[int]									=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__     :				int									=[3, 3, 3, 3]
                  SCREAMING_SNAKE_CASE__     :				Optional[int]									=[5, 5, 5, 5]
            elif "fl4" in model_name:
                  SCREAMING_SNAKE_CASE__     :				str									=[4, 4, 4, 4]
                  SCREAMING_SNAKE_CASE__     :				Optional[int]									=[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__     :				List[Any]									=[2, 2, 2, 2]

      if "tiny" in model_name:
            SCREAMING_SNAKE_CASE__     :				Dict									=9_6
      elif "small" in model_name:
            SCREAMING_SNAKE_CASE__     :				Dict									=9_6
      elif "base" in model_name:
            SCREAMING_SNAKE_CASE__     :				List[str]									=1_2_8
      elif "large" in model_name:
            SCREAMING_SNAKE_CASE__     :				Tuple									=1_9_2
      elif "xlarge" in model_name:
            SCREAMING_SNAKE_CASE__     :				Optional[int]									=2_5_6
      elif "huge" in model_name:
            SCREAMING_SNAKE_CASE__     :				str									=3_5_2

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

      SCREAMING_SNAKE_CASE__     :				Tuple									=json.load(open(hf_hub_download(UpperCamelCase__,	UpperCamelCase__,	repo_type='''dataset'''							),	'''r'''							)							)
      SCREAMING_SNAKE_CASE__     :				int									={int(UpperCamelCase__							): v for k, v in idalabel.items()}
      SCREAMING_SNAKE_CASE__     :				Any									={v: k for k, v in idalabel.items()}

      SCREAMING_SNAKE_CASE__     :				Optional[Any]									=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  _a(  UpperCamelCase__				:						Union[str, Any]							):




      '''simple docstring'''






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

      if name == "norm.weight":
            SCREAMING_SNAKE_CASE__     :				Dict									='''layernorm.weight'''
      if name == "norm.bias":
            SCREAMING_SNAKE_CASE__     :				Optional[int]									='''layernorm.bias'''

      if "head" in name:
            SCREAMING_SNAKE_CASE__     :				List[Any]									=name.replace('''head''',	'''classifier'''							)
      else:
            SCREAMING_SNAKE_CASE__     :				Any									='''focalnet.''' + name

      return name





def  _a(  UpperCamelCase__				:						str,	UpperCamelCase__				:						Any,	UpperCamelCase__				:						List[Any]=False							):




      '''simple docstring'''






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

      # rename keys
      for key in state_dict.copy().keys():
            SCREAMING_SNAKE_CASE__     :				Dict									=state_dict.pop(UpperCamelCase__							)
            SCREAMING_SNAKE_CASE__     :				Dict									=val

      SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=get_focalnet_config(UpperCamelCase__							)
      SCREAMING_SNAKE_CASE__     :				Optional[int]									=FocalNetForImageClassification(UpperCamelCase__							)
      model.eval()

      # load state dict
      model.load_state_dict(UpperCamelCase__							)

      # verify conversion
      SCREAMING_SNAKE_CASE__     :				List[Any]									='''http://images.cocodataset.org/val2017/000000039769.jpg'''

      SCREAMING_SNAKE_CASE__     :				str									=BitImageProcessor(
          do_resize=UpperCamelCase__,	size={'''shortest_edge''': 2_5_6},	resample=PILImageResampling.BILINEAR,	do_center_crop=UpperCamelCase__,	crop_size=2_2_4,	do_normalize=UpperCamelCase__,	image_mean=UpperCamelCase__,	image_std=UpperCamelCase__,	)
      SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=Image.open(requests.get(UpperCamelCase__,	stream=UpperCamelCase__							).raw							)
      SCREAMING_SNAKE_CASE__     :				Tuple									=processor(images=UpperCamelCase__,	return_tensors='''pt'''							)

      SCREAMING_SNAKE_CASE__     :				Optional[Any]									=transforms.Compose(
          [
              transforms.Resize(2_5_6							),
              transforms.CenterCrop(2_2_4							),
              transforms.ToTensor(),
              transforms.Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6],	std=[0.2_2_9, 0.2_2_4, 0.2_2_5]							),
          ]							)

      SCREAMING_SNAKE_CASE__     :				int									=image_transforms(UpperCamelCase__							).unsqueeze(0							)

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

      SCREAMING_SNAKE_CASE__     :				Optional[int]									=model(**UpperCamelCase__							)

      SCREAMING_SNAKE_CASE__     :				Optional[int]									=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__     :				List[Any]									=torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]							)
      elif model_name == "focalnet-tiny-lrf":
            SCREAMING_SNAKE_CASE__     :				Optional[int]									=torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]							)
      elif model_name == "focalnet-small":
            SCREAMING_SNAKE_CASE__     :				List[Any]									=torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]							)
      elif model_name == "focalnet-small-lrf":
            SCREAMING_SNAKE_CASE__     :				Optional[Any]									=torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]							)
      elif model_name == "focalnet-base":
            SCREAMING_SNAKE_CASE__     :				List[Any]									=torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]							)
      elif model_name == "focalnet-base-lrf":
            SCREAMING_SNAKE_CASE__     :				Any									=torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]							)
      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_					= 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)







'''simple docstring'''
class      __SCREAMING_SNAKE_CASE      :




 def __init__(      self	:							List[Any]					,      __lowercase	:							int      )	->						None:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=size
       SCREAMING_SNAKE_CASE__     :				List[Any]									=[0] * size
       SCREAMING_SNAKE_CASE__     :				str									=[0] * size




 @staticmethod
 def      __magic_name__	(      __lowercase	:							int      )	->						int:
       return index | (index + 1)




 @staticmethod
 def      __magic_name__	(      __lowercase	:							int      )	->						int:
       return (index & (index + 1)) - 1




 def      __magic_name__	(      self	:							Dict					,      __lowercase	:							int					,      __lowercase	:							int      )	->						None:
       SCREAMING_SNAKE_CASE__     :				List[str]									=value
       while index < self.size:
             SCREAMING_SNAKE_CASE__     :				Any									=self.get_prev(__lowercase      ) + 1
             if current_left_border == index:
                   SCREAMING_SNAKE_CASE__     :				List[str]									=value
             else:
                   SCREAMING_SNAKE_CASE__     :				str									=max(__lowercase					,      __lowercase					,      __lowercase      )
             SCREAMING_SNAKE_CASE__     :				Optional[Any]									=self.get_next(__lowercase      )




 def      __magic_name__	(      self	:							Optional[int]					,      __lowercase	:							int					,      __lowercase	:							int      )	->						int:
       right -= 1  # Because of right is exclusive
       SCREAMING_SNAKE_CASE__     :				str									=0
       while left <= right:
             SCREAMING_SNAKE_CASE__     :				Optional[Any]									=self.get_prev(__lowercase      )
             if left <= current_left:
                   SCREAMING_SNAKE_CASE__     :				List[Any]									=max(__lowercase					,      self.tree[right]      )
                   SCREAMING_SNAKE_CASE__     :				Any									=current_left
             else:
                   SCREAMING_SNAKE_CASE__     :				Optional[Any]									=max(__lowercase					,      self.arr[right]      )
                   right -= 1
       return result


if __name__ == "__main__":
  import doctest

  doctest.testmod()