

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()
lowerCAmelCase__	   = logging.get_logger(__name__)



def 					lowerCAmelCase__							(	SCREAMING_SNAKE_CASE_:							Dict	,       SCREAMING_SNAKE_CASE_:							List[Any]=False      )			->      int:




				'''simple docstring'''





				A__        =			[]

				# 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"
								A__        =			[(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_:							int	,       SCREAMING_SNAKE_CASE_:							Union[str, Any]	,       SCREAMING_SNAKE_CASE_:							str=False      )			->      Optional[Any]:




				'''simple docstring'''





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



def 					lowerCAmelCase__							(	SCREAMING_SNAKE_CASE_:							List[Any]      )			->      int:




				'''simple docstring'''





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



def 					lowerCAmelCase__							(	SCREAMING_SNAKE_CASE_:							Union[str, Any]	,       SCREAMING_SNAKE_CASE_:							str	,       SCREAMING_SNAKE_CASE_:							Tuple      )			->      Union[str, Any]:




				'''simple docstring'''





				A__        =			dct.pop(SCREAMING_SNAKE_CASE_      )
				A__        =			val



def 					lowerCAmelCase__							(	)			->      str:




				'''simple docstring'''





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



@torch.no_grad()
def 					lowerCAmelCase__							(	SCREAMING_SNAKE_CASE_:							List[Any]	,       SCREAMING_SNAKE_CASE_:							Optional[int]	,       SCREAMING_SNAKE_CASE_:							List[Any]=False      )			->      Dict:




				'''simple docstring'''





				A__        =			BitConfig(
				    global_padding="same"	,       layer_type="bottleneck"	,       depths=(3, 4, 9)	,       out_features=["stage3"]	,       embedding_dynamic_padding=SCREAMING_SNAKE_CASE_	,       )
				A__        =			ViTHybridConfig(backbone_config=SCREAMING_SNAKE_CASE_	,       image_size=3_8_4	,       num_labels=1_0_0_0      )
				A__        =			False

				# load original model from timm
				A__        =			timm.create_model(SCREAMING_SNAKE_CASE_	,       pretrained=SCREAMING_SNAKE_CASE_      )
				timm_model.eval()

				# load state_dict of original model, remove and rename some keys
				A__        =			timm_model.state_dict()
				if base_model:
								remove_classification_head_(SCREAMING_SNAKE_CASE_      )
				A__        =			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_      )

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

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

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

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

				A__        =			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()	,       )

				A__        =			prepare_img()
				A__        =			transform(SCREAMING_SNAKE_CASE_      ).unsqueeze(0      )
				A__        =			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():
								A__        =			model(SCREAMING_SNAKE_CASE_      )
								A__        =			outputs.logits

				print("Predicted class:"	,       logits.argmax(-1      ).item()      )
				if base_model:
								A__        =			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:
								A__        =			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__":
	lowerCAmelCase__	   = 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."""
	)

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




from typing import Union

import fire
import torch
from tqdm import tqdm



def 					lowerCAmelCase__							(	SCREAMING_SNAKE_CASE_:							str	,       SCREAMING_SNAKE_CASE_:							str = "cpu"	,       SCREAMING_SNAKE_CASE_:							Union[str, None] = None      )			->      None:




				'''simple docstring'''





				A__        =			torch.load(SCREAMING_SNAKE_CASE_	,       map_location=SCREAMING_SNAKE_CASE_      )
				for k, v in tqdm(state_dict.items()      ):
								if not isinstance(SCREAMING_SNAKE_CASE_	,       torch.Tensor      ):
												raise TypeError("FP16 conversion only works on paths that are saved state dicts, like pytorch_model.bin"      )
								A__        =			v.half()
				if save_path is None:  # overwrite src_path
								A__        =			src_path
				torch.save(SCREAMING_SNAKE_CASE_	,       SCREAMING_SNAKE_CASE_      )


if __name__ == "__main__":
	fire.Fire(convert)

