


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   ()->  Optional[Any]:




     """simple docstring"""

     snake_case_	=				'''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
     snake_case_	=				Image.open(requests.get(SCREAMING_SNAKE_CASE , stream=SCREAMING_SNAKE_CASE ).raw ).convert('''RGB''' )

     return image






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE )->  Any:




     """simple docstring"""

     snake_case_	=				[]
     # 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   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Tuple:




     """simple docstring"""

     snake_case_	=				dct.pop(SCREAMING_SNAKE_CASE )
     snake_case_	=				val






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Dict:




     """simple docstring"""

     for i in range(config.vision_config.num_hidden_layers ):
          # read in original q and v biases
          snake_case_	=				state_dict.pop(f'''visual_encoder.blocks.{i}.attn.q_bias''' )
          snake_case_	=				state_dict.pop(f'''visual_encoder.blocks.{i}.attn.v_bias''' )

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






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  List[str]:




     """simple docstring"""

     snake_case_	=				364 if '''coco''' in model_name else 224
     snake_case_	=				BlipaVisionConfig(image_size=SCREAMING_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 "opt-2.7b" in model_name:
          snake_case_	=				OPTConfig.from_pretrained('''facebook/opt-2.7b''' , eos_token_id=SCREAMING_SNAKE_CASE ).to_dict()
     elif "opt-6.7b" in model_name:
          snake_case_	=				OPTConfig.from_pretrained('''facebook/opt-6.7b''' , eos_token_id=SCREAMING_SNAKE_CASE ).to_dict()
     elif "t5-xl" in model_name:
          snake_case_	=				TaConfig.from_pretrained('''google/flan-t5-xl''' , dense_act_fn='''gelu''' , bos_token_id=1 ).to_dict()
     elif "t5-xxl" in model_name:
          snake_case_	=				TaConfig.from_pretrained('''google/flan-t5-xxl''' , dense_act_fn='''gelu''' , bos_token_id=1 ).to_dict()

     snake_case_	=				BlipaConfig(vision_config=SCREAMING_SNAKE_CASE , text_config=SCREAMING_SNAKE_CASE )

     return config, image_size






@torch.no_grad()
def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE=None , SCREAMING_SNAKE_CASE=False )->  List[Any]:




     """simple docstring"""

     snake_case_	=				(
         AutoTokenizer.from_pretrained('''facebook/opt-2.7b''' )
         if '''opt''' in model_name
         else AutoTokenizer.from_pretrained('''google/flan-t5-xl''' )
     )
     snake_case_	=				tokenizer('''\n''' , add_special_tokens=SCREAMING_SNAKE_CASE ).input_ids[0]
     snake_case_	,   snake_case_	=				get_blipa_config(SCREAMING_SNAKE_CASE , eos_token_id=SCREAMING_SNAKE_CASE )

     snake_case_	=				BlipaForConditionalGeneration(SCREAMING_SNAKE_CASE ).eval()

     snake_case_	=				{
         '''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'''),
     }

     snake_case_	,   snake_case_	=				model_name_to_original[model_name]

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

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

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

     # read in qv biases
     read_in_q_v_bias(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )

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

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

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

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

     original_model.to(SCREAMING_SNAKE_CASE )
     hf_model.to(SCREAMING_SNAKE_CASE )
     with torch.no_grad():
          if "opt" in model_name:
               snake_case_	=				original_model({'''image''': original_pixel_values, '''text_input''': ['''''']} ).logits
               snake_case_	=				hf_model(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE ).logits
          else:
               snake_case_	=				original_model(
                   {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']} ).logits
               snake_case_	=				input_ids.masked_fill(input_ids == tokenizer.pad_token_id , -100 )
               snake_case_	=				hf_model(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , labels=SCREAMING_SNAKE_CASE ).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":
          snake_case_	=				torch.tensor(
              [[-41.5850, -4.4_4_4_0, -8.9_9_2_2], [-47.4322, -5.9_1_4_3, -1.7_3_4_0]] , device=SCREAMING_SNAKE_CASE )
          assert torch.allclose(logits[0, :3, :3] , SCREAMING_SNAKE_CASE , atol=1E-4 )
     elif model_name == "blip2-flan-t5-xl-coco":
          snake_case_	=				torch.tensor(
              [[-57.0109, -9.8_9_6_7, -12.6280], [-68.6578, -12.7191, -10.5065]] , device=SCREAMING_SNAKE_CASE )
     else:
          # cast to same type
          snake_case_	=				logits.dtype
          assert torch.allclose(original_logits.to(SCREAMING_SNAKE_CASE ) , SCREAMING_SNAKE_CASE , atol=1E-2 )
     print('''Looks ok!''' )

     print('''Generating a caption...''' )
     snake_case_	=				''''''
     snake_case_	=				tokenizer(SCREAMING_SNAKE_CASE , return_tensors='''pt''' ).input_ids.to(SCREAMING_SNAKE_CASE )

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

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

     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)



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase      =   {"""configuration_xglm""": ["""XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """XGLMConfig"""]}

try:
   if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   ["""XGLMTokenizer"""]

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   ["""XGLMTokenizerFast"""]

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   [
       """XGLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """XGLMForCausalLM""",
       """XGLMModel""",
       """XGLMPreTrainedModel""",
   ]


try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   [
       """FlaxXGLMForCausalLM""",
       """FlaxXGLMModel""",
       """FlaxXGLMPreTrainedModel""",
   ]


try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   [
       """TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """TFXGLMForCausalLM""",
       """TFXGLMModel""",
       """TFXGLMPreTrainedModel""",
   ]


if TYPE_CHECKING:
   from .configuration_xglm import XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XGLMConfig

   try:
      if not is_sentencepiece_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_xglm import XGLMTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_xglm_fast import XGLMTokenizerFast

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_xglm import XGLM_PRETRAINED_MODEL_ARCHIVE_LIST, XGLMForCausalLM, XGLMModel, XGLMPreTrainedModel

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel, FlaxXGLMPreTrainedModel

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_xglm import (
          TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFXGLMForCausalLM,
          TFXGLMModel,
          TFXGLMPreTrainedModel,
      )


else:
   import sys

   UpperCAmelCase      =   _LazyModule(__name__, globals()["""__file__"""], _import_structure)