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				(      lowercase__  :			str ):







     '''simple docstring'''


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

     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__				:				Optional[int]			=		[4, 4, 4, 4]
               SCREAMING_SNAKE_CASE__				:				Tuple			=		[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__				:				Dict			=		[3, 3, 3, 3]
          else:
               SCREAMING_SNAKE_CASE__				:				Tuple			=		[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__				:				Any			=		1_28
     elif "large" in model_name:
          SCREAMING_SNAKE_CASE__				:				List[Any]			=		1_92
     elif "xlarge" in model_name:
          SCREAMING_SNAKE_CASE__				:				str			=		2_56
     elif "huge" in model_name:
          SCREAMING_SNAKE_CASE__				:				Dict			=		3_52

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

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

     SCREAMING_SNAKE_CASE__				:				List[Any]			=		FocalNetConfig(
         embed_dim=lowercase__     ,  depths=lowercase__     ,  focal_levels=lowercase__     ,  focal_windows=lowercase__     ,  use_conv_embed=lowercase__     ,  idalabel=lowercase__     ,  labelaid=lowercase__     ,  use_post_layernorm=lowercase__     ,  use_layerscale=lowercase__     ,  )

     return config







def     _a				(      lowercase__  :			List[Any] ):







     '''simple docstring'''


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

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

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

     return name







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

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

     SCREAMING_SNAKE_CASE__				:				Optional[int]			=		get_focalnet_config(lowercase__ )
     SCREAMING_SNAKE_CASE__				:				Any			=		FocalNetForImageClassification(lowercase__ )
     model.eval()

     # load state dict
     model.load_state_dict(lowercase__ )

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

     SCREAMING_SNAKE_CASE__				:				int			=		BitImageProcessor(
         do_resize=lowercase__     ,  size={'shortest_edge': 2_56}     ,  resample=PILImageResampling.BILINEAR     ,  do_center_crop=lowercase__     ,  crop_size=2_24     ,  do_normalize=lowercase__     ,  image_mean=lowercase__     ,  image_std=lowercase__     ,  )
     SCREAMING_SNAKE_CASE__				:				Tuple			=		Image.open(requests.get(lowercase__     ,  stream=lowercase__ ).raw )
     SCREAMING_SNAKE_CASE__				:				Any			=		processor(images=lowercase__     ,  return_tensors='pt' )

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

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

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

     SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		model(**lowercase__ )

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

     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__":
      SCREAMING_SNAKE_CASE__		:					Optional[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.",
      )

      SCREAMING_SNAKE_CASE__		:					List[Any]						=					parser.parse_args()
      convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




from __future__ import annotations







def     _a				(      lowercase__  :			list[int | float]     ,  lowercase__  :			int     ,  lowercase__  :			int ):







     '''simple docstring'''


     if len(lowercase__ ) == 0:
          raise ValueError('find_max() arg is an empty sequence' )
     if (
         left >= len(lowercase__ )
         or left < -len(lowercase__ )
         or right >= len(lowercase__ )
         or right < -len(lowercase__ )
     ):
          raise IndexError('list index out of range' )
     if left == right:
          return nums[left]
     SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		(left + right) >> 1  # the middle
     SCREAMING_SNAKE_CASE__				:				int			=		find_max(lowercase__     ,  lowercase__     ,  lowercase__ )  # find max in range[left, mid]
     SCREAMING_SNAKE_CASE__				:				Tuple			=		find_max(lowercase__     ,  mid + 1     ,  lowercase__ )  # find max in range[mid + 1, right]

     return left_max if left_max >= right_max else right_max


if __name__ == "__main__":
      import doctest

      doctest.testmod(verbose=True)



