



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



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_=False					):
			lowercase__   					=			[]

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



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



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					):
			lowercase__   					=			dct.pop(SCREAMING_SNAKE_CASE_					)
			lowercase__   					=			val



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



@torch.no_grad()
def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_=False					):
			lowercase__   					=			BitConfig(
			    global_padding="same"  ,					layer_type="bottleneck"  ,					depths=(3, 4, 9)  ,					out_features=["stage3"]  ,					embedding_dynamic_padding=SCREAMING_SNAKE_CASE_  ,					)
			lowercase__   					=			ViTHybridConfig(backbone_config=SCREAMING_SNAKE_CASE_  ,					image_size=384  ,					num_labels=1000					)
			lowercase__   					=			False

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

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

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

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

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

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

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

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

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

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









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


lowercase_						=       logging.get_logger(__name__)

lowercase_						=       {
    """facebook/nllb-moe-54B""": """https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json""",
}
class 			_snake_case     (   lowercase__):
				UpperCamelCase__			:   Tuple       						="""nllb-moe"""
				UpperCamelCase__			:   Dict       						=["""past_key_values"""]
				UpperCamelCase__			:   List[Any]       						={"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}




				def __init__(  self   :  List[str],   __lowercase   :  str=12_8112,   __lowercase   :  List[Any]=1024,   __lowercase   :  int=12,   __lowercase   :  Union[str, Any]=4096,   __lowercase   :  Union[str, Any]=16,   __lowercase   :  Tuple=12,   __lowercase   :  int=4096,   __lowercase   :  Any=16,   __lowercase   :  List[str]=0.05,   __lowercase   :  Optional[int]=0.05,   __lowercase   :  int=True,   __lowercase   :  Optional[Any]=True,   __lowercase   :  Optional[Any]="relu",   __lowercase   :  str=1024,   __lowercase   :  Optional[Any]=0.1,   __lowercase   :  List[Any]=0.1,   __lowercase   :  Union[str, Any]=0.0,   __lowercase   :  Tuple=0.02,   __lowercase   :  Optional[Any]=2,   __lowercase   :  Optional[int]=True,   __lowercase   :  List[Any]=False,   __lowercase   :  Any="float32",   __lowercase   :  Union[str, Any]=False,   __lowercase   :  Optional[int]=128,   __lowercase   :  Union[str, Any]=64,   __lowercase   :  Optional[int]=4,   __lowercase   :  Dict=4,   __lowercase   :  Any=0.001,   __lowercase   :  Optional[int]=0.001,   __lowercase   :  int="all",   __lowercase   :  Any=False,   __lowercase   :  Optional[int]=False,   __lowercase   :  Dict=1.0,   __lowercase   :  int=0.2,   __lowercase   :  str=1,   __lowercase   :  Optional[Any]=0,   __lowercase   :  Optional[int]=2,   __lowercase   :  Any=False,   **__lowercase   :  List[str],   ):
							lowercase__   					=			vocab_size
							lowercase__   					=			max_position_embeddings
							lowercase__   					=			d_model
							lowercase__   					=			encoder_ffn_dim
							lowercase__   					=			encoder_layers
							lowercase__   					=			encoder_attention_heads
							lowercase__   					=			decoder_ffn_dim
							lowercase__   					=			decoder_layers
							lowercase__   					=			decoder_attention_heads
							lowercase__   					=			dropout
							lowercase__   					=			attention_dropout
							lowercase__   					=			activation_dropout
							lowercase__   					=			activation_function
							lowercase__   					=			init_std
							lowercase__   					=			encoder_layerdrop
							lowercase__   					=			decoder_layerdrop
							lowercase__   					=			use_cache
							lowercase__   					=			encoder_layers
							lowercase__   					=			scale_embedding  # scale factor will be sqrt(d_model) if True
							lowercase__   					=			router_z_loss_coef
							lowercase__   					=			router_aux_loss_coef
							lowercase__   					=			decoder_sparse_step
							lowercase__   					=			encoder_sparse_step
							lowercase__   					=			num_experts
							lowercase__   					=			expert_capacity
							lowercase__   					=			router_bias
							if router_dtype not in ["float32", "float16", "bfloat16"]:
										raise ValueError(F'''`router_dtype` must be one of \'float32\', \'float16\' or \'bfloat16\', got {router_dtype}'''  )
							lowercase__   					=			router_dtype

							lowercase__   					=			router_ignore_padding_tokens
							lowercase__   					=			batch_prioritized_routing
							lowercase__   					=			second_expert_policy
							lowercase__   					=			normalize_router_prob_before_dropping
							lowercase__   					=			moe_eval_capacity_token_fraction
							lowercase__   					=			moe_token_dropout
							lowercase__   					=			output_router_logits
							super().__init__(
							    pad_token_id=__lowercase,   bos_token_id=__lowercase,   eos_token_id=__lowercase,   is_encoder_decoder=__lowercase,   decoder_start_token_id=__lowercase,   **__lowercase,   )




