




'''simple docstring'''

import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
a       :						List[Any]								=		logging.get_logger(__name__)


def __UpperCAmelCase	(   _UpperCAmelCase		:   List[Any]		,				_UpperCAmelCase		:   Any=False						)      ->      Optional[int]:
			__snake_case			=		[]
			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''')						)

			# projection layer + position embeddings
			rename_keys.extend(
			    [
			        ("cls_token", "vit.embeddings.cls_token"),
			        ("patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight"),
			        ("patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias"),
			        ("pos_embed", "vit.embeddings.position_embeddings"),
			    ]						)

			if base_model:
						# layernorm + pooler
						rename_keys.extend(
						    [
						        ("norm.weight", "layernorm.weight"),
						        ("norm.bias", "layernorm.bias"),
						    ]						)

						# if just the base model, we should remove "vit" from all keys that start with "vit"
						__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"),
						    ]						)

			return rename_keys


def __UpperCAmelCase	(   _UpperCAmelCase		:   Any		,				_UpperCAmelCase		:   Tuple		,				_UpperCAmelCase		:   List[Any]=False						)      ->      Any:



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


def __UpperCAmelCase	(   _UpperCAmelCase		:   Tuple						)      ->      Tuple:
			__snake_case			=		['''head.weight''', '''head.bias''']
			for k in ignore_keys:
						state_dict.pop(_UpperCamelCase		,				_UpperCamelCase						)


def __UpperCAmelCase	(   _UpperCAmelCase		:   Dict		,				_UpperCAmelCase		:   Any		,				_UpperCAmelCase		:   List[Any]						)      ->      Union[str, Any]:
			__snake_case			=		dct.pop(_UpperCamelCase						)
			__snake_case			=		val


def __UpperCAmelCase	(   )      ->      str:
			__snake_case			=		'''http://images.cocodataset.org/val2017/000000039769.jpg'''
			__snake_case			=		Image.open(requests.get(_UpperCamelCase		,				stream=_UpperCamelCase						).raw						)
			return im


@torch.no_grad()
def __UpperCAmelCase	(   _UpperCAmelCase		:   Any		,				_UpperCAmelCase		:   List[Any]		,				_UpperCAmelCase		:   Union[str, Any]=True						)      ->      Dict:
			__snake_case			=		ViTConfig()
			# patch_size
			if model_name[-1] == "8":
						__snake_case			=		8
			# set labels if required
			if not base_model:
						__snake_case			=		10_00
						__snake_case			=		'''huggingface/label-files'''
						__snake_case			=		'''imagenet-1k-id2label.json'''
						__snake_case			=		json.load(open(hf_hub_download(_UpperCamelCase		,				_UpperCamelCase		,				repo_type="dataset"						)		,				"r"						)						)
						__snake_case			=		{int(_UpperCamelCase						): v for k, v in idalabel.items()}
						__snake_case			=		idalabel
						__snake_case			=		{v: k for k, v in idalabel.items()}
			# size of the architecture
			if model_name in ["dino_vits8", "dino_vits16"]:
						__snake_case			=		3_84
						__snake_case			=		15_36
						__snake_case			=		12
						__snake_case			=		6

			# load original model from torch hub
			__snake_case			=		torch.hub.load("facebookresearch/dino:main"		,				_UpperCamelCase						)
			original_model.eval()

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

			# load HuggingFace model
			if base_model:
						__snake_case			=		ViTModel(_UpperCamelCase		,				add_pooling_layer=_UpperCamelCase						).eval()
			else:
						__snake_case			=		ViTForImageClassification(_UpperCamelCase						).eval()
			model.load_state_dict(_UpperCamelCase						)

			# Check outputs on an image, prepared by ViTImageProcessor
			__snake_case			=		ViTImageProcessor()
			__snake_case			=		image_processor(images=prepare_img()		,				return_tensors="pt"						)
			__snake_case			=		encoding['''pixel_values''']
			__snake_case			=		model(_UpperCamelCase						)

			if base_model:
						__snake_case			=		original_model(_UpperCamelCase						)
						assert torch.allclose(_UpperCamelCase		,				outputs.last_hidden_state[:, 0, :]		,				atol=1E-1						)
			else:
						__snake_case			=		original_model(_UpperCamelCase						)
						assert logits.shape == outputs.logits.shape
						assert torch.allclose(_UpperCamelCase		,				outputs.logits		,				atol=1E-3						)

			Path(_UpperCamelCase						).mkdir(exist_ok=_UpperCamelCase						)
			print(F'''Saving model {model_name} to {pytorch_dump_folder_path}'''						)
			model.save_pretrained(_UpperCamelCase						)
			print(F'''Saving image processor to {pytorch_dump_folder_path}'''						)
			image_processor.save_pretrained(_UpperCamelCase						)


