


'''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 _SCREAMING_SNAKE_CASE      (A       )       ->       int:






       """simple docstring"""







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

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

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

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

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

       lowercase__      			= FocalNetConfig(
           embed_dim=A  , depths=A  , focal_levels=A  , focal_windows=A  , use_conv_embed=A  , idalabel=A  , labelaid=A  , use_post_layernorm=A  , use_layerscale=A  , )

       return config

def _SCREAMING_SNAKE_CASE      (A       )       ->       Dict:






       """simple docstring"""







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

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

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

       return name

def _SCREAMING_SNAKE_CASE      (A  , A  , A=False       )       ->       Any:






       """simple docstring"""







       lowercase__      			= {
           '''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__      			= model_name_to_url[model_name]
       print('''Checkpoint URL: '''  , A       )
       lowercase__      			= torch.hub.load_state_dict_from_url(A  , map_location='''cpu'''       )['''model''']

       # rename keys
       for key in state_dict.copy().keys():
              lowercase__      			= state_dict.pop(A       )
              lowercase__      			= val

       lowercase__      			= get_focalnet_config(A       )
       lowercase__      			= FocalNetForImageClassification(A       )
       model.eval()

       # load state dict
       model.load_state_dict(A       )

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

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

       lowercase__      			= transforms.Compose(
           [
               transforms.Resize(256       ),
               transforms.CenterCrop(224       ),
               transforms.ToTensor(),
               transforms.Normalize(mean=[0.485, 0.456, 0.406]  , std=[0.229, 0.224, 0.225]       ),
           ]       )

       lowercase__      			= image_transforms(A       ).unsqueeze(0       )

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

       lowercase__      			= model(**A       )

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

       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__":
 lowerCamelCase							:   str  						=	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.',
 )

 lowerCamelCase							:   Union[str, Any]  						=	parser.parse_args()
 convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





'''simple docstring'''




import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowercase						: Optional[Any]						=						logging.get_logger(__name__)

_lowercase						: Tuple						=						{
    "BAAI/AltCLIP": "https://huggingface.co/BAAI/AltCLIP/resolve/main/config.json",
    # See all AltCLIP models at https://huggingface.co/models?filter=altclip
}





class     __magic_name__						(    _UpperCAmelCase):
			UpperCamelCase__				=       '''altclip_text_model'''







			def __init__(	self				:  Union[str, Any]				,			lowercase_				:  str=250002				,			lowercase_				:  Union[str, Any]=1024				,			lowercase_				:  Any=24				,			lowercase_				:  Union[str, Any]=16				,			lowercase_				:  Any=4096				,			lowercase_				:  Union[str, Any]="gelu"				,			lowercase_				:  Union[str, Any]=0.1				,			lowercase_				:  Union[str, Any]=0.1				,			lowercase_				:  int=514				,			lowercase_				:  Union[str, Any]=1				,			lowercase_				:  List[str]=0.02				,			lowercase_				:  Optional[int]=0.02				,			lowercase_				:  str=1E-05				,			lowercase_				:  List[str]=1				,			lowercase_				:  List[Any]=0				,			lowercase_				:  Dict=2				,			lowercase_				:  Union[str, Any]="absolute"				,			lowercase_				:  Any=True				,			lowercase_				:  Union[str, Any]=768				,			**lowercase_				:  Any				,			):
										super().__init__(pad_token_id=lowercase_				,			bos_token_id=lowercase_				,			eos_token_id=lowercase_				,			**lowercase_    )

										lowercase_					:			Union[str, Any]										=       vocab_size
										lowercase_					:			str										=       hidden_size
										lowercase_					:			Optional[Any]										=       num_hidden_layers
										lowercase_					:			int										=       num_attention_heads
										lowercase_					:			str										=       hidden_act
										lowercase_					:			List[str]										=       intermediate_size
										lowercase_					:			Optional[int]										=       hidden_dropout_prob
										lowercase_					:			str										=       attention_probs_dropout_prob
										lowercase_					:			str										=       max_position_embeddings
										lowercase_					:			List[str]										=       type_vocab_size
										lowercase_					:			Union[str, Any]										=       initializer_range
										lowercase_					:			List[Any]										=       initializer_factor
										lowercase_					:			str										=       layer_norm_eps
										lowercase_					:			Tuple										=       position_embedding_type
										lowercase_					:			List[Any]										=       use_cache
										lowercase_					:			Tuple										=       project_dim





