


import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers.testing_utils import require_vision
from transformers.utils import is_vision_available


if is_vision_available():
 from PIL import Image

 from transformers import (
     AutoProcessor,
     BertTokenizerFast,
     BlipImageProcessor,
     GPTaTokenizer,
     InstructBlipProcessor,
     PreTrainedTokenizerFast,
 )

@require_vision
class   A__ (   unittest.TestCase     ):







    """simple docstring"""

    def 		a_      (	self   ):
         snake_case					=    tempfile.mkdtemp()

         snake_case					=    BlipImageProcessor()
         snake_case					=    GPTaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-GPT2Model'''   )
         snake_case					=    BertTokenizerFast.from_pretrained('''hf-internal-testing/tiny-random-bert'''   )

         snake_case					=    InstructBlipProcessor(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE   )

         processor.save_pretrained(self.tmpdirname   )
    def 		a_      (	self      , **__snake_case   ):
         return AutoProcessor.from_pretrained(self.tmpdirname      , **_SCREAMING_SNAKE_CASE   ).tokenizer
    def 		a_      (	self      , **__snake_case   ):
         return AutoProcessor.from_pretrained(self.tmpdirname      , **_SCREAMING_SNAKE_CASE   ).image_processor
    def 		a_      (	self      , **__snake_case   ):
         return AutoProcessor.from_pretrained(self.tmpdirname      , **_SCREAMING_SNAKE_CASE   ).qformer_tokenizer
    def 		a_      (	self   ):
         shutil.rmtree(self.tmpdirname   )
    def 		a_      (	self   ):
         snake_case					=    [np.random.randint(2_5_5      , size=(3, 3_0, 4_0_0)      , dtype=np.uinta   )]

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

         return image_inputs
    def 		a_      (	self   ):
         snake_case					=    InstructBlipProcessor(
             tokenizer=self.get_tokenizer()      , image_processor=self.get_image_processor()      , qformer_tokenizer=self.get_qformer_tokenizer()      , )
         processor.save_pretrained(self.tmpdirname   )

         snake_case					=    self.get_tokenizer(bos_token='''(BOS)'''      , eos_token='''(EOS)'''   )
         snake_case					=    self.get_image_processor(do_normalize=_SCREAMING_SNAKE_CASE      , padding_value=1.0   )

         snake_case					=    InstructBlipProcessor.from_pretrained(
             self.tmpdirname      , bos_token='''(BOS)'''      , eos_token='''(EOS)'''      , do_normalize=_SCREAMING_SNAKE_CASE      , padding_value=1.0   )

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

         self.assertEqual(processor.image_processor.to_json_string()      , image_processor_add_kwargs.to_json_string()   )
         self.assertIsInstance(processor.image_processor      , _SCREAMING_SNAKE_CASE   )
         self.assertIsInstance(processor.qformer_tokenizer      , _SCREAMING_SNAKE_CASE   )
    def 		a_      (	self   ):
         snake_case					=    self.get_image_processor()
         snake_case					=    self.get_tokenizer()
         snake_case					=    self.get_qformer_tokenizer()

         snake_case					=    InstructBlipProcessor(
             tokenizer=_SCREAMING_SNAKE_CASE      , image_processor=_SCREAMING_SNAKE_CASE      , qformer_tokenizer=_SCREAMING_SNAKE_CASE   )

         snake_case					=    self.prepare_image_inputs()

         snake_case					=    image_processor(_SCREAMING_SNAKE_CASE      , return_tensors='''np'''   )
         snake_case					=    processor(images=_SCREAMING_SNAKE_CASE      , return_tensors='''np'''   )

         for key in input_feat_extract.keys():
              self.assertAlmostEqual(input_feat_extract[key].sum()      , input_processor[key].sum()      , delta=1E-2   )
    def 		a_      (	self   ):
         snake_case					=    self.get_image_processor()
         snake_case					=    self.get_tokenizer()
         snake_case					=    self.get_qformer_tokenizer()

         snake_case					=    InstructBlipProcessor(
             tokenizer=_SCREAMING_SNAKE_CASE      , image_processor=_SCREAMING_SNAKE_CASE      , qformer_tokenizer=_SCREAMING_SNAKE_CASE   )

         snake_case					=    "lower newer"

         snake_case					=    processor(text=_SCREAMING_SNAKE_CASE   )

         snake_case					=    tokenizer(_SCREAMING_SNAKE_CASE      , return_token_type_ids=_SCREAMING_SNAKE_CASE   )
         snake_case					=    qformer_tokenizer(_SCREAMING_SNAKE_CASE      , return_token_type_ids=_SCREAMING_SNAKE_CASE   )

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

         for key in encoded_tokens_qformer.keys():
              self.assertListEqual(encoded_tokens_qformer[key]      , encoded_processor['''qformer_''' + key]   )
    def 		a_      (	self   ):
         snake_case					=    self.get_image_processor()
         snake_case					=    self.get_tokenizer()
         snake_case					=    self.get_qformer_tokenizer()

         snake_case					=    InstructBlipProcessor(
             tokenizer=_SCREAMING_SNAKE_CASE      , image_processor=_SCREAMING_SNAKE_CASE      , qformer_tokenizer=_SCREAMING_SNAKE_CASE   )

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

         snake_case					=    processor(text=_SCREAMING_SNAKE_CASE      , images=_SCREAMING_SNAKE_CASE   )

         self.assertListEqual(
             list(inputs.keys()   )      , ['''input_ids''', '''attention_mask''', '''qformer_input_ids''', '''qformer_attention_mask''', '''pixel_values''']      , )

         # test if it raises when no input is passed
         with pytest.raises(_SCREAMING_SNAKE_CASE   ):
              processor()
    def 		a_      (	self   ):
         snake_case					=    self.get_image_processor()
         snake_case					=    self.get_tokenizer()
         snake_case					=    self.get_qformer_tokenizer()

         snake_case					=    InstructBlipProcessor(
             tokenizer=_SCREAMING_SNAKE_CASE      , image_processor=_SCREAMING_SNAKE_CASE      , qformer_tokenizer=_SCREAMING_SNAKE_CASE   )

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

         snake_case					=    processor.batch_decode(_SCREAMING_SNAKE_CASE   )
         snake_case					=    tokenizer.batch_decode(_SCREAMING_SNAKE_CASE   )

         self.assertListEqual(_SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE   )




    def 		a_      (	self   ):
         snake_case					=    self.get_image_processor()
         snake_case					=    self.get_tokenizer()
         snake_case					=    self.get_qformer_tokenizer()

         snake_case					=    InstructBlipProcessor(
             tokenizer=_SCREAMING_SNAKE_CASE      , image_processor=_SCREAMING_SNAKE_CASE      , qformer_tokenizer=_SCREAMING_SNAKE_CASE   )

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

         snake_case					=    processor(text=_SCREAMING_SNAKE_CASE      , images=_SCREAMING_SNAKE_CASE   )

         self.assertListEqual(
             list(inputs.keys()   )      , ['''input_ids''', '''attention_mask''', '''qformer_input_ids''', '''qformer_attention_mask''', '''pixel_values''']      , )








import argparse
import json

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

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_=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'''module.blocks.{i}.norm1.weight''', F'''vit.encoder.layer.{i}.layernorm_before.weight''')    )
          rename_keys.append((F'''module.blocks.{i}.norm1.bias''', F'''vit.encoder.layer.{i}.layernorm_before.bias''')    )
          rename_keys.append(
              (F'''module.blocks.{i}.attn.proj.weight''', F'''vit.encoder.layer.{i}.attention.output.dense.weight''')    )
          rename_keys.append((F'''module.blocks.{i}.attn.proj.bias''', F'''vit.encoder.layer.{i}.attention.output.dense.bias''')    )
          rename_keys.append((F'''module.blocks.{i}.norm2.weight''', F'''vit.encoder.layer.{i}.layernorm_after.weight''')    )
          rename_keys.append((F'''module.blocks.{i}.norm2.bias''', F'''vit.encoder.layer.{i}.layernorm_after.bias''')    )
          rename_keys.append((F'''module.blocks.{i}.mlp.fc1.weight''', F'''vit.encoder.layer.{i}.intermediate.dense.weight''')    )
          rename_keys.append((F'''module.blocks.{i}.mlp.fc1.bias''', F'''vit.encoder.layer.{i}.intermediate.dense.bias''')    )
          rename_keys.append((F'''module.blocks.{i}.mlp.fc2.weight''', F'''vit.encoder.layer.{i}.output.dense.weight''')    )
          rename_keys.append((F'''module.blocks.{i}.mlp.fc2.bias''', F'''vit.encoder.layer.{i}.output.dense.bias''')    )

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

     if base_model:
          # layernorm + pooler
          rename_keys.extend(
              [
                  ('''module.norm.weight''', '''layernorm.weight'''),
                  ('''module.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_							,UpperCamelCase_							,UpperCamelCase_=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'''module.blocks.{i}.attn.qkv.weight'''    )
          snake_case					=    state_dict.pop(F'''module.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_    ):



     """simple docstring"""






     snake_case					=    ['''head.weight''', '''head.bias''']
     for k in ignore_keys:
          state_dict.pop(UpperCamelCase_							,UpperCamelCase_    )
def       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    [
         '''module.fc.fc1.weight''',
         '''module.fc.fc1.bias''',
         '''module.fc.bn1.weight''',
         '''module.fc.bn1.bias''',
         '''module.fc.bn1.running_mean''',
         '''module.fc.bn1.running_var''',
         '''module.fc.bn1.num_batches_tracked''',
         '''module.fc.fc2.weight''',
         '''module.fc.fc2.bias''',
         '''module.fc.bn2.weight''',
         '''module.fc.bn2.bias''',
         '''module.fc.bn2.running_mean''',
         '''module.fc.bn2.running_var''',
         '''module.fc.bn2.num_batches_tracked''',
         '''module.fc.fc3.weight''',
         '''module.fc.fc3.bias''',
     ]
     for k in ignore_keys:
          state_dict.pop(UpperCamelCase_							,UpperCamelCase_    )
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    dct.pop(UpperCamelCase_    )
     snake_case					=    val
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    ViTMSNConfig()
     snake_case					=    10_00

     snake_case					=    '''datasets/huggingface/label-files'''
     snake_case					=    '''imagenet-1k-id2label.json'''
     snake_case					=    json.load(open(hf_hub_download(UpperCamelCase_							,UpperCamelCase_    )							,'''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()}

     if "s16" in checkpoint_url:
          snake_case					=    3_84
          snake_case					=    15_36
          snake_case					=    6
     elif "l16" in checkpoint_url:
          snake_case					=    10_24
          snake_case					=    40_96
          snake_case					=    24
          snake_case					=    16
          snake_case					=    0.1
     elif "b4" in checkpoint_url:
          snake_case					=    4
     elif "l7" in checkpoint_url:
          snake_case					=    7
          snake_case					=    10_24
          snake_case					=    40_96
          snake_case					=    24
          snake_case					=    16
          snake_case					=    0.1

     snake_case					=    ViTMSNModel(UpperCamelCase_    )

     snake_case					=    torch.hub.load_state_dict_from_url(UpperCamelCase_							,map_location='''cpu'''    )['''target_encoder''']

     snake_case					=    ViTImageProcessor(size=config.image_size    )

     remove_projection_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_							,base_model=UpperCamelCase_    )

     model.load_state_dict(UpperCamelCase_    )
     model.eval()

     snake_case					=    '''http://images.cocodataset.org/val2017/000000039769.jpg'''

     snake_case					=    Image.open(requests.get(UpperCamelCase_							,stream=UpperCamelCase_    ).raw    )
     snake_case					=    ViTImageProcessor(
         size=config.image_size							,image_mean=UpperCamelCase_							,image_std=UpperCamelCase_    )
     snake_case					=    image_processor(images=UpperCamelCase_							,return_tensors='''pt'''    )

     # forward pass
     torch.manual_seed(2    )
     snake_case					=    model(**UpperCamelCase_    )
     snake_case					=    outputs.last_hidden_state

     # The following Colab Notebook was used to generate these outputs:
     # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
     if "s16" in checkpoint_url:
          snake_case					=    torch.tensor([[-1.0915, -1.4876, -1.1809]]    )
     elif "b16" in checkpoint_url:
          snake_case					=    torch.tensor([[14.2889, -18.9045, 11.7281]]    )
     elif "l16" in checkpoint_url:
          snake_case					=    torch.tensor([[41.5028, -22.8681, 45.6475]]    )
     elif "b4" in checkpoint_url:
          snake_case					=    torch.tensor([[-4.3868, 5.2932, -0.4137]]    )
     else:
          snake_case					=    torch.tensor([[-0.1792, -0.6465, 2.4263]]    )

     # verify logits
     assert torch.allclose(last_hidden_state[:, 0, :3]							,UpperCamelCase_							,atol=1e-4    )

     print(F'''Saving model 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__":
 _SCREAMING_SNAKE_CASE      :			List[Any]      				=   argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--checkpoint_url",
     default="https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar",
     type=str,
     help="URL of the checkpoint you'd like to convert.",
 )
 parser.add_argument(
     "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
 )

 _SCREAMING_SNAKE_CASE      :			Any      				=   parser.parse_args()
 convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




