


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()
lowerCAmelCase       =							logging.get_logger(__name__)
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE=False     ):


		"""simple docstring"""




		lowercase__            =      []
		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"
				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'''),
				    ]     )

		return rename_keys
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE=False     ):


		"""simple docstring"""




		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 			_a (      SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		lowercase__            =      ['''head.weight''', '''head.bias''']
		for k in ignore_keys:
				state_dict.pop(SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     )
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		lowercase__            =      dct.pop(SCREAMING_SNAKE_CASE     )
		lowercase__            =      val
def 			_a (      ):


		"""simple docstring"""




		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 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE=True     ):


		"""simple docstring"""




		lowercase__            =      ViTConfig()
		# patch_size
		if model_name[-1] == "8":
				lowercase__            =      8
		# set labels if required
		if not base_model:
				lowercase__            =      10_00
				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()}
		# size of the architecture
		if model_name in ["dino_vits8", "dino_vits16"]:
				lowercase__            =      3_84
				lowercase__            =      15_36
				lowercase__            =      12
				lowercase__            =      6

		# load original model from torch hub
		lowercase__            =      torch.hub.load('''facebookresearch/dino:main'''						,				SCREAMING_SNAKE_CASE     )
		original_model.eval()

		# load state_dict of original model, remove and rename some keys
		lowercase__            =      original_model.state_dict()
		if base_model:
				remove_classification_head_(SCREAMING_SNAKE_CASE     )
		lowercase__            =      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:
				lowercase__            =      ViTModel(SCREAMING_SNAKE_CASE						,				add_pooling_layer=SCREAMING_SNAKE_CASE     ).eval()
		else:
				lowercase__            =      ViTForImageClassification(SCREAMING_SNAKE_CASE     ).eval()
		model.load_state_dict(SCREAMING_SNAKE_CASE     )

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

		if base_model:
				lowercase__            =      original_model(SCREAMING_SNAKE_CASE     )
				assert torch.allclose(SCREAMING_SNAKE_CASE						,				outputs.last_hidden_state[:, 0, :]						,				atol=1E-1     )
		else:
				lowercase__            =      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__":
			lowerCAmelCase       =							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)
			lowerCAmelCase       =							parser.parse_args()
			convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)






from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


lowerCAmelCase       =							logging.get_logger(__name__)
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		try:
				with open(SCREAMING_SNAKE_CASE						,				'''rb'''     ) as flax_state_f:
						lowercase__            =      from_bytes(SCREAMING_SNAKE_CASE						,				flax_state_f.read()     )
		except UnpicklingError as e:
				try:
						with open(SCREAMING_SNAKE_CASE     ) as f:
								if f.read().startswith('''version'''     ):
										raise OSError(
										    '''You seem to have cloned a repository without having git-lfs installed. Please'''
										    ''' install git-lfs and run `git lfs install` followed by `git lfs pull` in the'''
										    ''' folder you cloned.'''     )
								else:
										raise ValueError from e
				except (UnicodeDecodeError, ValueError):
						raise EnvironmentError(f'Unable to convert {model_file} to Flax deserializable object. '     )

		return load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     )
def 			_a (      SCREAMING_SNAKE_CASE						,				SCREAMING_SNAKE_CASE     ):


		"""simple docstring"""




		try:
				import torch  # noqa: F401
		except ImportError:
				logger.error(
				    '''Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see'''
				    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
				    ''' instructions.'''     )
				raise

		# check if we have bf16 weights
		lowercase__            =      flatten_dict(jax.tree_util.tree_map(lambda SCREAMING_SNAKE_CASE     : x.dtype == jnp.bfloataa						,				SCREAMING_SNAKE_CASE     )     ).values()
		if any(SCREAMING_SNAKE_CASE     ):
				# convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

				# and bf16 is not fully supported in PT yet.
				logger.warning(
				    '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
				    '''before loading those in PyTorch model.'''     )
				lowercase__            =      jax.tree_util.tree_map(
				    lambda SCREAMING_SNAKE_CASE     : params.astype(np.floataa     ) if params.dtype == jnp.bfloataa else params						,				SCREAMING_SNAKE_CASE     )

		lowercase__            =      ''''''

		lowercase__            =      flatten_dict(SCREAMING_SNAKE_CASE						,				sep='''.'''     )
		lowercase__            =      pt_model.state_dict()

		# keep track of unexpected & missing keys
		lowercase__            =      []
		lowercase__            =      set(pt_model_dict.keys()     )

		for flax_key_tuple, flax_tensor in flax_state_dict.items():
				lowercase__            =      flax_key_tuple.split('''.'''     )

				if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']
						lowercase__            =      jnp.transpose(SCREAMING_SNAKE_CASE						,				(3, 2, 0, 1)     )
				elif flax_key_tuple_array[-1] == "kernel":
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']
						lowercase__            =      flax_tensor.T
				elif flax_key_tuple_array[-1] == "scale":
						lowercase__            =      flax_key_tuple_array[:-1] + ['''weight''']

				if "time_embedding" not in flax_key_tuple_array:
						for i, flax_key_tuple_string in enumerate(SCREAMING_SNAKE_CASE     ):
								lowercase__            =      (
								    flax_key_tuple_string.replace('''_0'''						,				'''.0'''     )
								    .replace('''_1'''						,				'''.1'''     )
								    .replace('''_2'''						,				'''.2'''     )
								    .replace('''_3'''						,				'''.3'''     )
								    .replace('''_4'''						,				'''.4'''     )
								    .replace('''_5'''						,				'''.5'''     )
								    .replace('''_6'''						,				'''.6'''     )
								    .replace('''_7'''						,				'''.7'''     )
								    .replace('''_8'''						,				'''.8'''     )
								    .replace('''_9'''						,				'''.9'''     )
								)

				lowercase__            =      '''.'''.join(SCREAMING_SNAKE_CASE     )

				if flax_key in pt_model_dict:
						if flax_tensor.shape != pt_model_dict[flax_key].shape:
								raise ValueError(
								    f'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
								    f'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'     )
						else:
								# add weight to pytorch dict
								lowercase__            =      np.asarray(SCREAMING_SNAKE_CASE     ) if not isinstance(SCREAMING_SNAKE_CASE						,				np.ndarray     ) else flax_tensor
								lowercase__            =      torch.from_numpy(SCREAMING_SNAKE_CASE     )
								# remove from missing keys
								missing_keys.remove(SCREAMING_SNAKE_CASE     )
				else:
						# weight is not expected by PyTorch model
						unexpected_keys.append(SCREAMING_SNAKE_CASE     )

		pt_model.load_state_dict(SCREAMING_SNAKE_CASE     )

		# re-transform missing_keys to list
		lowercase__            =      list(SCREAMING_SNAKE_CASE     )

		if len(SCREAMING_SNAKE_CASE     ) > 0:
				logger.warning(
				    '''Some weights of the Flax model were not used when initializing the PyTorch model'''
				    f' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
				    f' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
				    ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
				    f' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
				    ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
				    ''' FlaxBertForSequenceClassification model).'''     )
		if len(SCREAMING_SNAKE_CASE     ) > 0:
				logger.warning(
				    f'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
				    f' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
				    ''' use it for predictions and inference.'''     )

		return pt_model


