



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 				_a				(   unittest.TestCase					):

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		str							      =     tempfile.mkdtemp()

						# fmt: off
						UpperCAmelCase_:		Optional[int]							      =     ["""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
						UpperCAmelCase_:		Any							      =     dict(zip(SCREAMING_SNAKE_CASE_,   range(len(SCREAMING_SNAKE_CASE_					)					)					)					)
						UpperCAmelCase_:		str							      =     ["""#version: 0.2""", """l o""", """lo w</w>""", """e r</w>""", """"""]
						UpperCAmelCase_:		List[Any]							      =     {"""unk_token""": """<unk>"""}

						UpperCAmelCase_:		List[str]							      =     os.path.join(self.tmpdirname,   VOCAB_FILES_NAMES["""vocab_file"""]					)
						UpperCAmelCase_:		Dict							      =     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(SCREAMING_SNAKE_CASE_					) + """\n"""					)
						with open(self.merges_file,   """w""",   encoding="""utf-8"""					) as fp:
									fp.write("""\n""".join(SCREAMING_SNAKE_CASE_					)					)

						UpperCAmelCase_:		List[str]							      =     {
						    """do_resize""": True,
						    """size""": 20,
						    """do_center_crop""": True,
						    """crop_size""": 18,
						    """do_normalize""": True,
						    """image_mean""": [0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3],
						    """image_std""": [0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1],
						}
						UpperCAmelCase_:		str							      =     os.path.join(self.tmpdirname,   SCREAMING_SNAKE_CASE_					)
						with open(self.image_processor_file,   """w""",   encoding="""utf-8"""					) as fp:
									json.dump(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self,   **SCREAMING_SNAKE_CASE_					)       ->   List[str]:
						return CLIPTokenizer.from_pretrained(self.tmpdirname,   **SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self,   **SCREAMING_SNAKE_CASE_					)       ->   Optional[int]:
						return CLIPTokenizerFast.from_pretrained(self.tmpdirname,   **SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self,   **SCREAMING_SNAKE_CASE_					)       ->   Tuple:
						return CLIPImageProcessor.from_pretrained(self.tmpdirname,   **SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   int:
						shutil.rmtree(self.tmpdirname					)

			def      __snake_case       (self					)       ->   int:
						UpperCAmelCase_:		Tuple							      =     [np.random.randint(255,   size=(3, 30, 400),   dtype=np.uinta					)]

						UpperCAmelCase_:		List[str]							      =     [Image.fromarray(np.moveaxis(SCREAMING_SNAKE_CASE_,   0,   -1					)					) for x in image_inputs]

						return image_inputs

			def      __snake_case       (self					)       ->   List[str]:
						UpperCAmelCase_:		List[str]							      =     self.get_tokenizer()
						UpperCAmelCase_:		int							      =     self.get_rust_tokenizer()
						UpperCAmelCase_:		int							      =     self.get_image_processor()

						UpperCAmelCase_:		int							      =     CLIPProcessor(tokenizer=SCREAMING_SNAKE_CASE_,   image_processor=SCREAMING_SNAKE_CASE_					)
						processor_slow.save_pretrained(self.tmpdirname					)
						UpperCAmelCase_:		Union[str, Any]							      =     CLIPProcessor.from_pretrained(self.tmpdirname,   use_fast=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Any							      =     CLIPProcessor(tokenizer=SCREAMING_SNAKE_CASE_,   image_processor=SCREAMING_SNAKE_CASE_					)
						processor_fast.save_pretrained(self.tmpdirname					)
						UpperCAmelCase_:		Union[str, Any]							      =     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,   SCREAMING_SNAKE_CASE_					)
						self.assertIsInstance(processor_fast.tokenizer,   SCREAMING_SNAKE_CASE_					)

						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,   SCREAMING_SNAKE_CASE_					)
						self.assertIsInstance(processor_fast.image_processor,   SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Dict:
						UpperCAmelCase_:		Tuple							      =     CLIPProcessor(tokenizer=self.get_tokenizer(),   image_processor=self.get_image_processor()					)
						processor.save_pretrained(self.tmpdirname					)

						UpperCAmelCase_:		Any							      =     self.get_tokenizer(bos_token="""(BOS)""",   eos_token="""(EOS)"""					)
						UpperCAmelCase_:		Union[str, Any]							      =     self.get_image_processor(do_normalize=SCREAMING_SNAKE_CASE_,   padding_value=1.0					)

						UpperCAmelCase_:		List[str]							      =     CLIPProcessor.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_					)

			def      __snake_case       (self					)       ->   List[Any]:
						UpperCAmelCase_:		Optional[Any]							      =     self.get_image_processor()
						UpperCAmelCase_:		Dict							      =     self.get_tokenizer()

						UpperCAmelCase_:		List[Any]							      =     CLIPProcessor(tokenizer=SCREAMING_SNAKE_CASE_,   image_processor=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		List[str]							      =     self.prepare_image_inputs()

						UpperCAmelCase_:		Tuple							      =     image_processor(SCREAMING_SNAKE_CASE_,   return_tensors="""np"""					)
						UpperCAmelCase_:		Optional[Any]							      =     processor(images=SCREAMING_SNAKE_CASE_,   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      __snake_case       (self					)       ->   Tuple:
						UpperCAmelCase_:		List[str]							      =     self.get_image_processor()
						UpperCAmelCase_:		str							      =     self.get_tokenizer()

						UpperCAmelCase_:		Any							      =     CLIPProcessor(tokenizer=SCREAMING_SNAKE_CASE_,   image_processor=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Optional[Any]							      =     """lower newer"""

						UpperCAmelCase_:		Optional[Any]							      =     processor(text=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Dict							      =     tokenizer(SCREAMING_SNAKE_CASE_					)

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

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		Tuple							      =     self.get_image_processor()
						UpperCAmelCase_:		Dict							      =     self.get_tokenizer()

						UpperCAmelCase_:		List[Any]							      =     CLIPProcessor(tokenizer=SCREAMING_SNAKE_CASE_,   image_processor=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Any							      =     """lower newer"""
						UpperCAmelCase_:		List[str]							      =     self.prepare_image_inputs()

						UpperCAmelCase_:		Optional[int]							      =     processor(text=SCREAMING_SNAKE_CASE_,   images=SCREAMING_SNAKE_CASE_					)

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

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

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		Optional[Any]							      =     self.get_image_processor()
						UpperCAmelCase_:		Dict							      =     self.get_tokenizer()

						UpperCAmelCase_:		str							      =     CLIPProcessor(tokenizer=SCREAMING_SNAKE_CASE_,   image_processor=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		str							      =     [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

						UpperCAmelCase_:		Optional[Any]							      =     processor.batch_decode(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Any							      =     tokenizer.batch_decode(SCREAMING_SNAKE_CASE_					)

						self.assertListEqual(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		List[Any]							      =     self.get_image_processor()
						UpperCAmelCase_:		str							      =     self.get_tokenizer()

						UpperCAmelCase_:		Optional[int]							      =     CLIPProcessor(tokenizer=SCREAMING_SNAKE_CASE_,   image_processor=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Union[str, Any]							      =     """lower newer"""
						UpperCAmelCase_:		Any							      =     self.prepare_image_inputs()

						UpperCAmelCase_:		Dict							      =     processor(text=SCREAMING_SNAKE_CASE_,   images=SCREAMING_SNAKE_CASE_					)

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











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_	():





			"""simple docstring"""


			UpperCAmelCase_:		List[Any]							      =     """https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
			UpperCAmelCase_:		List[Any]							      =     Image.open(requests.get(lowerCAmelCase__		,						stream=lowerCAmelCase__  ).raw  ).convert("""RGB"""  )

			return image







def  lowerCAmelCase_	(lowerCAmelCase__:		Any  ):





			"""simple docstring"""


			UpperCAmelCase_:		List[str]							      =     []
			# 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_	(lowerCAmelCase__:		Optional[int]		,						lowerCAmelCase__:		Dict		,						lowerCAmelCase__:		str  ):





			"""simple docstring"""


			UpperCAmelCase_:		int							      =     dct.pop(lowerCAmelCase__  )
			UpperCAmelCase_:		Optional[int]							      =     val







def  lowerCAmelCase_	(lowerCAmelCase__:		List[str]		,						lowerCAmelCase__:		str  ):





			"""simple docstring"""


			for i in range(config.vision_config.num_hidden_layers  ):
						# read in original q and v biases
						UpperCAmelCase_:		Optional[Any]							      =     state_dict.pop(F'visual_encoder.blocks.{i}.attn.q_bias'  )
						UpperCAmelCase_:		List[str]							      =     state_dict.pop(F'visual_encoder.blocks.{i}.attn.v_bias'  )

						# next, set bias in the state dict
						UpperCAmelCase_:		Dict							      =     torch.cat((q_bias, torch.zeros_like(lowerCAmelCase__		,						requires_grad=lowerCAmelCase__  ), v_bias)  )
						UpperCAmelCase_:		Optional[Any]							      =     qkv_bias







def  lowerCAmelCase_	(lowerCAmelCase__:		Optional[Any]		,						lowerCAmelCase__:		Union[str, Any]  ):





			"""simple docstring"""


			UpperCAmelCase_:		List[str]							      =     3_6_4 if """coco""" in model_name else 2_2_4
			UpperCAmelCase_:		List[str]							      =     BlipaVisionConfig(image_size=lowerCAmelCase__  ).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:
						UpperCAmelCase_:		List[Any]							      =     OPTConfig.from_pretrained("""facebook/opt-2.7b"""		,						eos_token_id=lowerCAmelCase__  ).to_dict()
			elif "opt-6.7b" in model_name:
						UpperCAmelCase_:		Union[str, Any]							      =     OPTConfig.from_pretrained("""facebook/opt-6.7b"""		,						eos_token_id=lowerCAmelCase__  ).to_dict()
			elif "t5-xl" in model_name:
						UpperCAmelCase_:		Dict							      =     TaConfig.from_pretrained("""google/flan-t5-xl"""		,						dense_act_fn="""gelu"""		,						bos_token_id=1  ).to_dict()
			elif "t5-xxl" in model_name:
						UpperCAmelCase_:		List[Any]							      =     TaConfig.from_pretrained("""google/flan-t5-xxl"""		,						dense_act_fn="""gelu"""		,						bos_token_id=1  ).to_dict()

			UpperCAmelCase_:		Dict							      =     BlipaConfig(vision_config=lowerCAmelCase__		,						text_config=lowerCAmelCase__  )

			return config, image_size







@torch.no_grad()
def  lowerCAmelCase_	(lowerCAmelCase__:		str		,						lowerCAmelCase__:		Optional[Any]=None		,						lowerCAmelCase__:		Any=False  ):





			"""simple docstring"""


			UpperCAmelCase_:		Any							      =     (
			    AutoTokenizer.from_pretrained("""facebook/opt-2.7b"""  )
			    if """opt""" in model_name
			    else AutoTokenizer.from_pretrained("""google/flan-t5-xl"""  )
			)
			UpperCAmelCase_:		Dict							      =     tokenizer("""\n"""		,						add_special_tokens=lowerCAmelCase__  ).input_ids[0]
			UpperCAmelCase_      ,							UpperCAmelCase_:		Optional[Any]							      =     get_blipa_config(lowerCAmelCase__		,						eos_token_id=lowerCAmelCase__  )

			UpperCAmelCase_:		Tuple							      =     BlipaForConditionalGeneration(lowerCAmelCase__  ).eval()

			UpperCAmelCase_:		Optional[Any]							      =     {
			    """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"""),
			}

			UpperCAmelCase_      ,							UpperCAmelCase_:		List[Any]							      =     model_name_to_original[model_name]

			# load original model
			print("""Loading original model..."""  )
			UpperCAmelCase_:		Any							      =     """cuda""" if torch.cuda.is_available() else """cpu"""
			UpperCAmelCase_      ,							UpperCAmelCase_      ,							UpperCAmelCase_:		int							      =     load_model_and_preprocess(
			    name=lowerCAmelCase__		,						model_type=lowerCAmelCase__		,						is_eval=lowerCAmelCase__		,						device=lowerCAmelCase__  )
			original_model.eval()
			print("""Done!"""  )

			# update state dict keys
			UpperCAmelCase_:		List[Any]							      =     original_model.state_dict()
			UpperCAmelCase_:		Union[str, Any]							      =     create_rename_keys(lowerCAmelCase__  )
			for src, dest in rename_keys:
						rename_key(lowerCAmelCase__		,						lowerCAmelCase__		,						lowerCAmelCase__  )

			# some keys can be renamed efficiently
			for key, val in state_dict.copy().items():
						UpperCAmelCase_:		List[str]							      =     state_dict.pop(lowerCAmelCase__  )
						if key.startswith("""Qformer.bert"""  ):
									UpperCAmelCase_:		Optional[Any]							      =     key.replace("""Qformer.bert"""		,						"""qformer"""  )
						if "attention.self" in key:
									UpperCAmelCase_:		int							      =     key.replace("""self"""		,						"""attention"""  )
						if "opt_proj" in key:
									UpperCAmelCase_:		Optional[Any]							      =     key.replace("""opt_proj"""		,						"""language_projection"""  )
						if "t5_proj" in key:
									UpperCAmelCase_:		Optional[int]							      =     key.replace("""t5_proj"""		,						"""language_projection"""  )
						if key.startswith("""opt"""  ):
									UpperCAmelCase_:		List[Any]							      =     key.replace("""opt"""		,						"""language"""  )
						if key.startswith("""t5"""  ):
									UpperCAmelCase_:		str							      =     key.replace("""t5"""		,						"""language"""  )
						UpperCAmelCase_:		int							      =     val

			# read in qv biases
			read_in_q_v_bias(lowerCAmelCase__		,						lowerCAmelCase__  )

			UpperCAmelCase_      ,							UpperCAmelCase_:		int							      =     hf_model.load_state_dict(lowerCAmelCase__		,						strict=lowerCAmelCase__  )
			assert len(lowerCAmelCase__  ) == 0
			assert unexpected_keys == ["qformer.embeddings.position_ids"]

			UpperCAmelCase_:		str							      =     load_demo_image()
			UpperCAmelCase_:		str							      =     vis_processors["""eval"""](lowerCAmelCase__  ).unsqueeze(0  ).to(lowerCAmelCase__  )
			UpperCAmelCase_:		Tuple							      =     tokenizer(["""\n"""]		,						return_tensors="""pt"""  ).input_ids.to(lowerCAmelCase__  )

			# create processor
			UpperCAmelCase_:		List[str]							      =     BlipImageProcessor(
			    size={"""height""": image_size, """width""": image_size}		,						image_mean=lowerCAmelCase__		,						image_std=lowerCAmelCase__  )
			UpperCAmelCase_:		Dict							      =     BlipaProcessor(image_processor=lowerCAmelCase__		,						tokenizer=lowerCAmelCase__  )
			UpperCAmelCase_:		Union[str, Any]							      =     processor(images=lowerCAmelCase__		,						return_tensors="""pt"""  ).pixel_values.to(lowerCAmelCase__  )

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

			original_model.to(lowerCAmelCase__  )
			hf_model.to(lowerCAmelCase__  )
			with torch.no_grad():
						if "opt" in model_name:
									UpperCAmelCase_:		Tuple							      =     original_model({"""image""": original_pixel_values, """text_input""": [""""""]}  ).logits
									UpperCAmelCase_:		Tuple							      =     hf_model(lowerCAmelCase__		,						lowerCAmelCase__  ).logits
						else:
									UpperCAmelCase_:		List[str]							      =     original_model(
									    {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}  ).logits
									UpperCAmelCase_:		List[Any]							      =     input_ids.masked_fill(input_ids == tokenizer.pad_token_id		,						-1_0_0  )
									UpperCAmelCase_:		List[Any]							      =     hf_model(lowerCAmelCase__		,						lowerCAmelCase__		,						labels=lowerCAmelCase__  ).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":
						UpperCAmelCase_:		Tuple							      =     torch.tensor(
						    [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]		,						device=lowerCAmelCase__  )
						assert torch.allclose(logits[0, :3, :3]		,						lowerCAmelCase__		,						atol=1e-4  )
			elif model_name == "blip2-flan-t5-xl-coco":
						UpperCAmelCase_:		int							      =     torch.tensor(
						    [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]		,						device=lowerCAmelCase__  )
			else:
						# cast to same type
						UpperCAmelCase_:		Union[str, Any]							      =     logits.dtype
						assert torch.allclose(original_logits.to(lowerCAmelCase__  )		,						lowerCAmelCase__		,						atol=1e-2  )
			print("""Looks ok!"""  )

			print("""Generating a caption..."""  )
			UpperCAmelCase_:		Tuple							      =     """"""
			UpperCAmelCase_:		int							      =     tokenizer(lowerCAmelCase__		,						return_tensors="""pt"""  ).input_ids.to(lowerCAmelCase__  )

			UpperCAmelCase_:		Tuple							      =     original_model.generate({"""image""": original_pixel_values}  )
			UpperCAmelCase_:		Tuple							      =     hf_model.generate(
			    lowerCAmelCase__		,						lowerCAmelCase__		,						do_sample=lowerCAmelCase__		,						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:"""		,						lowerCAmelCase__  )
			UpperCAmelCase_:		str							      =     input_ids.shape[1]
			UpperCAmelCase_:		int							      =     processor.batch_decode(outputs[:, prompt_length:]		,						skip_special_tokens=lowerCAmelCase__  )
			UpperCAmelCase_:		Dict							      =     [text.strip() for text in output_text]
			print("""HF generation:"""		,						lowerCAmelCase__  )

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

			if push_to_hub:
						processor.push_to_hub(F'nielsr/{model_name}'  )
						hf_model.push_to_hub(F'nielsr/{model_name}'  )


if __name__ == "__main__":
	a      :  Union[str, Any]			     =			argparse.ArgumentParser()
	a      :  int			     =			[
	    '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',
	)

	a      :  int			     =			parser.parse_args()

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






