'''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
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_     (		)    ->				Union[str, Any]:







 '''simple docstring'''
 _UpperCamelCase:     Optional[int]							= '''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
 _UpperCamelCase:     Tuple							= Image.open(requests.get(lowercase  ,					stream=lowercase		).raw		).convert('''RGB'''		)

 return image




def       lowerCAmelCase_     (		lowercase:							Tuple		)    ->				str:







 '''simple docstring'''
 _UpperCamelCase:     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_     (		lowercase:							Any  ,					lowercase:							str  ,					lowercase:							str		)    ->				Dict:







 '''simple docstring'''
 _UpperCamelCase:     Any							= dct.pop(lowercase		)
 _UpperCamelCase:     Optional[Any]							= val




def       lowerCAmelCase_     (		lowercase:							Union[str, Any]  ,					lowercase:							Tuple		)    ->				Optional[Any]:







 '''simple docstring'''
 for i in range(config.vision_config.num_hidden_layers		):
  # read in original q and v biases
  _UpperCamelCase:     Tuple							= state_dict.pop(F"""visual_encoder.blocks.{i}.attn.q_bias"""		)
  _UpperCamelCase:     Optional[int]							= state_dict.pop(F"""visual_encoder.blocks.{i}.attn.v_bias"""		)

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




def       lowerCAmelCase_     (		lowercase:							str  ,					lowercase:							Any		)    ->				Any:







 '''simple docstring'''
 _UpperCamelCase:     Dict							= 364 if '''coco''' in model_name else 224
 _UpperCamelCase:     Optional[Any]							= BlipaVisionConfig(image_size=lowercase		).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:     str							= OPTConfig.from_pretrained('''facebook/opt-2.7b'''  ,					eos_token_id=lowercase		).to_dict()
 elif "opt-6.7b" in model_name:
  _UpperCamelCase:     Union[str, Any]							= OPTConfig.from_pretrained('''facebook/opt-6.7b'''  ,					eos_token_id=lowercase		).to_dict()
 elif "t5-xl" in model_name:
  _UpperCamelCase:     str							= TaConfig.from_pretrained('''google/flan-t5-xl'''  ,					dense_act_fn='''gelu'''  ,					bos_token_id=1		).to_dict()
 elif "t5-xxl" in model_name:
  _UpperCamelCase:     Any							= TaConfig.from_pretrained('''google/flan-t5-xxl'''  ,					dense_act_fn='''gelu'''  ,					bos_token_id=1		).to_dict()

 _UpperCamelCase:     Union[str, Any]							= BlipaConfig(vision_config=lowercase  ,					text_config=lowercase		)

 return config, image_size




@torch.no_grad()
def       lowerCAmelCase_     (		lowercase:							str  ,					lowercase:							Tuple=None  ,					lowercase:							str=False		)    ->				Dict:







 '''simple docstring'''
 _UpperCamelCase:     List[Any]							= (
     AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''		)
     if '''opt''' in model_name
     else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''		)
 )
 _UpperCamelCase:     Tuple							= tokenizer('''\n'''  ,					add_special_tokens=lowercase		).input_ids[0]
 _UpperCamelCase:     Tuple							= get_blipa_config(lowercase  ,					eos_token_id=lowercase		)

 _UpperCamelCase:     List[Any]							= BlipaForConditionalGeneration(lowercase		).eval()

 _UpperCamelCase:     List[str]							= {
     '''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:     List[str]							= model_name_to_original[model_name]

 # load original model
 print('''Loading original model...'''		)
 _UpperCamelCase:     Dict							= '''cuda''' if torch.cuda.is_available() else '''cpu'''
 _UpperCamelCase:     Optional[Any]							= load_model_and_preprocess(
     name=lowercase  ,					model_type=lowercase  ,					is_eval=lowercase  ,					device=lowercase		)
 original_model.eval()
 print('''Done!'''		)

 # update state dict keys
 _UpperCamelCase:     Optional[int]							= original_model.state_dict()
 _UpperCamelCase:     Tuple							= create_rename_keys(lowercase		)
 for src, dest in rename_keys:
  rename_key(lowercase  ,					lowercase  ,					lowercase		)

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

 # read in qv biases
 read_in_q_v_bias(lowercase  ,					lowercase		)

 _UpperCamelCase:     Union[str, Any]							= hf_model.load_state_dict(lowercase  ,					strict=lowercase		)
 assert len(lowercase		) == 0
 assert unexpected_keys == ["qformer.embeddings.position_ids"]

 _UpperCamelCase:     Union[str, Any]							= load_demo_image()
 _UpperCamelCase:     str							= vis_processors['''eval'''](lowercase		).unsqueeze(0		).to(lowercase		)
 _UpperCamelCase:     Dict							= tokenizer(['''\n''']  ,					return_tensors='''pt'''		).input_ids.to(lowercase		)

 # create processor
 _UpperCamelCase:     Any							= BlipImageProcessor(
     size={'''height''': image_size, '''width''': image_size}  ,					image_mean=lowercase  ,					image_std=lowercase		)
 _UpperCamelCase:     Union[str, Any]							= BlipaProcessor(image_processor=lowercase  ,					tokenizer=lowercase		)
 _UpperCamelCase:     Any							= processor(images=lowercase  ,					return_tensors='''pt'''		).pixel_values.to(lowercase		)

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

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

 print('''Generating a caption...'''		)
 _UpperCamelCase:     Optional[Any]							= ''''''
 _UpperCamelCase:     List[Any]							= tokenizer(lowercase  ,					return_tensors='''pt'''		).input_ids.to(lowercase		)

 _UpperCamelCase:     List[Any]							= original_model.generate({'''image''': original_pixel_values}		)
 _UpperCamelCase:     Any							= hf_model.generate(
     lowercase  ,					lowercase  ,					do_sample=lowercase  ,					num_beams=5  ,					max_length=30  ,					min_length=1  ,					top_p=0.9  ,					repetition_penalty=1.0  ,					length_penalty=1.0  ,					temperature=1  ,					)
 print('''Original generation:'''  ,					lowercase		)
 _UpperCamelCase:     List[str]							= input_ids.shape[1]
 _UpperCamelCase:     int							= processor.batch_decode(outputs[:, prompt_length:]  ,					skip_special_tokens=lowercase		)
 _UpperCamelCase:     Tuple							= [text.strip() for text in output_text]
 print('''HF generation:'''  ,					lowercase		)

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

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


if __name__ == "__main__":
  UpperCAmelCase_			    =							argparse.ArgumentParser()
  UpperCAmelCase_			    =							[
      '''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''',
  )

  UpperCAmelCase_			    =							parser.parse_args()

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

from ..trainer import Trainer
from ..utils import logging


UpperCAmelCase_			    =							logging.get_logger(__name__)




class  __magic_name__  (		__a    ):





   """simple docstring"""




   def __init__(	self					:  List[Any]		,  _lowercase					:  int=None		,  **_lowercase					:  Optional[Any]							):



    """simple docstring"""



    warnings.warn(
        '''`SageMakerTrainer` is deprecated and will be removed in v5 of Transformers. You can use `Trainer` '''
        '''instead.'''		,  _lowercase		,  )
    super().__init__(args=_lowercase		,  **_lowercase							)