if __name__ == "__main__":
						a       :						Tuple								=		argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '''--model_name''',
						    default='''dino_vitb16''',
						    type=str,
						    help='''Name of the model trained with DINO 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(
						    '''--base_model''',
						    action='''store_true''',
						    help='''Whether to only convert the base model (no projection head weights).''',
						)

						parser.set_defaults(base_model=True)
						a       :						Dict								=		parser.parse_args()
						convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)







"""simple docstring"""




import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from ...utils import logging
from ..auto import CONFIG_MAPPING


A__ :       List[Any]				= logging.get_logger(__name__)

A__ :       Tuple				= {
    'Salesforce/instruct-blip-flan-t5': 'https://huggingface.co/Salesforce/instruct-blip-flan-t5/resolve/main/config.json',
}


class       __magic_name__							(							SCREAMING_SNAKE_CASE__			):
       UpperCamelCase_ =		'''instructblip_vision_model'''
       def __init__(  self			,					A_=1408			,					A_=6144			,					A_=39			,					A_=16			,					A_=224			,					A_=14			,					A_="gelu"			,					A_=1E-6			,					A_=0.0			,					A_=1E-10			,					A_=True			,					**A_			,					)  ->						Tuple:





         """simple docstring"""




         super().__init__(**A_							)

         _lowercase:      List[str]   			=  hidden_size
         _lowercase:      Tuple   			=  intermediate_size
         _lowercase:      List[str]   			=  num_hidden_layers
         _lowercase:      Union[str, Any]   			=  num_attention_heads
         _lowercase:      str   			=  patch_size
         _lowercase:      Union[str, Any]   			=  image_size
         _lowercase:      int   			=  initializer_range
         _lowercase:      List[str]   			=  attention_dropout
         _lowercase:      Tuple   			=  layer_norm_eps
         _lowercase:      Tuple   			=  hidden_act
         _lowercase:      Optional[int]   			=  qkv_bias



       @classmethod
       def 				lowercase_				(  cls			,					A_			,					**A_							)  ->						"PretrainedConfig":





         """simple docstring"""




         cls._set_token_in_kwargs(A_							)

         _lowercase			,					_lowercase:      str   			=  cls.get_config_dict(A_			,					**A_							)

         # get the vision config dict if we are loading from InstructBlipConfig
         if config_dict.get('''model_type'''							) == "instructblip":
           _lowercase:      Dict   			=  config_dict['''vision_config''']

         if "model_type" in config_dict and hasattr(cls			,					'''model_type'''							) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               f'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
               f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

         return cls.from_dict(A_			,					**A_							)


class       __magic_name__							(							SCREAMING_SNAKE_CASE__			):
       UpperCamelCase_ =		'''instructblip_qformer'''
       def __init__(  self			,					A_=3_0522			,					A_=768			,					A_=12			,					A_=12			,					A_=3072			,					A_="gelu"			,					A_=0.1			,					A_=0.1			,					A_=512			,					A_=0.02			,					A_=1E-12			,					A_=0			,					A_="absolute"			,					A_=2			,					A_=1408			,					**A_			,					)  ->						Optional[int]:





         """simple docstring"""




         super().__init__(pad_token_id=A_			,					**A_							)

         _lowercase:      Optional[Any]   			=  vocab_size
         _lowercase:      Optional[int]   			=  hidden_size
         _lowercase:      Dict   			=  num_hidden_layers
         _lowercase:      Union[str, Any]   			=  num_attention_heads
         _lowercase:      List[Any]   			=  hidden_act
         _lowercase:      Union[str, Any]   			=  intermediate_size
         _lowercase:      Union[str, Any]   			=  hidden_dropout_prob
         _lowercase:      Tuple   			=  attention_probs_dropout_prob
         _lowercase:      Optional[int]   			=  max_position_embeddings
         _lowercase:      Optional[int]   			=  initializer_range
         _lowercase:      Optional[Any]   			=  layer_norm_eps
         _lowercase:      Optional[Any]   			=  position_embedding_type
         _lowercase:      Tuple   			=  cross_attention_frequency
         _lowercase:      str   			=  encoder_hidden_size



       @classmethod
       def 				lowercase_				(  cls			,					A_			,					**A_							)  ->						"PretrainedConfig":





         """simple docstring"""




         cls._set_token_in_kwargs(A_							)

         _lowercase			,					_lowercase:      Tuple   			=  cls.get_config_dict(A_			,					**A_							)

         # get the qformer config dict if we are loading from InstructBlipConfig
         if config_dict.get('''model_type'''							) == "instructblip":
           _lowercase:      Tuple   			=  config_dict['''qformer_config''']

         if "model_type" in config_dict and hasattr(cls			,					'''model_type'''							) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               f'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
               f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

         return cls.from_dict(A_			,					**A_							)


class       __magic_name__							(							SCREAMING_SNAKE_CASE__			):
       UpperCamelCase_ =		'''instructblip'''
       UpperCamelCase_ =		True
       def __init__(  self			,					A_=None			,					A_=None			,					A_=None			,					A_=32			,					**A_							)  ->						str:





         """simple docstring"""




         super().__init__(**A_							)

         if vision_config is None:
           _lowercase:      str   			=  {}
           logger.info('''vision_config is None. initializing the InstructBlipVisionConfig with default values.'''							)

         if qformer_config is None:
           _lowercase:      List[str]   			=  {}
           logger.info('''qformer_config is None. Initializing the InstructBlipQFormerConfig with default values.'''							)

         if text_config is None:
           _lowercase:      Any   			=  {}
           logger.info('''text_config is None. Initializing the text config with default values (`OPTConfig`).'''							)

         _lowercase:      int   			=  InstructBlipVisionConfig(**A_							)
         _lowercase:      Optional[Any]   			=  InstructBlipQFormerConfig(**A_							)
         _lowercase:      Optional[Any]   			=  text_config['''model_type'''] if '''model_type''' in text_config else '''opt'''
         _lowercase:      Optional[Any]   			=  CONFIG_MAPPING[text_model_type](**A_							)

         _lowercase:      Tuple   			=  self.text_config.tie_word_embeddings
         _lowercase:      Dict   			=  self.text_config.is_encoder_decoder

         _lowercase:      List[Any]   			=  num_query_tokens
         _lowercase:      Any   			=  self.vision_config.hidden_size
         _lowercase:      List[Any]   			=  self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
         _lowercase:      Union[str, Any]   			=  1.0
         _lowercase:      Dict   			=  0.02
       @classmethod
       def 				lowercase_				(  cls			,					A_			,					A_			,					A_			,					**A_			,					)  ->						str:





         """simple docstring"""




         return cls(
             vision_config=vision_config.to_dict()			,					qformer_config=qformer_config.to_dict()			,					text_config=text_config.to_dict()			,					**A_			,					)



       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         _lowercase:      List[str]   			=  copy.deepcopy(self.__dict__							)
         _lowercase:      Tuple   			=  self.vision_config.to_dict()
         _lowercase:      Tuple   			=  self.qformer_config.to_dict()
         _lowercase:      int   			=  self.text_config.to_dict()
         _lowercase:      Any   			=  self.__class__.model_type
         return output
