




"""simple docstring"""






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_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD





def 							UpperCAmelCase		(		):
					_lowerCAmelCase:Any									=  '''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
					_lowerCAmelCase:int									=  Image.open(requests.get(snake_case      ,       stream=snake_case    ).raw    ).convert('''RGB'''    )

					return image





def 							UpperCAmelCase		(		snake_case			:					Dict    ):
					_lowerCAmelCase:Optional[int]									=  []
					# 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.embeddings.layernorm.weight''')    )
					rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.embeddings.layernorm.bias''')    )

					# fmt: on
					return rename_keys





def 							UpperCAmelCase		(		snake_case			:					List[str]      ,       snake_case			:					Optional[Any]      ,       snake_case			:					str    ):
					_lowerCAmelCase:int									=  dct.pop(snake_case    )
					_lowerCAmelCase:str									=  val





def 							UpperCAmelCase		(		snake_case			:					Any      ,       snake_case			:					Dict    ):
					for i in range(config.vision_config.num_hidden_layers    ):
										# read in original q and v biases
										_lowerCAmelCase:int									=  state_dict.pop(F'visual_encoder.blocks.{i}.attn.q_bias'    )
										_lowerCAmelCase:List[Any]									=  state_dict.pop(F'visual_encoder.blocks.{i}.attn.v_bias'    )

										# next, set bias in the state dict
										_lowerCAmelCase:List[Any]									=  torch.cat((q_bias, torch.zeros_like(snake_case      ,       requires_grad=snake_case    ), v_bias)    )
										_lowerCAmelCase:int									=  qkv_bias





def 							UpperCAmelCase		(		snake_case			:					Any    ):
					_lowerCAmelCase:List[str]									=  364 if '''coco''' in model_name else 224
					_lowerCAmelCase:Optional[Any]									=  InstructBlipVisionConfig(image_size=snake_case    ).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 "t5-xl" in model_name:
										_lowerCAmelCase:Any									=  TaConfig.from_pretrained('''google/flan-t5-xl'''      ,       dense_act_fn='''gelu'''      ,       bos_token_id=1    ).to_dict()
					elif "t5-xxl" in model_name:
										_lowerCAmelCase:Tuple									=  TaConfig.from_pretrained('''google/flan-t5-xxl'''      ,       dense_act_fn='''gelu'''      ,       bos_token_id=1    ).to_dict()
					elif "vicuna-7b" in model_name:
										_lowerCAmelCase:List[Any]									=  LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''      ,       vocab_size=32001    ).to_dict()
					elif "vicuna-13b" in model_name:
										_lowerCAmelCase:Union[str, Any]									=  LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''      ,       vocab_size=32001    ).to_dict()
					else:
										raise ValueError('''Model name not supported'''    )

					# the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
					_lowerCAmelCase:Dict									=  InstructBlipQFormerConfig(vocab_size=30523    ).to_dict()
					_lowerCAmelCase:Union[str, Any]									=  InstructBlipConfig(vision_config=snake_case      ,       text_config=snake_case      ,       qformer_config=snake_case    )

					return config, image_size





@torch.no_grad()
def 							UpperCAmelCase		(		snake_case			:					Any      ,       snake_case			:					int=None      ,       snake_case			:					Optional[Any]=False    ):
					_lowerCAmelCase:Any									=  AutoTokenizer.from_pretrained('''bert-base-uncased'''      ,       truncation_side='''left'''    )
					qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}    )

					if "t5" in model_name:
										_lowerCAmelCase:Dict									=  TaTokenizerFast.from_pretrained('''google/flan-t5-xl'''      ,       truncation_side='''left'''    )
					elif "vicuna" in model_name:
										# the following was used in the original implementation:
										# tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
										# tokenizer.add_special_tokens({"pad_token": "[PAD]"})
										# tokenizer.add_special_tokens({"bos_token": "</s>"})
										# tokenizer.add_special_tokens({"eos_token": "</s>"})
										# tokenizer.add_special_tokens({"unk_token": "</s>"})
										_lowerCAmelCase:Optional[int]									=  LlamaTokenizerFast.from_pretrained(
										    '''huggyllama/llama-7b'''      ,       truncation_side='''left'''      ,       bos_token='''</s>'''      ,       unk_token='''</s>'''    )
										tokenizer.add_special_tokens({'''pad_token''': '''[PAD]'''}    )

					_lowerCAmelCase							,   _lowerCAmelCase:Any									=  get_blipa_config(snake_case    )
					_lowerCAmelCase:Any									=  InstructBlipForConditionalGeneration(snake_case    ).eval()

					_lowerCAmelCase:Optional[int]									=  {
					    '''instructblip-vicuna-7b''': ('''blip2_vicuna_instruct''', '''vicuna7b'''),
					    '''instructblip-vicuna-13b''': ('''blip2_vicuna_instruct''', '''vicuna13b'''),
					    '''instructblip-flan-t5-xl''': ('''blip2_t5_instruct''', '''flant5xl'''),
					    '''instructblip-flan-t5-xxl''': ('''blip2_t5_instruct''', '''flant5xxl'''),
					}

					_lowerCAmelCase							,   _lowerCAmelCase:List[str]									=  model_name_to_original[model_name]

					# load original model
					print('''Loading original model...'''    )
					_lowerCAmelCase:str									=  '''cuda:1''' if torch.cuda.is_available() else '''cpu'''
					_lowerCAmelCase:int									=  '''cuda:2''' if torch.cuda.is_available() else '''cpu'''
					_lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase:str									=  load_model_and_preprocess(
					    name=snake_case      ,       model_type=snake_case      ,       is_eval=snake_case      ,       device=snake_case    )
					original_model.eval()
					print('''Done!'''    )

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

					# some keys can be renamed efficiently
					for key, val in state_dict.copy().items():
										_lowerCAmelCase:List[Any]									=  state_dict.pop(snake_case    )
										if key.startswith('''Qformer.bert'''    ):
															_lowerCAmelCase:Any									=  key.replace('''Qformer.bert'''      ,       '''qformer'''    )
										if "attention.self" in key:
															_lowerCAmelCase:Tuple									=  key.replace('''self'''      ,       '''attention'''    )
										if "llm_proj" in key:
															_lowerCAmelCase:Dict									=  key.replace('''llm_proj'''      ,       '''language_projection'''    )
										if "t5_proj" in key:
															_lowerCAmelCase:List[str]									=  key.replace('''t5_proj'''      ,       '''language_projection'''    )
										if key.startswith('''llm_model'''    ):
															_lowerCAmelCase:Union[str, Any]									=  key.replace('''llm_model'''      ,       '''language_model'''    )
										if key.startswith('''t5'''    ):
															_lowerCAmelCase:Optional[int]									=  key.replace('''t5'''      ,       '''language'''    )
										_lowerCAmelCase:Any									=  val

					# read in qv biases
					read_in_q_v_bias(snake_case      ,       snake_case    )

					# note: weights get loaded in torch.float32 by default
					hf_model.load_state_dict(snake_case      ,       strict=snake_case    )

					_lowerCAmelCase:str									=  load_demo_image()
					_lowerCAmelCase:Any									=  '''What is unusual about this image?'''

					# create processor
					_lowerCAmelCase:int									=  BlipImageProcessor(
					    size={'''height''': image_size, '''width''': image_size}      ,       image_mean=snake_case      ,       image_std=snake_case    )
					_lowerCAmelCase:int									=  InstructBlipProcessor(
					    image_processor=snake_case      ,       tokenizer=snake_case      ,       qformer_tokenizer=snake_case      ,       )
					_lowerCAmelCase:List[Any]									=  processor(images=snake_case      ,       text=snake_case      ,       return_tensors='''pt'''    ).to(snake_case    )

					# make sure processor creates exact same pixel values
					_lowerCAmelCase:Optional[int]									=  vis_processors['''eval'''](snake_case    ).unsqueeze(0    ).to(snake_case    )
					_lowerCAmelCase:Optional[int]									=  inputs.pixel_values
					assert torch.allclose(original_pixel_values.to(pixel_values.device    )      ,       snake_case    )

					original_model.to(snake_case    )
					hf_model.to(snake_case    )
					with torch.no_grad():
										if "vicuna" in model_name:
															_lowerCAmelCase:Dict									=  original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}    ).logits
															_lowerCAmelCase:Union[str, Any]									=  hf_model(**snake_case    ).logits
										else:
															_lowerCAmelCase:Dict									=  original_model(
															    {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}    ).logits
															_lowerCAmelCase:Optional[Any]									=  tokenizer('''\n'''      ,       return_tensors='''pt'''    ).input_ids.to(snake_case    )
															_lowerCAmelCase:Optional[int]									=  label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id      ,       -100    )
															_lowerCAmelCase:Union[str, Any]									=  hf_model(**snake_case      ,       labels=snake_case    ).logits

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

					# assert values
					assert original_logits.shape == logits.shape
					_lowerCAmelCase:str									=  1e-4 if '''vicuna''' in model_name else 1e-5
					assert torch.allclose(original_logits.to(logits.device    )      ,       snake_case      ,       atol=snake_case    )
					print('''Looks ok!'''    )

					print('''Generating with original model...'''    )
					_lowerCAmelCase:Dict									=  original_model.generate({'''image''': original_pixel_values, '''prompt''': prompt}      ,       num_beams=5    )

					# important: we need to cast the weights of the HF model to the appropriate type
					print('''Generating with HF model...'''    )
					_lowerCAmelCase:List[str]									=  hf_model.generate(
					    **snake_case      ,       do_sample=snake_case      ,       num_beams=5      ,       max_length=256      ,       min_length=1      ,       top_p=0.9      ,       repetition_penalty=1.5      ,       length_penalty=1.0      ,       temperature=1      ,       )
					if "vicuna" in model_name:
										# convert output id 0 to 2 (eos_token_id)
										# TODO add this in the generate method?
										_lowerCAmelCase:Any									=  2
					print('''Original generation:'''      ,       snake_case    )
					_lowerCAmelCase:str									=  processor.batch_decode(snake_case      ,       skip_special_tokens=snake_case    )
					_lowerCAmelCase:str									=  [text.strip() for text in output_text]
					print('''HF generation:'''      ,       snake_case    )

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

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


if __name__ == "__main__":
				UpperCamelCase__						=			argparse.ArgumentParser()
				UpperCamelCase__						=			[
				    '''instructblip-vicuna-7b''',
				    '''instructblip-vicuna-13b''',
				    '''instructblip-flan-t5-xl''',
				    '''instructblip-flan-t5-xxl''',
				]
				parser.add_argument(
				    '''--model_name''',
				    default='''instructblip-flan-t5-xl''',
				    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''',
				)

				UpperCamelCase__						=			parser.parse_args()

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












