





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				(    lowerCAmelCase_   )		->						List[str]:




       '''simple docstring'''






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

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

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

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

       _A=	json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_       ,					SCREAMING_SNAKE_CASE_       ,					repo_type='dataset'   )       ,					'r'   )   )
       _A=	{int(SCREAMING_SNAKE_CASE_   ): v for k, v in idalabel.items()}
       _A=	{v: k for k, v in idalabel.items()}

       _A=	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 							UpperCamelCase				(    lowerCAmelCase_   )		->						Optional[int]:




       '''simple docstring'''






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

       if name == "norm.weight":
              _A=	'layernorm.weight'
       if name == "norm.bias":
              _A=	'layernorm.bias'

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

       return name

def 							UpperCamelCase				(    lowerCAmelCase_       ,					lowerCAmelCase_       ,					lowerCAmelCase_=False   )		->						Dict:




       '''simple docstring'''






       _A=	{
           '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

       _A=	model_name_to_url[model_name]
       print('Checkpoint URL: '       ,					SCREAMING_SNAKE_CASE_   )
       _A=	torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_       ,					map_location='cpu'   )['model']

       # rename keys
       for key in state_dict.copy().keys():
              _A=	state_dict.pop(SCREAMING_SNAKE_CASE_   )
              _A=	val

       _A=	get_focalnet_config(SCREAMING_SNAKE_CASE_   )
       _A=	FocalNetForImageClassification(SCREAMING_SNAKE_CASE_   )
       model.eval()

       # load state dict
       model.load_state_dict(SCREAMING_SNAKE_CASE_   )

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

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

       _A=	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]   ),
           ]   )

       _A=	image_transforms(SCREAMING_SNAKE_CASE_   ).unsqueeze(0   )

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

       _A=	model(**SCREAMING_SNAKE_CASE_   )

       _A=	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":
              _A=	torch.tensor([0.2_166, -0.4_368, 0.2_191]   )
       elif model_name == "focalnet-tiny-lrf":
              _A=	torch.tensor([1.1_669, 0.0_125, -0.1_695]   )
       elif model_name == "focalnet-small":
              _A=	torch.tensor([0.4_917, -0.0_430, 0.1_341]   )
       elif model_name == "focalnet-small-lrf":
              _A=	torch.tensor([-0.2_588, -0.5_342, -0.2_331]   )
       elif model_name == "focalnet-base":
              _A=	torch.tensor([-0.1_655, -0.4_090, -0.1_730]   )
       elif model_name == "focalnet-base-lrf":
              _A=	torch.tensor([0.5_306, -0.0_483, -0.3_928]   )
       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__":
      UpperCAmelCase_								=							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.''',
      )

      UpperCAmelCase_								=							parser.parse_args()
      convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tensorflow_text_available, is_torch_available


UpperCAmelCase_								=							{
    '''configuration_ernie''': ['''ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ErnieConfig''', '''ErnieOnnxConfig'''],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCAmelCase_								=							[
          '''ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST''',
          '''ErnieForCausalLM''',
          '''ErnieForMaskedLM''',
          '''ErnieForMultipleChoice''',
          '''ErnieForNextSentencePrediction''',
          '''ErnieForPreTraining''',
          '''ErnieForQuestionAnswering''',
          '''ErnieForSequenceClassification''',
          '''ErnieForTokenClassification''',
          '''ErnieModel''',
          '''ErniePreTrainedModel''',
      ]

if TYPE_CHECKING:
      from .configuration_ernie import ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieConfig, ErnieOnnxConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_ernie import (
                ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST,
                ErnieForCausalLM,
                ErnieForMaskedLM,
                ErnieForMultipleChoice,
                ErnieForNextSentencePrediction,
                ErnieForPreTraining,
                ErnieForQuestionAnswering,
                ErnieForSequenceClassification,
                ErnieForTokenClassification,
                ErnieModel,
                ErniePreTrainedModel,
            )

else:
      import sys

      UpperCAmelCase_								=							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)