
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[int]:
     lowerCAmelCase_					:       Optional[int]    =		"""https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
     lowerCAmelCase_					:       Tuple    =		Image.open(requests.get(_lowercase					,   stream=_lowercase			).raw			).convert("""RGB"""			)

     return image






def 							_lowerCAmelCase (  _a   :					Dict			)	->			Tuple:
     lowerCAmelCase_					:       Dict    =		[]
     # 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 (  _a   :					Optional[Any]					,   _a   :					List[str]					,   _a   :					Tuple			)	->			List[Any]:
     lowerCAmelCase_					:       List[Any]    =		dct.pop(_lowercase			)
     lowerCAmelCase_					:       int    =		val






def 							_lowerCAmelCase (  _a   :					Optional[int]					,   _a   :					int			)	->			List[str]:







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

          # next, set bias in the state dict
          lowerCAmelCase_					:       Tuple    =		torch.cat((q_bias, torch.zeros_like(_lowercase					,   requires_grad=_lowercase			), v_bias)			)
          lowerCAmelCase_					:       Union[str, Any]    =		qkv_bias






def 							_lowerCAmelCase (  _a   :					Union[str, Any]					,   _a   :					Optional[Any]			)	->			Dict:
     lowerCAmelCase_					:       List[Any]    =		3_64 if """coco""" in model_name else 2_24
     lowerCAmelCase_					:       int    =		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:
          lowerCAmelCase_					:       Dict    =		OPTConfig.from_pretrained("""facebook/opt-2.7b"""					,   eos_token_id=_lowercase			).to_dict()
     elif "opt-6.7b" in model_name:
          lowerCAmelCase_					:       Optional[Any]    =		OPTConfig.from_pretrained("""facebook/opt-6.7b"""					,   eos_token_id=_lowercase			).to_dict()
     elif "t5-xl" in model_name:
          lowerCAmelCase_					:       List[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()

     lowerCAmelCase_					:       List[str]    =		BlipaConfig(vision_config=_lowercase					,   text_config=_lowercase			)

     return config, image_size






@torch.no_grad()
def 							_lowerCAmelCase (  _a   :					List[Any]					,   _a   :					int=None					,   _a   :					Optional[int]=False			)	->			Any:
     lowerCAmelCase_					:       List[str]    =		(
         AutoTokenizer.from_pretrained("""facebook/opt-2.7b"""			)
         if """opt""" in model_name
         else AutoTokenizer.from_pretrained("""google/flan-t5-xl"""			)
     )
     lowerCAmelCase_					:       str    =		tokenizer("""\n"""					,   add_special_tokens=_lowercase			).input_ids[0]
     lowerCAmelCase_   ,  lowerCAmelCase_					:       Dict    =		get_blipa_config(_lowercase					,   eos_token_id=_lowercase			)

     lowerCAmelCase_					:       Dict    =		BlipaForConditionalGeneration(_lowercase			).eval()

     lowerCAmelCase_					:       Union[str, 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"""),
     }

     lowerCAmelCase_   ,  lowerCAmelCase_					:       Union[str, Any]    =		model_name_to_original[model_name]

     # load original model
     print("""Loading original model..."""			)
     lowerCAmelCase_					:       Dict    =		"""cuda""" if torch.cuda.is_available() else """cpu"""
     lowerCAmelCase_   ,  lowerCAmelCase_   ,  lowerCAmelCase_					:       List[Any]    =		load_model_and_preprocess(
         name=_lowercase					,   model_type=_lowercase					,   is_eval=_lowercase					,   device=_lowercase			)
     original_model.eval()
     print("""Done!"""			)

     # update state dict keys
     lowerCAmelCase_					:       Tuple    =		original_model.state_dict()
     lowerCAmelCase_					:       Any    =		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():
          lowerCAmelCase_					:       Tuple    =		state_dict.pop(_lowercase			)
          if key.startswith("""Qformer.bert"""			):
               lowerCAmelCase_					:       List[str]    =		key.replace("""Qformer.bert"""					,   """qformer"""			)
          if "attention.self" in key:
               lowerCAmelCase_					:       Tuple    =		key.replace("""self"""					,   """attention"""			)
          if "opt_proj" in key:
               lowerCAmelCase_					:       Union[str, Any]    =		key.replace("""opt_proj"""					,   """language_projection"""			)
          if "t5_proj" in key:
               lowerCAmelCase_					:       int    =		key.replace("""t5_proj"""					,   """language_projection"""			)
          if key.startswith("""opt"""			):
               lowerCAmelCase_					:       Optional[int]    =		key.replace("""opt"""					,   """language"""			)
          if key.startswith("""t5"""			):
               lowerCAmelCase_					:       Dict    =		key.replace("""t5"""					,   """language"""			)
          lowerCAmelCase_					:       int    =		val

     # read in qv biases
     read_in_q_v_bias(_lowercase					,   _lowercase			)

     lowerCAmelCase_   ,  lowerCAmelCase_					:       List[Any]    =		hf_model.load_state_dict(_lowercase					,   strict=_lowercase			)
     assert len(_lowercase			) == 0
     assert unexpected_keys == ["qformer.embeddings.position_ids"]

     lowerCAmelCase_					:       List[str]    =		load_demo_image()
     lowerCAmelCase_					:       Optional[Any]    =		vis_processors["""eval"""](_lowercase			).unsqueeze(0			).to(_lowercase			)
     lowerCAmelCase_					:       Tuple    =		tokenizer(["""\n"""]					,   return_tensors="""pt"""			).input_ids.to(_lowercase			)

     # create processor
     lowerCAmelCase_					:       Dict    =		BlipImageProcessor(
         size={"""height""": image_size, """width""": image_size}					,   image_mean=_lowercase					,   image_std=_lowercase			)
     lowerCAmelCase_					:       Tuple    =		BlipaProcessor(image_processor=_lowercase					,   tokenizer=_lowercase			)
     lowerCAmelCase_					:       List[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:
               lowerCAmelCase_					:       Optional[Any]    =		original_model({"""image""": original_pixel_values, """text_input""": [""""""]}			).logits
               lowerCAmelCase_					:       Dict    =		hf_model(_lowercase					,   _lowercase			).logits
          else:
               lowerCAmelCase_					:       Tuple    =		original_model(
                   {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}			).logits
               lowerCAmelCase_					:       str    =		input_ids.masked_fill(input_ids == tokenizer.pad_token_id					,   -1_00			)
               lowerCAmelCase_					:       Dict    =		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":
          lowerCAmelCase_					:       List[str]    =		torch.tensor(
              [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]]					,   device=_lowercase			)
          assert torch.allclose(logits[0, :3, :3]					,   _lowercase					,   atol=1E-4			)
     elif model_name == "blip2-flan-t5-xl-coco":
          lowerCAmelCase_					:       int    =		torch.tensor(
              [[-57.0_109, -9.8_967, -12.6_280], [-68.6_578, -12.7_191, -10.5_065]]					,   device=_lowercase			)
     else:
          # cast to same type
          lowerCAmelCase_					:       int    =		logits.dtype
          assert torch.allclose(original_logits.to(_lowercase			)					,   _lowercase					,   atol=1E-2			)
     print("""Looks ok!"""			)

     print("""Generating a caption..."""			)
     lowerCAmelCase_					:       List[Any]    =		""""""
     lowerCAmelCase_					:       Tuple    =		tokenizer(_lowercase					,   return_tensors="""pt"""			).input_ids.to(_lowercase			)

     lowerCAmelCase_					:       str    =		original_model.generate({"""image""": original_pixel_values}			)
     lowerCAmelCase_					:       Dict    =		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			)
     lowerCAmelCase_					:       Optional[Any]    =		input_ids.shape[1]
     lowerCAmelCase_					:       Optional[Any]    =		processor.batch_decode(outputs[:, prompt_length:]					,   skip_special_tokens=_lowercase			)
     lowerCAmelCase_					:       Optional[Any]    =		[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_					:   List[Any] 							=						argparse.ArgumentParser()
      UpperCAmelCase_					:   List[Any] 							=						[
          """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_					:   int 							=						parser.parse_args()

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







import enum
import os
from hashlib import shaaaa
from typing import Optional

from .. import config
from .logging import get_logger


UpperCAmelCase_					:   int 							=						get_logger(__name__)





class      lowercase__    (							enum.Enum  ):
  __UpperCamelCase	     =  """all_checks"""
  __UpperCamelCase	     =  """basic_checks"""
  __UpperCamelCase	     =  """no_checks"""





class      lowercase__    (							__A  ):
  pass





class      lowercase__    (							__A  ):
  pass





class      lowercase__    (							__A  ):
  pass





class      lowercase__    (							__A  ):
  pass






def 							_lowerCAmelCase (  _a   :					Optional[dict]					,   _a   :					dict					,   _a   :					str=None			)	->			Optional[Any]:
     if expected_checksums is None:
          logger.info("""Unable to verify checksums."""			)
          return
     if len(set(_a			) - set(_a			)			) > 0:
          raise ExpectedMoreDownloadedFiles(str(set(_a			) - set(_a			)			)			)
     if len(set(_a			) - set(_a			)			) > 0:
          raise UnexpectedDownloadedFile(str(set(_a			) - set(_a			)			)			)
     lowerCAmelCase_					:       Union[str, Any]    =		[url for url in expected_checksums if expected_checksums[url] != recorded_checksums[url]]
     lowerCAmelCase_					:       Optional[Any]    =		""" for """ + verification_name if verification_name is not None else """"""
     if len(_a			) > 0:
          raise NonMatchingChecksumError(
              F'Checksums didn\'t match{for_verification_name}:\n'
              F'{bad_urls}\n'
              """Set `verification_mode='no_checks'` to skip checksums verification and ignore this error"""			)
     logger.info("""All the checksums matched successfully""" + for_verification_name			)





class      lowercase__    (							__A  ):
  pass





class      lowercase__    (							__A  ):
  pass





class      lowercase__    (							__A  ):
  pass





class      lowercase__    (							__A  ):
  pass






def 							_lowerCAmelCase (  _a   :					Optional[dict]					,   _a   :					dict			)	->			str:
     if expected_splits is None:
          logger.info("""Unable to verify splits sizes."""			)
          return
     if len(set(_a			) - set(_a			)			) > 0:
          raise ExpectedMoreSplits(str(set(_a			) - set(_a			)			)			)
     if len(set(_a			) - set(_a			)			) > 0:
          raise UnexpectedSplits(str(set(_a			) - set(_a			)			)			)
     lowerCAmelCase_					:       Optional[int]    =		[
         {"""expected""": expected_splits[name], """recorded""": recorded_splits[name]}
         for name in expected_splits
         if expected_splits[name].num_examples != recorded_splits[name].num_examples
     ]
     if len(_a			) > 0:
          raise NonMatchingSplitsSizesError(str(_a			)			)
     logger.info("""All the splits matched successfully."""			)






def 							_lowerCAmelCase (  _a   :					str					,   _a   :					bool = True			)	->			dict:
     if record_checksum:
          lowerCAmelCase_					:       int    =		shaaaa()
          with open(_a					,   """rb"""			) as f:
               for chunk in iter(lambda: f.read(1 << 20			)					,   B""""""			):
                    m.update(_a			)
               lowerCAmelCase_					:       List[str]    =		m.hexdigest()
     else:
          lowerCAmelCase_					:       Union[str, Any]    =		None
     return {"num_bytes": os.path.getsize(_a			), "checksum": checksum}







def 							_lowerCAmelCase (  _a   :					Tuple			)	->			Optional[int]:
     if dataset_size and config.IN_MEMORY_MAX_SIZE:
          return dataset_size < config.IN_MEMORY_MAX_SIZE
     else:
          return False
