import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


__UpperCamelCase	: Dict    		=  logging.get_logger(__name__)

# General docstring
__UpperCamelCase	: Optional[Any]    		=  'PoolFormerConfig'

# Base docstring
__UpperCamelCase	: List[str]    		=  'sail/poolformer_s12'
__UpperCamelCase	: Any    		=  [1, 512, 7, 7]

# Image classification docstring
__UpperCamelCase	: Dict    		=  'sail/poolformer_s12'
__UpperCamelCase	: int    		=  'tabby, tabby cat'

__UpperCamelCase	: Tuple    		=  [
    'sail/poolformer_s12',
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]
def    A			(					_lowercase	,						_lowercase = 0.0	,						_lowercase = False		):
						if drop_prob == 0.0 or not training:
												return input
						SCREAMING_SNAKE_CASE							:					int           =							1 - drop_prob
						SCREAMING_SNAKE_CASE							:					str           =							(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
						SCREAMING_SNAKE_CASE							:					int           =							keep_prob + torch.rand(_lowercase	,						dtype=input.dtype	,						device=input.device		)
						random_tensor.floor_()  # binarize
						SCREAMING_SNAKE_CASE							:					List[Any]           =							input.div(_lowercase		) * random_tensor
						return output





class       lowercase__		(     nn.Module):


						def __init__( self			:       Tuple					,	UpperCamelCase__			:       Optional[float] = None		):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					str           =							drop_prob


						def  __A ( self			:       Union[str, Any]					,	UpperCamelCase__			:       torch.Tensor		):





												'''simple docstring'''





												return drop_path(_a					,	self.drop_prob					,	self.training		)



						def  __A ( self			:       Dict		):





												'''simple docstring'''





												return "p={}".format(self.drop_prob		)






class       lowercase__		(     nn.Module):


						def __init__( self			:       Any					,	UpperCamelCase__			:       Dict					,	UpperCamelCase__			:       Optional[int]					,	UpperCamelCase__			:       Any					,	UpperCamelCase__			:       Optional[int]					,	UpperCamelCase__			:       Tuple					,	UpperCamelCase__			:       str=None		):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					str           =							patch_size if isinstance(_a					,	collections.abc.Iterable		) else (patch_size, patch_size)
												SCREAMING_SNAKE_CASE							:					Dict           =							stride if isinstance(_a					,	collections.abc.Iterable		) else (stride, stride)
												SCREAMING_SNAKE_CASE							:					Any           =							padding if isinstance(_a					,	collections.abc.Iterable		) else (padding, padding)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							nn.Convad(_a					,	_a					,	kernel_size=_a					,	stride=_a					,	padding=_a		)
												SCREAMING_SNAKE_CASE							:					Dict           =							norm_layer(_a		) if norm_layer else nn.Identity()



						def  __A ( self			:       List[str]					,	UpperCamelCase__			:       int		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					str           =							self.projection(_a		)
												SCREAMING_SNAKE_CASE							:					Tuple           =							self.norm(_a		)
												return embeddings






class       lowercase__		(     nn.GroupNorm):


						def __init__( self			:       Union[str, Any]					,	UpperCamelCase__			:       Optional[int]					,	**UpperCamelCase__			:       Dict		):





												'''simple docstring'''





												super().__init__(1					,	_a					,	**_a		)






class       lowercase__		(     nn.Module):


						def __init__( self			:       Optional[Any]					,	UpperCamelCase__			:       int		):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					Tuple           =							nn.AvgPoolad(_a					,	stride=1					,	padding=pool_size // 2					,	count_include_pad=_a		)



						def  __A ( self			:       Dict					,	UpperCamelCase__			:       Optional[int]		):





												'''simple docstring'''





												return self.pool(_a		) - hidden_states






class       lowercase__		(     nn.Module):


						def __init__( self			:       Any					,	UpperCamelCase__			:       List[str]					,	UpperCamelCase__			:       Any					,	UpperCamelCase__			:       int					,	UpperCamelCase__			:       int		):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					Dict           =							nn.Convad(_a					,	_a					,	1		)
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							nn.Convad(_a					,	_a					,	1		)
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							PoolFormerDropPath(_a		)
												if isinstance(config.hidden_act					,	_a		):
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							ACTaFN[config.hidden_act]
												else:
																		SCREAMING_SNAKE_CASE							:					Dict           =							config.hidden_act



						def  __A ( self			:       str					,	UpperCamelCase__			:       Tuple		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Any           =							self.conva(_a		)
												SCREAMING_SNAKE_CASE							:					str           =							self.act_fn(_a		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							self.drop(_a		)
												SCREAMING_SNAKE_CASE							:					List[str]           =							self.conva(_a		)
												SCREAMING_SNAKE_CASE							:					List[str]           =							self.drop(_a		)

												return hidden_states






class       lowercase__		(     nn.Module):


						def __init__( self			:       List[Any]					,	UpperCamelCase__			:       Any					,	UpperCamelCase__			:       Tuple					,	UpperCamelCase__			:       List[Any]					,	UpperCamelCase__			:       Tuple					,	UpperCamelCase__			:       Optional[Any]					,	UpperCamelCase__			:       List[Any]		):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					List[Any]           =							PoolFormerPooling(_a		)
												SCREAMING_SNAKE_CASE							:					str           =							PoolFormerOutput(_a					,	_a					,	_a					,	_a		)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							PoolFormerGroupNorm(_a		)
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							PoolFormerGroupNorm(_a		)

												# Useful for training neural nets
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							PoolFormerDropPath(_a		) if drop_path > 0.0 else nn.Identity()
												SCREAMING_SNAKE_CASE							:					Any           =							config.use_layer_scale
												if config.use_layer_scale:
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							nn.Parameter(
																		    config.layer_scale_init_value * torch.ones((_a)		)					,	requires_grad=_a		)
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							nn.Parameter(
																		    config.layer_scale_init_value * torch.ones((_a)		)					,	requires_grad=_a		)



						def  __A ( self			:       Optional[int]					,	UpperCamelCase__			:       Any		):





												'''simple docstring'''





												if self.use_layer_scale:
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							self.pooling(self.before_norm(_a		)		)
																		SCREAMING_SNAKE_CASE							:					str           =							self.layer_scale_a.unsqueeze(-1		).unsqueeze(-1		) * pooling_output
																		# First residual connection
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							hidden_states + self.drop_path(_a		)
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							()

																		SCREAMING_SNAKE_CASE							:					List[Any]           =							self.output(self.after_norm(_a		)		)
																		SCREAMING_SNAKE_CASE							:					str           =							self.layer_scale_a.unsqueeze(-1		).unsqueeze(-1		) * layer_output
																		# Second residual connection
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							hidden_states + self.drop_path(_a		)

																		SCREAMING_SNAKE_CASE							:					Dict           =							(output,) + outputs
																		return outputs

												else:
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							self.drop_path(self.pooling(self.before_norm(_a		)		)		)
																		# First residual connection
																		SCREAMING_SNAKE_CASE							:					Any           =							pooling_output + hidden_states
																		SCREAMING_SNAKE_CASE							:					List[str]           =							()

																		# Second residual connection inside the PoolFormerOutput block
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							self.drop_path(self.output(self.after_norm(_a		)		)		)
																		SCREAMING_SNAKE_CASE							:					str           =							hidden_states + layer_output

																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							(output,) + outputs
																		return outputs






class       lowercase__		(     nn.Module):


						def __init__( self			:       Tuple					,	UpperCamelCase__			:       Optional[Any]		):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					Any           =							config
												# stochastic depth decay rule
												SCREAMING_SNAKE_CASE							:					Dict           =							[x.item() for x in torch.linspace(0					,	config.drop_path_rate					,	sum(config.depths		)		)]

												# patch embeddings
												SCREAMING_SNAKE_CASE							:					Tuple           =							[]
												for i in range(config.num_encoder_blocks		):
																		embeddings.append(
																		    PoolFormerEmbeddings(
																		        patch_size=config.patch_sizes[i]					,	stride=config.strides[i]					,	padding=config.padding[i]					,	num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]					,	hidden_size=config.hidden_sizes[i]					,	)		)
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							nn.ModuleList(_a		)

												# Transformer blocks
												SCREAMING_SNAKE_CASE							:					Tuple           =							[]
												SCREAMING_SNAKE_CASE							:					int           =							0
												for i in range(config.num_encoder_blocks		):
																		# each block consists of layers
																		SCREAMING_SNAKE_CASE							:					str           =							[]
																		if i != 0:
																								cur += config.depths[i - 1]
																		for j in range(config.depths[i]		):
																								layers.append(
																								    PoolFormerLayer(
																								        _a					,	num_channels=config.hidden_sizes[i]					,	pool_size=config.pool_size					,	hidden_size=config.hidden_sizes[i]					,	intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio		)					,	drop_path=dpr[cur + j]					,	)		)
																		blocks.append(nn.ModuleList(_a		)		)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							nn.ModuleList(_a		)



						def  __A ( self			:       Tuple					,	UpperCamelCase__			:       List[str]					,	UpperCamelCase__			:       List[Any]=False					,	UpperCamelCase__			:       Dict=True		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[Any]           =							() if output_hidden_states else None

												SCREAMING_SNAKE_CASE							:					str           =							pixel_values
												for idx, layers in enumerate(zip(self.patch_embeddings					,	self.block		)		):
																		SCREAMING_SNAKE_CASE							:					List[str]           =							layers
																		# Get patch embeddings from hidden_states
																		SCREAMING_SNAKE_CASE							:					Any           =							embedding_layer(_a		)
																		# Send the embeddings through the blocks
																		for _, blk in enumerate(_a		):
																								SCREAMING_SNAKE_CASE							:					Tuple           =							blk(_a		)
																								SCREAMING_SNAKE_CASE							:					List[str]           =							layer_outputs[0]

																		if output_hidden_states:
																								SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							all_hidden_states + (hidden_states,)

												if not return_dict:
																		return tuple(v for v in [hidden_states, all_hidden_states] if v is not None		)

												return BaseModelOutputWithNoAttention(last_hidden_state=_a					,	hidden_states=_a		)






class       lowercase__		(     __SCREAMING_SNAKE_CASE):
						UpperCamelCase_  =	PoolFormerConfig
						UpperCamelCase_  =	"""poolformer"""
						UpperCamelCase_  =	"""pixel_values"""
						UpperCamelCase_  =	True


						def  __A ( self			:       int					,	UpperCamelCase__			:       List[str]		):





												'''simple docstring'''





												if isinstance(_a					,	(nn.Linear, nn.Convad)		):
																		module.weight.data.normal_(mean=0.0					,	std=self.config.initializer_range		)
																		if module.bias is not None:
																								module.bias.data.zero_()
												elif isinstance(_a					,	nn.LayerNorm		):
																		module.bias.data.zero_()
																		module.weight.data.fill_(1.0		)



						def  __A ( self			:       Any					,	UpperCamelCase__			:       Optional[int]					,	UpperCamelCase__			:       List[str]=False		):





												'''simple docstring'''





												if isinstance(_a					,	_a		):
																		SCREAMING_SNAKE_CASE							:					int           =							value






__UpperCamelCase	: Dict    		=  R'\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'

__UpperCamelCase	: Optional[Any]    		=  R'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`PoolFormerImageProcessor.__call__`] for details.\n'





@add_start_docstrings(
    """The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."""			,							__SCREAMING_SNAKE_CASE			,							)
class       lowercase__		(     __SCREAMING_SNAKE_CASE):


						def __init__( self			:       str					,	UpperCamelCase__			:       str		):





												'''simple docstring'''





												super().__init__(_a		)
												SCREAMING_SNAKE_CASE							:					List[str]           =							config

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							PoolFormerEncoder(_a		)

												# Initialize weights and apply final processing
												self.post_init()


						def  __A ( self			:       Optional[int]		):





												'''simple docstring'''





												return self.embeddings.patch_embeddings



						@add_start_docstrings_to_model_forward(_a		)
						@add_code_sample_docstrings(
						    checkpoint=_CHECKPOINT_FOR_DOC					,	output_type=_a					,	config_class=_CONFIG_FOR_DOC					,	modality='''vision'''					,	expected_output=_EXPECTED_OUTPUT_SHAPE					,	)
						def  __A ( self			:       List[str]					,	UpperCamelCase__			:       Optional[torch.FloatTensor] = None					,	UpperCamelCase__			:       Optional[bool] = None					,	UpperCamelCase__			:       Optional[bool] = None					,	):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Any           =							(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							return_dict if return_dict is not None else self.config.use_return_dict

												if pixel_values is None:
																		raise ValueError('''You have to specify pixel_values'''		)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							self.encoder(
												    _a					,	output_hidden_states=_a					,	return_dict=_a					,	)
												SCREAMING_SNAKE_CASE							:					int           =							encoder_outputs[0]

												if not return_dict:
																		return (sequence_output, None) + encoder_outputs[1:]

												return BaseModelOutputWithNoAttention(
												    last_hidden_state=_a					,	hidden_states=encoder_outputs.hidden_states					,	)






class       lowercase__		(     nn.Module):


						def __init__( self			:       Optional[int]					,	UpperCamelCase__			:       str		):





												'''simple docstring'''





												super().__init__()
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							nn.Linear(config.hidden_size					,	config.hidden_size		)



						def  __A ( self			:       str					,	UpperCamelCase__			:       Any		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							self.dense(_a		)
												return output






@add_start_docstrings(
    """
    PoolFormer Model transformer with an image classification head on top
    """			,							__SCREAMING_SNAKE_CASE			,							)
class       lowercase__		(     __SCREAMING_SNAKE_CASE):


						def __init__( self			:       int					,	UpperCamelCase__			:       List[str]		):





												'''simple docstring'''





												super().__init__(_a		)
												SCREAMING_SNAKE_CASE							:					List[str]           =							config.num_labels
												SCREAMING_SNAKE_CASE							:					str           =							PoolFormerModel(_a		)

												# Final norm
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							PoolFormerGroupNorm(config.hidden_sizes[-1]		)
												# Classifier head
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							(
												    nn.Linear(config.hidden_sizes[-1]					,	config.num_labels		) if config.num_labels > 0 else nn.Identity()
												)

												# Initialize weights and apply final processing
												self.post_init()



						@add_start_docstrings_to_model_forward(_a		)
						@add_code_sample_docstrings(
						    checkpoint=_IMAGE_CLASS_CHECKPOINT					,	output_type=_a					,	config_class=_CONFIG_FOR_DOC					,	expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT					,	)
						def  __A ( self			:       str					,	UpperCamelCase__			:       Optional[torch.FloatTensor] = None					,	UpperCamelCase__			:       Optional[torch.LongTensor] = None					,	UpperCamelCase__			:       Optional[bool] = None					,	UpperCamelCase__			:       Optional[bool] = None					,	):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					int           =							return_dict if return_dict is not None else self.config.use_return_dict

												SCREAMING_SNAKE_CASE							:					Tuple           =							self.poolformer(
												    _a					,	output_hidden_states=_a					,	return_dict=_a					,	)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							outputs[0]

												SCREAMING_SNAKE_CASE							:					Optional[int]           =							self.classifier(self.norm(_a		).mean([-2, -1]		)		)

												SCREAMING_SNAKE_CASE							:					Tuple           =							None
												if labels is not None:
																		if self.config.problem_type is None:
																								if self.num_labels == 1:
																														SCREAMING_SNAKE_CASE							:					Optional[Any]           =							"""regression"""
																								elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																														SCREAMING_SNAKE_CASE							:					List[Any]           =							"""single_label_classification"""
																								else:
																														SCREAMING_SNAKE_CASE							:					int           =							"""multi_label_classification"""

																		if self.config.problem_type == "regression":
																								SCREAMING_SNAKE_CASE							:					Optional[Any]           =							MSELoss()
																								if self.num_labels == 1:
																														SCREAMING_SNAKE_CASE							:					Optional[int]           =							loss_fct(logits.squeeze()					,	labels.squeeze()		)
																								else:
																														SCREAMING_SNAKE_CASE							:					List[Any]           =							loss_fct(_a					,	_a		)
																		elif self.config.problem_type == "single_label_classification":
																								SCREAMING_SNAKE_CASE							:					List[str]           =							CrossEntropyLoss()
																								SCREAMING_SNAKE_CASE							:					List[str]           =							loss_fct(logits.view(-1					,	self.num_labels		)					,	labels.view(-1		)		)
																		elif self.config.problem_type == "multi_label_classification":
																								SCREAMING_SNAKE_CASE							:					Dict           =							BCEWithLogitsLoss()
																								SCREAMING_SNAKE_CASE							:					List[str]           =							loss_fct(_a					,	_a		)

												if not return_dict:
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							(logits,) + outputs[2:]
																		return ((loss,) + output) if loss is not None else output

												return ImageClassifierOutputWithNoAttention(loss=_a					,	logits=_a					,	hidden_states=outputs.hidden_states		)






import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm import create_model
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import BitConfig, BitForImageClassification, BitImageProcessor
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCamelCase	: str    		=  logging.get_logger(__name__)
def    A			(					_lowercase		):
						SCREAMING_SNAKE_CASE							:					Any           =							'''huggingface/label-files'''
						SCREAMING_SNAKE_CASE							:					Any           =							'''imagenet-1k-id2label.json'''
						SCREAMING_SNAKE_CASE							:					Any           =							json.load(open(hf_hub_download(_lowercase	,						_lowercase	,						repo_type='''dataset'''		)	,						'''r'''		)		)
						SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							{int(_lowercase		): v for k, v in idalabel.items()}
						SCREAMING_SNAKE_CASE							:					Dict           =							{v: k for k, v in idalabel.items()}

						SCREAMING_SNAKE_CASE							:					str           =							'''std_conv''' if '''bit''' in model_name else False

						# note that when using BiT as backbone for ViT-hybrid checkpoints,
						# one needs to additionally set config.layer_type = "bottleneck", config.stem_type = "same",
						# config.conv_layer = "std_conv_same"
						SCREAMING_SNAKE_CASE							:					Optional[int]           =							BitConfig(
						    conv_layer=_lowercase	,						num_labels=1_000	,						idalabel=_lowercase	,						labelaid=_lowercase	,						)

						return config
def    A			(					_lowercase		):
						if "stem.conv" in name:
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							name.replace('''stem.conv'''	,						'''bit.embedder.convolution'''		)
						if "blocks" in name:
												SCREAMING_SNAKE_CASE							:					Tuple           =							name.replace('''blocks'''	,						'''layers'''		)
						if "head.fc" in name:
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							name.replace('''head.fc'''	,						'''classifier.1'''		)
						if name.startswith('''norm'''		):
												SCREAMING_SNAKE_CASE							:					str           =							'''bit.''' + name
						if "bit" not in name and "classifier" not in name:
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							'''bit.encoder.''' + name

						return name
def    A			(					):
						SCREAMING_SNAKE_CASE							:					Any           =							'''http://images.cocodataset.org/val2017/000000039769.jpg'''
						SCREAMING_SNAKE_CASE							:					List[str]           =							Image.open(requests.get(_lowercase	,						stream=_lowercase		).raw		)
						return im
@torch.no_grad()
def    A			(					_lowercase	,						_lowercase	,						_lowercase=False		):
						SCREAMING_SNAKE_CASE							:					List[Any]           =							get_config(_lowercase		)

						# load original model from timm
						SCREAMING_SNAKE_CASE							:					Optional[Any]           =							create_model(_lowercase	,						pretrained=_lowercase		)
						timm_model.eval()

						# load state_dict of original model
						SCREAMING_SNAKE_CASE							:					Optional[int]           =							timm_model.state_dict()
						for key in state_dict.copy().keys():
												SCREAMING_SNAKE_CASE							:					Dict           =							state_dict.pop(_lowercase		)
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							val.squeeze() if '''head''' in key else val

						# load HuggingFace model
						SCREAMING_SNAKE_CASE							:					str           =							BitForImageClassification(_lowercase		)
						model.eval()
						model.load_state_dict(_lowercase		)

						# create image processor
						SCREAMING_SNAKE_CASE							:					Optional[Any]           =							create_transform(**resolve_data_config({}	,						model=_lowercase		)		)
						SCREAMING_SNAKE_CASE							:					List[str]           =							transform.transforms

						SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							{
						    '''bilinear''': PILImageResampling.BILINEAR,
						    '''bicubic''': PILImageResampling.BICUBIC,
						    '''nearest''': PILImageResampling.NEAREST,
						}

						SCREAMING_SNAKE_CASE							:					Tuple           =							BitImageProcessor(
						    do_resize=_lowercase	,						size={'''shortest_edge''': timm_transforms[0].size}	,						resample=pillow_resamplings[timm_transforms[0].interpolation.value]	,						do_center_crop=_lowercase	,						crop_size={'''height''': timm_transforms[1].size[0], '''width''': timm_transforms[1].size[1]}	,						do_normalize=_lowercase	,						image_mean=timm_transforms[-1].mean.tolist()	,						image_std=timm_transforms[-1].std.tolist()	,						)

						SCREAMING_SNAKE_CASE							:					Any           =							prepare_img()
						SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							transform(_lowercase		).unsqueeze(0		)
						SCREAMING_SNAKE_CASE							:					Optional[int]           =							processor(_lowercase	,						return_tensors='''pt'''		).pixel_values

						# verify pixel values
						assert torch.allclose(_lowercase	,						_lowercase		)

						# verify logits
						with torch.no_grad():
												SCREAMING_SNAKE_CASE							:					Dict           =							model(_lowercase		)
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							outputs.logits

						print('''Logits:'''	,						logits[0, :3]		)
						print('''Predicted class:'''	,						model.config.idalabel[logits.argmax(-1		).item()]		)
						SCREAMING_SNAKE_CASE							:					List[Any]           =							timm_model(_lowercase		)
						assert timm_logits.shape == outputs.logits.shape
						assert torch.allclose(_lowercase	,						outputs.logits	,						atol=1e-3		)
						print('''Looks ok!'''		)

						if pytorch_dump_folder_path is not None:
												Path(_lowercase		).mkdir(exist_ok=_lowercase		)
												print(f"""Saving model {model_name} and processor to {pytorch_dump_folder_path}"""		)
												model.save_pretrained(_lowercase		)
												processor.save_pretrained(_lowercase		)

						if push_to_hub:
												print(f"""Pushing model {model_name} and processor to the hub"""		)
												model.push_to_hub(f"""ybelkada/{model_name}"""		)
												processor.push_to_hub(f"""ybelkada/{model_name}"""		)


if __name__ == "__main__":
					__UpperCamelCase	: Any    		=  argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '--model_name',
					    default='resnetv2_50x1_bitm',
					    type=str,
					    help='Name of the BiT timm 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 to the hub.',
					)

					__UpperCamelCase	: Optional[int]    		=  parser.parse_args()
					convert_bit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





