

"""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()
A_							:			Any											= logging.get_logger(__name__)

def 	A       (     snake_case__		,				snake_case__=False    ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__								=      []

			# 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"
						SCREAMING_SNAKE_CASE__								=      [(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 	A       (     snake_case__		,				snake_case__		,				snake_case__=False    ):




			'''simple docstring'''


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

def 	A       (     snake_case__    ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__								=      ["""head.weight""", """head.bias"""]
			for k in ignore_keys:
						state_dict.pop(snake_case__		,				snake_case__    )

def 	A       (     snake_case__		,				snake_case__		,				snake_case__    ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__								=      dct.pop(snake_case__    )
			SCREAMING_SNAKE_CASE__								=      val

def 	A       (     ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__								=      """http://images.cocodataset.org/val2017/000000039769.jpg"""
			SCREAMING_SNAKE_CASE__								=      Image.open(requests.get(snake_case__		,				stream=snake_case__    ).raw    )
			return im

@torch.no_grad()
def 	A       (     snake_case__		,				snake_case__		,				snake_case__=False    ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__								=      BitConfig(
			    global_padding="""same"""		,				layer_type="""bottleneck"""		,				depths=(3, 4, 9)		,				out_features=["""stage3"""]		,				embedding_dynamic_padding=snake_case__		,				)
			SCREAMING_SNAKE_CASE__								=      ViTHybridConfig(backbone_config=snake_case__		,				image_size=3_84		,				num_labels=10_00    )
			SCREAMING_SNAKE_CASE__								=      False

			# load original model from timm
			SCREAMING_SNAKE_CASE__								=      timm.create_model(snake_case__		,				pretrained=snake_case__    )
			timm_model.eval()

			# load state_dict of original model, remove and rename some keys
			SCREAMING_SNAKE_CASE__								=      timm_model.state_dict()
			if base_model:
						remove_classification_head_(snake_case__    )
			SCREAMING_SNAKE_CASE__								=      create_rename_keys(snake_case__		,				snake_case__    )
			for src, dest in rename_keys:
						rename_key(snake_case__		,				snake_case__		,				snake_case__    )
			read_in_q_k_v(snake_case__		,				snake_case__		,				snake_case__    )

			SCREAMING_SNAKE_CASE__								=      """huggingface/label-files"""
			SCREAMING_SNAKE_CASE__								=      """imagenet-1k-id2label.json"""
			SCREAMING_SNAKE_CASE__								=      json.load(open(hf_hub_download(snake_case__		,				snake_case__		,				repo_type="""dataset"""    )		,				"""r"""    )    )
			SCREAMING_SNAKE_CASE__								=      {int(snake_case__    ): v for k, v in idalabel.items()}
			SCREAMING_SNAKE_CASE__								=      idalabel
			SCREAMING_SNAKE_CASE__								=      {v: k for k, v in idalabel.items()}

			# load HuggingFace model
			if vit_name[-5:] == "in21k":
						SCREAMING_SNAKE_CASE__								=      ViTHybridModel(snake_case__    ).eval()
			else:
						SCREAMING_SNAKE_CASE__								=      ViTHybridForImageClassification(snake_case__    ).eval()
			model.load_state_dict(snake_case__    )

			# create image processor
			SCREAMING_SNAKE_CASE__								=      create_transform(**resolve_data_config({}		,				model=snake_case__    )    )
			SCREAMING_SNAKE_CASE__								=      transform.transforms

			SCREAMING_SNAKE_CASE__								=      {
			    """bilinear""": PILImageResampling.BILINEAR,
			    """bicubic""": PILImageResampling.BICUBIC,
			    """nearest""": PILImageResampling.NEAREST,
			}

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

			SCREAMING_SNAKE_CASE__								=      prepare_img()
			SCREAMING_SNAKE_CASE__								=      transform(snake_case__    ).unsqueeze(0    )
			SCREAMING_SNAKE_CASE__								=      processor(snake_case__		,				return_tensors="""pt"""    ).pixel_values

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

			# verify logits
			with torch.no_grad():
						SCREAMING_SNAKE_CASE__								=      model(snake_case__    )
						SCREAMING_SNAKE_CASE__								=      outputs.logits

			print("""Predicted class:"""		,				logits.argmax(-1    ).item()    )
			if base_model:
						SCREAMING_SNAKE_CASE__								=      timm_model.forward_features(snake_case__    )
						assert timm_pooled_output.shape == outputs.pooler_output.shape
						assert torch.allclose(snake_case__		,				outputs.pooler_output		,				atol=1e-3    )
			else:
						SCREAMING_SNAKE_CASE__								=      timm_model(snake_case__    )
						assert timm_logits.shape == outputs.logits.shape
						assert torch.allclose(snake_case__		,				outputs.logits		,				atol=1e-3    )
			print("""Looks ok!"""    )

			if pytorch_dump_folder_path is not None:
						Path(snake_case__    ).mkdir(exist_ok=snake_case__    )
						print(f"""Saving model {vit_name} to {pytorch_dump_folder_path}"""    )
						model.save_pretrained(snake_case__    )
						print(f"""Saving processor to {pytorch_dump_folder_path}"""    )
						processor.save_pretrained(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__":
						A_							:			Dict											= 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."
						)

						A_							:			str											= parser.parse_args()
						convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)





"""simple docstring"""




# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


A_							:			Dict											= {
    "configuration_efficientnet": [
        "EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "EfficientNetConfig",
        "EfficientNetOnnxConfig",
    ]
}

try:
						if not is_vision_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A_							:			List[Any]											= ["EfficientNetImageProcessor"]

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A_							:			Dict											= [
						    "EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST",
						    "EfficientNetForImageClassification",
						    "EfficientNetModel",
						    "EfficientNetPreTrainedModel",
						]

if TYPE_CHECKING:
						from .configuration_efficientnet import (
						    EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
						    EfficientNetConfig,
						    EfficientNetOnnxConfig,
						)

						try:
												if not is_vision_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .image_processing_efficientnet import EfficientNetImageProcessor

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_efficientnet import (
												    EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST,
												    EfficientNetForImageClassification,
												    EfficientNetModel,
												    EfficientNetPreTrainedModel,
												)

else:
						import sys

						A_							:			List[Any]											= _LazyModule(__name__, globals()["__file__"], _import_structure)


