


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        SCREAMING_SNAKE_CASE_						(			)	->		Any:

       """simple docstring"""

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

       return image







def        SCREAMING_SNAKE_CASE_						(			__A :			List[str]				)	->		Dict:

       """simple docstring"""

       a_					:		List[Any]							=  []
       # 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        SCREAMING_SNAKE_CASE_						(			__A :			List[str]  ,    __A :			Tuple  ,    __A :			Dict				)	->		Union[str, Any]:

       """simple docstring"""

       a_					:		List[str]							=  dct.pop(__A				)
       a_					:		Optional[Any]							=  val







def        SCREAMING_SNAKE_CASE_						(			__A :			str  ,    __A :			Optional[Any]				)	->		str:

       """simple docstring"""

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

              # next, set bias in the state dict
              a_					:		List[str]							=  torch.cat((q_bias, torch.zeros_like(__A  ,    requires_grad=__A				), v_bias)				)
              a_					:		Optional[int]							=  qkv_bias







def        SCREAMING_SNAKE_CASE_						(			__A :			Any  ,    __A :			List[str]				)	->		Any:

       """simple docstring"""

       a_					:		int							=  3_64 if 'coco' in model_name else 2_24
       a_					:		Union[str, Any]							=  BlipaVisionConfig(image_size=__A				).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:
              a_					:		Tuple							=  OPTConfig.from_pretrained('facebook/opt-2.7b'  ,    eos_token_id=__A				).to_dict()
       elif "opt-6.7b" in model_name:
              a_					:		str							=  OPTConfig.from_pretrained('facebook/opt-6.7b'  ,    eos_token_id=__A				).to_dict()
       elif "t5-xl" in model_name:
              a_					:		Optional[int]							=  TaConfig.from_pretrained('google/flan-t5-xl'  ,    dense_act_fn='gelu'  ,    bos_token_id=1				).to_dict()
       elif "t5-xxl" in model_name:
              a_					:		Any							=  TaConfig.from_pretrained('google/flan-t5-xxl'  ,    dense_act_fn='gelu'  ,    bos_token_id=1				).to_dict()

       a_					:		List[str]							=  BlipaConfig(vision_config=__A  ,    text_config=__A				)

       return config, image_size







