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







def 				SCREAMING_SNAKE_CASE__	(   __A			,				__A=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 				SCREAMING_SNAKE_CASE__	(   __A			,				__A			,				__A=False	)       ->	int:
	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 				SCREAMING_SNAKE_CASE__	(   __A	)       ->	List[str]:
	_snake_case   						=		['head.weight', 'head.bias']
	for k in ignore_keys:
		state_dict.pop(__SCREAMING_SNAKE_CASE			,				__SCREAMING_SNAKE_CASE	)







def 				SCREAMING_SNAKE_CASE__	(   __A			,				__A			,				__A	)       ->	Any:
	_snake_case   						=		dct.pop(__SCREAMING_SNAKE_CASE	)
	_snake_case   						=		val







def 				SCREAMING_SNAKE_CASE__	(   )       ->	Any:
	_snake_case   						=		'http://images.cocodataset.org/val2017/000000039769.jpg'
	_snake_case   						=		Image.open(requests.get(__SCREAMING_SNAKE_CASE			,				stream=__SCREAMING_SNAKE_CASE	).raw	)
	return im







@torch.no_grad()
def 				SCREAMING_SNAKE_CASE__	(   __A			,				__A			,				__A=True	)       ->	Optional[Any]:
	_snake_case   						=		ViTConfig()
	# patch_size
	if model_name[-1] == "8":
		_snake_case   						=		8
	# set labels if required
	if not base_model:
		_snake_case   						=		1_000
		_snake_case   						=		'huggingface/label-files'
		_snake_case   						=		'imagenet-1k-id2label.json'
		_snake_case   						=		json.load(open(hf_hub_download(__SCREAMING_SNAKE_CASE			,				__SCREAMING_SNAKE_CASE			,				repo_type='dataset'	)			,				'r'	)	)
		_snake_case   						=		{int(__SCREAMING_SNAKE_CASE	): 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   						=		384
		_snake_case   						=		1_536
		_snake_case   						=		12
		_snake_case   						=		6

	# load original model from torch hub
	_snake_case   						=		torch.hub.load('facebookresearch/dino:main'			,				__SCREAMING_SNAKE_CASE	)
	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_(__SCREAMING_SNAKE_CASE	)
	_snake_case   						=		create_rename_keys(__SCREAMING_SNAKE_CASE			,				base_model=__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	)

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

	# 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(__SCREAMING_SNAKE_CASE	)

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

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


if __name__ == "__main__":
					lowercase    :							Any      =					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)
					lowercase    :							Any      =					parser.parse_args()
					convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)



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


UpperCAmelCase	      = logging.get_logger(__name__)

UpperCAmelCase	      = {
    '''salesforce/blip2-opt-2.7b''': '''https://huggingface.co/salesforce/blip2-opt-2.7b/resolve/main/config.json''',
}







class        A_   (			__lowerCamelCase   ):




	'''simple docstring'''


	_UpperCamelCase     :     Optional[Any]							 =       """blip_2_vision_model"""





	def __init__(       self		,    snake_case=1408		,    snake_case=6144		,    snake_case=39		,    snake_case=16		,    snake_case=224		,    snake_case=14		,    snake_case="gelu"		,    snake_case=0.00_001		,    snake_case=0.0		,    snake_case=1E-10		,    snake_case=True		,    **snake_case		,    ):
		super().__init__(**snake_case       )

		lowercase       =    hidden_size
		lowercase       =    intermediate_size
		lowercase       =    num_hidden_layers
		lowercase       =    num_attention_heads
		lowercase       =    patch_size
		lowercase       =    image_size
		lowercase       =    initializer_range
		lowercase       =    attention_dropout
		lowercase       =    layer_norm_eps
		lowercase       =    hidden_act
		lowercase       =    qkv_bias





	@classmethod
	def 	SCREAMING_SNAKE_CASE__						(       cls		,    snake_case		,    **snake_case       ):
		cls._set_token_in_kwargs(snake_case       )

		lowercase  ,					lowercase       =    cls.get_config_dict(snake_case		,    **snake_case       )

		# get the vision config dict if we are loading from Blip2Config
		if config_dict.get('model_type'       ) == "blip-2":
			lowercase       =    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(snake_case		,    **snake_case       )







