



import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
       MODEL_MAPPING,
       LevitForImageClassification,
       LevitForImageClassificationWithTeacher,
       LevitModel,
   )
   from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import LevitImageProcessor



class      _UpperCAmelCase      (							A__  ):




     def 		snake_case_		(					self):
         A__					=       self.config_class(**self.inputs_dict)
         self.parent.assertTrue(hasattr(a__			,    '''hidden_sizes'''))
         self.parent.assertTrue(hasattr(a__			,    '''num_attention_heads'''))







class      _UpperCAmelCase      :




     def __init__(					self			,    a__			,    a__=1_3			,    a__=6_4			,    a__=3			,    a__=3			,    a__=2			,    a__=1			,    a__=1_6			,    a__=[1_2_8, 2_5_6, 3_8_4]			,    a__=[4, 6, 8]			,    a__=[2, 3, 4]			,    a__=[1_6, 1_6, 1_6]			,    a__=0			,    a__=[2, 2, 2]			,    a__=[2, 2, 2]			,    a__=0.0_2			,    a__=True			,    a__=True			,    a__=2			,    ):
         A__					=       parent
         A__					=       batch_size
         A__					=       image_size
         A__					=       num_channels
         A__					=       kernel_size
         A__					=       stride
         A__					=       padding
         A__					=       hidden_sizes
         A__					=       num_attention_heads
         A__					=       depths
         A__					=       key_dim
         A__					=       drop_path_rate
         A__					=       patch_size
         A__					=       attention_ratio
         A__					=       mlp_ratio
         A__					=       initializer_range
         A__					=       [
             ['''Subsample''', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
             ['''Subsample''', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
         ]
         A__					=       is_training
         A__					=       use_labels
         A__					=       num_labels
         A__					=       initializer_range




     def 		snake_case_		(					self):
         A__					=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

         A__					=       None
         if self.use_labels:
             A__					=       ids_tensor([self.batch_size]			,    self.num_labels)

         A__					=       self.get_config()
         return config, pixel_values, labels




     def 		snake_case_		(					self):
         return LevitConfig(
             image_size=self.image_size			,    num_channels=self.num_channels			,    kernel_size=self.kernel_size			,    stride=self.stride			,    padding=self.padding			,    patch_size=self.patch_size			,    hidden_sizes=self.hidden_sizes			,    num_attention_heads=self.num_attention_heads			,    depths=self.depths			,    key_dim=self.key_dim			,    drop_path_rate=self.drop_path_rate			,    mlp_ratio=self.mlp_ratio			,    attention_ratio=self.attention_ratio			,    initializer_range=self.initializer_range			,    down_ops=self.down_ops			,    )




     def 		snake_case_		(					self			,    a__			,    a__			,    a__):
         A__					=       LevitModel(config=a__)
         model.to(a__)
         model.eval()
         A__					=       model(a__)
         A__					=       (self.image_size, self.image_size)
         A__			,						A__					=       image_size[0], image_size[1]
         for _ in range(4):
             A__					=       floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1)
             A__					=       floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1)
         self.parent.assertEqual(
             result.last_hidden_state.shape			,    (self.batch_size, ceil(height / 4) * ceil(width / 4), self.hidden_sizes[-1])			,    )




     def 		snake_case_		(					self			,    a__			,    a__			,    a__):
         A__					=       self.num_labels
         A__					=       LevitForImageClassification(a__)
         model.to(a__)
         model.eval()
         A__					=       model(a__			,    labels=a__)
         self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.num_labels))




     def 		snake_case_		(					self):
         A__					=       self.prepare_config_and_inputs()
         A__			,						A__			,						A__					=       config_and_inputs
         A__					=       {'''pixel_values''': pixel_values}
         return config, inputs_dict







