





"""simple docstring"""





import argparse
import json
from pathlib import Path

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

from transformers import (
    BitConfig,
    ViTHybridConfig,
    ViTHybridForImageClassification,
    ViTHybridImageProcessor,
    ViTHybridModel,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
_UpperCamelCase					      =   logging.get_logger(__name__)



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				str     ,    SCREAMING_SNAKE_CASE							:				Optional[int]=False ):


							'''simple docstring'''







							__lowerCamelCase   :   int       =[]

							# fmt: off
							# stem:
							rename_keys.append(('''cls_token''', '''vit.embeddings.cls_token''') )
							rename_keys.append(('''pos_embed''', '''vit.embeddings.position_embeddings''') )

							rename_keys.append(('''patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight''') )
							rename_keys.append(('''patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias''') )

							# backbone
							rename_keys.append(('''patch_embed.backbone.stem.conv.weight''', '''vit.embeddings.patch_embeddings.backbone.bit.embedder.convolution.weight''') )
							rename_keys.append(('''patch_embed.backbone.stem.norm.weight''', '''vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.weight''') )
							rename_keys.append(('''patch_embed.backbone.stem.norm.bias''', '''vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.bias''') )

							for stage_idx in range(len(config.backbone_config.depths ) ):
														for layer_idx in range(config.backbone_config.depths[stage_idx] ):
																					rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv1.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv1.weight') )
																					rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.weight') )
																					rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.bias') )
																					rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv2.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv2.weight') )
																					rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.weight') )
																					rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.bias') )
																					rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv3.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv3.weight') )
																					rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.weight') )
																					rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.bias') )

														rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.conv.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.conv.weight') )
														rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.weight', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.weight') )
														rename_keys.append((F'patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.bias', F'vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.bias') )

							# transformer encoder
							for i in range(config.num_hidden_layers ):
														# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
														rename_keys.append((F'blocks.{i}.norm1.weight', F'vit.encoder.layer.{i}.layernorm_before.weight') )
														rename_keys.append((F'blocks.{i}.norm1.bias', F'vit.encoder.layer.{i}.layernorm_before.bias') )
														rename_keys.append((F'blocks.{i}.attn.proj.weight', F'vit.encoder.layer.{i}.attention.output.dense.weight') )
														rename_keys.append((F'blocks.{i}.attn.proj.bias', F'vit.encoder.layer.{i}.attention.output.dense.bias') )
														rename_keys.append((F'blocks.{i}.norm2.weight', F'vit.encoder.layer.{i}.layernorm_after.weight') )
														rename_keys.append((F'blocks.{i}.norm2.bias', F'vit.encoder.layer.{i}.layernorm_after.bias') )
														rename_keys.append((F'blocks.{i}.mlp.fc1.weight', F'vit.encoder.layer.{i}.intermediate.dense.weight') )
														rename_keys.append((F'blocks.{i}.mlp.fc1.bias', F'vit.encoder.layer.{i}.intermediate.dense.bias') )
														rename_keys.append((F'blocks.{i}.mlp.fc2.weight', F'vit.encoder.layer.{i}.output.dense.weight') )
														rename_keys.append((F'blocks.{i}.mlp.fc2.bias', F'vit.encoder.layer.{i}.output.dense.bias') )

							if base_model:
														# layernorm + pooler
														rename_keys.extend(
														    [
														        ('''norm.weight''', '''layernorm.weight'''),
														        ('''norm.bias''', '''layernorm.bias'''),
														        ('''pre_logits.fc.weight''', '''pooler.dense.weight'''),
														        ('''pre_logits.fc.bias''', '''pooler.dense.bias'''),
														    ] )

														# if just the base model, we should remove "vit" from all keys that start with "vit"
														__lowerCamelCase   :   Union[str, Any]       =[(pair[0], pair[1][4:]) if pair[1].startswith('''vit''' ) else pair for pair in rename_keys]
							else:
														# layernorm + classification head
														rename_keys.extend(
														    [
														        ('''norm.weight''', '''vit.layernorm.weight'''),
														        ('''norm.bias''', '''vit.layernorm.bias'''),
														        ('''head.weight''', '''classifier.weight'''),
														        ('''head.bias''', '''classifier.bias'''),
														    ] )
							# fmt: on

							return rename_keys



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				List[str]     ,    SCREAMING_SNAKE_CASE							:				Any     ,    SCREAMING_SNAKE_CASE							:				Any=False ):


							'''simple docstring'''







							for i in range(config.num_hidden_layers ):
														if base_model:
																					__lowerCamelCase   :   Any       =''''''
														else:
																					__lowerCamelCase   :   List[Any]       ='''vit.'''
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														__lowerCamelCase   :   str       =state_dict.pop(F'blocks.{i}.attn.qkv.weight' )
														__lowerCamelCase   :   Optional[Any]       =state_dict.pop(F'blocks.{i}.attn.qkv.bias' )
														# next, add query, keys and values (in that order) to the state dict
														__lowerCamelCase   :   int       =in_proj_weight[
														    : config.hidden_size, :
														]
														__lowerCamelCase   :   Dict       =in_proj_bias[: config.hidden_size]
														__lowerCamelCase   :   int       =in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														__lowerCamelCase   :   Any       =in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														__lowerCamelCase   :   Dict       =in_proj_weight[
														    -config.hidden_size :, :
														]
														__lowerCamelCase   :   List[str]       =in_proj_bias[-config.hidden_size :]



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Union[str, Any] ):


							'''simple docstring'''







							__lowerCamelCase   :   int       =['''head.weight''', '''head.bias''']
							for k in ignore_keys:
														state_dict.pop(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Optional[int]     ,    SCREAMING_SNAKE_CASE							:				str     ,    SCREAMING_SNAKE_CASE							:				Tuple ):


							'''simple docstring'''







							__lowerCamelCase   :   Any       =dct.pop(SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   Dict       =val



def 					lowerCAmelCase_      (   ):


							'''simple docstring'''







							__lowerCamelCase   :   str       ='''http://images.cocodataset.org/val2017/000000039769.jpg'''
							__lowerCamelCase   :   Dict       =Image.open(requests.get(SCREAMING_SNAKE_CASE     ,    stream=SCREAMING_SNAKE_CASE ).raw )
							return im



@torch.no_grad()
def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				List[str]     ,    SCREAMING_SNAKE_CASE							:				Any     ,    SCREAMING_SNAKE_CASE							:				Any=False ):


							'''simple docstring'''







							__lowerCamelCase   :   Union[str, Any]       =BitConfig(
							    global_padding='''same'''     ,    layer_type='''bottleneck'''     ,    depths=(3, 4, 9)     ,    out_features=['''stage3''']     ,    embedding_dynamic_padding=SCREAMING_SNAKE_CASE     ,    )
							__lowerCamelCase   :   Optional[int]       =ViTHybridConfig(backbone_config=SCREAMING_SNAKE_CASE     ,    image_size=384     ,    num_labels=1000 )
							__lowerCamelCase   :   Optional[int]       =False

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

							# load state_dict of original model, remove and rename some keys
							__lowerCamelCase   :   Optional[Any]       =timm_model.state_dict()
							if base_model:
														remove_classification_head_(SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   Any       =create_rename_keys(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )
							for src, dest in rename_keys:
														rename_key(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )
							read_in_q_k_v(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )

							__lowerCamelCase   :   str       ='''huggingface/label-files'''
							__lowerCamelCase   :   Optional[int]       ='''imagenet-1k-id2label.json'''
							__lowerCamelCase   :   Optional[int]       =json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    repo_type='''dataset''' )     ,    '''r''' ) )
							__lowerCamelCase   :   str       ={int(SCREAMING_SNAKE_CASE ): v for k, v in idalabel.items()}
							__lowerCamelCase   :   str       =idalabel
							__lowerCamelCase   :   int       ={v: k for k, v in idalabel.items()}

							# load HuggingFace model
							if vit_name[-5:] == "in21k":
														__lowerCamelCase   :   Tuple       =ViTHybridModel(SCREAMING_SNAKE_CASE ).eval()
							else:
														__lowerCamelCase   :   Union[str, Any]       =ViTHybridForImageClassification(SCREAMING_SNAKE_CASE ).eval()
							model.load_state_dict(SCREAMING_SNAKE_CASE )

							# create image processor
							__lowerCamelCase   :   int       =create_transform(**resolve_data_config({}     ,    model=SCREAMING_SNAKE_CASE ) )
							__lowerCamelCase   :   str       =transform.transforms

							__lowerCamelCase   :   List[Any]       ={
							    '''bilinear''': PILImageResampling.BILINEAR,
							    '''bicubic''': PILImageResampling.BICUBIC,
							    '''nearest''': PILImageResampling.NEAREST,
							}

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

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

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

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

							print('''Predicted class:'''     ,    logits.argmax(-1 ).item() )
							if base_model:
														__lowerCamelCase   :   Tuple       =timm_model.forward_features(SCREAMING_SNAKE_CASE )
														assert timm_pooled_output.shape == outputs.pooler_output.shape
														assert torch.allclose(SCREAMING_SNAKE_CASE     ,    outputs.pooler_output     ,    atol=1E-3 )
							else:
														__lowerCamelCase   :   List[str]       =timm_model(SCREAMING_SNAKE_CASE )
														assert timm_logits.shape == outputs.logits.shape
														assert torch.allclose(SCREAMING_SNAKE_CASE     ,    outputs.logits     ,    atol=1E-3 )
							print('''Looks ok!''' )

							if pytorch_dump_folder_path is not None:
														Path(SCREAMING_SNAKE_CASE ).mkdir(exist_ok=SCREAMING_SNAKE_CASE )
														print(F'Saving model {vit_name} to {pytorch_dump_folder_path}' )
														model.save_pretrained(SCREAMING_SNAKE_CASE )
														print(F'Saving processor to {pytorch_dump_folder_path}' )
														processor.save_pretrained(SCREAMING_SNAKE_CASE )

							if push_to_hub:
														print(F'Pushing model and processor to the hub {vit_name}' )
														model.push_to_hub(F'ybelkada/{vit_name}' )
														processor.push_to_hub(F'ybelkada/{vit_name}' )