class        A_   (			__lowerCamelCase   ):




	'''simple docstring'''


	_UpperCamelCase     :     int							 =       """blip_2_qformer"""





	def __init__(       self		,    snake_case=3_0522		,    snake_case=768		,    snake_case=12		,    snake_case=12		,    snake_case=3072		,    snake_case="gelu"		,    snake_case=0.1		,    snake_case=0.1		,    snake_case=512		,    snake_case=0.02		,    snake_case=1E-12		,    snake_case=0		,    snake_case="absolute"		,    snake_case=2		,    snake_case=1408		,    **snake_case		,    ):
		super().__init__(pad_token_id=snake_case		,    **snake_case       )

		lowercase       =    vocab_size
		lowercase       =    hidden_size
		lowercase       =    num_hidden_layers
		lowercase       =    num_attention_heads
		lowercase       =    hidden_act
		lowercase       =    intermediate_size
		lowercase       =    hidden_dropout_prob
		lowercase       =    attention_probs_dropout_prob
		lowercase       =    max_position_embeddings
		lowercase       =    initializer_range
		lowercase       =    layer_norm_eps
		lowercase       =    position_embedding_type
		lowercase       =    cross_attention_frequency
		lowercase       =    encoder_hidden_size





	@classmethod
	def 	SCREAMING_SNAKE_CASE__						(       cls		,    snake_case		,    **snake_case       ):
		cls._set_token_in_kwargs(snake_case       )

		lowercase  ,					lowercase       =    cls.get_config_dict(snake_case		,    **snake_case       )

		# get the qformer config dict if we are loading from Blip2Config
		if config_dict.get('model_type'       ) == "blip-2":
			lowercase       =    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(snake_case		,    **snake_case       )







class        A_   (			__lowerCamelCase   ):




	'''simple docstring'''


	_UpperCamelCase     :     List[Any]							 =       """blip-2"""
	_UpperCamelCase     :     str							 =       True





	def __init__(       self		,    snake_case=None		,    snake_case=None		,    snake_case=None		,    snake_case=32		,    **snake_case       ):
		super().__init__(**snake_case       )

		if vision_config is None:
			lowercase       =    {}
			logger.info('vision_config is None. initializing the Blip2VisionConfig with default values.'       )

		if qformer_config is None:
			lowercase       =    {}
			logger.info('qformer_config is None. Initializing the Blip2QFormerConfig with default values.'       )

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

		lowercase       =    BlipaVisionConfig(**snake_case       )
		lowercase       =    BlipaQFormerConfig(**snake_case       )
		lowercase       =    text_config['model_type'] if 'model_type' in text_config else 'opt'
		lowercase       =    CONFIG_MAPPING[text_model_type](**snake_case       )

		lowercase       =    self.text_config.tie_word_embeddings
		lowercase       =    self.text_config.is_encoder_decoder

		lowercase       =    num_query_tokens
		lowercase       =    self.vision_config.hidden_size
		lowercase       =    self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
		lowercase       =    1.0
		lowercase       =    0.02





	@classmethod
	def 	SCREAMING_SNAKE_CASE__						(       cls		,    snake_case		,    snake_case		,    snake_case		,    **snake_case		,    ):
		return cls(
		    vision_config=vision_config.to_dict()		,    qformer_config=qformer_config.to_dict()		,    text_config=text_config.to_dict()		,    **snake_case		,    )





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    copy.deepcopy(self.__dict__       )
		lowercase       =    self.vision_config.to_dict()
		lowercase       =    self.qformer_config.to_dict()
		lowercase       =    self.text_config.to_dict()
		lowercase       =    self.__class__.model_type
		return output

