






"""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     __lowercase					(    )					->					Union[str, Any]:
      __snake_case      :    int  			='''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
      __snake_case      :    Union[str, Any]  			=Image.open(requests.get(a       ,		stream=a ).raw ).convert('''RGB''' )

      return image



def     __lowercase					(    a :				Optional[int] )					->					Any:
      __snake_case      :    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.layernorm.weight''') )
      rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.layernorm.bias''') )

      # fmt: on
      return rename_keys



def     __lowercase					(    a :				str       ,		a :				Tuple       ,		a :				Optional[int] )					->					List[Any]:
      __snake_case      :    Union[str, Any]  			=dct.pop(a )
      __snake_case      :    Optional[Any]  			=val



def     __lowercase					(    a :				Optional[Any]       ,		a :				str )					->					Union[str, Any]:
      for i in range(config.vision_config.num_hidden_layers ):
            # read in original q and v biases
            __snake_case      :    Dict  			=state_dict.pop(f'''visual_encoder.blocks.{i}.attn.q_bias''' )
            __snake_case      :    Any  			=state_dict.pop(f'''visual_encoder.blocks.{i}.attn.v_bias''' )

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



def     __lowercase					(    a :				Optional[Any]       ,		a :				str )					->					Union[str, Any]:
      __snake_case      :    Optional[int]  			=364 if '''coco''' in model_name else 224
      __snake_case      :    List[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:
            __snake_case      :    Tuple  			=OPTConfig.from_pretrained('''facebook/opt-2.7b'''       ,		eos_token_id=a ).to_dict()
      elif "opt-6.7b" in model_name:
            __snake_case      :    Optional[Any]  			=OPTConfig.from_pretrained('''facebook/opt-6.7b'''       ,		eos_token_id=a ).to_dict()
      elif "t5-xl" in model_name:
            __snake_case      :    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:
            __snake_case      :    Union[str, Any]  			=TaConfig.from_pretrained('''google/flan-t5-xxl'''       ,		dense_act_fn='''gelu'''       ,		bos_token_id=1 ).to_dict()

      __snake_case      :    Optional[int]  			=BlipaConfig(vision_config=a       ,		text_config=a )

      return config, image_size



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

      __snake_case      :    Union[str, Any]  			=BlipaForConditionalGeneration(a ).eval()

      __snake_case      :    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'''),
      }

      __snake_case      :    List[Any]  			=model_name_to_original[model_name]

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

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

      # read in qv biases
      read_in_q_v_bias(a       ,		a )

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

      __snake_case      :    List[Any]  			=load_demo_image()
      __snake_case      :    Optional[int]  			=vis_processors['''eval'''](a ).unsqueeze(0 ).to(a )
      __snake_case      :    Dict  			=tokenizer(['''\n''']       ,		return_tensors='''pt''' ).input_ids.to(a )

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

      print('''Generating a caption...''' )
      __snake_case      :    Union[str, Any]  			=''''''
      __snake_case      :    Union[str, Any]  			=tokenizer(a       ,		return_tensors='''pt''' ).input_ids.to(a )

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

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










"""simple docstring"""





from typing import TYPE_CHECKING

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


UpperCamelCase_	:				Any										= {"""configuration_reformer""": ["""REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ReformerConfig"""]}

try:
    if not is_sentencepiece_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    UpperCamelCase_	:				List[str]										= ["""ReformerTokenizer"""]

try:
    if not is_tokenizers_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    UpperCamelCase_	:				List[str]										= ["""ReformerTokenizerFast"""]

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    UpperCamelCase_	:				Optional[int]										= [
        """REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """ReformerAttention""",
        """ReformerForMaskedLM""",
        """ReformerForQuestionAnswering""",
        """ReformerForSequenceClassification""",
        """ReformerLayer""",
        """ReformerModel""",
        """ReformerModelWithLMHead""",
        """ReformerPreTrainedModel""",
    ]


if TYPE_CHECKING:
    from .configuration_reformer import REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, ReformerConfig

    try:
        if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_reformer import ReformerTokenizer

    try:
        if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_reformer_fast import ReformerTokenizerFast

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_reformer import (
            REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
            ReformerAttention,
            ReformerForMaskedLM,
            ReformerForQuestionAnswering,
            ReformerForSequenceClassification,
            ReformerLayer,
            ReformerModel,
            ReformerModelWithLMHead,
            ReformerPreTrainedModel,
        )

else:
    import sys

    UpperCamelCase_	:				Tuple										= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