@require_torch
class      _UpperCAmelCase      (							A__      ,   A__      ,   unittest.TestCase  ):
     UpperCamelCase__		     =					(
         (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
         if is_torch_available()
         else ()
     )
     UpperCamelCase__		     =					(
         {
             '''feature-extraction''': LevitModel,
             '''image-classification''': (LevitForImageClassification, LevitForImageClassificationWithTeacher),
         }
         if is_torch_available()
         else {}
     )

     UpperCamelCase__		     =					False
     UpperCamelCase__		     =					False
     UpperCamelCase__		     =					False
     UpperCamelCase__		     =					False
     UpperCamelCase__		     =					False




     def 		snake_case_		(					self):
         A__					=       LevitModelTester(self)
         A__					=       ConfigTester(self			,    config_class=a__			,    has_text_modality=a__			,    hidden_size=3_7)




     def 		snake_case_		(					self):
         self.create_and_test_config_common_properties()
         self.config_tester.create_and_test_config_to_json_string()
         self.config_tester.create_and_test_config_to_json_file()
         self.config_tester.create_and_test_config_from_and_save_pretrained()
         self.config_tester.create_and_test_config_with_num_labels()
         self.config_tester.check_config_can_be_init_without_params()
         self.config_tester.check_config_arguments_init()




     def 		snake_case_		(					self):
         return




     @unittest.skip(reason='''Levit does not use inputs_embeds''')
     def 		snake_case_		(					self):
         pass




     @unittest.skip(reason='''Levit does not support input and output embeddings''')
     def 		snake_case_		(					self):
         pass




     @unittest.skip(reason='''Levit does not output attentions''')
     def 		snake_case_		(					self):
         pass




     def 		snake_case_		(					self):
         A__			,						A__					=       self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             A__					=       model_class(a__)
             A__					=       inspect.signature(model.forward)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             A__					=       [*signature.parameters.keys()]

             A__					=       ['''pixel_values''']
             self.assertListEqual(arg_names[:1]			,    a__)




     def 		snake_case_		(					self):
         def check_hidden_states_output(a__			,    a__			,    a__):
             A__					=       model_class(a__)
             model.to(a__)
             model.eval()

             with torch.no_grad():
                 A__					=       model(**self._prepare_for_class(a__			,    a__))

             A__					=       outputs.hidden_states

             A__					=       len(self.model_tester.depths) + 1
             self.assertEqual(len(a__)			,    a__)

             A__					=       (self.model_tester.image_size, self.model_tester.image_size)
             A__			,						A__					=       image_size[0], image_size[1]
             for _ in range(4):
                 A__					=       floor(
                     (
                         (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                         / self.model_tester.stride
                     )
                     + 1)
                 A__					=       floor(
                     (
                         (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                         / self.model_tester.stride
                     )
                     + 1)
             # verify the first hidden states (first block)
             self.assertListEqual(
                 list(hidden_states[0].shape[-2:])			,    [
                     height * width,
                     self.model_tester.hidden_sizes[0],
                 ]			,    )

         A__			,						A__					=       self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             A__					=       True
             check_hidden_states_output(a__			,    a__			,    a__)

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             A__					=       True

             check_hidden_states_output(a__			,    a__			,    a__)




     @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''')
     def 		snake_case_		(					self):
         pass




     def 		snake_case_		(					self			,    a__			,    a__			,    a__=False):
         A__					=       super()._prepare_for_class(a__			,    a__			,    return_labels=a__)

         if return_labels:
             if model_class.__name__ == "LevitForImageClassificationWithTeacher":
                 del inputs_dict["labels"]

         return inputs_dict




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__)




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*a__)




     def 		snake_case_		(					self):
         if not self.model_tester.is_training:
             return

         A__			,						A__					=       self.model_tester.prepare_config_and_inputs_for_common()
         A__					=       True

         for model_class in self.all_model_classes:
             # LevitForImageClassificationWithTeacher supports inference-only
             if (
                 model_class in get_values(a__)
                 or model_class.__name__ == "LevitForImageClassificationWithTeacher"
             ):
                 continue
             A__					=       model_class(a__)
             model.to(a__)
             model.train()
             A__					=       self._prepare_for_class(a__			,    a__			,    return_labels=a__)
             A__					=       model(**a__).loss
             loss.backward()




     def 		snake_case_		(					self):
         A__			,						A__					=       self.model_tester.prepare_config_and_inputs_for_common()
         if not self.model_tester.is_training:
             return

         A__					=       False
         A__					=       True

         for model_class in self.all_model_classes:
             if model_class in get_values(a__) or not model_class.supports_gradient_checkpointing:
                 continue
             # LevitForImageClassificationWithTeacher supports inference-only
             if model_class.__name__ == "LevitForImageClassificationWithTeacher":
                 continue
             A__					=       model_class(a__)
             model.gradient_checkpointing_enable()
             model.to(a__)
             model.train()
             A__					=       self._prepare_for_class(a__			,    a__			,    return_labels=a__)
             A__					=       model(**a__).loss
             loss.backward()




     def 		snake_case_		(					self):
         A__			,						A__					=       self.model_tester.prepare_config_and_inputs_for_common()

         A__					=       [
             {'''title''': '''multi_label_classification''', '''num_labels''': 2, '''dtype''': torch.float},
             {'''title''': '''single_label_classification''', '''num_labels''': 1, '''dtype''': torch.long},
             {'''title''': '''regression''', '''num_labels''': 1, '''dtype''': torch.float},
         ]

         for model_class in self.all_model_classes:
             if (
                 model_class
                 not in [
                     *get_values(a__),
                 ]
                 or model_class.__name__ == "LevitForImageClassificationWithTeacher"
             ):
                 continue

             for problem_type in problem_types:
                 with self.subTest(msg=F"Testing {model_class} with {problem_type['title']}"):
                     A__					=       problem_type['''title''']
                     A__					=       problem_type['''num_labels''']

                     A__					=       model_class(a__)
                     model.to(a__)
                     model.train()

                     A__					=       self._prepare_for_class(a__			,    a__			,    return_labels=a__)

                     if problem_type["num_labels"] > 1:
                         A__					=       inputs['''labels'''].unsqueeze(1).repeat(1			,    problem_type['''num_labels'''])

                     A__					=       inputs['''labels'''].to(problem_type['''dtype'''])

                     # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
                     # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
                     # they have the same size." which is a symptom something in wrong for the regression problem.
                     # See https://github.com/huggingface/transformers/issues/11780
                     with warnings.catch_warnings(record=a__) as warning_list:
                         A__					=       model(**a__).loss
                     for w in warning_list:
                         if "Using a target size that is different to the input size" in str(w.message):
                             raise ValueError(
                                 F"Something is going wrong in the regression problem: intercepted {w.message}")

                     loss.backward()




     @slow
     def 		snake_case_		(					self):
         for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             A__					=       LevitModel.from_pretrained(a__)
             self.assertIsNotNone(a__)







def 		lowerCAmelCase__     (       )->	str:
    A__					=       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						)
    return image



@require_torch
@require_vision
class      _UpperCAmelCase      (							unittest.TestCase  ):




     @cached_property
     def 		snake_case_		(					self):
         return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0])




     @slow
     def 		snake_case_		(					self):
         A__					=       LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]).to(
             a__)

         A__					=       self.default_image_processor
         A__					=       prepare_img()
         A__					=       image_processor(images=a__			,    return_tensors='''pt''').to(a__)

         # forward pass
         with torch.no_grad():
             A__					=       model(**a__)

         # verify the logits
         A__					=       torch.Size((1, 1_0_0_0))
         self.assertEqual(outputs.logits.shape			,    a__)

         A__					=       torch.tensor([1.0_4_4_8, -0.3_7_4_5, -1.8_3_1_7]).to(a__)

         self.assertTrue(torch.allclose(outputs.logits[0, :3]			,    a__			,    atol=1e-4))





import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD



def 		lowerCAmelCase__     (       )->	List[str]:
    A__					=       '''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
    A__					=       Image.open(requests.get(UpperCamelCase_			,							stream=UpperCamelCase_						).raw						).convert('''RGB'''						)

    return image



def 		lowerCAmelCase__     (       UpperCamelCase_					:							Dict						)->	Any:
    A__					=       []
    # fmt: off

    # vision encoder
    rename_keys.append(('''visual_encoder.cls_token''', '''vision_model.embeddings.class_embedding''')						)
    rename_keys.append(('''visual_encoder.pos_embed''', '''vision_model.embeddings.position_embedding''')						)
    rename_keys.append(('''visual_encoder.patch_embed.proj.weight''', '''vision_model.embeddings.patch_embedding.weight''')						)
    rename_keys.append(('''visual_encoder.patch_embed.proj.bias''', '''vision_model.embeddings.patch_embedding.bias''')						)
    rename_keys.append(('''ln_vision.weight''', '''vision_model.post_layernorm.weight''')						)
    rename_keys.append(('''ln_vision.bias''', '''vision_model.post_layernorm.bias''')						)

    for i in range(config.vision_config.num_hidden_layers						):
        rename_keys.append((f"visual_encoder.blocks.{i}.norm1.weight", f"vision_model.encoder.layers.{i}.layer_norm1.weight")						)
        rename_keys.append((f"visual_encoder.blocks.{i}.norm1.bias", f"vision_model.encoder.layers.{i}.layer_norm1.bias")						)
        rename_keys.append((f"visual_encoder.blocks.{i}.norm2.weight", f"vision_model.encoder.layers.{i}.layer_norm2.weight")						)
        rename_keys.append((f"visual_encoder.blocks.{i}.norm2.bias", f"vision_model.encoder.layers.{i}.layer_norm2.bias")						)
        rename_keys.append((f"visual_encoder.blocks.{i}.attn.qkv.weight", f"vision_model.encoder.layers.{i}.self_attn.qkv.weight")						)
        rename_keys.append((f"visual_encoder.blocks.{i}.attn.proj.weight", f"vision_model.encoder.layers.{i}.self_attn.projection.weight",)						)
        rename_keys.append((f"visual_encoder.blocks.{i}.attn.proj.bias", f"vision_model.encoder.layers.{i}.self_attn.projection.bias")						)
        rename_keys.append((f"visual_encoder.blocks.{i}.mlp.fc1.weight", f"vision_model.encoder.layers.{i}.mlp.fc1.weight")						)
        rename_keys.append((f"visual_encoder.blocks.{i}.mlp.fc1.bias", f"vision_model.encoder.layers.{i}.mlp.fc1.bias")						)
        rename_keys.append((f"visual_encoder.blocks.{i}.mlp.fc2.weight", f"vision_model.encoder.layers.{i}.mlp.fc2.weight")						)
        rename_keys.append((f"visual_encoder.blocks.{i}.mlp.fc2.bias", f"vision_model.encoder.layers.{i}.mlp.fc2.bias")						)

    # QFormer
    rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.weight''', '''qformer.layernorm.weight''')						)
    rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.layernorm.bias''')						)

    # fmt: on
    return rename_keys



def 		lowerCAmelCase__     (       UpperCamelCase_					:							int			,							UpperCamelCase_					:							Optional[int]			,							UpperCamelCase_					:							int						)->	List[Any]:
    A__					=       dct.pop(UpperCamelCase_						)
    A__					=       val



def 		lowerCAmelCase__     (       UpperCamelCase_					:							Union[str, Any]			,							UpperCamelCase_					:							Optional[int]						)->	Optional[int]:
    for i in range(config.vision_config.num_hidden_layers						):
        # read in original q and v biases
        A__					=       state_dict.pop(f"visual_encoder.blocks.{i}.attn.q_bias"						)
        A__					=       state_dict.pop(f"visual_encoder.blocks.{i}.attn.v_bias"						)

        # next, set bias in the state dict
        A__					=       torch.cat((q_bias, torch.zeros_like(UpperCamelCase_			,							requires_grad=UpperCamelCase_						), v_bias)						)
        A__					=       qkv_bias



def 		lowerCAmelCase__     (       UpperCamelCase_					:							List[Any]			,							UpperCamelCase_					:							List[Any]						)->	int:
    A__					=       3_6_4 if '''coco''' in model_name else 2_2_4
    A__					=       BlipaVisionConfig(image_size=UpperCamelCase_						).to_dict()

    # make sure the models have proper bos_token_id and eos_token_id set (important for generation)
    # seems like flan-T5 models don't have bos_token_id properly set?
    if "opt-2.7b" in model_name:
        A__					=       OPTConfig.from_pretrained('''facebook/opt-2.7b'''			,							eos_token_id=UpperCamelCase_						).to_dict()
    elif "opt-6.7b" in model_name:
        A__					=       OPTConfig.from_pretrained('''facebook/opt-6.7b'''			,							eos_token_id=UpperCamelCase_						).to_dict()
    elif "t5-xl" in model_name:
        A__					=       TaConfig.from_pretrained('''google/flan-t5-xl'''			,							dense_act_fn='''gelu'''			,							bos_token_id=1						).to_dict()
    elif "t5-xxl" in model_name:
        A__					=       TaConfig.from_pretrained('''google/flan-t5-xxl'''			,							dense_act_fn='''gelu'''			,							bos_token_id=1						).to_dict()

    A__					=       BlipaConfig(vision_config=UpperCamelCase_			,							text_config=UpperCamelCase_						)

    return config, image_size



@torch.no_grad()
def 		lowerCAmelCase__     (       UpperCamelCase_					:							Tuple			,							UpperCamelCase_					:							Optional[Any]=None			,							UpperCamelCase_					:							Any=False						)->	Optional[Any]:
    A__					=       (
        AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''						)
        if '''opt''' in model_name
        else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''						)
    )
    A__					=       tokenizer('''\n'''			,							add_special_tokens=UpperCamelCase_						).input_ids[0]
    A__			,						A__					=       get_blipa_config(UpperCamelCase_			,							eos_token_id=UpperCamelCase_						)

    A__					=       BlipaForConditionalGeneration(UpperCamelCase_						).eval()

    A__					=       {
        '''blip2-opt-2.7b''': ('''blip2_opt''', '''pretrain_opt2.7b'''),
        '''blip2-opt-6.7b''': ('''blip2_opt''', '''pretrain_opt6.7b'''),
        '''blip2-opt-2.7b-coco''': ('''blip2_opt''', '''caption_coco_opt2.7b'''),
        '''blip2-opt-6.7b-coco''': ('''blip2_opt''', '''caption_coco_opt6.7b'''),
        '''blip2-flan-t5-xl''': ('''blip2_t5''', '''pretrain_flant5xl'''),
        '''blip2-flan-t5-xl-coco''': ('''blip2_t5''', '''caption_coco_flant5xl'''),
        '''blip2-flan-t5-xxl''': ('''blip2_t5''', '''pretrain_flant5xxl'''),
    }

    A__			,						A__					=       model_name_to_original[model_name]

    # load original model
    print('''Loading original model...'''						)
    A__					=       '''cuda''' if torch.cuda.is_available() else '''cpu'''
    A__			,						A__			,						A__					=       load_model_and_preprocess(
        name=UpperCamelCase_			,							model_type=UpperCamelCase_			,							is_eval=UpperCamelCase_			,							device=UpperCamelCase_						)
    original_model.eval()
    print('''Done!'''						)

    # update state dict keys
    A__					=       original_model.state_dict()
    A__					=       create_rename_keys(UpperCamelCase_						)
    for src, dest in rename_keys:
        rename_key(UpperCamelCase_			,							UpperCamelCase_			,							UpperCamelCase_						)

    # some keys can be renamed efficiently
    for key, val in state_dict.copy().items():
        A__					=       state_dict.pop(UpperCamelCase_						)
        if key.startswith('''Qformer.bert'''						):
            A__					=       key.replace('''Qformer.bert'''			,							'''qformer'''						)
        if "attention.self" in key:
            A__					=       key.replace('''self'''			,							'''attention'''						)
        if "opt_proj" in key:
            A__					=       key.replace('''opt_proj'''			,							'''language_projection'''						)
        if "t5_proj" in key:
            A__					=       key.replace('''t5_proj'''			,							'''language_projection'''						)
        if key.startswith('''opt'''						):
            A__					=       key.replace('''opt'''			,							'''language'''						)
        if key.startswith('''t5'''						):
            A__					=       key.replace('''t5'''			,							'''language'''						)
        A__					=       val

    # read in qv biases
    read_in_q_v_bias(UpperCamelCase_			,							UpperCamelCase_						)

    A__			,						A__					=       hf_model.load_state_dict(UpperCamelCase_			,							strict=UpperCamelCase_						)
    assert len(UpperCamelCase_						) == 0
    assert unexpected_keys == ["qformer.embeddings.position_ids"]

    A__					=       load_demo_image()
    A__					=       vis_processors['''eval'''](UpperCamelCase_						).unsqueeze(0						).to(UpperCamelCase_						)
    A__					=       tokenizer(['''\n''']			,							return_tensors='''pt'''						).input_ids.to(UpperCamelCase_						)

    # create processor
    A__					=       BlipImageProcessor(
        size={'''height''': image_size, '''width''': image_size}			,							image_mean=UpperCamelCase_			,							image_std=UpperCamelCase_						)
    A__					=       BlipaProcessor(image_processor=UpperCamelCase_			,							tokenizer=UpperCamelCase_						)
    A__					=       processor(images=UpperCamelCase_			,							return_tensors='''pt'''						).pixel_values.to(UpperCamelCase_						)

    # make sure processor creates exact same pixel values
    assert torch.allclose(UpperCamelCase_			,							UpperCamelCase_						)

    original_model.to(UpperCamelCase_						)
    hf_model.to(UpperCamelCase_						)
    with torch.no_grad():
        if "opt" in model_name:
            A__					=       original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}						).logits
            A__					=       hf_model(UpperCamelCase_			,							UpperCamelCase_						).logits
        else:
            A__					=       original_model(
                {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}						).logits
            A__					=       input_ids.masked_fill(input_ids == tokenizer.pad_token_id			,							-1_0_0						)
            A__					=       hf_model(UpperCamelCase_			,							UpperCamelCase_			,							labels=UpperCamelCase_						).logits

    assert original_logits.shape == logits.shape
    print('''First values of original logits:'''			,							original_logits[0, :3, :3]						)
    print('''First values of HF logits:'''			,							logits[0, :3, :3]						)

    # assert values
    if model_name == "blip2-flan-t5-xl":
        A__					=       torch.tensor(
            [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]			,							device=UpperCamelCase_						)
        assert torch.allclose(logits[0, :3, :3]			,							UpperCamelCase_			,							atol=1E-4						)
    elif model_name == "blip2-flan-t5-xl-coco":
        A__					=       torch.tensor(
            [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]			,							device=UpperCamelCase_						)
    else:
        # cast to same type
        A__					=       logits.dtype
        assert torch.allclose(original_logits.to(UpperCamelCase_						)			,							UpperCamelCase_			,							atol=1E-2						)
    print('''Looks ok!'''						)

    print('''Generating a caption...'''						)
    A__					=       ''''''
    A__					=       tokenizer(UpperCamelCase_			,							return_tensors='''pt'''						).input_ids.to(UpperCamelCase_						)

    A__					=       original_model.generate({'''image''': original_pixel_values}						)
    A__					=       hf_model.generate(
        UpperCamelCase_			,							UpperCamelCase_			,							do_sample=UpperCamelCase_			,							num_beams=5			,							max_length=3_0			,							min_length=1			,							top_p=0.9			,							repetition_penalty=1.0			,							length_penalty=1.0			,							temperature=1			,							)
    print('''Original generation:'''			,							UpperCamelCase_						)
    A__					=       input_ids.shape[1]
    A__					=       processor.batch_decode(outputs[:, prompt_length:]			,							skip_special_tokens=UpperCamelCase_						)
    A__					=       [text.strip() for text in output_text]
    print('''HF generation:'''			,							UpperCamelCase_						)

    if pytorch_dump_folder_path is not None:
        processor.save_pretrained(UpperCamelCase_						)
        hf_model.save_pretrained(UpperCamelCase_						)

    if push_to_hub:
        processor.push_to_hub(f"nielsr/{model_name}"						)
        hf_model.push_to_hub(f"nielsr/{model_name}"						)


if __name__ == "__main__":
   _lowercase    	=		argparse.ArgumentParser()
   _lowercase    	=		[
       "blip2-opt-2.7b",
       "blip2-opt-6.7b",
       "blip2-opt-2.7b-coco",
       "blip2-opt-6.7b-coco",
       "blip2-flan-t5-xl",
       "blip2-flan-t5-xl-coco",
       "blip2-flan-t5-xxl",
   ]
   parser.add_argument(
       "--model_name",
       default="blip2-opt-2.7b",
       choices=choices,
       type=str,
       help="Path to hf config.json of model to convert",
   )
   parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
   parser.add_argument(
       "--push_to_hub",
       action="store_true",
       help="Whether to push the model and processor to the hub after converting",
   )

   _lowercase    	=		parser.parse_args()

   convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