class     __magic_name__						(    _UpperCAmelCase):
			UpperCamelCase__				=       '''altclip_vision_model'''







			def __init__(	self				:  Dict				,			lowercase_				:  Any=768				,			lowercase_				:  Dict=3072				,			lowercase_				:  Optional[Any]=512				,			lowercase_				:  Dict=12				,			lowercase_				:  Optional[int]=12				,			lowercase_				:  Optional[Any]=3				,			lowercase_				:  str=224				,			lowercase_				:  List[Any]=32				,			lowercase_				:  Union[str, Any]="quick_gelu"				,			lowercase_				:  Dict=1E-5				,			lowercase_				:  Optional[int]=0.0				,			lowercase_				:  Optional[Any]=0.02				,			lowercase_				:  Optional[Any]=1.0				,			**lowercase_				:  Dict				,			):
										super().__init__(**lowercase_    )

										lowercase_					:			Tuple										=       hidden_size
										lowercase_					:			Optional[Any]										=       intermediate_size
										lowercase_					:			Optional[Any]										=       projection_dim
										lowercase_					:			Tuple										=       num_hidden_layers
										lowercase_					:			List[Any]										=       num_attention_heads
										lowercase_					:			Any										=       num_channels
										lowercase_					:			Any										=       patch_size
										lowercase_					:			Dict										=       image_size
										lowercase_					:			Optional[Any]										=       initializer_range
										lowercase_					:			str										=       initializer_factor
										lowercase_					:			Any										=       attention_dropout
										lowercase_					:			Optional[int]										=       layer_norm_eps
										lowercase_					:			int										=       hidden_act







			@classmethod
			def 	SCREAMING_SNAKE_CASE_     (	cls				:  Optional[int]				,			lowercase_				:  Union[str, os.PathLike]				,			**lowercase_				:  Any    ):
										cls._set_token_in_kwargs(lowercase_    )

										lowercase_						,						lowercase_					:			str										=       cls.get_config_dict(lowercase_				,			**lowercase_    )

										# get the vision config dict if we are loading from AltCLIPConfig
										if config_dict.get("""model_type"""    ) == "altclip":
																	lowercase_					:			List[str]										=       config_dict["""vision_config"""]

										if "model_type" in config_dict and hasattr(cls				,			"""model_type"""    ) and config_dict["model_type"] != cls.model_type:
																	logger.warning(
																	    f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
																	    f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''    )

										return cls.from_dict(lowercase_				,			**lowercase_    )