@torch.no_grad()
def        SCREAMING_SNAKE_CASE_						(			__A :			Dict  ,    __A :			List[Any]=None  ,    __A :			Any=False				)	->		List[Any]:

       """simple docstring"""

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

       a_					:		Optional[int]							=  BlipaForConditionalGeneration(__A				).eval()

       a_					:		Dict							=  {
           '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'),
       }

       a_							,		a_					:		Dict							=  model_name_to_original[model_name]

       # load original model
       print('Loading original model...'				)
       a_					:		List[str]							=  'cuda' if torch.cuda.is_available() else 'cpu'
       a_							,		a_							,		a_					:		Union[str, Any]							=  load_model_and_preprocess(
           name=__A  ,    model_type=__A  ,    is_eval=__A  ,    device=__A				)
       original_model.eval()
       print('Done!'				)

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

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

       # read in qv biases
       read_in_q_v_bias(__A  ,    __A				)

       a_							,		a_					:		List[str]							=  hf_model.load_state_dict(__A  ,    strict=__A				)
       assert len(__A				) == 0
       assert unexpected_keys == ["qformer.embeddings.position_ids"]

       a_					:		Tuple							=  load_demo_image()
       a_					:		int							=  vis_processors['eval'](__A				).unsqueeze(0				).to(__A				)
       a_					:		str							=  tokenizer(['\n']  ,    return_tensors='pt'				).input_ids.to(__A				)

       # create processor
       a_					:		Any							=  BlipImageProcessor(
           size={'height': image_size, 'width': image_size}  ,    image_mean=__A  ,    image_std=__A				)
       a_					:		Optional[int]							=  BlipaProcessor(image_processor=__A  ,    tokenizer=__A				)
       a_					:		Tuple							=  processor(images=__A  ,    return_tensors='pt'				).pixel_values.to(__A				)

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

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

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

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

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

       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_      : Optional[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_      : Optional[Any]  					=       parser.parse_args()

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







'''simple docstring'''




from collections import UserDict
from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
							from PIL import Image

							from ..image_utils import load_image

if is_torch_available():
							from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING

if is_tf_available():
							from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
							from ..tf_utils import stable_softmax

a   :			Union[str, Any]					  =    logging.get_logger(__name__)

@add_end_docstrings(__magic_name__			)
class 			UpperCamelCase_ (							__magic_name__			):





	def __init__(  self		,		**A    )			->			List[str]:
								super().__init__(**A    )

								requires_backends(self		,		"""vision"""    )
								self.check_model_type(
								    TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
								    if self.framework == """tf"""
								    else MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING    )





	def __call__(  self		,		A		,		**A    )			->			Optional[Any]:
								return super().__call__(A		,		**A    )





	def  _lowercase(  self		,		**A    )			->			Optional[Any]:
								UpperCAmelCase     :		List[Any]							     =		{}
								if "candidate_labels" in kwargs:
															UpperCAmelCase     :		Dict							     =		kwargs["""candidate_labels"""]
								if "hypothesis_template" in kwargs:
															UpperCAmelCase     :		Optional[Any]							     =		kwargs["""hypothesis_template"""]

								return preprocess_params, {}, {}





	def  _lowercase(  self		,		A		,		A=None		,		A="This is a photo of {}."    )			->			Optional[Any]:
								UpperCAmelCase     :		int							     =		load_image(A    )
								UpperCAmelCase     :		List[str]							     =		self.image_processor(images=[image]		,		return_tensors=self.framework    )
								UpperCAmelCase     :		List[str]							     =		candidate_labels
								UpperCAmelCase     :		Tuple							     =		[hypothesis_template.format(A    ) for x in candidate_labels]
								UpperCAmelCase     :		Union[str, Any]							     =		self.tokenizer(A		,		return_tensors=self.framework		,		padding=A    )
								UpperCAmelCase     :		Union[str, Any]							     =		[text_inputs]
								return inputs





	def  _lowercase(  self		,		A    )			->			Optional[int]:
								UpperCAmelCase     :		List[Any]							     =		model_inputs.pop("""candidate_labels"""    )
								UpperCAmelCase     :		Optional[Any]							     =		model_inputs.pop("""text_inputs"""    )
								if isinstance(text_inputs[0]		,		A    ):
															UpperCAmelCase     :		Optional[Any]							     =		text_inputs[0]
								else:
															# Batching case.
															UpperCAmelCase     :		Any							     =		text_inputs[0][0]

								UpperCAmelCase     :		Dict							     =		self.model(**A		,		**A    )

								UpperCAmelCase     :		List[Any]							     =		{
								    """candidate_labels""": candidate_labels,
								    """logits""": outputs.logits_per_image,
								}
								return model_outputs






	def  _lowercase(  self		,		A    )			->			Union[str, Any]:
								UpperCAmelCase     :		int							     =		model_outputs.pop("""candidate_labels"""    )
								UpperCAmelCase     :		int							     =		model_outputs["""logits"""][0]
								if self.framework == "pt":
															UpperCAmelCase     :		Optional[int]							     =		logits.softmax(dim=-1    ).squeeze(-1    )
															UpperCAmelCase     :		Any							     =		probs.tolist()
															if not isinstance(A		,		A    ):
																						UpperCAmelCase     :		Any							     =		[scores]
								elif self.framework == "tf":
															UpperCAmelCase     :		List[str]							     =		stable_softmax(A		,		axis=-1    )
															UpperCAmelCase     :		Union[str, Any]							     =		probs.numpy().tolist()
								else:
															raise ValueError(f'''Unsupported framework: {self.framework}'''    )

								UpperCAmelCase     :		Any							     =		[
								    {"""score""": score, """label""": candidate_label}
								    for score, candidate_label in sorted(zip(A		,		A    )		,		key=lambda A    : -x[0]    )
								]
								return result