"""simple docstring"""






import unittest

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
				import jax.numpy as jnp

				from transformers.models.distilbert.modeling_flax_distilbert import (
				    FlaxDistilBertForMaskedLM,
				    FlaxDistilBertForMultipleChoice,
				    FlaxDistilBertForQuestionAnswering,
				    FlaxDistilBertForSequenceClassification,
				    FlaxDistilBertForTokenClassification,
				    FlaxDistilBertModel,
				)




class   a__			(       unittest.TestCase  ):





						def __init__(							self      :	Dict     ,a__      :	List[Any]     ,a__      :	Optional[int]=13     ,a__      :	Dict=7     ,a__      :	Any=True     ,a__      :	Tuple=True     ,a__      :	Union[str, Any]=True     ,a__      :	Dict=True     ,a__      :	Dict=99     ,a__      :	Optional[Any]=32     ,a__      :	Dict=5     ,a__      :	List[Any]=4     ,a__      :	Union[str, Any]=37     ,a__      :	str="gelu"     ,a__      :	Any=0.1     ,a__      :	Optional[int]=0.1     ,a__      :	Tuple=512     ,a__      :	Any=16     ,a__      :	Any=2     ,a__      :	Dict=0.02     ,a__      :	int=4     ,)						->       Union[str, Any]:




											"""simple docstring"""







											_lowerCAmelCase:Optional[int]									=  parent
											_lowerCAmelCase:List[str]									=  batch_size
											_lowerCAmelCase:Tuple									=  seq_length
											_lowerCAmelCase:List[Any]									=  is_training
											_lowerCAmelCase:Dict									=  use_attention_mask
											_lowerCAmelCase:Dict									=  use_token_type_ids
											_lowerCAmelCase:Any									=  use_labels
											_lowerCAmelCase:str									=  vocab_size
											_lowerCAmelCase:int									=  hidden_size
											_lowerCAmelCase:Union[str, Any]									=  num_hidden_layers
											_lowerCAmelCase:List[str]									=  num_attention_heads
											_lowerCAmelCase:Optional[int]									=  intermediate_size
											_lowerCAmelCase:str									=  hidden_act
											_lowerCAmelCase:Optional[int]									=  hidden_dropout_prob
											_lowerCAmelCase:Optional[int]									=  attention_probs_dropout_prob
											_lowerCAmelCase:Any									=  max_position_embeddings
											_lowerCAmelCase:Optional[int]									=  type_vocab_size
											_lowerCAmelCase:Optional[int]									=  type_sequence_label_size
											_lowerCAmelCase:Tuple									=  initializer_range
											_lowerCAmelCase:Dict									=  num_choices





						def       __UpperCamelCase     (							self      :	List[str])						->       int:




											"""simple docstring"""







											_lowerCAmelCase:int									=  ids_tensor([self.batch_size, self.seq_length]     ,self.vocab_size)

											_lowerCAmelCase:List[Any]									=  None
											if self.use_attention_mask:
																_lowerCAmelCase:str									=  random_attention_mask([self.batch_size, self.seq_length])

											_lowerCAmelCase:Dict									=  DistilBertConfig(
											    vocab_size=self.vocab_size     ,dim=self.hidden_size     ,n_layers=self.num_hidden_layers     ,n_heads=self.num_attention_heads     ,hidden_dim=self.intermediate_size     ,hidden_act=self.hidden_act     ,dropout=self.hidden_dropout_prob     ,attention_dropout=self.attention_probs_dropout_prob     ,max_position_embeddings=self.max_position_embeddings     ,initializer_range=self.initializer_range     ,tie_weights_=a__     ,)

											return config, input_ids, attention_mask







						def       __UpperCamelCase     (							self      :	Any)						->       Any:




											"""simple docstring"""







											_lowerCAmelCase:Optional[Any]									=  self.prepare_config_and_inputs()
											_lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase:Optional[Any]									=  config_and_inputs
											_lowerCAmelCase:Any									=  {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
											return config, inputs_dict




@require_flax
class   a__			(       UpperCamelCase_					,	unittest.TestCase  ):
						snake_case__             = (
						    (
						        FlaxDistilBertModel,
						        FlaxDistilBertForMaskedLM,
						        FlaxDistilBertForMultipleChoice,
						        FlaxDistilBertForQuestionAnswering,
						        FlaxDistilBertForSequenceClassification,
						        FlaxDistilBertForTokenClassification,
						        FlaxDistilBertForQuestionAnswering,
						    )
						    if is_flax_available()
						    else ()
						)





						def       __UpperCamelCase     (							self      :	str)						->       Dict:




											"""simple docstring"""







											_lowerCAmelCase:str									=  FlaxDistilBertModelTester(self)







						@slow
						def       __UpperCamelCase     (							self      :	List[str])						->       int:




											"""simple docstring"""







											for model_class_name in self.all_model_classes:
																_lowerCAmelCase:List[str]									=  model_class_name.from_pretrained('''distilbert-base-uncased''')
																_lowerCAmelCase:Tuple									=  model(np.ones((1, 1)))
																self.assertIsNotNone(a__)




@require_flax
class   a__			(       unittest.TestCase  ):





						@slow
						def       __UpperCamelCase     (							self      :	Union[str, Any])						->       Optional[Any]:




											"""simple docstring"""







											_lowerCAmelCase:Any									=  FlaxDistilBertModel.from_pretrained('''distilbert-base-uncased''')
											_lowerCAmelCase:Optional[int]									=  np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])
											_lowerCAmelCase:Tuple									=  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
											_lowerCAmelCase:List[str]									=  model(a__     ,attention_mask=a__)[0]
											_lowerCAmelCase:List[str]									=  (1, 11, 768)
											self.assertEqual(output.shape     ,a__)
											_lowerCAmelCase:int									=  np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]])

											self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]     ,a__     ,atol=1E-4))