if __name__ == "__main__":
							_UpperCamelCase					      =   argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--vit_name',
							    default='vit_base_r50_s16_384',
							    type=str,
							    help='Name of the hybrid ViT 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 upload the model to the HuggingFace hub.'
							)

							_UpperCamelCase					      =   parser.parse_args()
							convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)











"""simple docstring"""





import os
import sys
import tempfile

import torch

from .state import AcceleratorState
from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Any     ,    SCREAMING_SNAKE_CASE							:				List[str]=()     ,    SCREAMING_SNAKE_CASE							:				Union[str, Any]=None     ,    SCREAMING_SNAKE_CASE							:				Dict="no"     ,    SCREAMING_SNAKE_CASE							:				int="29500" ):


							'''simple docstring'''







							__lowerCamelCase   :   Optional[int]       =False
							__lowerCamelCase   :   Tuple       =False
							if any(key.startswith('''KAGGLE''' ) for key in os.environ.keys() ):
														__lowerCamelCase   :   Optional[int]       =True
							elif "IPython" in sys.modules:
														__lowerCamelCase   :   int       ='''google.colab''' in str(sys.modules['''IPython'''].get_ipython() )

							try:
														__lowerCamelCase   :   str       =PrecisionType(mixed_precision.lower() )
							except ValueError:
														raise ValueError(
														    F'Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.' )

							if (in_colab or in_kaggle) and (os.environ.get('''TPU_NAME'''     ,    SCREAMING_SNAKE_CASE ) is not None):
														# TPU launch
														import torch_xla.distributed.xla_multiprocessing as xmp

														if len(AcceleratorState._shared_state ) > 0:
																					raise ValueError(
																					    '''To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside '''
																					    '''your training function. Restart your notebook and make sure no cells initializes an '''
																					    '''`Accelerator`.''' )
														if num_processes is None:
																					__lowerCamelCase   :   int       =8

														__lowerCamelCase   :   Optional[Any]       =PrepareForLaunch(SCREAMING_SNAKE_CASE     ,    distributed_type='''TPU''' )
														print(F'Launching a training on {num_processes} TPU cores.' )
														xmp.spawn(SCREAMING_SNAKE_CASE     ,    args=SCREAMING_SNAKE_CASE     ,    nprocs=SCREAMING_SNAKE_CASE     ,    start_method='''fork''' )
							elif in_colab:
														# No need for a distributed launch otherwise as it's either CPU or one GPU.
														if torch.cuda.is_available():
																					print('''Launching training on one GPU.''' )
														else:
																					print('''Launching training on one CPU.''' )
														function(*SCREAMING_SNAKE_CASE )
							else:
														if num_processes is None:
																					raise ValueError(
																					    '''You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call.''' )

														if num_processes > 1:
																					# Multi-GPU launch
																					from torch.multiprocessing import start_processes
																					from torch.multiprocessing.spawn import ProcessRaisedException

																					if len(AcceleratorState._shared_state ) > 0:
																												raise ValueError(
																												    '''To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized '''
																												    '''inside your training function. Restart your notebook and make sure no cells initializes an '''
																												    '''`Accelerator`.''' )

																					if torch.cuda.is_initialized():
																												raise ValueError(
																												    '''To launch a multi-GPU training from your notebook, you need to avoid running any instruction '''
																												    '''using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA '''
																												    '''function.''' )

																					# torch.distributed will expect a few environment variable to be here. We set the ones common to each
																					# process here (the other ones will be set be the launcher).
																					with patch_environment(
																					    world_size=SCREAMING_SNAKE_CASE     ,    master_addr='''127.0.01'''     ,    master_port=SCREAMING_SNAKE_CASE     ,    mixed_precision=SCREAMING_SNAKE_CASE ):
																												__lowerCamelCase   :   List[str]       =PrepareForLaunch(SCREAMING_SNAKE_CASE     ,    distributed_type='''MULTI_GPU''' )
																												print(F'Launching training on {num_processes} GPUs.' )
																												try:
																																			start_processes(SCREAMING_SNAKE_CASE     ,    args=SCREAMING_SNAKE_CASE     ,    nprocs=SCREAMING_SNAKE_CASE     ,    start_method='''fork''' )
																												except ProcessRaisedException as e:
																																			if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
																																										raise RuntimeError(
																																										    '''CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. '''
																																										    '''This likely stems from an outside import causing issues once the `notebook_launcher()` is called. '''
																																										    '''Please review your imports and test them when running the `notebook_launcher()` to identify '''
																																										    '''which one is problematic.''' ) from e

														else:
																					# No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
																					if is_mps_available():
																												__lowerCamelCase   :   Optional[Any]       ='''1'''
																												print('''Launching training on MPS.''' )
																					elif torch.cuda.is_available():
																												print('''Launching training on one GPU.''' )
																					else:
																												print('''Launching training on CPU.''' )
																					function(*SCREAMING_SNAKE_CASE )



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				str     ,    SCREAMING_SNAKE_CASE							:				Optional[int]=()     ,    SCREAMING_SNAKE_CASE							:				Optional[Any]=2 ):


							'''simple docstring'''







							from torch.multiprocessing import start_processes

							with tempfile.NamedTemporaryFile() as tmp_file:
														# torch.distributed will expect a few environment variable to be here. We set the ones common to each
														# process here (the other ones will be set be the launcher).
														with patch_environment(
														    world_size=SCREAMING_SNAKE_CASE     ,    master_addr='''127.0.01'''     ,    master_port='''29500'''     ,    accelerate_mixed_precision='''no'''     ,    accelerate_debug_rdv_file=tmp_file.name     ,    accelerate_use_cpu='''yes'''     ,    ):
																					__lowerCamelCase   :   Tuple       =PrepareForLaunch(SCREAMING_SNAKE_CASE     ,    debug=SCREAMING_SNAKE_CASE )
																					start_processes(SCREAMING_SNAKE_CASE     ,    args=SCREAMING_SNAKE_CASE     ,    nprocs=SCREAMING_SNAKE_CASE     ,    start_method='''fork''' )