class     __magic_name__						(    _UpperCAmelCase):
			UpperCamelCase__				=       '''altclip'''
			UpperCamelCase__				=       True







			def __init__(	self				:  Optional[int]				,			lowercase_				:  Dict=None				,			lowercase_				:  List[Any]=None				,			lowercase_				:  Tuple=768				,			lowercase_				:  List[str]=2.65_92				,			**lowercase_				:  List[Any]    ):
										# If `_config_dict` exist, we use them for the backward compatibility.
										# We pop out these 2 attributes before calling `super().__init__` to avoid them being saved (which causes a lot
										# of confusion!).
										lowercase_					:			Dict										=       kwargs.pop("""text_config_dict"""				,			lowercase_    )
										lowercase_					:			str										=       kwargs.pop("""vision_config_dict"""				,			lowercase_    )

										super().__init__(**lowercase_    )

										# Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
										# `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
										# cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
										if text_config_dict is not None:
																	if text_config is None:
																								lowercase_					:			Dict										=       {}

																	# This is the complete result when using `text_config_dict`.
																	lowercase_					:			List[str]										=       AltCLIPTextConfig(**lowercase_    ).to_dict()

																	# Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
																	for key, value in _text_config_dict.items():
																								if key in text_config and value != text_config[key] and key not in ["transformers_version"]:
																															# If specified in `text_config_dict`
																															if key in text_config_dict:
																																						lowercase_					:			Optional[Any]										=       (
																																						    f'''`{key}` is found in both `text_config_dict` and `text_config` but with different values. '''
																																						    f'''The value `text_config_dict["{key}"]` will be used instead.'''
																																						)
																															# If inferred from default argument values (just to be super careful)
																															else:
																																						lowercase_					:			Tuple										=       (
																																						    f'''`text_config_dict` is provided which will be used to initialize `AltCLIPTextConfig`. The '''
																																						    f'''value `text_config["{key}"]` will be overriden.'''
																																						)
																															logger.warning(lowercase_    )

            # Update all values in `text_config` with the ones in `_text_config_dict`.
																	text_config.update(_text_config_dict    )

										if vision_config_dict is not None:
																	if vision_config is None:
																								lowercase_					:			int										=       {}

																	# This is the complete result when using `vision_config_dict`.
																	lowercase_					:			List[str]										=       AltCLIPVisionConfig(**lowercase_    ).to_dict()
																	# convert keys to string instead of integer
																	if "id2label" in _vision_config_dict:
																								lowercase_					:			List[str]										=       {
																								    str(lowercase_    ): value for key, value in _vision_config_dict["""id2label"""].items()
																								}

																	# Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
																	for key, value in _vision_config_dict.items():
																								if key in vision_config and value != vision_config[key] and key not in ["transformers_version"]:
																															# If specified in `vision_config_dict`
																															if key in vision_config_dict:
																																						lowercase_					:			Any										=       (
																																						    f'''`{key}` is found in both `vision_config_dict` and `vision_config` but with different '''
																																						    f'''values. The value `vision_config_dict["{key}"]` will be used instead.'''
																																						)
																															# If inferred from default argument values (just to be super careful)
																															else:
																																						lowercase_					:			List[str]										=       (
																																						    f'''`vision_config_dict` is provided which will be used to initialize `AltCLIPVisionConfig`. '''
																																						    f'''The value `vision_config["{key}"]` will be overriden.'''
																																						)
																															logger.warning(lowercase_    )

            # Update all values in `vision_config` with the ones in `_vision_config_dict`.
																	vision_config.update(_vision_config_dict    )

										if text_config is None:
																	lowercase_					:			int										=       {}
																	logger.info("""`text_config` is `None`. Initializing the `AltCLIPTextConfig` with default values."""    )

										if vision_config is None:
																	lowercase_					:			Optional[int]										=       {}
																	logger.info("""`vision_config` is `None`. initializing the `AltCLIPVisionConfig` with default values."""    )

										lowercase_					:			Optional[int]										=       AltCLIPTextConfig(**lowercase_    )
										lowercase_					:			Any										=       AltCLIPVisionConfig(**lowercase_    )

										lowercase_					:			List[Any]										=       projection_dim
										lowercase_					:			Optional[Any]										=       logit_scale_init_value
										lowercase_					:			int										=       1.0







			@classmethod
			def 	SCREAMING_SNAKE_CASE_     (	cls				:  Dict				,			lowercase_				:  AltCLIPTextConfig				,			lowercase_				:  AltCLIPVisionConfig				,			**lowercase_				:  Optional[int]    ):
										return cls(text_config=text_config.to_dict()				,			vision_config=vision_config.to_dict()				,			**lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										lowercase_					:			Union[str, Any]										=       copy.deepcopy(self.__dict__    )
										lowercase_					:			Optional[int]										=       self.text_config.to_dict()
										lowercase_					:			Any										=       self.vision_config.to_dict()
										lowercase_					:			List[str]										=       self.__class__.model_type
										return output





