

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







def 						lowerCamelCase__      ( __A    :Tuple	,__A    :Union[str, Any]=False		):


       """simple docstring"""




       __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 						lowerCamelCase__      ( __A    :Optional[Any]	,__A    :Dict	,__A    :Dict=False		):


       """simple docstring"""




       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 						lowerCamelCase__      ( __A    :List[str]		):


       """simple docstring"""




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







def 						lowerCamelCase__      ( __A    :int	,__A    :Any	,__A    :Any		):


       """simple docstring"""




       __snake_case       	=							dct.pop(__A		)
       __snake_case       	=							val







def 						lowerCamelCase__      ( ):


       """simple docstring"""




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







@torch.no_grad()
def 						lowerCamelCase__      ( __A    :str	,__A    :Tuple	,__A    :Dict=True		):


       """simple docstring"""




       __snake_case       	=							ViTConfig()
       # patch_size
       if model_name[-1] == "8":
              __snake_case       	=							8
       # set labels if required
       if not base_model:
              __snake_case       	=							1_0_0_0
              __snake_case       	=							"""huggingface/label-files"""
              __snake_case       	=							"""imagenet-1k-id2label.json"""
              __snake_case       	=							json.load(open(hf_hub_download(__A	,__A	,repo_type="""dataset"""		)	,"""r"""		)		)
              __snake_case       	=							{int(__A		): 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_8_4
              __snake_case       	=							1_5_3_6
              __snake_case       	=							1_2
              __snake_case       	=							6

       # load original model from torch hub
       __snake_case       	=							torch.hub.load("""facebookresearch/dino:main"""	,__A		)
       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_(__A		)
       __snake_case       	=							create_rename_keys(__A	,base_model=__A		)
       for src, dest in rename_keys:
              rename_key(__A	,__A	,__A		)
       read_in_q_k_v(__A	,__A	,__A		)

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

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

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

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


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




import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import CLIPTokenizer, CLIPTokenizerFast
from transformers.models.clip.tokenization_clip import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
    from PIL import Image

    from transformers import CLIPImageProcessor, CLIPProcessor

@require_vision
class     __snake_case    (				unittest.TestCase  ):
       """simple docstring"""







       def 			a		(		self   )      ->							Optional[int]:






              """simple docstring"""


              __snake_case       	=							tempfile.mkdtemp()

              # fmt: off
              __snake_case       	=							["""l""", """o""", """w""", """e""", """r""", """s""", """t""", """i""", """d""", """n""", """lo""", """l</w>""", """w</w>""", """r</w>""", """t</w>""", """low</w>""", """er</w>""", """lowest</w>""", """newer</w>""", """wider""", """<unk>""", """<|startoftext|>""", """<|endoftext|>"""]
              # fmt: on
              __snake_case       	=							dict(zip(_UpperCamelCase      , range(len(_UpperCamelCase   )   )   )   )
              __snake_case       	=							["""#version: 0.2""", """l o""", """lo w</w>""", """e r</w>""", """"""]
              __snake_case       	=							{"""unk_token""": """<unk>"""}

              __snake_case       	=							os.path.join(self.tmpdirname      , VOCAB_FILES_NAMES["""vocab_file"""]   )
              __snake_case       	=							os.path.join(self.tmpdirname      , VOCAB_FILES_NAMES["""merges_file"""]   )
              with open(self.vocab_file      , """w"""      , encoding="""utf-8"""   ) as fp:
                     fp.write(json.dumps(_UpperCamelCase   ) + """\n"""   )
              with open(self.merges_file      , """w"""      , encoding="""utf-8"""   ) as fp:
                     fp.write("""\n""".join(_UpperCamelCase   )   )

              __snake_case       	=							{
                  """do_resize""": True,
                  """size""": 20,
                  """do_center_crop""": True,
                  """crop_size""": 18,
                  """do_normalize""": True,
                  """image_mean""": [0.4814_5466, 0.457_8275, 0.4082_1073],
                  """image_std""": [0.2686_2954, 0.2613_0258, 0.2757_7711],
              }
              __snake_case       	=							os.path.join(self.tmpdirname      , _UpperCamelCase   )
              with open(self.image_processor_file      , """w"""      , encoding="""utf-8"""   ) as fp:
                     json.dump(_UpperCamelCase      , _UpperCamelCase   )


       def 			a		(		self      , **_UpperCamelCase   )      ->							Optional[int]:






              """simple docstring"""


              return CLIPTokenizer.from_pretrained(self.tmpdirname      , **_UpperCamelCase   )


       def 			a		(		self      , **_UpperCamelCase   )      ->							Union[str, Any]:






              """simple docstring"""


              return CLIPTokenizerFast.from_pretrained(self.tmpdirname      , **_UpperCamelCase   )


       def 			a		(		self      , **_UpperCamelCase   )      ->							List[str]:






              """simple docstring"""


              return CLIPImageProcessor.from_pretrained(self.tmpdirname      , **_UpperCamelCase   )


       def 			a		(		self   )      ->							int:






              """simple docstring"""


              shutil.rmtree(self.tmpdirname   )


       def 			a		(		self   )      ->							Optional[Any]:






              """simple docstring"""


              __snake_case       	=							[np.random.randint(2_55      , size=(3, 30, 4_00)      , dtype=np.uinta   )]

              __snake_case       	=							[Image.fromarray(np.moveaxis(_UpperCamelCase      , 0      , -1   )   ) for x in image_inputs]

              return image_inputs


       def 			a		(		self   )      ->							List[Any]:






              """simple docstring"""


              __snake_case       	=							self.get_tokenizer()
              __snake_case       	=							self.get_rust_tokenizer()
              __snake_case       	=							self.get_image_processor()

              __snake_case       	=							CLIPProcessor(tokenizer=_UpperCamelCase      , image_processor=_UpperCamelCase   )
              processor_slow.save_pretrained(self.tmpdirname   )
              __snake_case       	=							CLIPProcessor.from_pretrained(self.tmpdirname      , use_fast=_UpperCamelCase   )

              __snake_case       	=							CLIPProcessor(tokenizer=_UpperCamelCase      , image_processor=_UpperCamelCase   )
              processor_fast.save_pretrained(self.tmpdirname   )
              __snake_case       	=							CLIPProcessor.from_pretrained(self.tmpdirname   )

              self.assertEqual(processor_slow.tokenizer.get_vocab()      , tokenizer_slow.get_vocab()   )
              self.assertEqual(processor_fast.tokenizer.get_vocab()      , tokenizer_fast.get_vocab()   )
              self.assertEqual(tokenizer_slow.get_vocab()      , tokenizer_fast.get_vocab()   )
              self.assertIsInstance(processor_slow.tokenizer      , _UpperCamelCase   )
              self.assertIsInstance(processor_fast.tokenizer      , _UpperCamelCase   )

              self.assertEqual(processor_slow.image_processor.to_json_string()      , image_processor.to_json_string()   )
              self.assertEqual(processor_fast.image_processor.to_json_string()      , image_processor.to_json_string()   )
              self.assertIsInstance(processor_slow.image_processor      , _UpperCamelCase   )
              self.assertIsInstance(processor_fast.image_processor      , _UpperCamelCase   )


       def 			a		(		self   )      ->							List[str]:






              """simple docstring"""


              __snake_case       	=							CLIPProcessor(tokenizer=self.get_tokenizer()      , image_processor=self.get_image_processor()   )
              processor.save_pretrained(self.tmpdirname   )

              __snake_case       	=							self.get_tokenizer(bos_token="""(BOS)"""      , eos_token="""(EOS)"""   )
              __snake_case       	=							self.get_image_processor(do_normalize=_UpperCamelCase      , padding_value=1.0   )

              __snake_case       	=							CLIPProcessor.from_pretrained(
                  self.tmpdirname      , bos_token="""(BOS)"""      , eos_token="""(EOS)"""      , do_normalize=_UpperCamelCase      , padding_value=1.0   )

              self.assertEqual(processor.tokenizer.get_vocab()      , tokenizer_add_kwargs.get_vocab()   )
              self.assertIsInstance(processor.tokenizer      , _UpperCamelCase   )

              self.assertEqual(processor.image_processor.to_json_string()      , image_processor_add_kwargs.to_json_string()   )
              self.assertIsInstance(processor.image_processor      , _UpperCamelCase   )


       def 			a		(		self   )      ->							Any:






              """simple docstring"""


              __snake_case       	=							self.get_image_processor()
              __snake_case       	=							self.get_tokenizer()

              __snake_case       	=							CLIPProcessor(tokenizer=_UpperCamelCase      , image_processor=_UpperCamelCase   )

              __snake_case       	=							self.prepare_image_inputs()

              __snake_case       	=							image_processor(_UpperCamelCase      , return_tensors="""np"""   )
              __snake_case       	=							processor(images=_UpperCamelCase      , return_tensors="""np"""   )

              for key in input_image_proc.keys():
                     self.assertAlmostEqual(input_image_proc[key].sum()      , input_processor[key].sum()      , delta=1E-2   )


       def 			a		(		self   )      ->							str:






              """simple docstring"""


              __snake_case       	=							self.get_image_processor()
              __snake_case       	=							self.get_tokenizer()

              __snake_case       	=							CLIPProcessor(tokenizer=_UpperCamelCase      , image_processor=_UpperCamelCase   )

              __snake_case       	=							"""lower newer"""

              __snake_case       	=							processor(text=_UpperCamelCase   )

              __snake_case       	=							tokenizer(_UpperCamelCase   )

              for key in encoded_tok.keys():
                     self.assertListEqual(encoded_tok[key]      , encoded_processor[key]   )


       def 			a		(		self   )      ->							str:






              """simple docstring"""


              __snake_case       	=							self.get_image_processor()
              __snake_case       	=							self.get_tokenizer()

              __snake_case       	=							CLIPProcessor(tokenizer=_UpperCamelCase      , image_processor=_UpperCamelCase   )

              __snake_case       	=							"""lower newer"""
              __snake_case       	=							self.prepare_image_inputs()

              __snake_case       	=							processor(text=_UpperCamelCase      , images=_UpperCamelCase   )

              self.assertListEqual(list(inputs.keys()   )      , ["""input_ids""", """attention_mask""", """pixel_values"""]   )

              # test if it raises when no input is passed
              with pytest.raises(_UpperCamelCase   ):
                     processor()


       def 			a		(		self   )      ->							Any:






              """simple docstring"""


              __snake_case       	=							self.get_image_processor()
              __snake_case       	=							self.get_tokenizer()

              __snake_case       	=							CLIPProcessor(tokenizer=_UpperCamelCase      , image_processor=_UpperCamelCase   )

              __snake_case       	=							[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

              __snake_case       	=							processor.batch_decode(_UpperCamelCase   )
              __snake_case       	=							tokenizer.batch_decode(_UpperCamelCase   )

              self.assertListEqual(_UpperCamelCase      , _UpperCamelCase   )






       def 			a		(		self   )      ->							int:






              """simple docstring"""


              __snake_case       	=							self.get_image_processor()
              __snake_case       	=							self.get_tokenizer()

              __snake_case       	=							CLIPProcessor(tokenizer=_UpperCamelCase      , image_processor=_UpperCamelCase   )

              __snake_case       	=							"""lower newer"""
              __snake_case       	=							self.prepare_image_inputs()

              __snake_case       	=							processor(text=_UpperCamelCase      , images=_UpperCamelCase   )

              self.assertListEqual(list(inputs.keys()   )      , processor.model_input_names   )

